`ifndef ALU_H_SV
`define ALU_H_SV

`ifndef LU_H_SV
`define LU_H_SV
module lu #(parameter width = 2) (
    input logic [1:0] op,           // 2 bit operand
    input logic [2**width-1:0] A,  // width defined as 2^n 
	 input logic [2**width-1:0] B,
    output logic [2**width-1:0] Y);
	 
    always_comb begin
		case (op) // different operations of A and B based on operand
			2'b00:	Y =   A & B;
			2'b01:	Y =   A | B;
			2'b10:	Y =   A ^ B;
			2'b11:	Y = ~(A | B);
			default:Y = '0;
		endcase
    end
endmodule
`endif

`ifndef SEVEN_SEG_H_SV
`define SEVEN_SEG_H_SV
module seven_seg(input logic [3:0] num, input logic [0:0] rbi_inv, output logic [7:0] segments);
        always_comb begin
           case (num)
               4'b0000: segments = (rbi_inv) ? ~(8'b01111111) : ~(8'b00000000);
               4'b0001: segments = ~(8'b00001101);
               4'b0010: segments = ~(8'b10110111);
               4'b0011: segments = ~(8'b10011111);
               4'b0100: segments = ~(8'b11001101);
               4'b0101: segments = ~(8'b11011011);
               4'b0110: segments = ~(8'b11111011);
               4'b0111: segments = ~(8'b00001111);
               4'b1000: segments = ~(8'b11111111);
               4'b1001: segments = ~(8'b11001111);
               4'b1010: segments = ~(8'b11101111);
               4'b1011: segments = ~(8'b11111001);
               4'b1100: segments = ~(8'b01110011);
               4'b1101: segments = ~(8'b10111101);
               4'b1110: segments = ~(8'b11110011);
               4'b1111: segments = ~(8'b11100011);
               default: segments = ~(8'bxxxxxxxx);
           endcase     
        end
endmodule
`endif

`ifndef ASU_H_SV
`define ASU_H_SV

module asu #(parameter width = 2)(
	input logic [0:0] op,
	input logic[2**width-1:0] A, input logic[2**width-1:0] B,
   output logic[2**width-1:0] OUT_1, output logic[0:0] COUT, output logic[0:0] OUT_CIN
// Used to switch between demo mode and testbench mode
// TODO: Strip for assignment 2
`ifdef SYNTHESIS
	, output logic[7:0] out_seg_1);
`else
    );
`endif
	
    logic [2**width-1:0] operand;
    always_comb begin
        // A + (sub) ? B : -B
        // Adds an extra bit to all operands to calculate the carry bit. Not optimal for hardware synthesis
        //operand = (B ^ {(2**width){op}});
        operand = op == 1'b0 ? B : ~B;
        {COUT, OUT_1} = A + operand + {{(2**width-1){1'b0}}, op};
    end

    // Need to do this outside of the always_comb block because parameter-based bitslicing/indexing isn't supported?
    // Cin for final bit in adder to later be xor'ed with COUT to determine if a signed overflow occured for the MIPS architecture
    assign OUT_CIN = A[2**width-1] ^ operand[2**width-1] ^ OUT_1[2**width-1];

	// Include seven segment visual output for hardware builds only
    // TODO: Strip this for assignment 2 outright
    `ifdef SYNTHESIS 
	    seven_seg dut(OUT_1, 1'b1, out_seg_1);
    `endif
endmodule
`endif

// flags[0:0] = Cout
// flags[1:1] = Overflow
module alu #(parameter width=2)(input logic [3:0] op,
    input logic [2**width-1:0] A, input logic [2**width-1:0] B,
    output logic [2**width-1:0] OUT_1, output logic[1:0] flags
`ifdef SYNTHESIS
	, output logic [7:0] out_seg1);
`else
    );
`endif
    // intermediate logic for handling asu output
    logic [0:0] asu_op;
    logic [2**width-1:0] asu_out;
    logic [1:0] asu_flags;
    logic [7:0] asu_seven_seg_out;

    // NOTE: For demo only
    // TODO: Strip this for assignment 2
    `ifdef SYNTHESIS
        asu #(width) adder_unit(asu_op[op[1]], A, B, asu_out, asu_flags[0], asu_flags[1], asu_seven_seg_out);
    `else
        asu #(width) adder_unit(asu_op[op[1]], A, B, asu_out, asu_flags[0], asu_flags[1]);
    `endif

    logic [2**width-1:0] lu_out;
    lu #(width) logic_unit(op[1:0], A, B, lu_out);

    always_comb begin
        case (op) 
            4'b0000: begin OUT_1 = asu_out; flags = asu_flags; end
            4'b0001: begin OUT_1 = asu_out; flags = asu_flags; end
            4'b0010: begin OUT_1 = asu_out; flags = asu_flags; end
            4'b0011: begin OUT_1 = asu_out; flags = asu_flags; end
            4'b0100: begin OUT_1 = lu_out; flags = asu_flags; end
            4'b0101: begin OUT_1 = lu_out; flags = asu_flags; end
            4'b0110: begin OUT_1 = lu_out;  flags = asu_flags; end
            4'b0111: begin OUT_1 = lu_out; flags = asu_flags; end
            4'b1010: begin OUT_1 = (asu_out[2**width-1] ^ asu_flags[1]) & {{(width-1){1'b0}}, 1'b1}; flags = asu_flags; end
            4'b1011: begin OUT_1 = ~asu_flags[0] & {{(width-1){1'b0}}, 1'b1}; flags = asu_flags; end
            default: begin OUT_1 = {(2**width){1'b0}}; flags = asu_flags; end
        endcase
    end

    // NOTE: For demo only
    // TODO: Strip this for assignment 2
    `ifdef SYNTHESIS
	    seven_seg dut(OUT_1, 1'b1, out_seg1);
    `endif
endmodule
`endif

// Testbench

module alu_tb #(width=2)();
    localparam bitwidth = 2**width;

    // Guard against future me by testing the tests
    generate
        if (bitwidth < 4) begin : param_asserts
            initial begin
                $fatal(1, "Can not instantiate an ALU testbench with less than 4 bits (n=2)");
            end
        end
    endgenerate

    logic [3:0] op_in;
    logic [2**width-1:0] A_in;
    logic [2**width-1:0] B_in;
    logic [2**width-1:0] Y_out;
    logic [1:0] flags_out;

    localparam int NUM_TESTS = 28;
    logic [3:0] ops[NUM_TESTS];
    logic [2**width-1:0] A_cases[NUM_TESTS];
    logic [2**width-1:0] B_cases[NUM_TESTS];
    logic [2**width-1:0] Y_exp[NUM_TESTS];
    logic [1:0] flags_exp[NUM_TESTS];

    alu #(width) dut (op_in, A_in, B_in, Y_out, flags_out);

    integer error_count;
    initial begin
        ops[0] = 4'b0000;                       // sadd
        A_in[0] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        B_in[0] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[0] = {{(bitwidth-2){1'b0}}, 2'b10};  // 0
        flags_exp[0] = {2'b00};                 // ~Cout, ~Overflow

        ops[1] = 4'b0000;                       // sadd
        A_in[1] = {bitwidth{1'b1}};                // UINTMAX
        B_in[1] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[1] = {bitwidth{1'b0}};               // 0
        flags_exp[1] = {2'b01};                 // Cout, ~Overflow

        ops[2] = 4'b0000;                       // sadd
        A_in[2] = {{1'b0}, {(bitwidth-1){1'b1}}};  // INTMAX
        B_in[2] = {{(bitwidth-1){1'b0}} , {1'b1}}; // 1
        Y_exp[2] = {{1'b1}, {(bitwidth-1){1'b0}} };// -INTMAX
        flags_exp[2] = {2'b10};                 // ~Cout, Overflow

        ops[3] = 4'b0000;                       // sadd
        A_in[3] = {{1'b1}, {(bitwidth-1){1'b0}}};  // -INTMAX
        B_in[3] = {{1'b1}, {(bitwidth-1){1'b0}}};  // -INTMAX
        Y_exp[3] = {bitwidth{1'b0}};               // 0
        flags_exp[3] = {2'b11};                 // Cout, Overflow

        ops[4] = 4'b0001;                       // uadd
        A_in[4] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        B_in[4] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[4] = {{(bitwidth-2){1'b0}}, 2'b10};  // 0
        flags_exp[4] = {2'b00};                 // ~Cout, ~Overflow

        ops[5] = 4'b0001;                       // uadd
        A_in[5] = {bitwidth{1'b1}};                // UINTMAX
        B_in[5] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[5] = {bitwidth{1'b0}};               // 0
        flags_exp[5] = {2'b01};                 // Cout, ~Overflow

        ops[6] = 4'b0001;                       // uadd
        A_in[6] = {{1'b0}, {(bitwidth-1){1'b1}}};  // INTMAX
        B_in[6] = {{(bitwidth-1){1'b0}} , {1'b1}}; // 1
        Y_exp[6] = {{1'b1}, {(bitwidth-1){1'b0}} };// -INTMAX
        flags_exp[6] = {2'b10};                 // ~Cout, Overflow

        ops[7] = 4'b0001;                       // uadd
        A_in[7] = {{1'b1}, {(bitwidth-1){1'b0}}};  // -INTMAX
        B_in[7] = {{1'b1}, {(bitwidth-1){1'b0}}};  // -INTMAX
        Y_exp[7] = {bitwidth{1'b0}};               // 0
        flags_exp[7] = {2'b11};                 // Cout, Overflow

        ops[8] = 4'b0010;                       // ssub
        A_in[8] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        B_in[8] = {{(bitwidth-2){1'b0}}, 2'b10};   // 2
        Y_exp[8] = {bitwidth{1'b1}};               // -1
        flags_exp[8] = {2'b00};                 // ~Cout, ~Overflow

        ops[9] = 4'b0010;                       // ssub
        A_in[9] = {1'b0, {(bitwidth-1){1'b1}}};    // INTMAX
        B_in[9] = {{(bitwidth-1){1'b0}}, 1'b1};    // -1
        Y_exp[9] = {1'b1, {(bitwidth-1){1'b0}}};   // -INTMAX
        flags_exp[9] = {2'b10};                 // ~Cout, Overflow

        ops[10] = 4'b0010;                       // ssub
        A_in[10] = {{(bitwidth-2){1'b0}}, 2'b10};   // 2
        B_in[10] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[10] = {{(bitwidth-1){1'b0}}, 1'b1};   // 1
        flags_exp[10] = {2'b01};                 // Cout, ~Overflow

        ops[11] = 4'b0010;                       // ssub
        A_in[11] = {1'b1, {(bitwidth-1){1'b0}}};    // -INTMAX
        B_in[11] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[11] = {1'b0, {(bitwidth-1){1'b1}}};   // INTMAX
        flags_exp[11] = {2'b11};                 // Cout, Overflow

        ops[12] = 4'b0011;                       // ssub
        A_in[12] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        B_in[12] = {{(bitwidth-2){1'b0}}, 2'b10};   // 2
        Y_exp[12] = {bitwidth{1'b1}};               // -1
        flags_exp[12] = {2'b00};                 // ~Cout, ~Overflow

        ops[13] = 4'b0011;                       // ssub
        A_in[13] = {1'b0, {(bitwidth-1){1'b1}}};    // INTMAX
        B_in[13] = {{(bitwidth-1){1'b0}}, 1'b1};    // -1
        Y_exp[13] = {1'b1, {(bitwidth-1){1'b0}}};   // -INTMAX
        flags_exp[13] = {2'b10};                 // ~Cout, Overflow

        ops[14] = 4'b0011;                       // ssub
        A_in[14] = {{(bitwidth-2){1'b0}}, 2'b10};   // 2
        B_in[14] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[14] = {{(bitwidth-1){1'b0}}, 1'b1};   // 1
        flags_exp[14] = {2'b01};                 // Cout, ~Overflow

        ops[15] = 4'b0011;                       // ssub
        A_in[15] = {1'b1, {(bitwidth-1){1'b0}}};    // -INTMAX
        B_in[15] = {{(bitwidth-1){1'b0}}, 1'b1};    // 1
        Y_exp[15] = {1'b0, {(bitwidth-1){1'b1}}};   // INTMAX
        flags_exp[15] = {2'b11};                 // Cout, Overflow

        ops[16] = 4'b0100; // and
        A_in[16] = {{(bitwidth-4){1'b0}}, 4'b1100};    // c
        B_in[16] = {{(bitwidth-4){1'b0}}, 4'b1010};    // a
        Y_exp[16] = {{(bitwidth-4){1'b0}}, 4'b1000};   // 8
        flags_exp[16] = {2'bxx};                    // Don't care

        ops[17] = 4'b0101; // or
        A_in[17] = {{(bitwidth-4){1'b0}}, 4'b1100};    // c
        B_in[17] = {{(bitwidth-4){1'b0}}, 4'b1010};    // a
        Y_exp[17] = {{(bitwidth-4){1'b0}}, 4'b1110};   // e
        flags_exp[17] = {2'bxx};                    // Don't care

        ops[18] = 4'b0110; // xor
        A_in[18] = {{(bitwidth-4){1'b0}}, 4'b1100};    // c
        B_in[18] = {{(bitwidth-4){1'b0}}, 4'b1010};    // a
        Y_exp[18] = {{(bitwidth-4){1'b0}}, 4'b0110};   // 6
        flags_exp[18] = {2'bxx};                    // Don't care

        ops[19] = 4'b0111; // nor
        A_in[19] = {{(bitwidth-4){1'b0}}, 4'b1100};    // c
        B_in[19] = {{(bitwidth-4){1'b0}}, 4'b1010};    // a
        Y_exp[19] = {{(bitwidth-4){1'b0}}, 4'b0001};   // 1
        flags_exp[19] = {2'bxx};                    // Don't care

        ops[20] = 4'b1010;                                      // sslt: 1 < 2 (No overflow)
        A_in[20] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        B_in[20] = {{(bitwidth-2){1'b0}}, 2'b10};                  // 2
        Y_exp[20] = {{(bitwidth-1){1'b0}}, 1'b1};                  // 1 (True)
        flags_exp[20] = {2'b00};                                // ~Cout, ~Overflow

        ops[21] = 4'b1010;                                      // sslt: 2 < 1 (No overflow)
        A_in[21] = {{(bitwidth-2){1'b0}}, 2'b10};                  // 2
        B_in[21] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        Y_exp[21] = {bitwidth{1'b0}};                              // 0 (False)
        flags_exp[21] = {2'b10};                                // Cout, ~Overflow

        ops[22] = 4'b1010;                                      // sslt: -INTMAX < 1 (Overflow test)
        A_in[22] = {1'b1, {(bitwidth-1){1'b0}}};                   // -INTMAX
        B_in[22] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        Y_exp[22] = {{(bitwidth-1){1'b0}}, 1'b1};                  // 1 (True)
        flags_exp[22] = {2'b11};                                // Cout, Overflow

        ops[23] = 4'b1010;                                      // sslt: INTMAX < -1 (Overflow test)
        A_in[23] = {1'b0, {(bitwidth-1){1'b1}}};                   // INTMAX
        B_in[23] = {bitwidth{1'b1}};                               // -1
        Y_exp[23] = {bitwidth{1'b0}};                              // 0 (False)
        flags_exp[23] = {2'b01};                                // ~Cout, Overflow

        ops[24] = 4'b1011;                                      // uslt: 1 < 2 
        A_in[24] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        B_in[24] = {{(bitwidth-2){1'b0}}, 2'b10};                  // 2
        Y_exp[24] = {{(bitwidth-1){1'b0}}, 1'b1};                  // 1 (True)
        flags_exp[24] = {2'b00};                                // ~Cout, ~Overflow

        ops[25] = 4'b1011;                                      // uslt: 2 < 1 
        A_in[25] = {{(bitwidth-2){1'b0}}, 2'b10};                  // 2
        B_in[25] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        Y_exp[25] = {bitwidth{1'b0}};                              // 0 (False)
        flags_exp[25] = {2'b10};                                // Cout, ~Overflow

        ops[26] = 4'b1011;                                      // uslt: 1 < UINTMAX
        A_in[26] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        B_in[26] = {bitwidth{1'b1}};                               // UINTMAX
        Y_exp[26] = {{(bitwidth-1){1'b0}}, 1'b1};                  // 1 (True)
        flags_exp[26] = {2'b00};                                // ~Cout, ~Overflow

        ops[27] = 4'b1011;                                      // uslt: UINTMAX < 1
        A_in[27] = {bitwidth{1'b1}};                               // UINTMAX
        B_in[27] = {{(bitwidth-1){1'b0}}, 1'b1};                   // 1
        Y_exp[27] = {bitwidth{1'b0}};                              // 0 (False)
        flags_exp[27] = {2'b10};                                // Cout, ~Overflow

        error_count = 0;
        $display("Beginning ALU tests");
        for (integer N = 0; N < $size(ops); ++N) begin
            op_in = ops[N];
            A_in = A_cases[N];
            B_in = B_cases[N];
            #1;
            if (Y_out != Y_exp[N]) begin
                $display("\t\x1b[38;5;9m[%0d] ALU Instruction 0b%b produced 0b%b_%b but 0b%b_%b was expected\x1b[m", N, op_in, flags_out, Y_out, flags_exp[N], Y_exp[N]);
                error_count = error_count + 1;
            end
        end

        if (error_count == 0) $display("\x1b[38;5;10mAll ALU tests passing\x1b[m");
        else $display("\x1b[38;5;9m%0d ALU test(s) failing\x1b[m", error_count);
    end
endmodule
