module alu (
    input logic [31:0] a_i,
    input logic [31:0] b_i,
    input logic [ 1:0] alu_control_i,

    output logic alu_zero_o,
    output logic [31:0] alu_result_o
);

  logic [31:0] result;

  always_comb begin
    unique case (alu_control_i)
      // add
      2'b00: result = a_i + b_i;
      // sub
      2'b01: result = a_i - b_i;
      // sll
      // can only shift by 32 pos
      2'b10: result = a_i << b_i[4:0];
      // srl
      // can only shift by 32 pos
      2'b11: result = a_i >> b_i[4:0];
    endcase

    alu_zero_o = (result == 1'b0);
  end
endmodule

module execute (
    input logic clk_i,
    input logic alu_src_i,
    input logic [31:0] imm_ext_i,
    input logic [1:0] alu_control_i,
    input logic [31:0] current_pc_i,
    input logic [31:0] rs1_data_i,
    input logic [31:0] rs2_data_i,

    output logic alu_zero_o,
    output logic [31:0] alu_result_o,
    output logic [31:0] pc_target_o
);

  logic [31:0] a;
  logic [31:0] b;

  assign a = rs1_data_i;
  assign b = alu_src_i ? imm_ext_i : rs2_data_i;
  assign pc_target_o = current_pc_i + imm_ext_i;

  alu alu (
      .a_i(a),
      .b_i(b),
      .alu_control_i(alu_control_i),
      .alu_zero_o(alu_zero_o),
      .alu_result_o(alu_result_o)
  );
endmodule