package vexiiriscv.misc

import spinal.core._
import spinal.core.fiber.Retainer
import spinal.lib._
import spinal.lib.fsm._
import spinal.lib.misc.plugin.FiberPlugin
import vexiiriscv.Global._
import vexiiriscv.execute.{CsrAccessPlugin, CsrListFilter, CsrRamPlugin, CsrRamService, ExecuteLanePlugin}
import vexiiriscv.riscv._
import vexiiriscv.riscv.Riscv._
import vexiiriscv._
import vexiiriscv.decode.Decode
import vexiiriscv.decode.Decode.{INSTRUCTION_SLICE_COUNT, INSTRUCTION_SLICE_COUNT_WIDTH, INSTRUCTION_WIDTH}
import vexiiriscv.execute.lsu.{LsuL1Service, LsuService}
import vexiiriscv.fetch.{Fetch, FetchL1Service, InitService, PcService}
import vexiiriscv.memory.AddressTranslationService
import vexiiriscv.prediction.{HistoryPlugin, Prediction}
import vexiiriscv.schedule.Ages

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer


case class TrapSpec(bus : Flow[Trap], age : Int, subAge : Int)
case class Trap(laneAgeWidth : Int, full : Boolean) extends Bundle{
  val exception = Bool()
  val tval = TVAL()
  val tval2 = TVAL()
  val code = CODE()
  val arg = TRAP_ARG()
  val laneAge = full generate UInt(laneAgeWidth bits)
  val hartId = full generate HART_ID()

  def toRaw(): Trap = {
    val r = new Trap(laneAgeWidth, false)
    r.assignSomeByName(this)
    r
  }
}

case class extendedInterruptPriority() extends Bundle {
  /*
   * A high-priority interrupt means its default priority is higher than
   * external interrupt.
   */
  val high = Bool()
  val external = Bool()
  val local = UInt(TRAP_IPRIO_WIDTH bits)
  val default = UInt(InterruptInfo.defaultOrderWidth bits)
}

object extendedInterruptPriority {
  implicit def priorityAsUInt(p: extendedInterruptPriority) = p.asBits.asUInt
}

case class InterruptState(idWidth: Int, priorityWidth: Int) extends Bundle {
  val id = UInt(idWidth bits)
  val priority = UInt(priorityWidth bits)
  val privilege = PrivilegeMode.TYPE()
  val valid = Bool()
}

trait CauseUser{
  def getCauseWidthMin() : Int
}

/**
 * fetch (page fault, access fault)
 * decode (illegal)
 * execute
 * - miss aligned load/store/branch
 * - page fault, access fault
 * - ecall, ebreak
 *
 * To do a trap request :
 * - Flush all youngers instructions
 * - Store cause / tval
 * - Rise trap flag in the pipe (this will disable side-effects)
 */
trait TrapService extends Area{
  val trapLock = Retainer()
  val traps = ArrayBuffer[TrapSpec]()

  // age => main priority, bigger win, meaning it a stage further down the pipeline
  // laneAge => On the same pipeline stage, specify at run time local age (in the stage over multiple lanes). smaller => win
  // subAge => Static priority between trap ports with same age and same laneAge, bigger win
  def newTrap(age: Int, laneAgeWidth: Int, subAge : Int = 0): Flow[Trap] = {
    traps.addRet(TrapSpec(Flow(Trap(laneAgeWidth, true)), age, subAge)).bus
  }

  def trapHandelingAt : Int
  val trapPendings = ArrayBuffer[Bits]()
  def newTrapPending() = trapPendings.addRet(Bits(Global.HART_COUNT bits))
}

case class TrapPending() extends Bundle{
  val hartId = Global.HART_ID()
}


object TrapReason{
  val INTERRUPT = 0
  val PRIV_RET = 1 //Used for MRET, SRET
  val FENCE_I = 2
  val DEBUG_TRIGGER = CSR.MCAUSE_ENUM.BREAKPOINT // Need to match to reduce hardware
  val REDO = 4 // Allows the retry the current instruction
  val NEXT = 5 // Allows to fence the current instruction against all the future one by flushing the younger instructions and resuming on the next instruction.
  val SFENCE_VMA = 6
  val MMU_REFILL = 7
  val WFI = 8
  val HFENCE_GMA = 9
  val SMMU_REFILL = 10
}

/**
 * Trap args are used in a few cases to provide additional information.
 * Mainly, when using TrapReason MMU_REFILL, to specify which TLB storage need to be refilled.
 */
object TrapArg{
  val LOAD = 0
  val STORE = 1
  val FETCH = 2
  /* Special Value for FETCH refill request from LSU (HLVX instruction) */
  val FETCH_LSU = 3
}

/**
 * Mainly, this plugin implement a state-machine which handle regular traps (interrupt/exception), but also a set of
 * special "hardware" trap used by the CPU to handle special cases as instruction retry/fences, MMU refill, ...
 *
 * Also, as VexiiRiscv implements a few large CSR directly into a shared memory (mepc, mtvec, ...), the TrapPlugin state-machine handles
 * the hardware read/write with those CSR (durring trap, mret, ...).
 */
class TrapPlugin(val trapAt : Int, val recordHtinst : Boolean) extends FiberPlugin with TrapService {
  override def trapHandelingAt: Int = trapAt

  def askWake(hartId : Int) = api.harts(hartId).askWake := True

  val api = during build new Area{
    val harts = for(hartId <- 0 until HART_COUNT) yield new Area{
      // Mainly used to interface with the RISC-V debugger implemented in the PrivilegedPlugin
      val redo = False // Instruction need to retried.
      val askWake = False // ex : wake the CPU as an interrupt is pending
      val rvTrap = False // Instruction got a trap
      val fsmBusy = Bool() //TrapPlugin FSM is doing work, hold on

      val holdPrivChange = False
    }
  }

  val logic = during setup new Area{
    val priv = host[PrivilegedPlugin]
    val cap = host[CsrAccessPlugin]
    val pp = host[PipelineBuilderPlugin]
    val fl1p = host.get[FetchL1Service]
    val lsu = host.get[LsuService]
    val lsul1 = host.get[LsuL1Service]
    val pcs = host[PcService]
    val hp = host.get[HistoryPlugin]
    val ats = host.find[AddressTranslationService](!_.isShadowMmu)
    val sats = host.find[AddressTranslationService](_.isShadowMmu)
    val withRam = host.get[CsrRamService].nonEmpty
    val crs = withRam generate host[CsrRamService]
    val invalidationLocks = retains(fl1p.map(_.invalidationRetainer).toList ++ lsu.map(_.invalidationRetainer))
    val buildBefore = retains(List(pp.elaborationLock, pcs.elaborationLock, cap.csrLock, ats.portsLock, sats.portsLock) ++ hp.map(_.elaborationLock))
    val ramPortRetainers = withRam generate crs.portLock()
    awaitBuild()

    val fetchL1Invalidate = fl1p.nonEmpty generate (0 until HART_COUNT).map(hartId => fl1p.get.newInvalidationPort())
    val lsuL1Invalidate = lsu.nonEmpty generate (0 until HART_COUNT).map(hartId => lsu.get.newInvalidationPort())

    invalidationLocks.release()

    val trapArgWidths = ArrayBuffer[Int](3)
    if(ats.mayNeedRedo) trapArgWidths += 3+ats.getStorageIdWidth()
    /*
     * A standalone field is required, because there is no guarantee
     * that ats.storageId == sats.storageId
     */
    if(sats.mayNeedRedo) trapArgWidths += 3+ats.getStorageIdWidth()+sats.getStorageIdWidth()
    TRAP_ARG_WIDTH.set(trapArgWidths.max)

    trapLock.await()

    val initHold = Bool()

    val harts = for(hartId <- 0 until HART_COUNT) yield new Area{
      val csr = priv.logic.harts(hartId)

      // Used to read/write the CSR stored in RAM as mepc, mtvec, ... in the state machine
      val crsPorts = withRam generate new Area{
        val read = crs.ramReadPort(CsrRamService.priority.TRAP)
        read.valid := False
        read.address.assignDontCare()

        val write = crs.ramWritePort(CsrRamService.priority.TRAP)
        write.valid := False
        write.address.assignDontCare()
        write.data.assignDontCare()
      }

      def privilegeMux[T <: Data](payloads: mutable.LinkedHashMap[Int, T], privilege: SInt): T = {
        val ret = CombInit(payloads(PrivilegeMode.M))
        switch(privilege) {
          for ((mode, data) <- payloads) {
            is(mode) {
              ret := data
            }
          }
        }
        ret
      }

      //Process interrupt request, code and privilege
      val interrupt = new Area {
        val privilegeAllowInterrupts = mutable.LinkedHashMap[Int, Bool]()
        var privilegs: List[Int] = Nil
        privilegs :+= PrivilegeMode.M
        privilegeAllowInterrupts += PrivilegeMode.M -> (csr.m.status.mie || !csr.withMachinePrivilege)

        if (priv.p.withSupervisor) {
          privilegs = PrivilegeMode.S :: privilegs
          privilegeAllowInterrupts += PrivilegeMode.S -> ((csr.s.status.sie && !csr.withMachinePrivilege) || !csr.withSupervisorPrivilege)
        }

        if (priv.p.withUserTrap) {
          privilegs = PrivilegeMode.U :: privilegs
          ??? // privilegeAllowInterrupts += 1 -> ((ustatus.UIE && !setup.supervisorPrivilege))
        }

        if (priv.p.withHypervisor) {
          privilegs = PrivilegeMode.VS :: privilegs
          privilegeAllowInterrupts += PrivilegeMode.VS -> ((csr.vs.status.sie && csr.withGuestPrivilege) || !csr.withVirtualSupervisorPrivilege)
        }

        val privilegeTriggers = privilegs.map(p => new Area {
          val interrupts = csr.spec.interrupt.filter(
            i => (i.privilege <= p) && //EX : Machine timer interrupt can't go into
            privilegs.filter(_ > p).forall(e => i.delegators.exists(_.privilege == e)) // EX : Supervisor timer need to have machine mode delegator
          ).map(i => {
            val int = InterruptState(CODE_WIDTH, TRAP_IPRIO_WIDTH)
            val delegUpOn = i.delegators.filter(_.privilege > p).map(_.enable).fold(True)(_ && _)
            val delegDownOff = !i.delegators.filter(_.privilege <= p).map(_.enable).orR

            int.id := i.id
            int.priority := i.priority
            int.privilege := p
            int.valid := i.cond && delegUpOn && delegDownOff

            int
          })

          val result = RegNext(interrupts.reduceBalancedTree((a, b) => {
            val priorityCheck = (a.priority < b.priority)
            val takeA = !b.valid || (a.valid && (priorityCheck))
            takeA ? a | b
          }))

          val triggerId = B(result.id).andMask(result.valid).resized
          val triggerPriority = result.priority.andMask(result.valid).resized
          p match {
            case PrivilegeMode.M  => {
              csr.m.candidate.interrupt := triggerId
              csr.m.candidate.priority  := triggerPriority
            }
            case PrivilegeMode.S  => {
              csr.s.candidate.interrupt := triggerId
              csr.s.candidate.priority  := triggerPriority
            }
            case PrivilegeMode.VS => {
              csr.vs.candidate.interrupt := triggerId
              csr.vs.candidate.priority  := triggerPriority
            }
          }

          val injects = csr.spec.injectedInterrupt.filter(i => i.privilege == p)
          val injectInterrupts = injects.nonEmpty generate injects.map(i => {
            val int = InterruptState(CODE_WIDTH, widthOf(extendedInterruptPriority()))
            int.id := i.id.resized
            int.priority := i.priority.resized
            int.privilege := p
            int.valid := i.cond

            int
          })

          val injectResult = injects.nonEmpty generate RegNext(injectInterrupts.reduceBalancedTree((a, b) => {
            val priorityCheck = a.priority < b.priority
            val takeA = !b.valid || (a.valid && (priorityCheck))
            takeA ? a | b
          }))

          val triggered = injects.nonEmpty.mux(injectResult, result)

          val effectiveInterrupt = B(triggered.id).andMask(triggered.valid).resized
          val effectivePriority = UInt(8 bits)
          val withLocalOrder = widthOf(triggered.priority) > InterruptInfo.defaultOrderWidth
          val localPriority = withLocalOrder.mux(triggered.priority.drop(InterruptInfo.defaultOrderWidth).resize(TRAP_IPRIO_WIDTH).asUInt, triggered.priority)
          val defaultPriority = triggered.priority.resize(InterruptInfo.defaultOrderWidth)
          val priorityBuilder = WhenBuilder()
          priorityBuilder.when(!triggered.valid) {
            effectivePriority := 0
          }
          if (TRAP_IPRIO_WIDTH.get > 8) priorityBuilder.when(localPriority > 255) {
            effectivePriority := 255
          }
          priorityBuilder.when(localPriority === 0) {
            val higher = defaultPriority < InterruptInfo.defaultOrder.indexOf((p == PrivilegeMode.M).mux(11, 9))
            effectivePriority := Mux(higher, U(0), U(255))
          }
          priorityBuilder.otherwise {
            effectivePriority := triggered.priority.resized
          }
        })

        val privilegeIndexedTriggers = privilegs.zip(privilegeTriggers)

        /* Link xTOPI register */
        privilegeIndexedTriggers.foreach{case (p, i) => p match {
            case PrivilegeMode.M  => csr.m.topi.interrupt := i.effectiveInterrupt
            case PrivilegeMode.S  => csr.s.topi.interrupt := i.effectiveInterrupt
            case PrivilegeMode.VS => {
              csr.vs.topi.interrupt := i.effectiveInterrupt
              csr.vs.topi.priority  := i.effectivePriority
            }
          }
        }

        val result = privilegeIndexedTriggers.map{case (p, i) => {
          val int = InterruptState(CODE_WIDTH, 0)
          val triggered = i.triggered

          int.id := triggered.id
          int.priority.assignDontCare()
          int.privilege := triggered.privilege
          int.valid := triggered.valid && privilegeAllowInterrupts(p)

          int
        }}.reduceBalancedTree((a, b) => {
          val takeA = !b.valid || (a.valid && (a.privilege > b.privilege))
          takeA ? a | b
        })

        val valid = CombInit(result.valid)
        val code = CombInit(result.id.asBits)
        val targetPrivilege = CombInit(result.privilege)

        if (priv.p.withDebug) {
          valid clearWhen (csr.debug.dcsr.step && !csr.debug.dcsr.stepie)
          valid setWhen (csr.debug.doHalt)
        }

        val validBuffer = RegNext(valid) init(False)
        val pendingInterrupt = validBuffer && priv.api.harts(hartId).allowInterrupts
        csr.int.pending setWhen(pendingInterrupt)

        when(privilegeIndexedTriggers.map(_._2.triggered.valid).orR){
          askWake(hartId)
        }
      }


      // Capture the state of the oldest pending trap requests (from the pipelines)
      val trap = new Area {
        val pending = new Area {
          val requests = traps.map(e => new AgedArbiterUp(e.bus.valid && e.bus.hartId === hartId, e.bus.payload.toRaw(), e.age, e.bus.laneAge, e.subAge))
          val arbiter = new AgedArbiter(requests)
          val state = arbiter.down.toReg
          val pc = Reg(PC)
          val uop = priv.p.withHypervisor generate Reg(Decode.UOP)
          val uopIsCompressed = priv.p.withHypervisor generate Reg(Decode.UOP_IS_COMPRESSED)
          val history = hp.nonEmpty generate Reg(Prediction.BRANCH_HISTORY)
          val slices = Reg(UInt(INSTRUCTION_SLICE_COUNT_WIDTH+1 bits))

          val xret = new Area {
            val sourcePrivilege = state.arg(2 downto 0).asSInt
            val targetPrivilegeMapping = mutable.LinkedHashMap[Int, SInt](
              PrivilegeMode.M -> ((csr.m.status.mpp === PrivilegeMode.M).mux(False, csr.m.status.mpv) ## csr.m.status.mpp).asSInt
            )

            if (priv.p.withSupervisor) targetPrivilegeMapping += PrivilegeMode.S ->
              (priv.p.withHypervisor.mux(csr.h.status.spv, U"0") ## U"0" ## csr.s.status.spp).asSInt

            if (priv.p.withHypervisor) targetPrivilegeMapping += PrivilegeMode.VS -> (U"10" ## csr.vs.status.spp).asSInt
            val targetPrivilege = privilegeMux(targetPrivilegeMapping, sourcePrivilege)
          }
        }

        val exception = new Area {
          val exceptionTargetPrivilegeUncapped = PrivilegeMode(PrivilegeMode.M)
          val code = CombInit(pending.state.code)
          switch(code) {
            for (s <- csr.spec.exception) {
              is(s.id) {
                var exceptionPrivilegs: List[Int] = List(PrivilegeMode.M)

                if (priv.p.withSupervisor) {
                  exceptionPrivilegs = PrivilegeMode.S :: exceptionPrivilegs
                }

                if (priv.p.withHypervisor) {
                  exceptionPrivilegs = PrivilegeMode.VS :: exceptionPrivilegs
                }

                while (exceptionPrivilegs.length != 1) {
                  val p = exceptionPrivilegs.head
                  if (exceptionPrivilegs.tail.forall(e => s.delegators.exists(_.privilege == e))) {
                    val delegUpOn = s.delegators.filter(_.privilege > p).map(_.enable).fold(True)(_ && _)
                    val delegDownOff = !s.delegators.filter(_.privilege <= p).map(_.enable).orR
                    when(delegUpOn && delegDownOff) {
                      exceptionTargetPrivilegeUncapped := p
                    }
                  }
                  exceptionPrivilegs = exceptionPrivilegs.tail
                }
              }
            }
          }
          // The lowest privilege for a trap is S mode
          val targetPrivilege = (csr.privilege | S(1)).max(exceptionTargetPrivilegeUncapped)
        }

        // Detect when an instruction which did a trap reach the commit point.
        val trigger = new Area {
          val lanes = host.list[ExecuteLanePlugin] //This is a bit pessimistic, as a given lane may not be able to trap, but the synthesis should be able to take care of that.
          csr.commitMask := B(for (self <- lanes; sn = self.execute(trapAt+1).down) yield sn.isFiring && sn(COMMIT))
          val oh = B(for (self <- lanes; sn = self.execute(trapAt).down) yield sn.isFiring && sn(TRAP))
          val valid = oh.orR
          val reader = lanes.map(_.execute(trapAt).down).reader(oh, true)
          when(valid) {
            pending.pc := reader(_(PC))
            if(priv.p.withHypervisor) {
              pending.uop := reader(_(Decode.UOP))
              pending.uopIsCompressed := reader(_(Decode.UOP_IS_COMPRESSED))
            }
            if(hp.nonEmpty) pending.history := reader(_(Prediction.BRANCH_HISTORY))
            pending.slices := reader(_(INSTRUCTION_SLICE_COUNT)).resize(INSTRUCTION_SLICE_COUNT_WIDTH+1)+1
          }
        }

        val whitebox = new Area {
          val trap = False
          val interrupt = Bool().assignDontCare()
          val code = Global.CODE().assignDontCare()
        }

        val historyPort = hp.nonEmpty generate hp.get.newPort(Integer.MAX_VALUE, 0)
        if(hp.nonEmpty) {
          historyPort.valid := False
          historyPort.history := pending.history
        }

        val pcPort = pcs.newJumpInterface(Ages.TRAP, 0, 0)
        pcPort.valid := False
        pcPort.fault := False
        pcPort.hartId := hartId
        pcPort.pc.assignDontCare()
        // This state machine is kinda the heart of the CPU.
        val fsm = new StateMachine {
          val RESET = makeInstantEntry()
          val RUNNING, COMPUTE = new State()
          val TRAP_EPC, TRAP_TVAL, TRAP_HTVAL, TRAP_HTINST, TRAP_TVEC, TRAP_WAIT, TRAP_APPLY = new State()
          val XRET_EPC, XRET_APPLY = new State()
          val ATS_RSP = ats.mayNeedRedo generate new State()
          val SATS_RSP = sats.mayNeedRedo generate new State()
          val JUMP = new State()
          val LSU_FLUSH = lsu.nonEmpty generate new State()
          val FETCH_FLUSH = fl1p.nonEmpty generate new State()
          val ENTER_DEBUG_WAIT, ENTER_DEBUG, DPC_READ, RESUME = (priv.p.withDebug) generate new State()

          val inflightTrap = trapPendings.map(_(hartId)).orR
          val holdPort = pcs.newHoldPort(hartId)
          holdPort := inflightTrap || !isActive(RUNNING)
          api.harts(hartId).fsmBusy := !isActive(RUNNING)

          val wfi = False //Whitebox

          val buffer = new Area {
            val sampleIt = False
            def sample[T <: Data](that : T) : T = RegNextWhen(that, sampleIt)
            val i = new Area {
              val valid = sample(interrupt.valid)
              val code = sample(interrupt.code)
              val targetPrivilege = sample(interrupt.targetPrivilege)
            }
            val trap = new Area{
              val interrupt = pending.state.code === TrapReason.INTERRUPT && i.valid
              val targetPrivilege = interrupt.mux(i.targetPrivilege, exception.targetPrivilege)
              val tval = pending.state.tval.andMask(!interrupt)
              val code = interrupt.mux(i.code, pending.state.code)
              val tval2 = priv.p.withHypervisor generate Reg(TVAL)
              val pseudoUop = priv.p.withHypervisor generate Reg(Decode.UOP)
            }
            val state = new Area {
              val guestMemoryAccess = priv.p.withHypervisor generate Reg(Bool())
            }
          }

          val resetToRunConditions = ArrayBuffer[Bool](!initHold)

          // Interface with the address translation service (MMU)
          val atsPorts = ats.mayNeedRedo generate new Area{
            val refill = ats.newRefillPort()
            val isGuestRefill = pending.state.arg(2) || PrivilegeMode.isGuest(priv.getPrivilege(hartId)) || (csr.m.status.mprv && csr.m.status.mpv)
            refill.cmd.valid := False
            refill.cmd.indirect := isGuestRefill
            refill.cmd.permission.read := !pending.state.arg(1)
            refill.cmd.permission.write := pending.state.arg(0, 2 bits) === TrapArg.STORE
            refill.cmd.permission.execute := pending.state.arg(1)
            refill.cmd.storageEnable := True
            refill.cmd.address := pending.state.tval.asUInt
            refill.cmd.storageId := pending.state.arg(3, ats.getStorageIdWidth() bits).asUInt
            refill.rsp.ready := False

            val invalidate = ats.newInvalidationPort()
            invalidate.cmd.valid := False
            invalidate.cmd.hartId := hartId
            invalidate.cmd.address := pending.state.tval.asUInt.resized
            invalidate.cmd.asid := pending.state.tval2.resized
            if (invalidate.p.requestGuest) invalidate.cmd.guest := pending.state.arg(0)
            invalidate.cmd.anyAddress := pending.state.arg(1)
            if (invalidate.cmd.withAnyAsid) invalidate.cmd.anyAsid := pending.state.arg(2)

            val invalidated = RegInit(False) setWhen(invalidate.cmd.fire)
            invalidate.cmd.valid setWhen(!invalidated)
            resetToRunConditions += invalidated
            invalidate.cmd.force := !invalidated
          }

          val satsPorts = sats.mayNeedRedo generate new Area{
            val refill = sats.newRefillPort()
            refill.cmd.valid := False
            refill.cmd.indirect := False
            refill.cmd.permission.read := !pending.state.arg(1)
            refill.cmd.permission.write := pending.state.arg(0, 2 bits) === TrapArg.STORE
            refill.cmd.permission.execute := pending.state.arg(1)
            refill.cmd.storageEnable := True
            refill.cmd.address := (pending.state.tval2.asUInt << 2).resized
            refill.cmd.storageId := pending.state.arg(3+ats.getStorageIdWidth(), sats.getStorageIdWidth() bits).asUInt
            refill.rsp.ready := False

            val invalidate = sats.newInvalidationPort()
            invalidate.cmd.valid := False
            invalidate.cmd.hartId := hartId

            val invalidated = RegInit(False) setWhen(invalidate.cmd.fire)
            invalidate.cmd.valid setWhen(!invalidated)
            resetToRunConditions += invalidated
          }

          // Used to wait until everybody is ready after reset
          RESET.whenIsActive{
            when(resetToRunConditions.andR){
              goto(RUNNING)
            }
          }

          // While CPU is running code
          RUNNING.whenIsActive {
            when(trigger.valid) {
              buffer.sampleIt := True
              if (priv.p.withHypervisor) {
                buffer.trap.tval2 := 0
                buffer.trap.pseudoUop := 0
                buffer.state.guestMemoryAccess := PrivilegeMode.isGuest(csr.privilege) || (csr.m.status.mprv && csr.m.status.mpv)
              }
              goto(COMPUTE)
            }
            if (priv.p.withDebug) when(!csr.hartRunning && csr.debug.doResume) {
              goto(DPC_READ)
            }
          }

          val jumpTarget = Reg(PC)
          val jumpOffset = UInt(INSTRUCTION_SLICE_COUNT_WIDTH+1 bits)
          jumpTarget := pending.pc + (jumpOffset << Fetch.SLICE_RANGE_LOW)
          jumpOffset := pending.slices.andMask(
            List(
              TrapReason.NEXT,
              TrapReason.FENCE_I,
              TrapReason.SFENCE_VMA,
              TrapReason.HFENCE_GMA,
              TrapReason.WFI
            ).map(pending.state.code === _).orR
          )

          val writeGVA = List(
            CSR.MCAUSE_ENUM.BREAKPOINT,
            CSR.MCAUSE_ENUM.FETCH_MISSALIGNED,
            CSR.MCAUSE_ENUM.LOAD_MISALIGNED,
            CSR.MCAUSE_ENUM.STORE_MISALIGNED,
            CSR.MCAUSE_ENUM.INSTRUCTION_ACCESS_FAULT,
            CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT,
            CSR.MCAUSE_ENUM.STORE_ACCESS_FAULT,
            CSR.MCAUSE_ENUM.INSTRUCTION_PAGE_FAULT,
            CSR.MCAUSE_ENUM.LOAD_PAGE_FAULT,
            CSR.MCAUSE_ENUM.STORE_PAGE_FAULT,
            CSR.MCAUSE_ENUM.INSTRUCTION_GUEST_PAGE_FAULT,
            CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT,
            CSR.MCAUSE_ENUM.STORE_GUEST_PAGE_FAULT
          ).map(pending.state.code === _).orR

          if (fl1p.nonEmpty) fetchL1Invalidate(hartId).cmd.valid := False
          if (lsu.nonEmpty) lsuL1Invalidate(hartId).cmd.valid := False
          val trapEnterDebug = RegInit(False)
          val triggerEbreak = (priv.p.debugTriggers == 0).mux(False, !pending.state.exception && pending.state.code === TrapReason.DEBUG_TRIGGER && csr.trigger.slots.reader(pending.state.tval.asUInt.resized)(_.tdata1.doEbreak))
          val triggerEbreakReg = Reg(Bool())
          // Got a trap, need to figure out exactly what to do.
          COMPUTE.whenIsActive{
            triggerEbreakReg := triggerEbreak
            if(priv.p.debugTriggers > 0 ) {
              pending.state.exception setWhen(triggerEbreak) //Patch to reduce logic in next stages
              when(!pending.state.exception && pending.state.code === TrapReason.DEBUG_TRIGGER) {
                csr.trigger.slots.onSel(U(pending.state.tval).resized) { slot =>
                  slot.tdata1.hit := True
                }
              }
            }

            trapEnterDebug := False
            if(hp.nonEmpty) historyPort.valid := True
            when(pending.state.exception || triggerEbreak || buffer.trap.interrupt) {
              goto(TRAP_TVAL)
              if(priv.p.withDebug){
                when(!csr.debugMode) {
                  val doIt = False
                  when(pending.state.exception && exception.code === CSR.MCAUSE_ENUM.BREAKPOINT || triggerEbreak) {
                    doIt setWhen(csr.privilege === PrivilegeMode.M && csr.debug.dcsr.ebreakm)
                    if (priv.p.withUser) doIt setWhen (csr.privilege === PrivilegeMode.U && csr.debug.dcsr.ebreaku)
                    if (priv.p.withSupervisor) doIt setWhen (csr.privilege === PrivilegeMode.S && csr.debug.dcsr.ebreaks)
                  }
                  doIt setWhen(buffer.trap.interrupt && csr.debug.doHalt)
                  when(doIt){
                    trapEnterDebug := True
                    goto(TRAP_EPC)
                  }
                } otherwise {
                  goto(ENTER_DEBUG_WAIT)
                }
              }
              if(priv.p.withHypervisor) {
                buffer.state.guestMemoryAccess.setWhen(writeGVA && pending.state.arg(2))

                val writeTval2 = List(
                  CSR.MCAUSE_ENUM.INSTRUCTION_GUEST_PAGE_FAULT,
                  CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT,
                  CSR.MCAUSE_ENUM.STORE_GUEST_PAGE_FAULT
                ).map(pending.state.code === _).orR
                when(writeTval2 && pending.state.exception && (pending.state.arg(2) || buffer.state.guestMemoryAccess)) {
                  buffer.trap.tval2 := pending.state.tval2.resized
                }

                val writeHtinst = List(
                  CSR.MCAUSE_ENUM.LOAD_MISALIGNED,
                  CSR.MCAUSE_ENUM.STORE_MISALIGNED,
                  CSR.MCAUSE_ENUM.INSTRUCTION_ACCESS_FAULT,
                  CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT,
                  CSR.MCAUSE_ENUM.STORE_ACCESS_FAULT,
                  CSR.MCAUSE_ENUM.LOAD_PAGE_FAULT,
                  CSR.MCAUSE_ENUM.STORE_PAGE_FAULT,
                  CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT,
                  CSR.MCAUSE_ENUM.STORE_GUEST_PAGE_FAULT
                ).map(pending.state.code === _).orR
                if (recordHtinst) when(writeHtinst) {
                  val transformer = RvhTransformer(
                    uop = pending.uop,
                    compressed = pending.uopIsCompressed,
                    rvfd = RVF || RVD,
                    rva = RVZaamo || RVZalrsc,
                    xlen = XLEN.get,
                  )
                  buffer.trap.pseudoUop := transformer.inst
                }
              }
            } otherwise {
              switch(pending.state.code) {
                is(TrapReason.INTERRUPT) { //Got a sporadic interrupt => resume
                  assert(!buffer.i.valid)
                  pcPort.valid := True
                  pcPort.pc := pending.pc
                  goto(RUNNING)
                }
                is(TrapReason.PRIV_RET) {
                  when(!api.harts(hartId).holdPrivChange){
                    goto(XRET_EPC)
                  }
                }
                is(TrapReason.FENCE_I) {
                  (lsul1.nonEmpty) match {
                    case true => goto(LSU_FLUSH)
                    case false => fl1p.nonEmpty match {
                      case true => goto(FETCH_FLUSH)
                      case false => goto(JUMP)
                    }
                  }
                }
                is(TrapReason.REDO) {
                  api.harts(hartId).redo := True
                  goto(JUMP)
                }
                is(TrapReason.NEXT) {
                  goto(JUMP)
                }
                is(TrapReason.WFI) {
                  wfi := True
                  when(api.harts(hartId).askWake) {
                    goto(JUMP)
                  }
                }
                is(TrapReason.SFENCE_VMA) {
                  if(ats.mayNeedRedo) {
                    atsPorts.invalidate.cmd.valid := True
                    when(atsPorts.invalidate.cmd.ready) {
                      goto(JUMP)
                    }
                  } else {
                    goto(JUMP)
                  }
                }
                if(ats.mayNeedRedo) is(TrapReason.MMU_REFILL) {
                  atsPorts.refill.cmd.valid := True
                  when(atsPorts.refill.cmd.ready){
                    goto(ATS_RSP)
                  }
                }
                is(TrapReason.HFENCE_GMA) {
                  if(sats.mayNeedRedo) {
                    satsPorts.invalidate.cmd.valid := True
                    when(satsPorts.invalidate.cmd.ready) {
                      goto(JUMP)
                    }
                  } else {
                    goto(JUMP)
                  }
                }
                if(sats.mayNeedRedo) is(TrapReason.SMMU_REFILL) {
                  satsPorts.refill.cmd.valid := True
                  when(satsPorts.refill.cmd.ready){
                    goto(SATS_RSP)
                  }
                }
                if(priv.p.debugTriggers > 0 ) is(TrapReason.DEBUG_TRIGGER) {
                  trapEnterDebug := True
                  goto(TRAP_EPC)
                }
                default {
                  assert(False, "Unexpected trap reason")
                }
              }
            }
          }

          if(lsu.nonEmpty) LSU_FLUSH.whenIsActive{
            lsuL1Invalidate(hartId).cmd.valid := True
            when(lsuL1Invalidate(hartId).cmd.ready) {
              fl1p.nonEmpty match {
                case true => goto(FETCH_FLUSH)
                case false => goto(JUMP)
              }
            }
          }

          if (fl1p.nonEmpty) FETCH_FLUSH.whenIsActive {
            fetchL1Invalidate(hartId).cmd.valid := True
            when(fetchL1Invalidate(hartId).cmd.ready) {
              goto(JUMP)
            }
          }

          if(ats.mayNeedRedo) ATS_RSP.whenIsActive {
            when(atsPorts.refill.rsp.valid) {
              when(atsPorts.refill.rsp.guestFault || atsPorts.refill.rsp.pageFault || atsPorts.refill.rsp.accessFault) {
                atsPorts.refill.rsp.ready := True
                pending.state.exception := True
                if (priv.p.withHypervisor) when(atsPorts.isGuestRefill) {
                  buffer.trap.tval2 := atsPorts.refill.rsp.address.dropLow(2).asBits.resized
                  when (atsPorts.refill.rsp.guestFault) {
                    buffer.trap.pseudoUop := (XLEN.get == 32).mux(0x00002000, 0x00003000)
                  }
                }
                switch(atsPorts.refill.rsp.guestFault ## atsPorts.refill.rsp.accessFault ## pending.state.arg(1 downto 0)) {
                  def add(k : Int, v : Int) = is(k){pending.state.code := v}
                  add(TrapArg.FETCH | 8, CSR.MCAUSE_ENUM.INSTRUCTION_GUEST_PAGE_FAULT)
                  add(TrapArg.LOAD  | 8, CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT)
                  add(TrapArg.STORE | 8, CSR.MCAUSE_ENUM.STORE_GUEST_PAGE_FAULT)
                  add(TrapArg.FETCH | 4, CSR.MCAUSE_ENUM.INSTRUCTION_ACCESS_FAULT)
                  add(TrapArg.LOAD  | 4, CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT)
                  add(TrapArg.STORE | 4, CSR.MCAUSE_ENUM.STORE_ACCESS_FAULT)
                  add(TrapArg.FETCH    , CSR.MCAUSE_ENUM.INSTRUCTION_PAGE_FAULT)
                  add(TrapArg.LOAD     , CSR.MCAUSE_ENUM.LOAD_PAGE_FAULT)
                  add(TrapArg.STORE    , CSR.MCAUSE_ENUM.STORE_PAGE_FAULT)
                  add(TrapArg.FETCH_LSU, CSR.MCAUSE_ENUM.LOAD_PAGE_FAULT)
                  add(TrapArg.FETCH_LSU | 4, CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT)
                  add(TrapArg.FETCH_LSU | 8, CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT)
                }
                goto(TRAP_TVAL)
              } otherwise {
                if(sats.mayNeedRedo) {
                  when (atsPorts.isGuestRefill) {
                    satsPorts.refill.cmd.address := atsPorts.refill.rsp.address.resized
                    satsPorts.refill.cmd.valid := True
                    when(satsPorts.refill.cmd.ready){
                      atsPorts.refill.rsp.ready := True
                      goto(SATS_RSP)
                    }
                  } otherwise {
                    api.harts(hartId).redo := True
                    atsPorts.refill.rsp.ready := True
                    goto(JUMP)
                  }
                } else {
                  api.harts(hartId).redo := True
                  atsPorts.refill.rsp.ready := True
                  goto(JUMP) // improvement: shave one cycle
                }
              }
            }
          }

          if(sats.mayNeedRedo) SATS_RSP.whenIsActive {
            buffer.state.guestMemoryAccess.set()
            when(satsPorts.refill.rsp.valid) {
              api.harts(hartId).redo := True
              satsPorts.refill.rsp.ready := True
              goto(JUMP) // improvement: shave one cycle
              when(satsPorts.refill.rsp.pageFault || satsPorts.refill.rsp.accessFault) {
                pending.state.exception := True
                buffer.trap.tval2 := satsPorts.refill.rsp.address.dropLow(2).asBits.resized
                switch(satsPorts.refill.rsp.accessFault ## pending.state.arg(1 downto 0)){
                  def add(k : Int, v : Int) = is(k){pending.state.code := v}
                  add(TrapArg.FETCH | 4, CSR.MCAUSE_ENUM.INSTRUCTION_ACCESS_FAULT)
                  add(TrapArg.LOAD  | 4, CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT)
                  add(TrapArg.STORE | 4, CSR.MCAUSE_ENUM.STORE_ACCESS_FAULT)
                  add(TrapArg.FETCH    , CSR.MCAUSE_ENUM.INSTRUCTION_GUEST_PAGE_FAULT)
                  add(TrapArg.LOAD     , CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT)
                  add(TrapArg.STORE    , CSR.MCAUSE_ENUM.STORE_GUEST_PAGE_FAULT)
                  add(TrapArg.FETCH_LSU | 4, CSR.MCAUSE_ENUM.LOAD_ACCESS_FAULT)
                  add(TrapArg.FETCH_LSU, CSR.MCAUSE_ENUM.LOAD_GUEST_PAGE_FAULT)
                }
                goto(TRAP_TVAL)
              }
            }
          }

          TRAP_TVAL.whenIsActive {
            crsPorts.write.valid := True
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.tval.getAddress()
            )
            if (priv.p.withSupervisor) addressMapping += PrivilegeMode.S -> csr.s.tval.getAddress()
            if (priv.p.withHypervisor) addressMapping += PrivilegeMode.VS -> csr.vs.tval.getAddress()
            crsPorts.write.address := privilegeMux(addressMapping, buffer.trap.targetPrivilege)
            crsPorts.write.data := Global.expendPc(buffer.trap.tval.asUInt, XLEN).asBits
            when(triggerEbreakReg){
              crsPorts.write.data := Global.expendPc(pending.pc, XLEN).asBits
            }

            when(crsPorts.write.ready) {
              if (priv.p.withHypervisor) {
                goto(TRAP_HTVAL)
              } else {
                goto(TRAP_EPC)
              }
            }
          }

          val isGuestTarget = PrivilegeMode.isGuest(buffer.trap.targetPrivilege)
          if (priv.p.withHypervisor) TRAP_HTVAL.whenIsActive {
            crsPorts.write.valid := !isGuestTarget
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.tval2.getAddress(),
              PrivilegeMode.S -> csr.h.tval.getAddress()
            )
            crsPorts.write.address := privilegeMux(addressMapping, buffer.trap.targetPrivilege)
            crsPorts.write.data := buffer.trap.tval2.resized

            when(crsPorts.write.ready || isGuestTarget) {
              goto(TRAP_HTINST)
            }
          }

          if (priv.p.withHypervisor) TRAP_HTINST.whenIsActive {
            val writeUop = buffer.trap.pseudoUop

            crsPorts.write.valid := !isGuestTarget
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.tinst.getAddress(),
              PrivilegeMode.S -> csr.h.tinst.getAddress()
            )
            crsPorts.write.address := privilegeMux(addressMapping, buffer.trap.targetPrivilege)
            crsPorts.write.data := writeUop.resized

            when(crsPorts.write.ready || isGuestTarget) {
              goto(TRAP_EPC)
            }
          }

          TRAP_EPC.whenIsActive {
            crsPorts.write.valid := True
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.epc.getAddress()
            )
            if (priv.p.withSupervisor) addressMapping += PrivilegeMode.S -> csr.s.epc.getAddress()
            if (priv.p.withHypervisor) addressMapping += PrivilegeMode.VS -> csr.vs.epc.getAddress()
            crsPorts.write.address := privilegeMux(addressMapping, buffer.trap.targetPrivilege)
            if (priv.p.withDebug) when(trapEnterDebug) {  crsPorts.write.address := csr.debug.dpc.getAddress() }
            crsPorts.write.data := Global.expendPc(pending.pc, XLEN).asBits
            when(crsPorts.write.ready) {
              goto(TRAP_TVEC)
              if (priv.p.withDebug) when(trapEnterDebug) {
                goto(ENTER_DEBUG_WAIT)
              }
            }
          }

          val readed = RegNextWhen(crsPorts.read.data, crsPorts.read.valid && crsPorts.read.ready)
          TRAP_TVEC.whenIsActive {
            crsPorts.read.valid := True
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.tvec.getAddress()
            )
            if (priv.p.withSupervisor) addressMapping += PrivilegeMode.S -> csr.s.tvec.getAddress()
            if (priv.p.withHypervisor) addressMapping += PrivilegeMode.VS -> csr.vs.tvec.getAddress()
            crsPorts.read.address := privilegeMux(addressMapping, buffer.trap.targetPrivilege)
            when(crsPorts.read.ready) {
              goto(TRAP_WAIT)
            }
          }

          TRAP_WAIT.whenIsActive{
            when(!api.harts(hartId).holdPrivChange){
              goto(TRAP_APPLY)
            }
          }

          TRAP_APPLY.whenIsActive{
            api.harts(hartId).rvTrap := True
            pcPort.valid := True
            pcPort.pc := U(readed).resized //PC RESIZED

            csr.privilege := buffer.trap.targetPrivilege
            switch(buffer.trap.targetPrivilege) {
              is(PrivilegeMode.M) {
                csr.m.status.mie := False
                csr.m.status.mpie := csr.m.status.mie
                if (priv.p.withUser) csr.m.status.mpp := csr.privilege(1 downto 0).asUInt
                if (priv.p.withHypervisor) csr.m.status.gva := writeGVA && buffer.state.guestMemoryAccess
                if (priv.p.withHypervisor) csr.m.status.mpv := PrivilegeMode.isGuest(csr.privilege)

                csr.m.cause.code := buffer.trap.code
                csr.m.cause.interrupt := buffer.trap.interrupt
              }
              priv.p.withSupervisor generate is(PrivilegeMode.S) {
                csr.s.status.sie := False
                csr.s.status.spie := csr.s.status.sie
                if (priv.p.withUser) csr.s.status.spp := csr.privilege(0, 1 bits).asUInt
                if (priv.p.withHypervisor) csr.h.status.spv := PrivilegeMode.isGuest(csr.privilege)
                if (priv.p.withHypervisor) csr.h.status.gva := writeGVA && buffer.state.guestMemoryAccess
                if (priv.p.withHypervisor) when(PrivilegeMode.isGuest(csr.privilege)) {
                  csr.h.status.spvp := csr.privilege(0, 1 bits).asBool
                }

                csr.s.cause.code := buffer.trap.code
                csr.s.cause.interrupt := buffer.trap.interrupt
              }
              priv.p.withHypervisor generate is(PrivilegeMode.VS) {
                csr.vs.status.sie := False
                csr.vs.status.spie := csr.vs.status.sie
                csr.vs.status.spp := csr.privilege(0, 1 bits).asUInt

                csr.vs.cause.code := buffer.trap.code.mux(
                  B(10) -> B(9),
                  B(6)  -> B(5),
                  B(2)  -> B(1),
                  default -> buffer.trap.code
                )
                csr.vs.cause.interrupt := buffer.trap.interrupt
              }
            }

            whitebox.trap := True
            whitebox.interrupt := buffer.trap.interrupt
            whitebox.code :=  buffer.trap.code

            goto(RUNNING)
          }

          if(priv.p.withDebug) {
            ENTER_DEBUG_WAIT.whenIsActive{
              when(!api.harts(hartId).holdPrivChange){
                goto(ENTER_DEBUG)
              }
            }

            csr.debug.bus.exception := False
            csr.debug.bus.ebreak := False
            ENTER_DEBUG.whenIsActive{
              csr.hartRunning := False
              when(!csr.debugMode) {
                csr.debug.dcsr.cause := 0
                when(csr.debug.dcsr.step){ csr.debug.dcsr.cause := 4 }
                when(csr.debug.bus.haltReq) { csr.debug.dcsr.cause := 3 }
                when(pending.state.exception && exception.code === CSR.MCAUSE_ENUM.BREAKPOINT) { csr.debug.dcsr.cause := 1 }
                when(!pending.state.exception && exception.code === TrapReason.DEBUG_TRIGGER) { csr.debug.dcsr.cause := 2 }
                csr.debug.dcsr.prv := csr.privilege
              } otherwise {
                csr.debug.bus.exception := pending.state.exception && exception.code =/= CSR.MCAUSE_ENUM.BREAKPOINT
                csr.debug.bus.ebreak    := pending.state.exception && exception.code === CSR.MCAUSE_ENUM.BREAKPOINT
              }
              csr.privilege := PrivilegeMode.M
              goto(RUNNING)
            }

            DPC_READ.whenIsActive {
              crsPorts.read.valid := True
              crsPorts.read.address := csr.debug.dpc.getAddress
              when(crsPorts.read.ready) {
                goto(RESUME)
              }
            }
            RESUME.whenIsActive {
              pcPort.valid := True
              pcPort.pc := U(readed).resized //PC RESIZED
              csr.privilege := csr.debug.dcsr.prv
              csr.hartRunning := True
              csr.debug.bus.resume.rsp.valid := True
              goto(RUNNING)
            }
          }

          val xretPrivilege = pending.state.arg(2 downto 0).asSInt
          XRET_EPC.whenIsActive{
            crsPorts.read.valid := True
            val addressMapping = mutable.LinkedHashMap[Int, UInt](
              PrivilegeMode.M -> csr.m.epc.getAddress()
            )
            if (priv.p.withSupervisor) addressMapping += PrivilegeMode.S -> csr.s.epc.getAddress()
            if (priv.p.withHypervisor) addressMapping += PrivilegeMode.VS -> csr.vs.epc.getAddress()
            crsPorts.read.address := privilegeMux(addressMapping, xretPrivilege)
            when(crsPorts.read.ready) {
              goto(XRET_APPLY)
            }
          }

          XRET_APPLY.whenIsActive{
            pcPort.valid := True
            pcPort.pc := U(readed).resized //PC RESIZED

            csr.privilege := pending.xret.targetPrivilege
            csr.xretAwayFromMachine setWhen (pending.xret.targetPrivilege =/= PrivilegeMode.M)
            switch(pending.state.arg(2 downto 0).asSInt) {
              is(PrivilegeMode.M) {
                if (priv.p.withUser) csr.m.status.mpp := 0
                if (priv.p.withHypervisor) csr.m.status.mpv := False
                csr.m.status.mie := csr.m.status.mpie
                csr.m.status.mpie := True
              }
              priv.p.withSupervisor generate is(PrivilegeMode.S) {
                csr.s.status.spp := U"0"
                if (priv.p.withHypervisor) csr.h.status.spv := False
                csr.s.status.sie := csr.s.status.spie
                csr.s.status.spie := True
              }
              priv.p.withHypervisor generate is(PrivilegeMode.VS) {
                csr.vs.status.spp := U"0"
                csr.vs.status.sie := csr.vs.status.spie
                csr.vs.status.spie := True
              }
            }
            goto(RUNNING)
          }

          JUMP.whenIsActive {
            pcPort.valid := True
            pcPort.pc := jumpTarget //PC RESIZED
            goto(RUNNING)
          }
        }
      }
    }

    ramPortRetainers.release()
    buildBefore.release()

    initHold := host.list[InitService].map(_.initHold()).orR
  }
}
