`include "opcode.vh"

module branch_ctrl (
    input [31:0] inst,
    input BrEq,
    input BrLt,
    output branch
);

    wire [4:0] opcode = inst[6:2];
    wire [2:0] funct3 = inst[14:12];

    assign branch = (opcode == `OPC_BRANCH_5) && (funct3[2] ? BrLt : BrEq) ^ (funct3[0]);
    
endmodule

module imm_ctrl (
    input [31:0] inst,
    output [2:0] ImmSel
);

    wire [4:0] opcode = inst[6:2];

    reg [2:0] ImmSelOut;
    always @(*) begin
        case (opcode)
            `OPC_STORE_5:  ImmSelOut = 3'b001;
            `OPC_AUIPC_5:  ImmSelOut = 3'b011;
            `OPC_LUI_5:    ImmSelOut = 3'b011;
            `OPC_JAL_5:    ImmSelOut = 3'b100;
            `OPC_BRANCH_5: ImmSelOut = 3'b010;
            default:       ImmSelOut = 3'b000;
        endcase
    end
    assign ImmSel = ImmSelOut;

endmodule

module exec_mem_ctrl (
    input [31:0] inst,
    output BrUn,
    output [2:0] ImmSel,
    output ASel,
    output BSel,
    output [3:0] ALUSel,
    output [3:0] MemRW
);

    wire [4:0] opcode = inst[6:2];
    wire [2:0] funct3 = inst[14:12];
    
    assign BrUn = funct3[1];

    reg [2:0] ImmSelOut;
    always @(*) begin
        case (opcode)
            `OPC_STORE_5:  ImmSelOut = 3'b001;
            `OPC_AUIPC_5:  ImmSelOut = 3'b011;
            `OPC_LUI_5:    ImmSelOut = 3'b011;
            `OPC_JAL_5:    ImmSelOut = 3'b100;
            `OPC_BRANCH_5: ImmSelOut = 3'b010;
            default:       ImmSelOut = 3'b000;
        endcase
    end
    assign ImmSel = ImmSelOut;

    assign ASel = (opcode == `OPC_BRANCH_5) | (opcode == `OPC_AUIPC_5) | (opcode == `OPC_JAL_5);
    assign BSel = ~(opcode == `OPC_ARI_RTYPE_5);

    reg [3:0] ALUSelOut;
    always @(*) begin
        case (opcode)
            `OPC_ARI_RTYPE_5: ALUSelOut = {inst[30], funct3};
            `OPC_ARI_ITYPE_5: begin
                case (funct3)
                    `FNC_SRL_SRA: ALUSelOut = {inst[30], funct3};
                    `FNC_SLL:     ALUSelOut = {inst[30], funct3};
                    default:      ALUSelOut = {1'd0, funct3};
                endcase
            end
            `OPC_LUI_5: ALUSelOut = 4'b1111;
            default: ALUSelOut = 4'd0;
        endcase
    end
    assign ALUSel = ALUSelOut;

    reg [3:0] MemRWOut;
    always @(*) begin
        case ({opcode, funct3})
            default:                MemRWOut = 4'b0000;
            {`OPC_STORE_5, 3'b000}: MemRWOut = 4'b0001; // SB
            {`OPC_STORE_5, 3'b001}: MemRWOut = 4'b0011; // SH
            {`OPC_STORE_5, 3'b010}: MemRWOut = 4'b1111; // SW
        endcase
    end
    assign MemRW = MemRWOut;

endmodule

module writeback_ctrl (
    input [31:0] inst,
    output [1:0] WBSel,
    output RegWEn 
);

    wire [4:0] opcode = inst[6:2];

    reg [1:0] WBSelOut;
    always @(*) begin
        case (opcode)
            `OPC_LOAD_5:      WBSelOut = 2'b00; // Mem
            `OPC_ARI_RTYPE_5: WBSelOut = 2'b01; // ALU
            `OPC_ARI_ITYPE_5: WBSelOut = 2'b01; // ALU
            `OPC_AUIPC_5:     WBSelOut = 2'b01; // ALU
            `OPC_LUI_5:       WBSelOut = 2'b01; // ALU
            `OPC_JALR_5:      WBSelOut = 2'b10; // PC + 4
            `OPC_JAL_5:       WBSelOut = 2'b10; // PC + 4
            default:          WBSelOut = 2'b00;
        endcase
    end
    assign WBSel = WBSelOut;

    reg RegWEnOut;
    always @(*) begin
        case (opcode)
            `OPC_BRANCH_5: RegWEnOut = 1'b0;
            `OPC_STORE_5:  RegWEnOut = 1'b0;
            default:       RegWEnOut = 1'b1;
        endcase
    end 
    assign RegWEn = RegWEnOut;

endmodule

module csr_ctrl (
    input [31:0] inst,
    output CSRSel,
    output CSRWEn
);

    wire [4:0] opcode = inst[6:2];
    wire [2:0] funct3 = inst[14:12];

    assign CSRSel = funct3[2];
    assign CSRWEn = (opcode == 5'b11100) && (inst[31:20] == 12'h51e);

endmodule