package xiangshan.frontend.ifu

import chisel3._
import chisel3.util._
import org.chipsalliance.cde.config.Parameters
import xiangshan.frontend.PrunedAddr

class InstrCompact(implicit p: Parameters) extends IfuModule {
  class InstrCompactIO extends IfuBundle {
    class InstrCompactReq(implicit p: Parameters) extends IfuBundle {
      val fetchBlockSelect        = Vec(FetchBlockInstNum, Bool())
      val twoFetchInstrOffset     = Vec(FetchBlockInstNum, UInt(FetchBlockInstOffsetWidth.W))
      val twoFetchPcLower         = Vec(FetchBlockInstNum, UInt((PcCutPoint + 1).W))
      val instrCountBeforeCurrent = Vec(FetchBlockInstNum + 1, UInt(log2Ceil(FetchBlockInstNum + 1).W))
      val rawInstrValid           = Vec(FetchBlockInstNum, Bool())
      val rawIsRvc                = Vec(FetchBlockInstNum, Bool())
    }
    val req:  InstrCompactReq    = Input(new InstrCompactReq)
    val resp: InstrCompactBundle = Output(new InstrCompactBundle(FetchBlockInstNum))
  }
  val io: InstrCompactIO = IO(new InstrCompactIO)

  private val instrCountBeforeCurrent = io.req.instrCountBeforeCurrent
  private val rawInstrValid           = io.req.rawInstrValid
  private val rawIsRvc                = io.req.rawIsRvc
  private val fetchBlockSelect        = io.req.fetchBlockSelect
  private val twoFetchPcLower         = io.req.twoFetchPcLower

  private val twoFetchBlockIndex = VecInit.tabulate(FetchBlockInstNum)(i =>
    twoFetchPcLower(i)(log2Ceil(ICacheLineBytes) - 1, 1)
  )
  private val twoFetchInstrOffset = VecInit(
    io.req.twoFetchInstrOffset :+ 0.U(FetchBlockInstOffsetWidth.W)
  )

  private val instrSelectLowIndex   = WireDefault(VecInit.fill(FetchBlockInstNum)(true.B))
  private val instrSelectFetchBlock = WireDefault(VecInit.fill(FetchBlockInstNum)(false.B))

  private val instrPcLowerResult = WireDefault(VecInit.fill(FetchBlockInstNum)(0.U((PcCutPoint + 1).W)))
  private val instrIsRvc         = WireDefault(VecInit.fill(FetchBlockInstNum)(false.B))
  private val instrEndOffset     = WireDefault(VecInit.fill(FetchBlockInstNum)(0.U(FetchBlockInstOffsetWidth.W)))
  private val instrIndexEntry    = Wire(Vec(FetchBlockInstNum, new InstrIndexEntry))

  // Fetch PC and index info for valid instructions based on their positions.
  instrIndexEntry.zipWithIndex.foreach {
    case (instrIndex, idx) =>
      val instrRange = idx until Math.min(2 * idx + 2, FetchBlockInstNum)

      val validOH = instrRange.map {
        i => rawInstrValid(i) & (instrCountBeforeCurrent(i) === idx.U)
      }

      val index         = instrRange.map(twoFetchBlockIndex(_))
      val select        = instrRange.map(fetchBlockSelect(_))
      val pcLowerResult = instrRange.map(twoFetchPcLower(_))
      val isRvc         = instrRange.map(rawIsRvc(_))
      val instrOffset   = instrRange.map(i => Mux(rawIsRvc(i), twoFetchInstrOffset(i), twoFetchInstrOffset(i + 1)))
      // FIXME: This is wrong when 2-taken is enabled

      instrIndex.valid           := validOH.reduce(_ || _)
      instrIndex.value           := Mux1H(validOH, index)
      instrSelectFetchBlock(idx) := Mux1H(validOH, select)
      instrPcLowerResult(idx)    := Mux1H(validOH, pcLowerResult)
      instrIsRvc(idx)            := Mux1H(validOH, isRvc)
      instrEndOffset(idx)        := Mux1H(validOH, instrOffset)
  }

  io.resp.instrIndex     := instrIndexEntry
  io.resp.instrIsRvc     := instrIsRvc
  io.resp.selectBlock    := instrSelectFetchBlock
  io.resp.instrPcLower   := instrPcLowerResult
  io.resp.instrEndOffset := instrEndOffset
}
