package FiveStage

import chisel3._
import chisel3.experimental.MultiIOModule

class IFBarrier extends MultiIOModule {
  val io = IO(
    new Bundle {
      val PCin = Input(UInt(32.W))
      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 := Mux(io.stall, PC, io.PCin)
  io.PCout := PC

  val instruction = Reg(new Instruction)
  val replay = RegInit(Bool(), false.B)
  val flushRemaining = RegInit(UInt(2.W), 0.U)
  flushRemaining := Mux(
    io.flush,
    1.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
    )
  )
}