diff --git a/src/main/scala/CPU.scala b/src/main/scala/CPU.scala
index bc6ffe1..230c3af 100644
--- a/src/main/scala/CPU.scala
+++ b/src/main/scala/CPU.scala
@@ -54,6 +54,25 @@ class CPU extends MultiIOModule {
   /**
     TODO: Your code here
     */
+  def forward(data: UInt, addr: UInt, useForward: Bool, mem: Forwarding, wb: Forwarding, id: Forwarding): UInt = {
+    Mux(
+      !useForward,
+      data,
+      Mux(
+        mem.valid && mem.writeAddr === addr,
+        mem.writeData,
+        Mux(
+          wb.valid && wb.writeAddr === addr,
+          wb.writeData,
+          Mux(
+            id.valid && id.writeAddr === addr,
+            id.writeData,
+            data,
+          )
+        )
+      )
+    )
+  }
 
   IFBarrier.PCin := IF.io.PC
   IFBarrier.instructionIn := IF.io.instruction
@@ -61,74 +80,65 @@ class CPU extends MultiIOModule {
   ID.io.instruction := IFBarrier.instructionOut
   ID.io.pc := IFBarrier.PCout
 
-  IDBarrier.op1in := ID.io.op1
-  IDBarrier.op2in := ID.io.op2
-  IDBarrier.isOp1RValue := ID.io.isOp1RValue
-  IDBarrier.isOp2RValue := ID.io.isOp2RValue
-  IDBarrier.r1ValueIn := ID.io.r1Value
-  IDBarrier.r2ValueIn := ID.io.r2Value
-  IDBarrier.r1AddressIn := ID.io.r1Address
-  IDBarrier.r2AddressIn := ID.io.r2Address
-  IDBarrier.ALUopIn := ID.io.ALUOp
-  IDBarrier.returnAddrIn := ID.io.returnAddr
-  IDBarrier.jumpIn := ID.io.jump
-  IDBarrier.branchTypeIn := ID.io.branchType
-  IDBarrier.writeEnableIn := ID.io.writeEnableOut
-  IDBarrier.writeAddrIn := ID.io.writeAddrOut
-  IDBarrier.memWriteIn := ID.io.memWrite
-  IDBarrier.memReadIn := ID.io.memRead
+  IDBarrier.in.op1 := ID.io.op1
+  IDBarrier.in.op2 := ID.io.op2
+  IDBarrier.in.isOp1RValue := ID.io.isOp1RValue
+  IDBarrier.in.isOp2RValue := ID.io.isOp2RValue
+  IDBarrier.in.r1Value := ID.io.r1Value
+  IDBarrier.in.r2Value := ID.io.r2Value
+  IDBarrier.in.r1Address := ID.io.r1Address
+  IDBarrier.in.r2Address := ID.io.r2Address
+  IDBarrier.in.ALUop := ID.io.ALUOp
+  IDBarrier.in.returnAddr := ID.io.returnAddr
+  IDBarrier.in.jump := ID.io.jump
+  IDBarrier.in.branchType := ID.io.branchType
+  IDBarrier.in.writeEnable := ID.io.writeEnableOut
+  IDBarrier.in.writeAddr := ID.io.writeAddrOut
+  IDBarrier.in.memWrite := ID.io.memWrite
+  IDBarrier.in.memRead := ID.io.memRead
 
-  EX.io.op1 := IDBarrier.op1out
-  EX.io.op2 := IDBarrier.op2out
-  EX.io.ALUOp := IDBarrier.ALUopOut
-  EX.io.branchType := IDBarrier.branchTypeOut
-  EX.io.rs1ValueIn := IDBarrier.r1ValueOut.asSInt()
-  EX.io.rs2ValueIn := IDBarrier.r2ValueOut.asSInt()
+  EX.io.op1 := forward(IDBarrier.out.op1.asUInt(), IDBarrier.out.r1Address, IDBarrier.out.isOp1RValue, mem = MEMBarrier.forwardMem, wb = MEMBarrier.forwardWb, id = MEMBarrier.forwardId).asSInt()
+  EX.io.op2 := forward(IDBarrier.out.op2.asUInt(), IDBarrier.out.r2Address, IDBarrier.out.isOp2RValue, mem = MEMBarrier.forwardMem, wb = MEMBarrier.forwardWb, id = MEMBarrier.forwardId).asSInt()
+  EX.io.ALUOp := IDBarrier.out.ALUop
+  EX.io.branchType := IDBarrier.out.branchType
+  EX.io.rs1ValueIn := forward(IDBarrier.out.r1Value, IDBarrier.out.r1Address, true.B, mem = MEMBarrier.forwardMem, wb = MEMBarrier.forwardWb, id = MEMBarrier.forwardId).asSInt()
+  EX.io.rs2ValueIn := forward(IDBarrier.out.r2Value, IDBarrier.out.r2Address, true.B, mem = MEMBarrier.forwardMem, wb = MEMBarrier.forwardWb, id = MEMBarrier.forwardId).asSInt()
 
-  EXBarrier.r2ValueIn := EX.io.rs2ValueOut.asUInt()
-  EXBarrier.ALUResultIn := EX.io.ALUResult.asUInt()
-  EXBarrier.branchIn := EX.io.branch
-  EXBarrier.jumpIn := IDBarrier.jumpOut
-  EXBarrier.returnAddrIn := IDBarrier.returnAddrOut
-  EXBarrier.writeEnableIn := IDBarrier.writeEnableOut
-  EXBarrier.writeAddrIn := IDBarrier.writeAddrOut
-  EXBarrier.memWriteIn := IDBarrier.memWriteOut
-  EXBarrier.memReadIn := IDBarrier.memReadOut
+  EXBarrier.in.r2Value := EX.io.rs2ValueOut.asUInt()
+  EXBarrier.in.ALUResult := EX.io.ALUResult.asUInt()
+  EXBarrier.in.branch := EX.io.branch
+  EXBarrier.in.jump := IDBarrier.out.jump
+  EXBarrier.in.returnAddr := IDBarrier.out.returnAddr
+  EXBarrier.in.writeEnable := IDBarrier.out.writeEnable
+  EXBarrier.in.writeAddr := IDBarrier.out.writeAddr
+  EXBarrier.in.memWrite := IDBarrier.out.memWrite
+  EXBarrier.in.memRead := IDBarrier.out.memRead
 
-  MEM.io.ALUResult := EXBarrier.ALUResultOut
-  MEM.io.jump := EXBarrier.jumpOut
-  MEM.io.returnAddr := EXBarrier.returnAddrOut
-  MEM.io.writeMem := EXBarrier.memWriteOut
-  MEM.io.readMem := EXBarrier.memReadOut
-  MEM.io.writeData := EXBarrier.r2ValueOut
+  MEM.io.ALUResult := EXBarrier.out.ALUResult
+  MEM.io.jump := EXBarrier.out.jump
+  MEM.io.returnAddr := EXBarrier.out.returnAddr
+  MEM.io.writeMem := EXBarrier.out.memWrite
+  MEM.io.readMem := EXBarrier.out.memRead
+  MEM.io.writeData := EXBarrier.out.r2Value
 
-  MEMBarrier.memRead := EXBarrier.memReadOut
-  MEMBarrier.dataIn := MEM.io.dataOut
-  MEMBarrier.writeEnableIn := EXBarrier.writeEnableOut
-  MEMBarrier.writeAddrIn := EXBarrier.writeAddrOut
+  MEMBarrier.memRead := EXBarrier.out.memRead
+  MEMBarrier.in.data := MEM.io.dataOut
+  MEMBarrier.in.writeEnable := EXBarrier.out.writeEnable
+  MEMBarrier.in.writeAddr := EXBarrier.out.writeAddr
 
   // Write back
-  ID.io.writeData := MEMBarrier.dataOut
-  ID.io.writeEnableIn := MEMBarrier.writeEnableOut
-  ID.io.writeAddrIn := MEMBarrier.writeAddrOut
+  ID.io.writeData := MEMBarrier.out.data
+  ID.io.writeEnableIn := MEMBarrier.out.writeEnable
+  ID.io.writeAddrIn := MEMBarrier.out.writeAddr
 
   // Branching
-  IF.io.branch := EXBarrier.branchOut
-  IF.io.branchAddress := EXBarrier.ALUResultOut
-
-  // Forwarding
-  IDBarrier.forwardExData := EXBarrier.forwardExData
-  IDBarrier.forwardEx := EXBarrier.forwardEx
-  IDBarrier.forwardExAddr := EXBarrier.forwardExAddr
-
-  IDBarrier.forwardMemData := MEMBarrier.forwardMemData
-  IDBarrier.forwardMem := MEMBarrier.forwardMem
-  IDBarrier.forwardMemAddr := MEMBarrier.forwardMemAddr
-
-  IDBarrier.forwardWbData := MEMBarrier.forwardWbData
-  IDBarrier.forwardWb := MEMBarrier.forwardWb
-  IDBarrier.forwardWbAddr := MEMBarrier.forwardWbAddr
+  IF.io.branch := EXBarrier.out.branch
+  IF.io.branchAddress := EXBarrier.out.ALUResult
 
   // Stall
-  IF.io.stall := ID.io.stall
+  IF.io.stall := IDBarrier.stall
+  IFBarrier.stall := IDBarrier.stall
+
+  // Flush
+  IFBarrier.flush := EXBarrier.flush
 }
diff --git a/src/main/scala/EXBarrier.scala b/src/main/scala/EXBarrier.scala
index 9ba6cfc..222f20c 100644
--- a/src/main/scala/EXBarrier.scala
+++ b/src/main/scala/EXBarrier.scala
@@ -3,69 +3,29 @@ package FiveStage
 import chisel3._
 import chisel3.experimental.MultiIOModule
 
+class EXBarrierIO extends Bundle {
+  val ALUResult = UInt(32.W)
+  val returnAddr = UInt(32.W)
+  val r2Value = UInt(32.W)
+  val writeAddr = UInt(5.W)
+  val writeEnable = Bool()
+  val memRead = Bool()
+  val memWrite = Bool()
+  val branch = Bool()
+  val jump = Bool()
+}
+
 class EXBarrier extends MultiIOModule {
   val io = IO(
     new Bundle {
-      val ALUResultIn = Input(UInt(32.W))
-      val ALUResultOut = Output(UInt(32.W))
-      val returnAddrIn = Input(UInt(32.W))
-      val returnAddrOut = Output(UInt(32.W))
-      val r2ValueIn = Input(UInt(32.W))
-      val r2ValueOut = Output(UInt(32.W))
-      val writeAddrIn = Input(UInt(5.W))
-      val writeAddrOut = Output(UInt(5.W))
-      val writeEnableIn = Input(Bool())
-      val writeEnableOut = Output(Bool())
-      val memReadIn = Input(Bool())
-      val memReadOut = Output(Bool())
-      val memWriteIn = Input(Bool())
-      val memWriteOut = Output(Bool())
-      val branchIn = Input(Bool())
-      val branchOut = Output(Bool())
-      val jumpIn = Input(Bool())
-      val jumpOut = Output(Bool())
-      val forwardEx = Output(Bool())
-      val forwardExAddr = Output(UInt(5.W))
-      val forwardExData = Output(UInt(32.W))
+      val in = Input(new EXBarrierIO)
+      val out = Output(new EXBarrierIO)
+      val flush = Output(Bool())
     })
 
-  val ALUResult = RegInit(UInt(32.W), 0.U)
-  ALUResult := io.ALUResultIn
-  io.ALUResultOut := ALUResult
+  val delay = Reg(new EXBarrierIO)
+  delay := io.in
+  io.out := delay
 
-  val returnAddr = RegInit(UInt(32.W), 0.U)
-  returnAddr := io.returnAddrIn
-  io.returnAddrOut := returnAddr
-
-  val r2Value = RegInit(UInt(32.W), 0.U)
-  r2Value := io.r2ValueIn
-  io.r2ValueOut := r2Value
-
-  val writeAddr = RegInit(UInt(5.W), 0.U)
-  writeAddr := io.writeAddrIn
-  io.writeAddrOut := writeAddr
-
-  val writeEnable = RegInit(Bool(), false.B)
-  writeEnable := io.writeEnableIn
-  io.writeEnableOut := writeEnable
-
-  val memRead = RegInit(Bool(), false.B)
-  memRead := io.memReadIn
-  io.memReadOut := memRead
-
-  val memWrite = RegInit(Bool(), false.B)
-  memWrite := io.memWriteIn
-  io.memWriteOut := memWrite
-
-  val branch = RegInit(Bool(), false.B)
-  branch := io.branchIn
-  io.branchOut := branch
-
-  val jump = RegInit(Bool(), false.B)
-  jump := io.jumpIn
-  io.jumpOut := jump
-
-  io.forwardEx := io.writeEnableIn && !io.memReadIn
-  io.forwardExAddr := io.writeAddrIn
-  io.forwardExData := io.ALUResultIn
+  io.flush := io.in.branch
 }
diff --git a/src/main/scala/ID.scala b/src/main/scala/ID.scala
index 8942aa8..1027e70 100644
--- a/src/main/scala/ID.scala
+++ b/src/main/scala/ID.scala
@@ -39,7 +39,6 @@ class InstructionDecode extends MultiIOModule {
       val branchType = Output(UInt(3.W))
       val jump = Output(Bool())
       val returnAddr = Output(UInt(32.W))
-      val stall = Output(Bool())
     }
   )
 
@@ -86,18 +85,13 @@ class InstructionDecode extends MultiIOModule {
   io.r2Address := registers.io.readAddress2
 
   io.ALUOp := decoder.ALUop
-  io.branchType := decoder.branchType
   io.writeAddrOut := decoder.instruction.registerRd
 
-  val stallDelay = RegInit(Bool(), false.B)
-  val stall = Mux(stallDelay, false.B, decoder.controlSignals.memRead)
-  io.stall :=  stall
-  stallDelay := stall
-
-  io.jump := Mux(stallDelay, false.B, decoder.controlSignals.jump)
+  io.jump := decoder.controlSignals.jump
+  io.branchType := decoder.branchType
   io.returnAddr := io.pc + 4.U
 
-  io.writeEnableOut := Mux(stallDelay, false.B, decoder.controlSignals.regWrite)
-  io.memRead := Mux(stallDelay, false.B, decoder.controlSignals.memRead)
-  io.memWrite := Mux(stallDelay, false.B, decoder.controlSignals.memWrite)
+  io.writeEnableOut := decoder.controlSignals.regWrite
+  io.memRead := decoder.controlSignals.memRead
+  io.memWrite := decoder.controlSignals.memWrite
 }
diff --git a/src/main/scala/IDBarrier.scala b/src/main/scala/IDBarrier.scala
index c297af2..c924233 100644
--- a/src/main/scala/IDBarrier.scala
+++ b/src/main/scala/IDBarrier.scala
@@ -3,157 +3,35 @@ package FiveStage
 import chisel3._
 import chisel3.experimental.MultiIOModule
 
+class IDBarrierIO extends Bundle {
+  val op1 = SInt(32.W)
+  val isOp1RValue = Bool()
+  val op2 = SInt(32.W)
+  val isOp2RValue = Bool()
+  val r1Value = UInt(32.W)
+  val r1Address = UInt(5.W)
+  val r2Value = UInt(32.W)
+  val r2Address = UInt(5.W)
+  val returnAddr = UInt(32.W)
+  val jump = Bool()
+  val ALUop = UInt(4.W)
+  val branchType = UInt(3.W)
+  val writeAddr = UInt(5.W)
+  val writeEnable = Bool()
+  val memRead = Bool()
+  val memWrite = Bool()
+}
+
 class IDBarrier extends MultiIOModule {
   val io = IO(
     new Bundle {
-      val op1in = Input(SInt(32.W))
-      val op1out = Output(SInt(32.W))
-      val isOp1RValue = Input(Bool())
-      val op2in = Input(SInt(32.W))
-      val op2out = Output(SInt(32.W))
-      val isOp2RValue = Input(Bool())
-      val r1ValueIn = Input(UInt(32.W))
-      val r1ValueOut = Output(UInt(32.W))
-      val r1AddressIn = Input(UInt(5.W))
-      val r1AddressOut = Output(UInt(5.W))
-      val r2ValueIn = Input(UInt(32.W))
-      val r2ValueOut = Output(UInt(32.W))
-      val r2AddressIn = Input(UInt(5.W))
-      val r2AddressOut = Output(UInt(5.W))
-      val returnAddrIn = Input(UInt(32.W))
-      val returnAddrOut = Output(UInt(32.W))
-      val jumpIn = Input(Bool())
-      val jumpOut = Output(Bool())
-      val ALUopIn = Input(UInt(4.W))
-      val ALUopOut = Output(UInt(4.W))
-      val branchTypeIn = Input(UInt(3.W))
-      val branchTypeOut = Output(UInt(3.W))
-      val writeAddrIn = Input(UInt(5.W))
-      val writeAddrOut = Output(UInt(5.W))
-      val writeEnableIn = Input(Bool())
-      val writeEnableOut = Output(Bool())
-      val memReadIn = Input(Bool())
-      val memReadOut = Output(Bool())
-      val memWriteIn = Input(Bool())
-      val memWriteOut = Output(Bool())
-
-      val forwardEx = Input(Bool())
-      val forwardExAddr = Input(UInt(5.W))
-      val forwardExData = Input(UInt(32.W))
-      val forwardMem = Input(Bool())
-      val forwardMemAddr = Input(UInt(5.W))
-      val forwardMemData = Input(UInt(32.W))
-      val forwardWb = Input(Bool())
-      val forwardWbAddr = Input(UInt(5.W))
-      val forwardWbData = Input(UInt(32.W))
+      val in = Input(new IDBarrierIO)
+      val out = Output(new IDBarrierIO)
+      val stall = Output(Bool())
     })
 
-  val isOp1RValue = RegInit(Bool(), false.B)
-  isOp1RValue := io.isOp1RValue
-  val isOp2RValue = RegInit(Bool(), false.B)
-  isOp2RValue := io.isOp2RValue
-
-  val op1 = RegInit(SInt(32.W), 0.S)
-  op1 := Mux(
-    io.forwardEx && io.isOp1RValue && io.r1AddressIn === io.forwardExAddr,
-    io.forwardExData.asSInt(),
-    Mux(
-      io.forwardMem && io.isOp1RValue && io.r1AddressIn === io.forwardMemAddr,
-      io.forwardMemData.asSInt(),
-      Mux(
-        io.forwardWb && io.isOp1RValue && io.r1AddressIn === io.forwardWbAddr,
-        io.forwardWbData.asSInt(),
-        io.op1in.asSInt(),
-      ),
-    ),
-  )
-  io.op1out := op1
-
-  val op2 = RegInit(SInt(32.W), 0.S)
-  op2 := Mux(
-    io.forwardEx && io.isOp2RValue && io.r2AddressIn === io.forwardExAddr,
-    io.forwardExData.asSInt(),
-    Mux(
-      io.forwardMem && io.isOp2RValue && io.r2AddressIn === io.forwardMemAddr,
-      io.forwardMemData.asSInt(),
-      Mux(
-        io.forwardWb && io.isOp2RValue && io.r2AddressIn === io.forwardWbAddr,
-        io.forwardWbData.asSInt(),
-        io.op2in,
-      ),
-    ),
-  )
-  io.op2out := op2
-
-  val r1Value = RegInit(UInt(32.W), 0.U)
-  r1Value := Mux(
-    io.forwardEx && io.r1AddressIn === io.forwardExAddr,
-    io.forwardExData,
-    Mux(
-      io.forwardMem && io.isOp1RValue && io.r1AddressIn === io.forwardMemAddr,
-      io.forwardMemData,
-      Mux(
-        io.forwardWb && io.isOp1RValue && io.r1AddressIn === io.forwardWbAddr,
-        io.forwardWbData,
-        io.r1ValueIn,
-      ),
-    ),
-  )
-  io.r1ValueOut := r1Value
-
-  val r1Address = RegInit(UInt(5.W), 0.U)
-  r1Address :=  io.r1AddressIn
-  io.r1AddressOut := r1Address
-
-  val r2Value = RegInit(UInt(32.W), 0.U)
-  r2Value := Mux(
-    io.forwardEx && io.r2AddressIn === io.forwardExAddr,
-    io.forwardExData,
-    Mux(
-      io.forwardMem && io.r2AddressIn === io.forwardMemAddr,
-      io.forwardMemData,
-      Mux(
-        io.forwardWb && io.r2AddressIn === io.forwardWbAddr,
-        io.forwardWbData,
-        io.r2ValueIn,
-      ),
-    ),
-  )
-  io.r2ValueOut := r2Value
-
-  val r2Address = RegInit(UInt(5.W), 0.U)
-  r2Address := io.r2AddressIn
-  io.r2AddressOut := r2Address
-
-  val returnAddr = RegInit(UInt(32.W), 0.U)
-  returnAddr := io.returnAddrIn
-  io.returnAddrOut := returnAddr
-
-  val jump = RegInit(UInt(32.W), 0.U)
-  jump := io.jumpIn
-  io.jumpOut := jump
-
-  val ALUop = RegInit(UInt(4.W), 0.U)
-  ALUop := io.ALUopIn
-  io.ALUopOut := ALUop
-
-  val branchType = RegInit(UInt(5.W), 0.U)
-  branchType := io.branchTypeIn
-  io.branchTypeOut := branchType
-
-  val writeAddr = RegInit(UInt(5.W), 0.U)
-  writeAddr := io.writeAddrIn
-  io.writeAddrOut := writeAddr
-
-  val writeEnable = RegInit(Bool(), false.B)
-  writeEnable := io.writeEnableIn
-  io.writeEnableOut := writeEnable
-
-  val memRead = RegInit(Bool(), false.B)
-  memRead := io.memReadIn
-  io.memReadOut := memRead
-
-  val memWrite = RegInit(Bool(), false.B)
-  memWrite := io.memWriteIn
-  io.memWriteOut := memWrite
+  val delay = Reg(new IDBarrierIO)
+  delay := io.in
+  io.out := delay
+  io.stall := io.out.memRead
 }
diff --git a/src/main/scala/IF.scala b/src/main/scala/IF.scala
index 35b0842..b4fd54a 100644
--- a/src/main/scala/IF.scala
+++ b/src/main/scala/IF.scala
@@ -48,7 +48,7 @@ class InstructionFetch extends MultiIOModule {
     */
   io.PC := PC
   IMEM.io.instructionAddress := PC
-  PC := Mux(io.stall, PC, Mux(io.branch, io.branchAddress, PC + 4.U))
+  PC := Mux(io.branch, io.branchAddress, Mux(io.stall, PC, PC + 4.U))
 
   val instruction = Wire(new Instruction)
   instruction := IMEM.io.instruction.asTypeOf(new Instruction)
diff --git a/src/main/scala/IFBarrier.scala b/src/main/scala/IFBarrier.scala
index 0523874..49555ea 100644
--- a/src/main/scala/IFBarrier.scala
+++ b/src/main/scala/IFBarrier.scala
@@ -10,11 +10,37 @@ class IFBarrier extends MultiIOModule {
       val PCout = Output(UInt(32.W))
       val instructionIn = Input(new Instruction)
       val instructionOut = Output(new Instruction)
+      val stall = Input(Bool())
+      val flush = Input(Bool())
     })
 
   val PC = RegInit(UInt(32.W), 0.U)
-  PC := io.PCin
+  PC := Mux(io.stall, PC, io.PCin)
   io.PCout := PC
 
-  io.instructionOut := io.instructionIn
+  val instruction = Reg(new Instruction)
+  val replay = RegInit(Bool(), false.B)
+  val flushRemaining = RegInit(UInt(2.W), 0.U)
+  flushRemaining := Mux(
+    io.flush,
+    2.U,
+    Mux(
+      flushRemaining === 0.U,
+      0.U,
+      flushRemaining - 1.U
+    )
+  )
+
+  replay := io.stall
+  instruction := io.instructionIn
+
+  io.instructionOut := Mux(
+    io.stall || io.flush || flushRemaining > 0.U,
+    Instruction.NOP,
+    Mux(
+      replay,
+      instruction,
+      io.instructionIn
+    )
+  )
 }
diff --git a/src/main/scala/MEMBarrier.scala b/src/main/scala/MEMBarrier.scala
index ed36d1a..bef66ad 100644
--- a/src/main/scala/MEMBarrier.scala
+++ b/src/main/scala/MEMBarrier.scala
@@ -2,44 +2,50 @@ package FiveStage
 import chisel3._
 import chisel3.experimental.MultiIOModule
 
+class MEMBarrierIO extends Bundle {
+  val data = UInt(32.W)
+  val writeAddr = UInt(5.W)
+  val writeEnable = Bool()
+}
+
 class MEMBarrier extends MultiIOModule {
   val io = IO(
     new Bundle {
-      val dataIn = Input(UInt(32.W))
-      val dataOut = Output(UInt(32.W))
-      val writeAddrIn = Input(UInt(5.W))
-      val writeAddrOut = Output(UInt(5.W))
-      val writeEnableIn = Input(Bool())
-      val writeEnableOut = Output(Bool())
+      val in = Input(new MEMBarrierIO)
+      val out = Output(new MEMBarrierIO)
       val memRead = Input(Bool())
-      val forwardMem = Output(Bool())
-      val forwardMemAddr = Output(UInt(5.W))
-      val forwardMemData = Output(UInt(32.W))
-      val forwardWb = Output(Bool())
-      val forwardWbAddr = Output(UInt(5.W))
-      val forwardWbData = Output(UInt(32.W))
+      val forwardMem = Output(new Forwarding)
+      val forwardWb = Output(new Forwarding)
+      val forwardId = Output(new Forwarding)
     })
 
   val memRead = RegInit(Bool(), false.B)
   memRead := io.memRead
 
-  val data = RegInit(UInt(32.W), 0.U)
-  data := io.dataIn
-  io.dataOut := Mux(memRead, io.dataIn, data)
+  val delay = Reg(new MEMBarrierIO)
+  val delay2 = Reg(new MEMBarrierIO)
 
-  val writeAddr = RegInit(UInt(5.W), 0.U)
-  writeAddr := io.writeAddrIn
-  io.writeAddrOut := writeAddr
+  delay := io.in
+  io.out := delay
+  io.out.data := Mux(memRead, io.in.data, delay.data)
+  delay2 := io.out
 
-  val writeEnable = RegInit(Bool(), false.B)
-  writeEnable := io.writeEnableIn
-  io.writeEnableOut := writeEnable
+  io.forwardMem.write := io.in.writeEnable && !io.memRead
+  io.forwardMem.writeAddr := io.in.writeAddr
+  io.forwardMem.writeData := io.in.data
 
-  io.forwardMem := io.writeEnableIn && !io.memRead
-  io.forwardMemAddr := io.writeAddrIn
-  io.forwardMemData := io.dataIn
+  io.forwardWb.write := io.out.writeEnable
+  io.forwardWb.writeAddr := io.out.writeAddr
+  io.forwardWb.writeData := io.out.data
 
-  io.forwardWb := writeEnable
-  io.forwardWbAddr := writeAddr
-  io.forwardWbData := Mux(memRead, io.dataIn, data)
+  io.forwardId.write := delay2.writeEnable
+  io.forwardId.writeAddr := delay2.writeAddr
+  io.forwardId.writeData := delay2.data
+}
+
+class Forwarding extends Bundle {
+  val writeAddr = UInt(5.W)
+  val writeData = UInt(32.W)
+  val write = Bool()
+  def valid = write && (writeAddr =/= 0.U)
 }
diff --git a/src/test/resources/tests/addload.s b/src/test/resources/tests/addload.s
new file mode 100644
index 0000000..ea26489
--- /dev/null
+++ b/src/test/resources/tests/addload.s
@@ -0,0 +1,5 @@
+main:
+    addi    a5,a5,2
+    lw    a4,-36(s0)
+    add    a5,a4,a5
+    done
\ No newline at end of file
diff --git a/src/test/resources/tests/branch.s b/src/test/resources/tests/branch.s
index 0761396..15f6937 100644
--- a/src/test/resources/tests/branch.s
+++ b/src/test/resources/tests/branch.s
@@ -4,4 +4,7 @@ main:
 loop:
     addi x2, x2, 1
     blt x2, x1, loop
+    nop
+    nop
+    nop
     done
\ No newline at end of file
diff --git a/src/test/resources/tests/fucked.s b/src/test/resources/tests/fucked.s
new file mode 100644
index 0000000..2661c5e
--- /dev/null
+++ b/src/test/resources/tests/fucked.s
@@ -0,0 +1,18 @@
+main:
+    addi x2, x2, 0
+    j loop
+    nop
+    nop
+
+
+loop:
+    addi x1, x1, 32
+    blt x1, x2, loop
+    nop
+    nop
+    nop
+    nop
+    j end
+
+end:
+    done
diff --git a/src/test/resources/tests/jump.s b/src/test/resources/tests/jump.s
index cb18ee4..a72a9df 100644
--- a/src/test/resources/tests/jump.s
+++ b/src/test/resources/tests/jump.s
@@ -1,6 +1,8 @@
 main:
     jal x1, end
     addi x1, x1, 0
+    nop
+    nop
     done
 
 end:
diff --git a/src/test/scala/Manifest.scala b/src/test/scala/Manifest.scala
index eaf8823..d10ed8a 100644
--- a/src/test/scala/Manifest.scala
+++ b/src/test/scala/Manifest.scala
@@ -19,7 +19,7 @@ import LogParser._
 
 object Manifest {
 
-  val singleTest = "load3.s"
+  val singleTest = "branch.s"
 
   val nopPadded = false