AES.sv

/************************************************************************
AES Decryption Core Logic

Dong Kai Wang, Fall 2017

For use with ECE 385 Experiment 9
University of Illinois ECE Department
************************************************************************/

module AES (
    input    logic CLK,
    input  logic RESET,
    input  logic AES_START,
    output logic AES_DONE,
    input  logic [127:0] AES_KEY,
    input  logic [127:0] AES_MSG_ENC,
    output logic [127:0] AES_MSG_DEC
);



    // Internal logic variables
    logic latch_prevent;
    logic [1407:0] KeySchedule;
    logic [127:0] mux_output; // input for temp_state
    logic [127:0] temp_state_out; // output from temp_state reggyboi
    logic [127:0] shiftyboi_out; // output from shiftyboi, goes into mux
    logic [127:0] subbyboi_out; // output from subbyboi, goes into mux - this is formed from concatenating all the parts
    logic [127:0] roundyboi_out; // output from roundyboi, goes into mux
    logic [127:0] mixyboi_out; // output from mixyboi, goes into mux - this is formed from concatenating all the parts after 4 states.
    logic [127:0] curr_roundkey; //input for roundyboi's key param

    logic [7:0] subbyboi_1_out, subbyboi_2_out, subbyboi_3_out, subbyboi_4_out, subbyboi_5_out, subbyboi_6_out, subbyboi_7_out,
                    subbyboi_8_out, subbyboi_9_out, subbyboi_10_out, subbyboi_11_out, subbyboi_12_out, subbyboi_13_out, subbyboi_14_out,
                    subbyboi_15_out, subbyboi_16_out;

    // InvMixColumn stuff
    logic [31:0]    InvMC_reg, // Gets assigned to the relevant bits from state based on current state
                        InvMC_out_1, InvMC_out_2, InvMC_out_3, InvMC_out_4, // Gets the current state's mixyboi_piece, these will all be
                                                                                             // concatenated afterwards.
                        mixyboi_piece;

    // Select signals for the function select mux or load enables for the InvMixColumns registers, all assigned in the state machine.
    logic Enc_msg_sel, InvARK_sel, InvSR_sel, InvSB_sel, InvMC_sel, InvMC_out_1_en, InvMC_out_2_en, InvMC_out_3_en, InvMC_out_4_en;

    // Counter declaration
    logic [4:0] Counter;
    logic [4:0] Next_counter;

    // * ALL MODULE DECLARATIONS OF THINGS LIKE FUNCTION CALLS ARE AT THE BOTTOM OF THIS MODULE *

    // Assign state register to output. This will make the output wrong till the algorithm completes obviously
    assign AES_MSG_DEC = temp_state_out;

    enum logic [6:0] {
                                // Waiting for start signal
                                Wait,

                                // Start signal has been pressed
                                KeyExp_1,
                                KeyExp_2,
                                KeyExp_3,
                                KeyExp_4,
                                KeyExp_5,
                                KeyExp_6,
                                KeyExp_7, // this oughta be enough lolol

                                // pre-loop add round key
                                InvARK_0,

                                // 9 times loop
                                InvSR_1,
                                InvSB_1_1,
                                InvSB_1_2,
                                InvARK_1,
                                InvMC_1_1,
                                InvMC_1_2,
                                InvMC_1_3,
                                InvMC_1_4,
                                InvMC_1_5,

                                InvSR_2,
                                InvSB_2_1,
                                InvSB_2_2,
                                InvARK_2,
                                InvMC_2_1,
                                InvMC_2_2,
                                InvMC_2_3,
                                InvMC_2_4,
                                InvMC_2_5,

                                InvSR_3,
                                InvSB_3_1,
                                InvSB_3_2,
                                InvARK_3,
                                InvMC_3_1,
                                InvMC_3_2,
                                InvMC_3_3,
                                InvMC_3_4,
                                InvMC_3_5,

                                InvSR_4,
                                InvSB_4_1,
                                InvSB_4_2,
                                InvARK_4,
                                InvMC_4_1,
                                InvMC_4_2,
                                InvMC_4_3,
                                InvMC_4_4,
                                InvMC_4_5,

                                InvSR_5,
                                InvSB_5_1,
                                InvSB_5_2,
                                InvARK_5,
                                InvMC_5_1,
                                InvMC_5_2,
                                InvMC_5_3,
                                InvMC_5_4,
                                InvMC_5_5,

                                InvSR_6,
                                InvSB_6_1,
                                InvSB_6_2,
                                InvARK_6,
                                InvMC_6_1,
                                InvMC_6_2,
                                InvMC_6_3,
                                InvMC_6_4,
                                InvMC_6_5,

                                InvSR_7,
                                InvSB_7_1,
                                InvSB_7_2,
                                InvARK_7,
                                InvMC_7_1,
                                InvMC_7_2,
                                InvMC_7_3,
                                InvMC_7_4,
                                InvMC_7_5,

                                InvSR_8,
                                InvSB_8_1,
                                InvSB_8_2,
                                InvARK_8,
                                InvMC_8_1,
                                InvMC_8_2,
                                InvMC_8_3,
                                InvMC_8_4,
                                InvMC_8_5,

                                InvSR_9,
                                InvSB_9_1,
                                InvSB_9_2,
                                InvARK_9,
                                InvMC_9_1,
                                InvMC_9_2,
                                InvMC_9_3,
                                InvMC_9_4,
                                InvMC_9_5,

                                // post-loop final steps
                                InvSR_Fin,
                                InvSB_Fin_1,
                                InvSB_Fin_2,
                                InvARK_Fin,

                                // Final states
                                Done_1,
                                Done_2
                          }   State, Next_state;   // Internal state logic

    always_ff @ (posedge CLK)
    begin
        if (RESET)
            begin
                State <= Wait;
                Counter <= 5'b0;
            end
        else
            begin
                State <= Next_state;
                Counter <= Next_counter;

            end
    end

    always_comb
    begin
        // Default next state is staying at current state
        Next_state = State;
        Next_counter = Counter;

        // Default controls signal values
        AES_DONE = 1'b0;
        InvARK_sel = 1'b0;
        InvSR_sel = 1'b0;
        InvSB_sel = 1'b0;
        InvMC_sel = 1'b0;
        Enc_msg_sel = 1'b0;
        latch_prevent = 1'b1;

        // Default register vals for InvMixColumns
        InvMC_out_1_en = 1'b0;
        InvMC_out_2_en = 1'b0;
        InvMC_out_3_en = 1'b0;
        InvMC_out_4_en = 1'b0;
        InvMC_reg = 32'b0;
        mixyboi_out = 128'b0;

        // Assign next state
        unique case (State)
            Wait    :
                if (AES_START)
                    begin
                        Next_state = KeyExp_1;
                        Next_counter = 4'b0;
                    end
            KeyExp_1  :
                Next_state = KeyExp_2;
            KeyExp_2  :
                Next_state = KeyExp_3;
            KeyExp_3  :
                Next_state = KeyExp_4;
            KeyExp_4  :
                Next_state = KeyExp_5;
            KeyExp_5     :
                Next_state = KeyExp_6;
            KeyExp_6     :
                Next_state = KeyExp_7;
            KeyExp_7  :
                Next_state = InvARK_0;

            InvARK_0  : // Counter should be 0 here.
                Next_state = InvSR_1;

            InvSR_1   :
                begin
                    Next_state = InvSB_1_1;
                    Next_counter = Counter + 1'b1; // necessary to have a counter for ARK
                end
            InvSB_1_1 :
                Next_state = InvSB_1_2;
            InvSB_1_2 :
                Next_state = InvARK_1;
            InvARK_1  :
                Next_state = InvMC_1_1;
            InvMC_1_1 :
                Next_state = InvMC_1_2;
            InvMC_1_2 :
                Next_state = InvMC_1_3;
            InvMC_1_3 :
                Next_state = InvMC_1_4;
            InvMC_1_4 :
                Next_state = InvMC_1_5;
            InvMC_1_5 :
                Next_state = InvSR_2;

            InvSR_2   :
                begin
                    Next_state = InvSB_2_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_2_1 :
                Next_state = InvSB_2_2;
            InvSB_2_2 :
                Next_state = InvARK_2;
            InvARK_2  :
                Next_state = InvMC_2_1;
            InvMC_2_1 :
                Next_state = InvMC_2_2;
            InvMC_2_2 :
                Next_state = InvMC_2_3;
            InvMC_2_3 :
                Next_state = InvMC_2_4;
            InvMC_2_4 :
                Next_state = InvMC_2_5;
            InvMC_2_5 :
                Next_state = InvSR_3;

            InvSR_3   :
                begin
                    Next_state = InvSB_3_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_3_1 :
                Next_state = InvSB_3_2;
            InvSB_3_2 :
                Next_state = InvARK_3;
            InvARK_3  :
                Next_state = InvMC_3_1;
            InvMC_3_1 :
                Next_state = InvMC_3_2;
            InvMC_3_2 :
                Next_state = InvMC_3_3;
            InvMC_3_3 :
                Next_state = InvMC_3_4;
            InvMC_3_4 :
                Next_state = InvMC_3_5;
            InvMC_3_5 :
                Next_state = InvSR_4;

            InvSR_4   :
                begin
                    Next_state = InvSB_4_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_4_1 :
                Next_state = InvSB_4_2;
            InvSB_4_2 :
                Next_state = InvARK_4;
            InvARK_4  :
                Next_state = InvMC_4_1;
            InvMC_4_1 :
                Next_state = InvMC_4_2;
            InvMC_4_2 :
                Next_state = InvMC_4_3;
            InvMC_4_3 :
                Next_state = InvMC_4_4;
            InvMC_4_4 :
                Next_state = InvMC_4_5;
            InvMC_4_5 :
                Next_state = InvSR_5;

            InvSR_5   :
                begin
                    Next_state = InvSB_5_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_5_1   :
                Next_state = InvSB_5_2;
            InvSB_5_2   :
                Next_state = InvARK_5;
            InvARK_5  :
                Next_state = InvMC_5_1;
            InvMC_5_1 :
                Next_state = InvMC_5_2;
            InvMC_5_2 :
                Next_state = InvMC_5_3;
            InvMC_5_3 :
                Next_state = InvMC_5_4;
            InvMC_5_4 :
                Next_state = InvMC_5_5;
            InvMC_5_5 :
                Next_state = InvSR_6;

            InvSR_6   :
                begin
                    Next_state = InvSB_6_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_6_1 :
                Next_state = InvSB_6_2;
            InvSB_6_2 :
                Next_state = InvARK_6;
            InvARK_6  :
                Next_state = InvMC_6_1;
            InvMC_6_1 :
                Next_state = InvMC_6_2;
            InvMC_6_2 :
                Next_state = InvMC_6_3;
            InvMC_6_3 :
                Next_state = InvMC_6_4;
            InvMC_6_4 :
                Next_state = InvMC_6_5;
            InvMC_6_5 :
                Next_state = InvSR_7;

            InvSR_7   :
                begin
                    Next_state = InvSB_7_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_7_1 :
                Next_state = InvSB_7_2;
            InvSB_7_2 :
                Next_state = InvARK_7;
            InvARK_7  :
                Next_state = InvMC_7_1;
            InvMC_7_1 :
                Next_state = InvMC_7_2;
            InvMC_7_2 :
                Next_state = InvMC_7_3;
            InvMC_7_3 :
                Next_state = InvMC_7_4;
            InvMC_7_4 :
                Next_state = InvMC_7_5;
            InvMC_7_5 :
                Next_state = InvSR_8;

            InvSR_8   :
                begin
                    Next_state = InvSB_8_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_8_1 :
                Next_state = InvSB_8_2;
            InvSB_8_2 :
                Next_state = InvARK_8;
            InvARK_8  :
                Next_state = InvMC_8_1;
            InvMC_8_1 :
                Next_state = InvMC_8_2;
            InvMC_8_2 :
                Next_state = InvMC_8_3;
            InvMC_8_3 :
                Next_state = InvMC_8_4;
            InvMC_8_4 :
                Next_state = InvMC_8_5;
            InvMC_8_5 :
                Next_state = InvSR_9;

            InvSR_9   :
                begin
                    Next_state = InvSB_9_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_9_1   :
                Next_state = InvSB_9_2;
            InvSB_9_2   :
                Next_state = InvARK_9;
            InvARK_9  :
                Next_state = InvMC_9_1;
            InvMC_9_1 :
                Next_state = InvMC_9_2;
            InvMC_9_2 :
                Next_state = InvMC_9_3;
            InvMC_9_3 :
                Next_state = InvMC_9_4;
            InvMC_9_4 :
                Next_state = InvMC_9_5;
            InvMC_9_5 :
                Next_state = InvSR_Fin;

            InvSR_Fin :
                begin
                    Next_state = InvSB_Fin_1;
                    Next_counter = Counter + 1'b1;
                end
            InvSB_Fin_1 :
                Next_state = InvSB_Fin_2;
            InvSB_Fin_2 :
                Next_state = InvARK_Fin;
            InvARK_Fin:
                Next_state = Done_1;
            Done_1      :
                Next_state = Done_2;
            Done_2      :
                if (!AES_START)
                    Next_state = Wait;
            default : ;
        endcase

        case (State)
            Wait            :
                latch_prevent = 1'b1;

            KeyExp_1    :
                Enc_msg_sel = 1'b1;
            KeyExp_2    :
                latch_prevent = 1'b1;
            KeyExp_3    :
                latch_prevent = 1'b1;
            KeyExp_4    :
                latch_prevent = 1'b1;
            KeyExp_5    :
                latch_prevent = 1'b1;
            KeyExp_6        :
                latch_prevent = 1'b1;
            KeyExp_7        :
                latch_prevent = 1'b1;

            InvARK_0  :
                InvARK_sel = 1'b1;

            InvSR_1   :
                InvSR_sel = 1'b1;
            InvSB_1_1 :
                InvSB_sel = 1'b1;
            InvSB_1_2 :
                latch_prevent = 1'b1;
            InvARK_1  :
                InvARK_sel = 1'b1;
            InvMC_1_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_1_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_1_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_1_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_1_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_2   :
                InvSR_sel = 1'b1;
            InvSB_2_1 :
                InvSB_sel = 1'b1;
            InvSB_2_2 :
                latch_prevent = 1'b1;
            InvARK_2  :
                InvARK_sel = 1'b1;
            InvMC_2_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_2_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_2_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_2_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_2_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_3   :
                InvSR_sel = 1'b1;
            InvSB_3_1 :
                InvSB_sel = 1'b1;
            InvSB_3_2 :
                latch_prevent = 1'b1;
            InvARK_3  :
                InvARK_sel = 1'b1;
            InvMC_3_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_3_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_3_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_3_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_3_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_4   :
                InvSR_sel = 1'b1;
            InvSB_4_1 :
                InvSB_sel = 1'b1;
            InvSB_4_2 :
                latch_prevent = 1'b1;
            InvARK_4  :
                InvARK_sel = 1'b1;
            InvMC_4_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_4_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_4_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_4_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_4_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_5   :
                InvSR_sel = 1'b1;
            InvSB_5_1 :
                InvSB_sel = 1'b1;
            InvSB_5_2 :
                latch_prevent = 1'b1;
            InvARK_5  :
                InvARK_sel = 1'b1;
            InvMC_5_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_5_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_5_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_5_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_5_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_6   :
                InvSR_sel = 1'b1;
            InvSB_6_1 :
                InvSB_sel = 1'b1;
            InvSB_6_2 :
                latch_prevent = 1'b1;
            InvARK_6  :
                InvARK_sel = 1'b1;
            InvMC_6_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_6_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_6_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_6_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_6_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_7   :
                InvSR_sel = 1'b1;
            InvSB_7_1 :
                InvSB_sel = 1'b1;
            InvSB_7_2 :
                latch_prevent = 1'b1;
            InvARK_7  :
                InvARK_sel = 1'b1;
            InvMC_7_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_7_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_7_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_7_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_7_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_8   :
                InvSR_sel = 1'b1;
            InvSB_8_1 :
                InvSB_sel = 1'b1;
            InvSB_8_2 :
                latch_prevent = 1'b1;
            InvARK_8  :
                InvARK_sel = 1'b1;
            InvMC_8_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_8_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_8_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_8_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_8_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1};
                end

            InvSR_9   :
                InvSR_sel = 1'b1;
            InvSB_9_1 :
                InvSB_sel = 1'b1;
            InvSB_9_2 :
                latch_prevent = 1'b1;
            InvARK_9  :
                InvARK_sel = 1'b1;
            InvMC_9_1 :
                begin
                    InvMC_reg = temp_state_out[31:0];
                    InvMC_out_1_en = 1'b1;
                end
            InvMC_9_2 :
                begin
                    InvMC_reg = temp_state_out[63:32];
                    InvMC_out_2_en = 1'b1;
                end
            InvMC_9_3 :
                begin
                    InvMC_reg = temp_state_out[95:64];
                    InvMC_out_3_en = 1'b1;
                end
            InvMC_9_4 :
                begin
                    InvMC_reg = temp_state_out[127:96];
                    InvMC_out_4_en = 1'b1;
                end
            InvMC_9_5 :
                begin
                    InvMC_sel = 1'b1;
                    mixyboi_out = {InvMC_out_4, InvMC_out_3, InvMC_out_2, InvMC_out_1}; // concatenating it all together from most to least
                                                                                                              // significant!
                end

            InvSR_Fin :
                InvSR_sel = 1'b1;
            InvSB_Fin_1 :
                InvSB_sel = 1'b1;
            InvSB_Fin_2 :
                latch_prevent = 1'b1;
            InvARK_Fin:
                InvARK_sel = 1'b1;

            Done_1    :
                latch_prevent = 1'b1;
            Done_2   :
                AES_DONE = 1'b1;
        endcase
    end

    // Inverse Function Declarations
    InvShiftRows shiftyboi(.data_in(temp_state_out), .data_out(shiftyboi_out));

    InvMixColumns mixyboi(.in(InvMC_reg), .out(mixyboi_piece));

    InvAddRoundKey roundyboi(.inny_boi(temp_state_out), .key(curr_roundkey), .outty_boi(roundyboi_out));

    // This dumbass function only takes 8 bits at a time, need to pass all 128 of temp_state in so i guess i have to have 16 of these shits
    InvSubBytes subbyboi_1(.clk(CLK), .in(temp_state_out[7:0]), .out(subbyboi_1_out));
    InvSubBytes subbyboi_2(.clk(CLK), .in(temp_state_out[15:8]), .out(subbyboi_2_out));
    InvSubBytes subbyboi_3(.clk(CLK), .in(temp_state_out[23:16]), .out(subbyboi_3_out));
    InvSubBytes subbyboi_4(.clk(CLK), .in(temp_state_out[31:24]), .out(subbyboi_4_out));
    InvSubBytes subbyboi_5(.clk(CLK), .in(temp_state_out[39:32]), .out(subbyboi_5_out));
    InvSubBytes subbyboi_6(.clk(CLK), .in(temp_state_out[47:40]), .out(subbyboi_6_out));
    InvSubBytes subbyboi_7(.clk(CLK), .in(temp_state_out[55:48]), .out(subbyboi_7_out));
    InvSubBytes subbyboi_8(.clk(CLK), .in(temp_state_out[63:56]), .out(subbyboi_8_out));
    InvSubBytes subbyboi_9(.clk(CLK), .in(temp_state_out[71:64]), .out(subbyboi_9_out));
    InvSubBytes subbyboi_10(.clk(CLK), .in(temp_state_out[79:72]), .out(subbyboi_10_out));
    InvSubBytes subbyboi_11(.clk(CLK), .in(temp_state_out[87:80]), .out(subbyboi_11_out));
    InvSubBytes subbyboi_12(.clk(CLK), .in(temp_state_out[95:88]), .out(subbyboi_12_out));
    InvSubBytes subbyboi_13(.clk(CLK), .in(temp_state_out[103:96]), .out(subbyboi_13_out));
    InvSubBytes subbyboi_14(.clk(CLK), .in(temp_state_out[111:104]), .out(subbyboi_14_out));
    InvSubBytes subbyboi_15(.clk(CLK), .in(temp_state_out[119:112]), .out(subbyboi_15_out));
    InvSubBytes subbyboi_16(.clk(CLK), .in(temp_state_out[127:120]), .out(subbyboi_16_out));
    assign subbyboi_out = {subbyboi_16_out, subbyboi_15_out, subbyboi_14_out, subbyboi_13_out, subbyboi_12_out, subbyboi_11_out, subbyboi_10_out,
                                  subbyboi_9_out, subbyboi_8_out, subbyboi_7_out, subbyboi_6_out, subbyboi_5_out, subbyboi_4_out,
                                  subbyboi_3_out, subbyboi_2_out, subbyboi_1_out};

    // Key Expansion declaration
    KeyExpansion expandyboi(.clk(CLK), .Cipherkey(AES_KEY), .KeySchedule);

    // Mux for picking which function is the input to temp_state
    one_hot_6 data_reg_mux (
                            .Din1(roundyboi_out), // roundyboi output
                            .Din2(shiftyboi_out), // shiftyboi output
                            .Din3(subbyboi_out), // subbyboi output
                            .Din4(mixyboi_out), // mixyboi output
                            .Din5(AES_MSG_ENC), // original encoded message
                            .Din6(temp_state_out), // No function, give temp_state its old value.
                            .sel1(InvARK_sel),
                            .sel2(InvSR_sel),
                            .sel3(InvSB_sel),
                            .sel4(InvMC_sel),
                            .sel5(Enc_msg_sel),
                            .Dout(mux_output) // writes to state reg
                          );

    // Mux for picking the right roundkey for InvAddRoundKey
    mux11 roundkey_select (
                                        .Din1(KeySchedule[127:0]),
                                        .Din2(KeySchedule[255:128]),
                                        .Din3(KeySchedule[383:256]),
                                        .Din4(KeySchedule[511:384]),
                                        .Din5(KeySchedule[639:512]),
                                        .Din6(KeySchedule[767:640]),
                                        .Din7(KeySchedule[895:768]),
                                        .Din8(KeySchedule[1023:896]),
                                        .Din9(KeySchedule[1151:1024]),
                                        .Din10(KeySchedule[1279:1152]),
                                        .Din11(KeySchedule[1407:1280]),
                                        .Counter,
                                        .Dout(curr_roundkey)
                                        );

    // data register declaration
    reg_128 temp_state(
                            .Clk(CLK),
                            .Reset(RESET),
                            .D(mux_output),
                            .Dout(temp_state_out)
                            );

    // InvMC_reg 1-4 declarations here
    reg_32_ld mixcolreg1(
                                .Clk(CLK),
                                .Reset(RESET),
                                .Load(InvMC_out_1_en),
                                .D(mixyboi_piece),
                                .Dout(InvMC_out_1)
                                );
    reg_32_ld mixcolreg2(
                                .Clk(CLK),
                                .Reset(RESET),
                                .Load(InvMC_out_2_en),
                                .D(mixyboi_piece),
                                .Dout(InvMC_out_2)
                                );
    reg_32_ld mixcolreg3(
                                .Clk(CLK),
                                .Reset(RESET),
                                .Load(InvMC_out_3_en),
                                .D(mixyboi_piece),
                                .Dout(InvMC_out_3)
                                );
    reg_32_ld mixcolreg4(
                                .Clk(CLK),
                                .Reset(RESET),
                                .Load(InvMC_out_4_en),
                                .D(mixyboi_piece),
                                .Dout(InvMC_out_4)
                                );

endmodule


// OTHER MODULE DEFINITIONS

// one hot mux, used for picking current function
module one_hot_6 (input logic [127:0] Din1,
                     input logic [127:0] Din2,
                     input logic [127:0] Din3,
                     input logic [127:0] Din4,
                     input logic [127:0] Din5,
                     input logic [127:0] Din6, // this comes from the data register itself
                     input logic sel1, sel2, sel3, sel4, sel5,
                     output logic [127:0] Dout
                    );
    always_comb begin
        if (sel1)
            Dout = Din1;
        else if (sel2)
            Dout = Din2;
        else if (sel3)
            Dout = Din3;
        else if (sel4)
            Dout = Din4;
        else if (sel5)
            Dout = Din5;
        else
            Dout = Din6; //no function is being done, give the data register its old value.
    end
endmodule

// Used for the state register
module reg_128 (input  logic Clk, Reset,
              input  logic [127:0]  D,
              output logic [127:0]  Dout);

    always_ff @ (posedge Clk)
    begin
         if (Reset) // notice, this is a sycnrhonous reset, which is recommended on the FPGA
              Dout <= 128'b0;
         else
              Dout <= D;
    end
endmodule

// one hot mux with 11 signals, used for InvAddRoundKey's key parameter
module mux11 (
                         input logic [127:0] Din1,
                         input logic [127:0] Din2,
                         input logic [127:0] Din3,
                         input logic [127:0] Din4,
                         input logic [127:0] Din5,
                         input logic [127:0] Din6,
                         input logic [127:0] Din7,
                         input logic [127:0] Din8,
                         input logic [127:0] Din9,
                         input logic [127:0] Din10,
                         input logic [127:0] Din11,
                         input logic [4:0] Counter,
                         output logic [127:0] Dout
                      );
    always_comb begin
        case (Counter)
            4'd0 :
                Dout = Din1;
            4'd1 :
                Dout = Din2;
            4'd2 :
                Dout = Din3;
            4'd3 :
                Dout = Din4;
            4'd4 :
                Dout = Din5;
            4'd5 :
                Dout = Din6;
            4'd6 :
                Dout = Din7;
            4'd7 :
                Dout = Din8;
            4'd8 :
                Dout = Din9;
            4'd9 :
                Dout = Din10;
            4'd10 :
                Dout = Din11;
            default :
                Dout = 128'bX;
        endcase
    end
endmodule

// This other register is able to load, this is useful for the mixcolumns registers that need to hold their values after a state
module reg_32_ld (
                        input logic Clk, Reset, Load,
                        input  logic [31:0]  D,
                        output logic [31:0]  Dout
                        );

    always_ff @ (posedge Clk)
    begin
         if (Reset) //notice, this is a sycnrhonous reset, which is recommended on the FPGA
              Dout <= 32'b0;
         else if (Load)
              Dout <= D;
    end
endmodule