/*
CSR Unit
This module is responsible for handling traps. It contains the CSRs.
The module utilizes a finite state machine to realize the task.
*/

//bits in the CSRs
`define mstatus_mie mstatus[3]
`define mstatus_mpie mstatus[7]
`define mie_meie mie[11]
`define mip_meip mip[11]
`define mie_mtie mie[7]
`define mip_mtip mip[7]
`define mie_msie mie[3]
`define mip_msip mip[3]

module csr_unit(input clk_i,
                input reset_i,
                input [31:0] pc_i,
                input [11:0] csr_r_addr_i, //CSR read address
                input [11:0] csr_w_addr_i, //CSR write address
                input [31:0] csr_reg_i, //CSR input
                input csr_wen_i, //CSR write enable
                input meip_i, //machine external interrupt
                input mtip_i, //machine timer interrupt
                input msip_i, //machine software interrupt
                input [15:0] fast_irq_i, //fast interrupts
                input take_branch_i,
                input mem_wen_i, //MEM stage write enable signal
                input ex_dummy_i,
                input mem_dummy_i,
                input mret_id_i,
                input mret_wb_i,
                input misaligned_ex,
                input instr_access_fault_i, illegal_instr_i, instr_addr_misaligned_i, ecall_i, ebreak_i, data_err_i,

                output reg [31:0] csr_reg_o,
                output [31:0] irq_addr_o, mepc_o,
                output mux1_ctrl_o, mux2_ctrl_o,
                output reg ack_o,
                output csr_if_flush_o, csr_id_flush_o, csr_ex_flush_o, csr_mem_flush_o);

//state encoding
parameter STAND_BY = 0;
parameter S1 = 1;

//state register for the FSM
reg STATE;

//CSRs
reg [31:0] mstatus, mie, mip, mcause, mtvec, mepc, mscratch;

//pipeline flush signals generated by the CSR unit.
//the pipeline is flushed whenever an interrupt occurs.
wire csr_if_flush, csr_id_flush, csr_ex_flush, csr_mem_flush;

//interrupt handler addresses for different interrupt handling modes
wire [31:0] direct_mode_addr, vector_mode_addr;

//Priority Encoder index
reg [4:0] fast_irq_index;
//Priority Encoder Valid output
reg PE_valid;

wire pending_irq, pending_exception;
wire [31:0] masked_irq;

assign direct_mode_addr = mtvec;
assign vector_mode_addr = mcause[31] ? {mtvec[31:1],1'b0} + (mcause << 2) : {mtvec[31:1],1'b0};

assign masked_irq = mie & mip & {32{`mstatus_mie}};
assign pending_exception = (illegal_instr_i | instr_addr_misaligned_i | ecall_i | ebreak_i) & ~take_branch_i;
assign pending_irq = masked_irq != 32'b0;

assign csr_if_flush = (`mstatus_mie & pending_irq) | (STATE == S1) | (mret_id_i & ~take_branch_i) | pending_exception;
assign csr_id_flush = csr_ex_flush | (`mstatus_mie & pending_irq) | pending_exception;
assign csr_ex_flush = csr_mem_flush | (`mstatus_mie & pending_irq & !ex_dummy_i & !misaligned_ex) | instr_addr_misaligned_i;
assign csr_mem_flush = (`mstatus_mie & pending_irq & mem_wen_i & !mem_dummy_i) | instr_access_fault_i;

//outputs
assign irq_addr_o = mtvec[0] ? vector_mode_addr : direct_mode_addr;

assign mux1_ctrl_o = mret_id_i & ~take_branch_i;
assign mux2_ctrl_o = !((STATE == S1) | (mret_id_i & ~take_branch_i));

assign csr_if_flush_o = csr_if_flush;
assign csr_id_flush_o = csr_id_flush;
assign csr_ex_flush_o = csr_ex_flush;
assign csr_mem_flush_o = csr_mem_flush;

assign mepc_o = mepc;

//state transitions are done on the rising edge
always @(posedge clk_i or negedge reset_i)
begin
    if(!reset_i)
    begin
        STATE <= STAND_BY;
        ack_o <= 1'b0;
        mcause <= 32'd0;
    end

    else
    begin
        case(STATE)

            STAND_BY:
            begin
                if(!csr_wen_i)
                begin
                    if(csr_w_addr_i[11:0] == 12'h342) //0x342 - mcause
                        mcause <= csr_reg_i;
                end

                if(masked_irq[31:16] != 16'b0) //fast interrupts have the highest priority
                begin
                    STATE <= S1;
                    mcause[31] <= 1'b1;
                    mcause[30:0] <= {26'b0,fast_irq_index};
                end

                else
                begin
                    if(`mstatus_mie & `mie_meie & `mip_meip) //external interrupts have the second highest priority
                    begin
                        STATE <= S1;
                        ack_o <= 1'b1;
                        mcause[31] <= 1'b1;
                        mcause[30:0] <= 31'd11;
                    end

                    else if(`mstatus_mie & `mie_msie & `mip_msip) //software interrupts have the third highest priority
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b1;
                        mcause[30:0] <= 31'd3;
                    end

                    else if(`mstatus_mie & `mie_mtie & `mip_mtip) //timer interrupts have the fourth highest priority
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b1;
                        mcause[30:0] <= 31'd7;
                    end
                    else if(instr_access_fault_i) //exceptions have the lowest priority
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd1;
                    end
                    else if(instr_addr_misaligned_i & !take_branch_i)
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd0;
                    end
                    else if(illegal_instr_i & !take_branch_i)
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd2;
                    end
                    else if(ecall_i & !take_branch_i)
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd11;
                    end
                    else if(ebreak_i & !take_branch_i)
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd3;
                    end
                    else if(data_err_i & !mem_wen_i) //store access fault
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd7;
                    end
                    else if(data_err_i & mem_wen_i) //load access fault
                    begin
                        STATE <= S1;
                        mcause[31] <= 1'b0;
                        mcause[30:0] <= 31'd5;
                    end
                end
            end

            S1:
            begin
                STATE <= STAND_BY;
                ack_o <= 1'b0;
            end
        endcase
    end
end

always @(posedge clk_i or negedge reset_i)
begin
    if(!reset_i)
        csr_reg_o <= 32'b0;

    else
    begin
        if(csr_r_addr_i == 12'h300) //0x300 - mstatus
            csr_reg_o <= mstatus;

        else if(csr_r_addr_i[11:0] == 12'h304) //0x304 - mie
            csr_reg_o <= mie;

        else if(csr_r_addr_i[11:0] == 12'h305) //0x305 - mtvec
            csr_reg_o <= mtvec;

        else if(csr_r_addr_i[11:0] == 12'h340) //0x340 - mscratch
            csr_reg_o <= mscratch;

        else if(csr_r_addr_i[11:0] == 12'h341) //0x341 - mepc
            csr_reg_o <= {mepc[31:2],2'b0};

        else if(csr_r_addr_i[11:0] == 12'h342) //0x342 - mcause
            csr_reg_o <= mcause;

        else if(csr_r_addr_i[11:0] == 12'h344) //0x344 - mip
            csr_reg_o <= mip;

        else
            csr_reg_o <= 32'd0;
    end
end

//Priority Encoder for fast interrupts.
always @(*)
begin
    fast_irq_index = 5'd15;
    PE_valid = 1'b0;
    while(fast_irq_index != 5'd31 && PE_valid != 1'b1)
    begin
        fast_irq_index = fast_irq_index + 5'd1;
        PE_valid = masked_irq[fast_irq_index];
    end
end

integer i;
always @(posedge clk_i or negedge reset_i)
begin
    if(!reset_i)
        mip <= 32'b0;
    else
    begin
        `mip_meip <= meip_i; //meip bit is set by the interrupt controller
        `mip_mtip <= mtip_i; //timer interrupt bit
        `mip_msip <= msip_i; //software interrupt bit

        for (i = 16; i<32; i=i+1)
        begin
            if(masked_irq[i] == 1'b1 && i == {27'b0,fast_irq_index})
                mip[i] <= fast_irq_i[i-16];
            else if(mip[i] == 1'b0)
                mip[i] <= fast_irq_i[i-16];
        end
    end
end

//CSR assignments
always @(posedge clk_i or negedge reset_i)
begin
    if(!reset_i)
    begin
        mepc <= 32'b0;
        mie <= 32'b0;
        mscratch <= 32'b0;
        mtvec <= 32'b0;
        //unused fields are hardwired to 0
        mstatus[31:13] <= 19'b0; mstatus[10:0] <= 11'b0;
        //mstatus.mpp
        mstatus[12:11] <= 2'b11;
    end

    else
    begin
        if(!csr_wen_i)
        begin
            if(mret_wb_i)
            begin
                `mstatus_mie <= `mstatus_mpie;
                `mstatus_mpie <= 1'b1;
            end

            else if(csr_w_addr_i[11:0] == 12'h300) //0x300 - mstatus
            begin
                `mstatus_mie <= csr_reg_i[3];
                `mstatus_mpie <= csr_reg_i[7];
            end

            else if(csr_w_addr_i[11:0] == 12'h304) //0x304 - mie
            begin
                `mie_meie <= csr_reg_i[11];
                `mie_mtie <= csr_reg_i[7];
                `mie_msie <= csr_reg_i[3];
                mie[31:16] <= csr_reg_i[31:16];
            end

            else if(csr_w_addr_i[11:0] == 12'h305) //0x305 - mtvec
                mtvec <= csr_reg_i;

            else if(csr_w_addr_i[11:0] == 12'h340) //0x340 - mscratch
                mscratch <= csr_reg_i;

            else if(csr_w_addr_i[11:0] == 12'h341) //0x341 - mepc
                mepc <= csr_reg_i;
        end

        else
        begin
            case(STATE)
                S1:
                begin
                    mepc <= pc_i;
                    `mstatus_mpie <= `mstatus_mie;
                    `mstatus_mie <= 1'b0;
                end
            endcase
        end
    end
end

endmodule
