// Sumbission for ELEC4720 Assignment 1 Question 5
// Oliver Boutros (c3428958)
// Riley King-Saunders (cc3430382)

`ifndef SHIFTER_H_SV
`define SHIFTER_H_SV
module shifter #(width=2)(input logic [2:0] op,
    input logic[2**width-1:0] A, input logic[2**width-1:0] B, input logic[width-1:0] C, 
    output logic[2**width-1:0] OUT_1);
// A and B are 2^n bits wide and C is n bits wide as according to the docs
    always_comb begin
        case (op)
            3'b000: OUT_1 = A << C;
            3'b001: OUT_1 = A >> C;
            3'b011: OUT_1 = $signed(A) >>> C;
            3'b100: OUT_1 = A << B;
            3'b101: OUT_1 = A >> B;
            3'b111: OUT_1 = $signed(A) >>> B;
            default:OUT_1 = A;
        endcase
    end    
endmodule
`endif

// Testbench

module shifter_tb #(width=4)();
    logic [2:0] op;
    logic [2**width-1:0] Y;
    logic [2**width-1:0] A_in;
    logic [2**width-1:0] B_in;
    logic [width-1:0] C_in;

    logic [2:0] ops[6];
    logic [2**width-1:0] A[6];
    logic [2**width-1:0] B[6];
    logic [width-1:0] C[6];
    logic [2**width-1:0] Y_exp[6];

    shifter #(width) dut(op, A_in,B_in,C_in,Y);

    integer error_count;
    initial begin
        // Test left shift by C
        ops[0] = 3'b000;
        A[0] = {{(2**width-1){1'b0}}, {1'b1}};
        C[0] = {{(width-1){1'b0}}, {1'b1}};
        B[0] = {2**width{1'b0}};
        Y_exp[0] = {{(2**width-2){1'b0}}, {2'b10}};

        // Test left shift by B
        ops[1] = 3'b100;
        A[1] = {{(2**width-1){1'b0}}, {1'b1}};
        C[1] = {{(width-1){1'b0}}, {1'b1}};
        B[1] = {{(2**width-2){1'b0}}, 2'b10};
        Y_exp[1] = {{(2**width-3){1'b0}}, {3'b100}};

        // Test Logical right shift by C
        ops[2] = 3'b001;
        A[2] = {{(2**width-2){1'b0}}, {2'b10}};
        C[2] = {{(width-1){1'b0}}, {1'b1}};
        B[2] = {2**width{1'b0}};
        Y_exp[2] = {{(2**width-1){1'b0}}, {1'b1}};

        // Test Logical right shift by B
        ops[3] = 3'b101;
        A[3] = {{(2**width-1){1'b0}}, {2'b10}};
        C[3] = {{(width-1){1'b0}}, {1'b1}};
        B[3] = {{(2**width-2){1'b0}}, 2'b10};
        Y_exp[3] = {{(2**width-1){1'b0}}, {1'b0}};

        // Test arithmetic shift right by C
        ops[4] = 3'b011;
        A[4] = {{(2**width-1){1'b1}}, 1'b0};   // -2
        C[4] = {{(width-1){1'b0}}, 1'b1};   // 1
        B[4] = {{(2**width-2){1'b0}}, 2'b10};  // 2
        Y_exp[4] = {2**width{1'b1}};           // -1

        // Test arithmetic shift right by B
        ops[5] = 3'b111;
        A[5] = {{(2**width-2){1'b1}}, 2'b00};   // -4
        C[5] = {{(width-1){1'b0}}, 1'b1};   // 1
        B[5] = {{(2**width-2){1'b0}}, 2'b10};  // 2
        Y_exp[5] = {2**width{1'b1}};           // -1

        error_count = 0;
        $display("Beginning Shifter tests");
        for (integer N = 0; N < $size(ops); ++N) begin
            op = ops[N];
            A_in = A[N];
            B_in = B[N];
            C_in = C[N];
            #1;
            if (Y != Y_exp[N]) begin
                $display("\t\x1b[38;5;9mShifter(%0d) Op 0b%b(0b%b, 0b%b, 0b%b) produced 0b%b but 0b%b was expected\x1b[m", N, op, A_in, B_in, C_in, Y, Y_exp[N]);
                error_count = error_count + 1;
            end
        end
        if (error_count == 0) $display("\x1b[38;5;10mAll Shifter tests passing\x1b[m");
        else $display("\x1b[38;5;9m%0d Shifter test(s) failing\x1b[m", error_count);
    end
endmodule