`timescale 1ns / 1ps
module InstFetch (
    input clk,
    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 adder_out;
    assign adder_out = PC + 4;

    always @(posedge clk) begin
        if (need_stall) 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
endmodule