`timescale 1ns / 1ps
`include "cpu.h"

module mem_stage(
    input                          clk           ,
    input                          reset         ,
    input                          ms_flush,
    
    input                          ws_allowin    ,
    output                         ms_allowin    ,
    
    //from es
    input                          es_to_ms_valid,
    input  [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus  ,
    //to es
    output [`MS_TO_ES_BUS_WD -1:0] ms_to_es_bus  ,
    //to fw
    output [`MS_TO_FW_BUS_WD -1:0] ms_to_fw_bus  ,
    //to ws
    output                         ms_to_ws_valid,
    output [`MS_TO_WS_BUS_WD -1:0] ms_to_ws_bus  ,
    
    output [`BR_BUS_WD       -1:0] br_bus        ,
    //from data-sram
    input  [31                 :0] data_ram_rdata
    );
    
    reg         ms_valid;
    wire        ms_ready_go;

    reg  [`ES_TO_MS_BUS_WD -1:0] es_to_ms_bus_r;
    wire [ 5:0] ms_Branch    ;
    wire [ 1:0] ms_Jump      ;
    wire        ms_MemtoReg  ;
    wire        ms_MemWr     ;
    wire        ms_RegWr     ;
    wire [ 4:0] ms_Load      ;
    wire [ 4:0] ms_rd        ;
    wire [ 4:0] ms_rs1       ;
    wire [ 4:0] ms_rs2       ;
    wire [31:0] ms_alu_result;
    wire [31:0] ms_pc        ;
    wire [31:0] br_target    ;
    wire        ms_Carry     ;
    wire        ms_Sign      ;
    wire        ms_Overflow  ;
    wire        ms_Zero      ;
    
    wire        br_taken     ;
    
    assign { ms_Carry        ,  //119:119
             ms_Sign         ,  //118:118
             ms_Overflow     ,  //117:117
             ms_Zero         ,  //116:116
             br_target       ,  //115:84
             ms_Load         ,  //83 :79
             ms_Branch       ,  //78 :73
             ms_Jump         ,  //72 :71
             ms_MemtoReg     ,  //70 :70
             ms_RegWr        ,  //69 :69
             ms_rd           ,  //68 :64
             ms_alu_result   ,  //63 :32
             ms_pc              //31 :0
           } = es_to_ms_bus_r;
    
    wire [31:0] mem_result;
    
    assign br_bus       = { br_taken,   //32:32
                            br_target   //31:0
                            };
                            
    assign ms_to_es_bus = { ms_alu_result};
    assign ms_to_fw_bus = { ms_rd , ms_RegWr};
    
    assign ms_to_ws_bus = { ms_MemtoReg    ,  //102:102
                            ms_RegWr       ,  //101:101
                            mem_result     ,  //100:69
                            ms_rd          ,  //68 :64
                            ms_alu_result  ,  //63 :32
                            ms_pc             //31 :0
                            };
    
    assign ms_ready_go    = 1'b1;
    assign ms_allowin     = !ms_valid || ms_ready_go && ws_allowin;
    assign ms_to_ws_valid = ms_valid && ms_ready_go;
    always @(posedge clk) begin
        if (reset) begin
            ms_valid <= 1'b0;
        end
        else if (ms_allowin) begin
            ms_valid <= es_to_ms_valid;
        end
    
        if (reset || ms_flush) begin
            es_to_ms_bus_r <= 0;
        end
        else if (es_to_ms_valid && ms_allowin) begin
            es_to_ms_bus_r <= es_to_ms_bus;
        end
//        else begin
//            es_to_ms_bus_r <= 0;
//        end
    end
    
    assign mem_result = (ms_Load[0] || ms_Load[3]) ? ((ms_alu_result[1:0] == 2'b00) ? {{24{ms_Load[3] ? data_ram_rdata[ 7] : 1'b0 }}, data_ram_rdata[ 7:0]       }:
                                                      (ms_alu_result[1:0] == 2'b01) ? {{16{ms_Load[3] ? data_ram_rdata[ 7] : 1'b0 }}, data_ram_rdata[ 7:0],  8'b0}:
                                                      (ms_alu_result[1:0] == 2'b10) ? {{ 8{ms_Load[3] ? data_ram_rdata[ 7] : 1'b0 }}, data_ram_rdata[ 7:0], 16'b0}:
                                                                                      {                                               data_ram_rdata[ 7:0], 24'b0}
                                                      ) :
                        (ms_Load[1] || ms_Load[4]) ? ((ms_alu_result[1:0] == 2'b00) ? {{16{ms_Load[4] ? data_ram_rdata[15] : 1'b0 }}, data_ram_rdata[15:0]       }:
                                                                                      {                                               data_ram_rdata[15:0], 16'b0}
                                                      ) :
                         ms_Load[2]                ? (                                                                                data_ram_rdata
                                                      ) :                                                                                                   32'b0;
    
    
    assign br_taken = ( ms_Branch[0] &  ms_Zero                            ) 
                    | ( ms_Branch[1] & ~ms_Zero                            )
                    | ( ms_Branch[2] & (ms_Sign != ms_Overflow)            )
                    | ( ms_Branch[3] & (ms_Zero | (ms_Sign == ms_Overflow)))
                    | ( ms_Branch[4] &  ms_Carry                           )
                    | ( ms_Branch[5] & (ms_Zero | ~ms_Carry               ))
                    | (^ms_Jump                                            );

endmodule
