`timescale 1ns / 1ps
module InstFetch (
    input clk,
    input reset,
    input [31:0] branch_target,
    input [31:0] jump_target,
    input [31:0] jump_register_target,
    input [1:0] PC_jump,
    input PC_branch,
    input need_stall,
    output [31:0] fetched_instruction,
    output [31:0] PC_plus_4
);

    reg [31:0] PC;
    InstructionMemory instruction_memory (
        .address(PC),
        .instruction(fetched_instruction)
    );

    wire [31:0] adder_out;
    assign adder_out = PC + 4;
    assign PC_plus_4 = adder_out;

    always @(posedge clk) begin
        if (reset) begin
            PC <= 32'h00000000;
        end else begin
            // Note: if the inst is branch followed by a load word, the stall of the load word
            // will prevent the branch from jump
            if (need_stall == 1'h1 && PC_branch == 1'h0) begin
                PC <= PC;
            end else begin
                if (PC_branch) begin
                    PC <= branch_target;
                end else begin
                    case (PC_jump)
                        2'b00:   PC <= adder_out;
                        2'b01:   PC <= jump_target;
                        2'b10:   PC <= jump_register_target;
                        default: PC <= adder_out;
                    endcase
                end
            end
        end
    end
endmodule
