diff --git a/modules/alu.sv b/modules/alu.sv
index 87f8633..4bc9398 100644
--- a/modules/alu.sv
+++ b/modules/alu.sv
@@ -21,7 +21,6 @@ module alu #(parameter width=2)(input logic [3:0] op,
     lu #(width) logic_unit(op[1:0], A, B, lu_out);
 
     always_comb begin
-        asu_flags[2] = (A == B);
         case (op) 
             4'b0000: begin OUT_1 = asu_out; flags = asu_flags; end
             4'b0001: begin OUT_1 = asu_out; flags = asu_flags; end
@@ -35,6 +34,7 @@ module alu #(parameter width=2)(input logic [3:0] op,
             4'b1011: begin OUT_1 = ~asu_flags[0] & {{(width-1){1'b0}}, 1'b1}; flags = asu_flags; end
             default: begin OUT_1 = {(2**width){1'b0}}; flags = asu_flags; end
         endcase
+        asu_flags[2] = OUT_1 == 0;
     end
 endmodule
 `endif
diff --git a/modules/assignment2.sv b/modules/assignment2.sv
index 6a8d9c5..e9dd7e0 100644
--- a/modules/assignment2.sv
+++ b/modules/assignment2.sv
@@ -61,6 +61,9 @@ module assignment2 #(width=5, ramsz=9, romsz=10) (
 	instruction_counter #(width) PC (clk, pc_we, pc_addr_write, pc_addr);
 	
 	rom #(romsz, 2**width) ROM (pc_addr, instruction);
+
+	logic [2**width-1:0] rd_mem;
+	ram #(ramsz, 2**width) RAM (clk, rd_alu, B1, rt, rd_mem);
 	
 	instruction_classifier INSTRUCTION_CLASSIFIER(instruction[31:26], R, B1, J, B2, I, F, M);
 	
@@ -70,9 +73,9 @@ module assignment2 #(width=5, ramsz=9, romsz=10) (
 		en_mul, 
 		en_alu
 	);
-    
+
 	register_file #(width, 5) REGS (clk,
-		rd_we, instruction[15:11],
+		rd_we, (F) ? instruction[15:11] : instruction[20:16],
 		rs, instruction[25:21],
 		rt, instruction[20:16]
    );
@@ -84,16 +87,13 @@ module assignment2 #(width=5, ramsz=9, romsz=10) (
     //  - PCBranch (rel jump of 4-byte words)
 
 	logic [2**width-1:0] rd_alu;
-	alu #(width) ALU (instruction[3:0], rs, rt, rd_alu, flags[1:0]);
-
-	logic [2**width-1:0] rd_shifter;
-	shifter #(width) SHIFTER (instruction[2:0], rt, rs[4:0], instruction[10:6], rd_shifter);
-    
-
-	assign rd = (I | (R & en_alu)) ? rd_alu : ( (R & en_shifter) ? rd_shifter : {2**width{1'b1}});
-	assign rd_we = M; // TODO: Verify if diagram is in order
-
+	logic [2**width-1:0] rt_alu;
+	alu #(width) ALU (instruction[3:0], rs, rt_alu, rd_alu, flags[1:0]);
+    assign rt_alu = (I) ? (rt_imm) : (rt);
 
+	assign rd_we = M;
+	assign pc_we = B1 & flags[2];
+	assign pc_addr_write <= (rt_imm << 2) + pc_addr;
 
 
 
diff --git a/modules/register_file.sv b/modules/register_file.sv
index 43da9cb..5d7bdc7 100644
--- a/modules/register_file.sv
+++ b/modules/register_file.sv
@@ -5,7 +5,8 @@ module register_file #(parameter width=5, parameter nregs=5)(
     input logic [0:0] clk, input logic[0:0] we,
     input logic [nregs-1:0] WA, input logic [2**width-1:0] WD,
     input logic [nregs-1:0] RA1, output logic [2**width-1:0] RD1,
-    input logic [nregs-1:0] RA2, output logic [2**width-1:0] RD2
+    input logic [nregs-1:0] RA2, output logic [2**width-1:0] RD2,
+    input logic [nregs-1:0] RA3, output logic [2**width-1:0] RD3
 );
     localparam BW = 2**width;
     logic [2**width-1:0] regs[2**nregs];
@@ -22,5 +23,6 @@ module register_file #(parameter width=5, parameter nregs=5)(
 
     assign RD1 = regs[RA1];
     assign RD2 = regs[RA2];
+    assign RD3 = regs[RA3];
 endmodule
 `endif