diff --git a/src/main/scala/Barriers.scala b/src/main/scala/Barriers.scala
index b010faa..7f4581b 100644
--- a/src/main/scala/Barriers.scala
+++ b/src/main/scala/Barriers.scala
@@ -5,92 +5,78 @@ import chisel3.util.ShiftRegister
 
 
 class IFBarrier extends Module {
+  class IFBundle extends Bundle {
+    val PC = UInt()
+    val instruction = new Instruction
+  }
+
   val io = IO(new Bundle {
-    val in = Input(new Bundle {
-      val PC = UInt()
-      val instruction = new Instruction
-    })
-    val stage2 = Output(new Bundle {
-      val instruction = new Instruction
-    })
-    val stage3 = Output(new Bundle {
-      val PC = UInt()
-    })
-    val stage5 = Output(new Bundle {
-      val PC = UInt()
-      val instruction = new Instruction
-    })
+    val in = Input(new IFBundle)
+    val out = Output(new IFBundle)
   })
 
-  io.stage2 := io.in
-  io.stage3 := ShiftRegister(io.in, 1)
-  io.stage5 := ShiftRegister(io.in, 3)
+  io.out.instruction := io.in.instruction
+  io.out.PC := ShiftRegister(io.in.PC, 1)
 }
 
 
 class IDBarrier extends Module {
+  class IDBundle extends Bundle {
+    val PC = UInt()
+    val instruction = new Instruction
+    val controlSignals = new ControlSignals
+    val branchType     = UInt(3.W)
+    val reg1  = UInt(32.W)
+    val reg2  = UInt(32.W)
+    val imm   = UInt(32.W)
+    val ALUop = UInt(4.W)
+  }
+
   val io = IO(new Bundle {
-    val in = Input(new Bundle {
-      val controlSignals = new ControlSignals
-      val branchType     = UInt(3.W)
-      val reg1  = UInt(32.W)
-      val reg2  = UInt(32.W)
-      val imm   = UInt(32.W)
-      val ALUop = UInt(4.W)
-    })
-    val stage3 = Output(new Bundle {
-      val controlSignals = new ControlSignals
-      val branchType     = UInt(3.W)
-      val reg1  = UInt(32.W)
-      val reg2  = UInt(32.W)
-      val imm   = UInt(32.W)
-      val ALUop = UInt(4.W)
-    })
-    val stage4 = Output(new Bundle {
-      val controlSignals = new ControlSignals
-      val reg2  = UInt(32.W)
-    })
-    val stage5 = Output(new Bundle {
-      val controlSignals = new ControlSignals
-      val reg2  = UInt(32.W)
-    })
+    val in = Input(new IDBundle)
+    val out = Output(new IDBundle)
   })
 
-  io.stage3 := io.in
-  io.stage4 := ShiftRegister(io.in, 1)
-  io.stage5 := ShiftRegister(io.in, 2)
+  io.out := ShiftRegister(io.in, 1)
 }
 
 
 class EXBarrier extends Module {
+  class EXBundle extends Bundle {
+    val PC = UInt()
+    val instruction = new Instruction
+    val controlSignals = new ControlSignals
+    val reg2  = UInt(32.W)
+    val result = UInt(32.W)
+    val branch = Bool()
+  }
+
   val io = IO(new Bundle {
-    val in = Input(new Bundle {
-      val result = UInt(32.W)
-      val branch = Bool()
-    })
-    val stage4 = Output(new Bundle {
-      val result = UInt(32.W)
-      val branch = Bool()
-    })
-    val stage5 = Output(new Bundle {
-      val result = UInt(32.W)
-    })
+    val in = Input(new EXBundle)
+    val out = Output(new EXBundle)
   })
 
-  io.stage4 := io.in
-  io.stage5 := ShiftRegister(io.in, 1)
+  io.out := ShiftRegister(io.in, 1)
 }
 
 
 class MEMBarrier extends Module {
+  class MEMBundle extends Bundle {
+    val PC = UInt()
+    val instruction = new Instruction
+    val controlSignals = new ControlSignals
+    val result = UInt(32.W)
+    val dataOut = UInt(32.W)
+  }
+
   val io = IO(new Bundle {
-    val in = Input(new Bundle {
-      val dataOut = UInt(32.W)
-    })
-    val stage5 = Output(new Bundle {
-      val dataOut = UInt(32.W)
-    })
+    val in = Input(new MEMBundle)
+    val out = Output(new MEMBundle)
   })
 
-  io.stage5 := io.in
+  io.out.PC             := ShiftRegister(io.in.PC, 1)
+  io.out.instruction    := ShiftRegister(io.in.instruction, 1)
+  io.out.controlSignals := ShiftRegister(io.in.controlSignals, 1)
+  io.out.result         := ShiftRegister(io.in.result, 1)
+  io.out.dataOut        := io.in.dataOut
 }
diff --git a/src/main/scala/CPU.scala b/src/main/scala/CPU.scala
index 8caf1f4..3ed99da 100644
--- a/src/main/scala/CPU.scala
+++ b/src/main/scala/CPU.scala
@@ -55,53 +55,59 @@ class CPU extends MultiIOModule {
     TODO: Your code here
     */
   // Stage 1
-
-  // The IF gets the instruction, but it won't be ready until the next cycle
-  //printf(p"S1: PC=${Hexadecimal(IF.io.PC)} || ")
   IFBarrier.in := IF.io
 
   // Stage 2
-  ID.io.instruction := IFBarrier.stage2.instruction
-  IDBarrier.in      := ID.io
+  ID.io.instruction            := IFBarrier.out.instruction
+  IDBarrier.in.instruction     := IFBarrier.out.instruction
+  IDBarrier.in.PC              := IFBarrier.out.PC
+  IDBarrier.in.controlSignals  := ID.io.controlSignals
+  IDBarrier.in.branchType      := ID.io.branchType
+  IDBarrier.in.reg1            := ID.io.reg1
+  IDBarrier.in.reg2            := ID.io.reg2
+  IDBarrier.in.imm             := ID.io.imm
+  IDBarrier.in.ALUop           := ID.io.ALUop
 
   // Stage 3
-  EX.io.PC             := IFBarrier.stage3.PC
-  EX.io.controlSignals := IDBarrier.stage3.controlSignals
-  EX.io.branchType     := IDBarrier.stage3.branchType
-  EX.io.ALUop          := IDBarrier.stage3.ALUop
-  EX.io.reg1           := IDBarrier.stage3.reg1
-  EX.io.reg2           := IDBarrier.stage3.reg2
-  EX.io.imm            := IDBarrier.stage3.imm
-
-  EXBarrier.in := EX.io
-  //printf(p"S3: ALUop=${ID.io.ALUop}, reg1=${ID.io.reg1.asSInt}, reg2=${ID.io.reg2.asSInt}, imm=${ID.io.imm.asSInt} || ")
+  EX.io.PC                     := IDBarrier.out.PC
+  EX.io.controlSignals         := IDBarrier.out.controlSignals
+  EX.io.branchType             := IDBarrier.out.branchType
+  EX.io.reg1                   := IDBarrier.out.reg1
+  EX.io.reg2                   := IDBarrier.out.reg2
+  EX.io.imm                    := IDBarrier.out.imm
+  EX.io.ALUop                  := IDBarrier.out.ALUop
+
+  EXBarrier.in.PC              := IDBarrier.out.PC
+  EXBarrier.in.instruction     := IDBarrier.out.instruction
+  EXBarrier.in.controlSignals  := IDBarrier.out.controlSignals
+  EXBarrier.in.reg2            := IDBarrier.out.reg2
+  EXBarrier.in.result          := EX.io.result
+  EXBarrier.in.branch          := EX.io.branch
 
   // Stage 4
-  MEM.io.dataIn      := IDBarrier.stage4.reg2
-  MEM.io.writeEnable := IDBarrier.stage4.controlSignals.memWrite
-  MEM.io.dataAddress := EXBarrier.stage4.result
-
-  IF.io.jumpEnable := EXBarrier.stage4.branch
-  IF.io.jumpAddr   := EXBarrier.stage4.result
+  MEM.io.dataIn                := EXBarrier.out.reg2
+  MEM.io.writeEnable           := EXBarrier.out.controlSignals.memWrite
+  MEM.io.dataAddress           := EXBarrier.out.result
 
-  MEMBarrier.in := MEM.io
+  IF.io.jumpEnable             := EXBarrier.out.branch
+  IF.io.jumpAddr               := EXBarrier.out.result
 
-  //printf(p"S4: res=${EX.io.result} || ")
+  MEMBarrier.in.PC             := EXBarrier.out.PC
+  MEMBarrier.in.instruction    := EXBarrier.out.instruction
+  MEMBarrier.in.controlSignals := EXBarrier.out.controlSignals
+  MEMBarrier.in.result         := EXBarrier.out.result
+  MEMBarrier.in.dataOut        := MEM.io.dataOut
 
   // Stage 5
-  ID.io.writeEnable := IDBarrier.stage5.controlSignals.regWrite
-  ID.io.writeAddr   := IFBarrier.stage5.instruction.registerRd
-  ID.io.writeData   := EXBarrier.stage5.result
+  ID.io.writeEnable            := MEMBarrier.out.controlSignals.regWrite
+  ID.io.writeAddr              := MEMBarrier.out.instruction.registerRd
+  ID.io.writeData              := MEMBarrier.out.result
 
-  when (IDBarrier.stage5.controlSignals.memToReg) {
-    ID.io.writeData := MEMBarrier.stage5.dataOut
+  when (MEMBarrier.out.controlSignals.memToReg) {
+    ID.io.writeData := MEMBarrier.out.dataOut
   }
 
-  when (IDBarrier.stage5.controlSignals.jump) {
-    ID.io.writeData := IFBarrier.stage5.PC + 4.U
+  when (MEMBarrier.out.controlSignals.jump) {
+    ID.io.writeData := MEMBarrier.out.PC + 4.U
   }
-
-  //printf(p"S5: WB=${ID.io.writeData}")
-
-  //printf("\n\n")
 }
diff --git a/src/main/scala/Execute.scala b/src/main/scala/Execute.scala
index 3b2bcf6..464885b 100644
--- a/src/main/scala/Execute.scala
+++ b/src/main/scala/Execute.scala
@@ -31,6 +31,6 @@ class Execute extends Module {
   comparator.data1      := io.reg1
   comparator.data2      := io.reg2
 
-  io.result := ShiftRegister(alu.result, 1)
-  io.branch := ShiftRegister(io.controlSignals.jump | (comparator.result & io.controlSignals.branch), 1)
+  io.result := alu.result
+  io.branch := io.controlSignals.jump | (comparator.result & io.controlSignals.branch)
 }
diff --git a/src/main/scala/ID.scala b/src/main/scala/ID.scala
index 6394a61..1711be9 100644
--- a/src/main/scala/ID.scala
+++ b/src/main/scala/ID.scala
@@ -38,14 +38,6 @@ class InstructionDecode extends MultiIOModule {
   val registers = Module(new Registers)
   val decoder   = Module(new Decoder).io
 
-  val reg1  = RegInit(UInt(32.W), 0.U)
-  val reg2  = RegInit(UInt(32.W), 0.U)
-  val imm   = RegInit(UInt(32.W), 0.U)
-  val ALUop = RegInit(UInt(4.W), 0.U)
-  val branchType = RegInit(UInt(3.W), 0.U)
-  val controlSignals = Reg(new ControlSignals)
-
-
   /**
     * Setup. You should not change this code
     */
@@ -65,19 +57,12 @@ class InstructionDecode extends MultiIOModule {
   registers.io.writeAddress := io.writeAddr
   registers.io.writeData    := io.writeData
 
-  ALUop := decoder.ALUop
-  io.ALUop := ALUop
-  controlSignals := decoder.controlSignals
-  io.controlSignals := controlSignals
-  branchType := decoder.branchType
-  io.branchType := branchType
-
-  reg1 := registers.io.readData1
-  reg2 := registers.io.readData2
+  io.ALUop := decoder.ALUop
+  io.controlSignals := decoder.controlSignals
+  io.branchType := decoder.branchType
 
-  io.reg1 := reg1
-  io.reg2 := reg2
-  io.imm  := imm
+  io.reg1 := registers.io.readData1
+  io.reg2 := registers.io.readData2
 
 
   val immTypeMap = Array(
@@ -89,7 +74,7 @@ class InstructionDecode extends MultiIOModule {
     ImmFormat.SHAMT -> 0.S, // TODO: Implement SHAMT
   )
 
-  imm := MuxLookup(
+  io.imm := MuxLookup(
     decoder.immType,
     0.S,
     immTypeMap,