package constellation.router

import chisel3._
import chisel3.util._

import org.chipsalliance.cde.config.{Field, Parameters}
import constellation.channel._
import constellation.routing.{FlowRoutingBundle}
import constellation.noc.{HasNoCParams}

class OutputCreditAlloc extends Bundle {
  val alloc = Bool()
  val tail = Bool()
}

class OutputChannelStatus(implicit val p: Parameters) extends Bundle with HasNoCParams {
  val occupied = Bool()
  def available = !occupied
  val flow = new FlowRoutingBundle
}

class OutputChannelAlloc(implicit val p: Parameters) extends Bundle with HasNoCParams {
  val alloc = Bool()
  val flow = new FlowRoutingBundle
}

class AbstractOutputUnitIO(
  val inParams: Seq[ChannelParams],
  val ingressParams: Seq[IngressChannelParams],
  val cParam: BaseChannelParams
)(implicit val p: Parameters) extends Bundle with HasRouterInputParams {
  val nodeId = cParam.srcId
  val nVirtualChannels = cParam.nVirtualChannels
  val in = Flipped(Vec(cParam.srcSpeedup, Valid(new Flit(cParam.payloadBits))))
  val credit_available = Output(Vec(nVirtualChannels, Bool()))
  val channel_status = Output(Vec(nVirtualChannels, new OutputChannelStatus))
  val allocs = Input(Vec(nVirtualChannels, new OutputChannelAlloc))
  val credit_alloc = Input(Vec(nVirtualChannels, new OutputCreditAlloc))
}

abstract class AbstractOutputUnit(
  val inParams: Seq[ChannelParams],
  val ingressParams: Seq[IngressChannelParams],
  val cParam: BaseChannelParams
)(implicit val p: Parameters) extends Module with HasRouterInputParams with HasNoCParams {
  val nodeId = cParam.srcId

  def io: AbstractOutputUnitIO
}

class OutputUnit(inParams: Seq[ChannelParams], ingressParams: Seq[IngressChannelParams], cParam: ChannelParams)
  (implicit p: Parameters) extends AbstractOutputUnit(inParams, ingressParams, cParam)(p) {

  class OutputUnitIO extends AbstractOutputUnitIO(inParams, ingressParams, cParam) {
    val out = new Channel(cParam.asInstanceOf[ChannelParams])
  }
  val io = IO(new OutputUnitIO)

  class OutputState(val bufferSize: Int) extends Bundle {
    val occupied = Bool()
    val c = UInt(log2Up(1+bufferSize).W)
    val flow = new FlowRoutingBundle
  }

  (states zip io.channel_status).map { case (s,a) =>
    a.occupied := s.occupied
    a.flow := s.flow
  }
  io.out.flit := io.in

  states.zipWithIndex.map { case (s,i) => if (cParam.virtualChannelParams(i).traversable) {
    when (io.out.vc_free(i)) {
      assert(s.occupied)
      s.occupied := false.B
    }
  } }

  (states zip io.allocs).zipWithIndex.map { case ((s,a),i) => if (cParam.virtualChannelParams(i).traversable) {
    when (a.alloc) {
      s.occupied := true.B
      s.flow := a.flow
    }
  } }

  (io.credit_available zip states).zipWithIndex.map { case ((c,s),i) =>
    c := s.c =/= 0.U //|| (io.out.credit_return.valid && io.out.credit_return.bits === i.U)
  }

  states.zipWithIndex.map { case (s,i) =>
    val free = io.out.credit_return(i)
    val alloc = io.credit_alloc(i).alloc
    if (cParam.virtualChannelParams(i).traversable) {
      s.c := s.c +& free - alloc
    }
  }

  when (reset.asBool) {
    states.foreach(_.occupied := false.B)
    states.foreach(s => s.c := s.bufferSize.U)
  }
}

val states = Reg(MixedVec(cParam.virtualChannelParams.map { u => new OutputState(u.bufferSize) }))


package constellation.router

import chisel3._
import chisel3.util._
import chisel3.util.random.{LFSR}

import org.chipsalliance.cde.config.{Field, Parameters}
import freechips.rocketchip.util._

import constellation.channel._
import constellation.routing.{ChannelRoutingInfo, FlowRoutingBundle}

// Allocates 1 VC per cycle
abstract class SingleVCAllocator(vP: VCAllocatorParams)(implicit p: Parameters) extends VCAllocator(vP)(p) {
  // get single input
  val mask = RegInit(0.U(allInParams.size.W))
  val in_arb_reqs = Wire(Vec(allInParams.size, MixedVec(allOutParams.map { u => Vec(u.nVirtualChannels, Bool()) })))
  val in_arb_vals = Wire(Vec(allInParams.size, Bool()))
  val in_arb_filter = PriorityEncoderOH(Cat(in_arb_vals.asUInt, in_arb_vals.asUInt & ~mask))
  val in_arb_sel = (in_arb_filter(allInParams.size-1,0) | (in_arb_filter >> allInParams.size))
  when (in_arb_vals.orR) {
    mask := Mux1H(in_arb_sel, (0 until allInParams.size).map { w => ~(0.U((w+1).W)) })
  }

  for (i <- 0 until allInParams.size) {
    (0 until allOutParams.size).map { m =>
      (0 until allOutParams(m).nVirtualChannels).map { n =>
        in_arb_reqs(i)(m)(n) := io.req(i).bits.vc_sel(m)(n) && !io.channel_status(m)(n).occupied
      }
    }

    in_arb_vals(i) := io.req(i).valid && in_arb_reqs(i).map(_.orR).toSeq.orR
  }

  // Input arbitration
  io.req.foreach(_.ready := false.B)

  val in_flow = Mux1H(in_arb_sel, io.req.map(_.bits.flow).toSeq)
  val in_vc = Mux1H(in_arb_sel, io.req.map(_.bits.in_vc).toSeq)
  val in_vc_sel = Mux1H(in_arb_sel, in_arb_reqs)
  in_alloc := Mux(in_arb_vals.orR,
    inputAllocPolicy(in_flow, in_vc_sel, OHToUInt(in_arb_sel), in_vc, io.req.map(_.fire).toSeq.orR),
    0.U.asTypeOf(in_alloc))

  // send allocation to inputunits
  for (i <- 0 until allInParams.size) {
    io.req(i).ready := in_arb_sel(i)
    for (m <- 0 until allOutParams.size) {
      (0 until allOutParams(m).nVirtualChannels).map { n =>
        io.resp(i).vc_sel(m)(n) := in_alloc(m)(n)
      }
    }
    assert(PopCount(io.resp(i).vc_sel.asUInt) <= 1.U)
  }

  // send allocation to output units
  for (i <- 0 until allOutParams.size) {
    (0 until allOutParams(i).nVirtualChannels).map { j =>
      io.out_allocs(i)(j).alloc := in_alloc(i)(j)
      io.out_allocs(i)(j).flow := in_flow
    }
  }
}

val in_alloc = Wire(MixedVec(allOutParams.map { u => Vec(u.nVirtualChannels, Bool()) }))


package constellation.routing

import scala.math.{min, max, pow}
import scala.collection.mutable.HashMap
import org.chipsalliance.cde.config.{Parameters}
import scala.collection.immutable.ListMap
import constellation.topology.{PhysicalTopology, Mesh2DLikePhysicalTopology, HierarchicalTopology, CustomTopology}

import scala.collection.mutable

/** Routing and channel allocation policy
 *
 * @param f function that takes in source channel, next channel, and packet routing info.
 *          Returns True if packet can acquire/proceed to this next channel, False if not. An example
 *          is the bidirectionalLine method; see its docstring.
 * @param isEscape function that takes in ChannelRoutingInfo and a virtual network ID. Returns True if the
 *                 channel represented by ChannelRoutingInfo is an escape channel and False if it is not.
 *                 By default, we ignore the escape-channel-based deadlock-free properties, and just check
 *                 for deadlock-freedom by verifying lack of a cyclic dependency.
 */
// BEGIN: RoutingRelation
abstract class RoutingRelation(topo: PhysicalTopology) {
  // Child classes must implement these
  def rel       (srcC: ChannelRoutingInfo,
                 nxtC: ChannelRoutingInfo,
                 flow: FlowRoutingInfo): Boolean

  def isEscape  (c: ChannelRoutingInfo,
                 vNetId: Int): Boolean = true

  def getNPrios (src: ChannelRoutingInfo): Int = 1

  def getPrio   (srcC: ChannelRoutingInfo,
                 nxtC: ChannelRoutingInfo,
                 flow: FlowRoutingInfo): Int = 0

  // END: RoutingRelation

  private val memoize = new HashMap[(ChannelRoutingInfo, ChannelRoutingInfo, FlowRoutingInfo), Boolean]()
  def apply(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Boolean = {
    require(srcC.dst == nxtC.src)
    val key = (srcC, nxtC, flow)
    memoize.getOrElseUpdate(key, rel(srcC, nxtC, flow))
  }
}
// END: RoutingRelation

/** Given a deadlock-prone routing relation and a routing relation representing the network's escape
  *  channels, returns a routing relation that adds the escape channels to the deadlock-prone relation.
  *
  *  @param escapeRouter routing relation representing the network's escape channels
  *  @param normalrouter deadlock-prone routing relation
  *  @param nEscapeChannels number of escape channels
  */
object EscapeChannelRouting {
  def apply(
    escapeRouter: PhysicalTopology => RoutingRelation,
    normalRouter: PhysicalTopology => RoutingRelation,
    nEscapeChannels: Int = 1) = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    val escape = escapeRouter(topo)
    val normal = normalRouter(topo)
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      if (srcC.src == -1) {
        if (nxtC.vc >= nEscapeChannels) { // if ingress and jumping into normal channel, use normal relation
          normal(srcC, nxtC.copy(vc=nxtC.vc-nEscapeChannels, n_vc=nxtC.n_vc-nEscapeChannels), flow)
        } else { // else use ingress and jump into escape channel
          escape(srcC, nxtC.copy(n_vc=nEscapeChannels), flow)
        }
      } else if (nxtC.vc < nEscapeChannels) {
        escape(srcC, nxtC.copy(n_vc=nEscapeChannels), flow)
      } else if (srcC.vc >= nEscapeChannels && nxtC.vc >= nEscapeChannels) {
        normal(
          srcC.copy(vc=srcC.vc-nEscapeChannels, n_vc=srcC.n_vc-nEscapeChannels),
          nxtC.copy(vc=nxtC.vc-nEscapeChannels, n_vc=nxtC.n_vc-nEscapeChannels),
          flow)
      } else {
        false
      }
    }
    override def isEscape(c: ChannelRoutingInfo, v: Int) = {
      c.isIngress || c.isEgress || c.vc < nEscapeChannels
    }
    override def getNPrios(c: ChannelRoutingInfo): Int = {
      escape.getNPrios(c) + normal.getNPrios(c)
    }
    override def getPrio(src: ChannelRoutingInfo, a: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
      val aEscape = isEscape(a, flow.vNetId)
      val nSrc = if (src.isIngress) {
        src
      } else if (isEscape(src, flow.vNetId)) {
        src.copy(vc=src.vc, n_vc=nEscapeChannels)
      } else {
        src.copy(vc=src.vc-nEscapeChannels, n_vc=src.n_vc-nEscapeChannels)
      }
      if (aEscape) {
        escape.getPrio(nSrc, a.copy(n_vc=nEscapeChannels), flow) + normal.getNPrios(src)
      } else {
        normal.getPrio(nSrc, a.copy(vc=a.vc-nEscapeChannels, n_vc=a.n_vc-nEscapeChannels), flow)
      }
    }
  }
}

/**
  * Allows all routing. Probably not very useful
  */
object AllLegalRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = true
  }
}

object UnidirectionalLineRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      nxtC.dst <= flow.egressNode
    }
    override def getNPrios(src: ChannelRoutingInfo) = 4
    override def getPrio(src: ChannelRoutingInfo, a: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
      max(src.dst - a.dst + getNPrios(src), 0)
    }
  }
}

object BidirectionalLineRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      if (nxtC.src < nxtC.dst) flow.egressNode >= nxtC.dst else flow.egressNode <= nxtC.dst
    }
  }
}

object UnidirectionalTorus1DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      if (srcC.src == -1)  {
        nxtC.vc != 0
      } else if (srcC.vc == 0) {
        nxtC.vc == 0
      } else if (nxtC.src == topo.nNodes - 1) {
        nxtC.vc < srcC.vc
      } else {
        nxtC.vc <= srcC.vc && (nxtC.vc != 0 || flow.egressNode > nxtC.src)
      }
    }
    override def getNPrios(src: ChannelRoutingInfo) = 2
    override def getPrio(src: ChannelRoutingInfo, nxt: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
      if (nxt.vc == 0 || nxt.vc == 1) 1 else 0
    }
    override def isEscape(c: ChannelRoutingInfo, v: Int) = {
      c.isIngress || c.isEgress || c.vc < 2
    }
  }
}

object BidirectionalTorus1DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      if (srcC.src == -1)  {
        nxtC.vc != 0
      } else if (srcC.vc == 0) {
        nxtC.vc == 0
      } else if ((nxtC.dst + topo.nNodes - nxtC.src) % topo.nNodes == 1) {
        if (nxtC.src == topo.nNodes - 1) {
          nxtC.vc < srcC.vc
        } else {
          nxtC.vc <= srcC.vc && (nxtC.vc != 0 || flow.egressNode > nxtC.src)
        }
      } else if ((nxtC.src + topo.nNodes - nxtC.dst) % topo.nNodes == 1) {
        if (nxtC.src == 0) {
          nxtC.vc < srcC.vc
        } else {
          nxtC.vc <= srcC.vc && (nxtC.vc != 0 || flow.egressNode < nxtC.src)
        }
      } else {
        false
      }
    }
  }
}

object BidirectionalTorus1DShortestRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    val base = BidirectionalTorus1DDatelineRouting()(topo)
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      val cwDist = (flow.egressNode + topo.nNodes - nxtC.src) % topo.nNodes
      val ccwDist = (nxtC.src + topo.nNodes - flow.egressNode) % topo.nNodes
      val distSel = if (cwDist < ccwDist) {
        (nxtC.dst + topo.nNodes - nxtC.src) % topo.nNodes == 1
      } else if (cwDist > ccwDist) {
        (nxtC.src + topo.nNodes - nxtC.dst) % topo.nNodes == 1
      } else {
        true
      }
      distSel && base(srcC, nxtC, flow)
    }
  }
}

object BidirectionalTorus1DRandomRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    val base = BidirectionalTorus1DDatelineRouting()(topo)
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      val sel = if (srcC.src == -1) {
        true
      } else if ((nxtC.src + topo.nNodes - srcC.src) % topo.nNodes == 1) {
        (nxtC.dst + topo.nNodes - nxtC.src) % topo.nNodes == 1
      } else {
        (nxtC.src + topo.nNodes - nxtC.dst) % topo.nNodes == 1
      }
      sel && base(srcC, nxtC, flow)
    }
  }
}

object ButterflyRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: constellation.topology.Butterfly => new RoutingRelation(topo) {
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY) = (nxtC.dst / topo.height, nxtC.dst % topo.height)
        val (nodeX, nodeY) = (nxtC.src / topo.height, nxtC.src % topo.height)
        val (dstX, dstY) = (flow.egressNode / topo.height, flow.egressNode % topo.height)
        if (dstX <= nodeX) {
          false
        } else if (nodeX == topo.nFly-1) {
          true
        } else {
          val dsts = (nxtX until topo.nFly-1).foldRight((0 until topo.height).map { i => Seq(i) }) {
            case (i,l) => (0 until topo.height).map { s => topo.channels(i).filter(_._1 == s).map { case (_,d) =>
              l(d)
            }.flatten }
          }
          dsts(nxtY).contains(flow.egressNode % topo.height)
        }
      }
    }
  }
}

object BidirectionalTreeRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: constellation.topology.BidirectionalTree => new RoutingRelation(topo) {
      /** Returns a boolean indicating whether src is an ancestor node of dst */
      def isAncestor(src: Int, dst: Int): Boolean = {
        if (src == dst) {
          true
        } else if (src > dst) {
          false
        } else {
          ((topo.dAry * src + 1) to (topo.dAry * src + topo.dAry)).foldLeft(false)((sofar, nextChild) => sofar || isAncestor(nextChild, dst))
        }
      }
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        if (isAncestor(nxtC.src, flow.egressNode)) {
          isAncestor(nxtC.dst, flow.egressNode) && (nxtC.dst >= nxtC.src)
        } else {
          isAncestor(nxtC.dst, nxtC.src)
        }
      }
    }
  }
}

object Mesh2DDimensionOrderedRouting {
  def apply(firstDim: Int = 0) = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)

        if (firstDim == 0) {
          if (dstX != nodeX) {
            (if (nodeX < nxtX) dstX >= nxtX else dstX <= nxtX) && nxtY == nodeY
          } else {
            (if (nodeY < nxtY) dstY >= nxtY else dstY <= nxtY) && nxtX == nodeX
          }
        } else {
          if (dstY != nodeY) {
            (if (nodeY < nxtY) dstY >= nxtY else dstY <= nxtY) && nxtX == nodeX
          } else {
            (if (nodeX < nxtX) dstX >= nxtX else dstX <= nxtX) && nxtY == nodeY
          }
        }
      }
    }
  }
}

// WARNING: Not deadlock free
object Mesh2DMinimalRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)

        val xR = (if (nodeX < nxtX) dstX >= nxtX else if (nodeX > nxtX) dstX <= nxtX else nodeX == nxtX)
        val yR = (if (nodeY < nxtY) dstY >= nxtY else if (nodeY > nxtY) dstY <= nxtY else nodeY == nxtY)
        xR && yR
      }
    }
  }
}

object Mesh2DWestFirstRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val base = Mesh2DMinimalRouting()(topo)
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)

        if (dstX < nodeX) {
          nxtX == nodeX - 1
        } else {
          base(srcC, nxtC, flow)
        }
      }
    }
  }
}

object Mesh2DNorthLastRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val base = Mesh2DMinimalRouting()(topo)
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)

        val minimal = base(srcC, nxtC, flow)
        if (dstY > nodeY && dstX != nodeX) {
          minimal && nxtY != nodeY + 1
        } else if (dstY > nodeY) {
          nxtY == nodeY + 1
        } else {
          minimal
        }
      }
    }
  }
}

object Mesh2DEscapeRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => EscapeChannelRouting(
      escapeRouter=Mesh2DDimensionOrderedRouting(),
      normalRouter=Mesh2DMinimalRouting())(topo)
  }
}

object UnidirectionalTorus2DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val baseX = UnidirectionalTorus1DDatelineRouting()(constellation.topology.UnidirectionalTorus1D(topo.nX))
      val baseY = UnidirectionalTorus1DDatelineRouting()(constellation.topology.UnidirectionalTorus1D(topo.nY))
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {

        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)
        val (srcX, srcY)   = (srcC.src % topo.nX , srcC.src / topo.nX)

        val turn = nxtX != srcX && nxtY != srcY
        if (srcC.src == -1 || turn) {
          nxtC.vc != 0
        } else if (srcX == nxtX) {
          baseY(
            srcC.copy(src=srcY, dst=nodeY),
            nxtC.copy(src=nodeY, dst=nxtY),
            flow.copy(egressNode=dstY)
          )
        } else if (srcY == nxtY) {
          baseX(
            srcC.copy(src=srcX, dst=nodeX),
            nxtC.copy(src=nodeX, dst=nxtX),
            flow.copy(egressNode=dstX)
          )
        } else {
          false
        }
      }
      override def isEscape(c: ChannelRoutingInfo, v: Int) = {
        c.isIngress || c.isEgress || c.vc < 2
      }
    }
  }
}

object BidirectionalTorus2DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val baseX = UnidirectionalTorus1DDatelineRouting()(constellation.topology.UnidirectionalTorus1D(topo.nX))
      val baseY = UnidirectionalTorus1DDatelineRouting()(constellation.topology.UnidirectionalTorus1D(topo.nY))
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)
        val (srcX, srcY)   = (srcC.src % topo.nX , srcC.src / topo.nX)

        if (srcC.src == -1) {
          nxtC.vc != 0
        } else if (nodeX == nxtX) {
          baseY(
            srcC.copy(src=srcY, dst=nodeY),
            nxtC.copy(src=nodeY, dst=nxtY),
            flow.copy(egressNode=dstY)
          )
        } else if (nodeY == nxtY) {
          baseX(
            srcC.copy(src=srcX, dst=nodeX),
            nxtC.copy(src=nodeX, dst=nxtX),
            flow.copy(egressNode=dstX)
          )
        } else {
          false
        }
      }
    }
  }
}

object DimensionOrderedUnidirectionalTorus2DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val base = UnidirectionalTorus2DDatelineRouting()(topo)
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)
        val (srcX, srcY)   = (srcC.src % topo.nX , srcC.src / topo.nX)

        def sel = if (dstX != nodeX) {
          nxtY == nodeY
        } else {
          nxtX == nodeX
        }
        base(srcC, nxtC, flow) && sel
      }
    }
  }
}

object DimensionOrderedBidirectionalTorus2DDatelineRouting {
  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: Mesh2DLikePhysicalTopology => new RoutingRelation(topo) {
      val baseX = BidirectionalTorus1DShortestRouting()(constellation.topology.BidirectionalTorus1D(topo.nX))
      val baseY = BidirectionalTorus1DShortestRouting()(constellation.topology.BidirectionalTorus1D(topo.nY))
      val base = BidirectionalTorus2DDatelineRouting()(topo)
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {

        val (nxtX, nxtY)   = (nxtC.dst % topo.nX , nxtC.dst / topo.nX)
        val (nodeX, nodeY) = (nxtC.src % topo.nX , nxtC.src / topo.nX)
        val (dstX, dstY)   = (flow.egressNode % topo.nX , flow.egressNode / topo.nX)
        val (srcX, srcY)   = (srcC.src % topo.nX , srcC.src / topo.nX)

        val xdir = baseX(
          srcC.copy(src=(if (srcC.src == -1) -1 else srcX), dst=nodeX),
          nxtC.copy(src=nodeX, dst=nxtX),
          flow.copy(egressNode=dstX)
        ) && nxtX != nodeX
        val ydir = baseY(
          srcC.copy(src=(if (srcC.src == -1) -1 else srcY), dst=nodeY),
          nxtC.copy(src=nodeY, dst=nxtY),
          flow.copy(egressNode=dstY)
        ) && nxtY != nodeY

        val sel = if (dstX != nodeX) xdir else ydir

        sel && base(srcC, nxtC, flow)
      }
    }
  }
}

object TerminalRouterRouting {
  def apply(baseRouting: PhysicalTopology => RoutingRelation) = (topo: PhysicalTopology) => topo match {
    case topo: constellation.topology.TerminalRouter => new RoutingRelation(topo) {
      val base = baseRouting(topo.base)
      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        if (srcC.isIngress && topo.isBase(nxtC.dst)) {
          true
        } else if (nxtC.dst == flow.egressNode) {
          true
        } else if (topo.isBase(nxtC.src) &&
                   topo.isBase(nxtC.dst) &&
                   nxtC.src % topo.base.nNodes != flow.egressNode % topo.base.nNodes) {
          if (topo.isTerminal(srcC.src)) {
            base(
              srcC.copy(src=(-1), dst=srcC.dst % topo.base.nNodes, vc=0, n_vc=1),
              nxtC.copy(src=nxtC.src % topo.base.nNodes, dst=nxtC.dst % topo.base.nNodes),
              flow.copy(egressNode=flow.egressNode % topo.base.nNodes))
          } else {
            base(
              srcC.copy(src=srcC.src % topo.base.nNodes, dst=srcC.dst % topo.base.nNodes),
              nxtC.copy(src=nxtC.src % topo.base.nNodes, dst=nxtC.dst % topo.base.nNodes),
              flow.copy(egressNode=flow.egressNode % topo.base.nNodes))
          }
        } else {
          false
        }
      }
      override def getNPrios(c: ChannelRoutingInfo): Int = {
        if (topo.isBase(c.dst) && topo.isBase(c.src)){
          base.getNPrios(c.copy(src=c.src % topo.base.nNodes, dst=c.dst % topo.base.nNodes))
        } else {
          1
        }
      }
      override def getPrio(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
        if (topo.isBase(srcC.dst) && topo.isBase(nxtC.dst)) {
          if (topo.isTerminal(srcC.src)) {
            base.getPrio(
              srcC.copy(src=(-1)                       , dst=srcC.dst % topo.base.nNodes, vc=0, n_vc=1),
              nxtC.copy(src=nxtC.src % topo.base.nNodes, dst=nxtC.dst % topo.base.nNodes),
              flow.copy(egressNode=flow.egressNode % topo.base.nNodes))
          } else {
            base.getPrio(
              srcC.copy(src=srcC.src % topo.base.nNodes, dst=srcC.dst % topo.base.nNodes),
              nxtC.copy(src=nxtC.src % topo.base.nNodes, dst=nxtC.dst % topo.base.nNodes),
              flow.copy(egressNode=flow.egressNode % topo.base.nNodes))
          }
        } else {
          0
        }
      }

      override def isEscape(c: ChannelRoutingInfo, v: Int) = {
        !topo.isBase(c.src) || !topo.isBase(c.dst) || base.isEscape(
          c.copy(src=c.src % topo.base.nNodes, dst=c.dst % topo.base.nNodes), v)
      }
    }
  }
}

// The below tables implement support for virtual subnetworks in a variety of ways
// NOTE: The topology must have sufficient virtual channels for these to work correctly
// TODO: Write assertions to check this

/** Produces a system with support for n virtual subnetworks. The output routing relation ensures that
  * the virtual subnetworks do not block each other. This is accomplished by allocating a portion of the
  * virtual channels to each subnetwork; all physical resources are shared between virtual subnetworks
  * but virtual channels and buffers are not shared.
  */
object NonblockingVirtualSubnetworksRouting {
  def apply(f: PhysicalTopology => RoutingRelation, n: Int, nDedicatedChannels: Int) = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def trueVIdToVirtualVId(vId: Int) = if (vId < n * nDedicatedChannels) {
      vId % nDedicatedChannels
    } else {
      nDedicatedChannels + vId - n * nDedicatedChannels
    }
    def lower(c: ChannelRoutingInfo) = c.copy(
      vc = trueVIdToVirtualVId(c.vc),
      n_vc = if (c.isIngress) 1 else c.n_vc - n * nDedicatedChannels + nDedicatedChannels
    )
    val base = f(topo)
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      val able = (nxtC.vc >= n * nDedicatedChannels) || ((nxtC.vc / nDedicatedChannels) == flow.vNetId)
      able && base(lower(srcC), lower(nxtC), flow.copy(vNetId=0))
    }
    override def isEscape(c: ChannelRoutingInfo, v: Int) = {
      base.isEscape(lower(c), 0)
    }
  }
}

object BlockingVirtualSubnetworksRouting {
  def apply(f: PhysicalTopology => RoutingRelation, n: Int, nDedicated: Int = 1) = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    val base = f(topo)
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      val lNxtV = nxtC.vc - flow.vNetId * nDedicated
      val lSrcV = srcC.vc - flow.vNetId * nDedicated
      val r = if (srcC.isIngress && lNxtV >= 0) {
        base(srcC,
          nxtC.copy(n_vc = nxtC.n_vc - flow.vNetId * nDedicated, vc = lNxtV),
          flow.copy(vNetId=0)
        )
      } else if (lNxtV < 0 || lSrcV < 0) {
        false
      } else {
        base(srcC.copy(n_vc = srcC.n_vc - flow.vNetId * nDedicated, vc = lSrcV),
          nxtC.copy(n_vc = nxtC.n_vc - flow.vNetId * nDedicated, vc = lNxtV),
          flow.copy(vNetId=0)
        )
      }
      r
    }
    override def isEscape(c: ChannelRoutingInfo, v: Int) = {
      c.vc >= v * nDedicated && base.isEscape(c.copy(vc=c.vc - v * nDedicated, n_vc=c.n_vc - v * nDedicated), 0)
    }
  }
}

object HierarchicalRouting {
  def apply(
    baseRouting: PhysicalTopology => RoutingRelation,
    childRouting: Seq[PhysicalTopology => RoutingRelation]
  ) = (topo: PhysicalTopology) => topo match {
    case topo: HierarchicalTopology => new RoutingRelation(topo) {
      val base = baseRouting(topo.base)
      val children = (childRouting zip topo.children).map { case (l,r) => l(r.topo) }
      def copyChannel(c: ChannelRoutingInfo, o: Int) = {
        val nSrc = if (c.src < topo.base.nNodes) -1 else c.src - o
        val nDst = if (c.dst == -1) -1 else c.dst - o
        assert(nSrc >= -1 && nDst >= -1, s"$c $o")
        c.copy(src=nSrc, dst=nDst)
      }

      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
        val thisBase = topo.isBase(srcC.dst)
        val flowBase = topo.isBase(flow.egressNode)
        val nextBase = topo.isBase(nxtC.dst)
        (thisBase, flowBase, nextBase) match {
          case (true , true , true ) => base(
            srcC.copy(src=if (srcC.src >= topo.base.nNodes) -1 else srcC.src),
            nxtC, flow
          )
          case (true , true , false) => false
          case (true , false, true ) => {
            val src = topo.children(topo.childId(flow.egressNode)).src
            val atDst = src == srcC.dst
            !atDst && base(
              srcC.copy(src=if (srcC.src >= topo.base.nNodes) -1 else srcC.src),
              nxtC,
              flow.copy(egressNode=src)
            )
          }
          case (true , false, false) => topo.childId(nxtC.dst) == topo.childId(flow.egressNode)
          case (false, true , true ) => true
          case (false, true , false) => {
            val id = topo.childId(srcC.dst)
            children(id)(
              srcC=copyChannel(srcC, topo.offsets(id)),
              nxtC=copyChannel(nxtC, topo.offsets(id)),
              flow=flow.copy(egressNode=topo.children(id).dst)
            ) && topo.children(id).dst != srcC.dst - topo.offsets(id)
          }
          case (false, false, true ) => topo.childId(srcC.dst) != topo.childId(flow.egressNode)
          case (false, false, false) => {
            val fid = topo.childId(flow.egressNode)
            val sid = topo.childId(srcC.dst)
            val fdst = if (fid == sid) {
              flow.egressNode - topo.offsets(sid)
            } else {
              topo.children(sid).dst
            }
            val at_exit = (fid != sid &&
              (srcC.dst - topo.offsets(sid) == topo.children(sid).dst))
            children(sid)(
              srcC=copyChannel(srcC, topo.offsets(sid)),
              nxtC=copyChannel(nxtC, topo.offsets(sid)),
              flow=flow.copy(egressNode=fdst)
            ) && !at_exit
          }
        }
      }
      override def getNPrios(c: ChannelRoutingInfo): Int = {
        if (topo.isBase(c.dst) && topo.isBase(c.src)){
          base.getNPrios(c)
        } else if (topo.childId(c.src) == topo.childId(c.dst) && !c.isEgress) {
          children(topo.childId(c.src)).getNPrios(c.copy(
            src=if (c.src >= topo.base.nNodes) -1 else c.src))
        } else {
          1
        }
      }
      override def getPrio(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
        val thisBase = topo.isBase(srcC.dst)
        val nextBase = topo.isBase(nxtC.dst)
        val id = topo.childId(srcC.dst)
        (thisBase, nextBase) match {
          case (true, true) => base.getPrio(
            srcC.copy(src=if (srcC.src >= topo.base.nNodes) -1 else srcC.src),
            nxtC,
            flow.copy(egressNode=topo.children(topo.childId(flow.egressNode)).src))
          case (false, false) => {
            if (topo.childId(srcC.dst) == topo.childId(nxtC.dst)) {
              children(id).getPrio(
                copyChannel(srcC, topo.offsets(id)),
                copyChannel(nxtC, topo.offsets(id)),
                flow.copy(egressNode=topo.children(id).dst))
            } else {
              0
            }
          }
          case _ => 0
        }
      }

      override def isEscape(c: ChannelRoutingInfo, v: Int) = {
        val srcBase = topo.isBase(c.src)
        val dstBase = topo.isBase(c.dst)
          (srcBase, dstBase) match {
          case (true, true) => base.isEscape(c, v)
          case (true, false) => true
          case (false, true) => true
          case (false, false) => {
            val sid = topo.childId(c.dst)
            children(sid).isEscape(copyChannel(c, topo.offsets(sid)), v)
          }
        }
      }
    }
  }
}

/**
 * Pure shortest-path routing strategy with some VC-based load balancing.
 *
 * This routing relation computes single-path shortest routes between all node pairs
 * using BFS and restricts flow transitions to only those that follow the precomputed
 * shortest-path next hop. No deadlock avoidance mechanism is enforced.
 *
 * If multiple virtual channels (VCs) are available, we can assist the
 * PrioritizingVCAllocator by assigning priorities deterministically via a hashing
 * of (flow.ingressNode, flow.egressNode) to help with better VC utilization.
 * However VC transitions are not enforced and per-hop flows are allowed to use any VC
 * as long as the edge is on the shortest path.
 *
 * This routing strategy is best for minimal-hop designs, and
 * scenarios where higher-level deadlock prevention is handled
 * externally such as through EscapeChannelRouting.
 *
 * maxVCs: The number of virtual channels available per physical channel.
 *         This must be >= 1. Defaults to 1 (no VC usage).
 *
 */

object ShortestPathRouting {
  def apply(maxVCs: Int = 1) = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    require(maxVCs >= 1, s"[ShortestPathRouting] maxVCs must be >= 1 (got $maxVCs)")

    // Build adjacency list
    val adjList: Map[Int, Seq[Int]] = (0 until topo.nNodes).map { src =>
      src -> (0 until topo.nNodes).filter(dst => topo.topo(src, dst))
    }.toMap

    // Precompute shortest-path next hops using BFS
    val nextHop: Map[(Int, Int), Int] = {
      val paths = for (src <- 0 until topo.nNodes; dst <- 0 until topo.nNodes if src != dst) yield {
        val next = bfsNextHop(src, dst)
        (src, dst) -> next
      }
      paths.collect { case ((s, d), Some(n)) => (s, d) -> n }.toMap
    }

    // BFS shortest-path helper
    private def bfsNextHop(src: Int, dst: Int): Option[Int] = {
      val visited = scala.collection.mutable.Set[Int]()
      val queue = scala.collection.mutable.Queue[(Int, List[Int])]()
      queue.enqueue((src, List()))
      while (queue.nonEmpty) {
        val (node, path) = queue.dequeue()
        if (node == dst) return path.headOption
        if (!visited.contains(node)) {
          visited += node
          adjList(node).foreach { neighbor =>
            queue.enqueue((neighbor, if (path.isEmpty) List(neighbor) else path))
          }
        }
      }
      None
    }

    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Boolean = {
      val flowKey = (flow.ingressNode, flow.egressNode)

      if (srcC.src == -1) {
        // Ingress: only inject into first hop of shortest path
        val expected = nextHop.get((flow.ingressNode, flow.egressNode))
        val legal = nxtC.src == flow.ingressNode && expected.contains(nxtC.dst)
        return legal
      }

      if (srcC.dst == flow.egressNode) return false

      val expected = nextHop.get((srcC.dst, flow.egressNode))
      val legal = expected.contains(nxtC.dst)
      legal
    }

    // VC configuration

    override def getPrio(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
      (flow.ingressNode * 31 + flow.egressNode) % maxVCs
    }

    override def isEscape(c: ChannelRoutingInfo, vNetId: Int): Boolean = c.isIngress || c.isEgress

    println(s"[ShortestPathRouting] Initialized with maxVCs = $maxVCs")
  }
}

object EscapeIDOrderRouting {
  def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo) = {
      if (srcC.src == -1) {
        true // allow ingress to inject anywhere
      } else {
        (srcC.dst < nxtC.dst) && topo.topo(srcC.dst, nxtC.dst)
      }
    }
  }
}

/*
 * CustomLayeredRouting: This is a a generalized deadlock-free routing policy for arbitrary
 * topologies.
 *
 * This routing relation assigns each flow (src, dst) a dedicated virtual channel (VC) layer
 * such that all paths in the same layer do not introduce cyclic dependencies. This is done by:
 *   1. Finding all shortest paths between flow pairs
 *   2. Deduplicating and pruning strict subpaths
 *   3. Packing flow paths into the minimum number of layers while preserving DAG constraints
 *   4. Assigning a VC layer to each flow based on its path
 *
 * The rel() function enforces this policy by:
 *   - Ensuring packets flow only along the assigned VC
 *   - Ensuring only legal hops in the assigned flow path are taken
 *   - Ensuring the hop is within the assigned VC layer
 */

object CustomLayeredRouting {

  def apply() = (topo: PhysicalTopology) => topo match {
    case topo: CustomTopology => new RoutingRelation(topo) {
      val edgeList = topo.edgeList.map(e => (e.src, e.dst))
      val graph = new CustomGraph(topo.nNodes, edgeList)
      val allSSPs = graph.generateSSPs()

      val deduped = graph.deduplicateSSPs(allSSPs)
      val pruned = graph.pruneSubpaths(deduped)
      val packed = graph.packLayersMinimal(pruned)

      println(s"[CustomLayeredRouting] Required VCs: ${packed.length}")

      // VC assignment for all flows based on first matching layer
      val allFlowAssignments: Map[(Int, Int), Int] = allSSPs.map { case ((src, dst), path) =>
        val assignedLayer = packed.indexWhere(layer => path.toSet.subsetOf(layer.flatMap(_._2).toSet))
        ((src, dst), assignedLayer)
      }.toMap

      // Build nextHop map per VC (layer)
      val nextHopMap: Map[Int, Map[(Int, Int), Int]] = packed.zipWithIndex.map { case (layer, vc) =>
        val nhs = mutable.Map[(Int, Int), Int]()

        for (((src, dst), path) <- layer) {
          val nodes = path.map(_._1) :+ path.last._2
          for (i <- 0 until nodes.length - 1) {
            nhs((nodes(i), dst)) = nodes(i + 1)
          }
        }

        vc -> nhs.toMap
      }.toMap

      val flowPaths: Map[(Int, Int), List[(Int, Int)]] = allSSPs.toMap
      val vcToEdges = packed.map(_.flatMap(_._2).toSet)

      def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Boolean = {
        val flowKey = (flow.ingressNode, flow.egressNode)

        // Step 1: Lookup assigned VC
        val assignedVC = allFlowAssignments.getOrElse(flowKey, {
          return false
        })

        // Step 2: Lookup full path and verify layer correctness
        val flowPath = flowPaths.getOrElse(flowKey, {
          return false
        })

        val flowPathEdges = flowPath.toSet
        val assignedEdgeSet = vcToEdges(assignedVC)

        val fullPathOk = flowPathEdges.subsetOf(assignedEdgeSet)
        if (!fullPathOk) {
          return false
        }

        // Step 3: Current and next node
        val curNode  = if (srcC.src == -1) flow.ingressNode else srcC.dst
        val nextNode = nxtC.dst
        val edge     = (curNode, nextNode)

        // Step 4: Legal check
        val edgeOk = flowPathEdges.contains(edge)
        val vcOk   = nxtC.vc == assignedVC && (srcC.src == -1 || srcC.vc == assignedVC)

        val legal = edgeOk && vcOk

        legal
      }

      override def isEscape(c: ChannelRoutingInfo, v: Int): Boolean =
        c.isIngress || c.isEgress

      val guAssignedLayer = allFlowAssignments
      val nVirtualLayers = vcToEdges.length

      override def getNPrios(c: ChannelRoutingInfo): Int = nVirtualLayers

      override def getPrio(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = 0

    }
  }

}

/**
 * ShortestPathGeneralizedDatelineRouting implements a deadlock-free routing relation
 * based on identifying cycle-breaking datelines and the presence of virtual channels.
 *
 * This routing strategy:
 * 1. Analyzes the given topology using DatelineAnalyzer to:
 *    - Identify cycles in the channel dependency graph (CDG)
 *    - Select a minimal set of "dateline" edges to break those cycles
 *    - Compute the minimum number of virtual channels (VCs) needed to support
 *      all shortest paths without deadlock
 *
 * 2. Recomputes all-pairs shortest paths and annotates each hop with the
 *    required VC level. VC level increases only at dateline crossings,
 *    and remains constant otherwise. This leads to an enforced acyclic traversal
 *    of the dependency graph.
 *
 * 3. rel()) only permits transitions which:
 *    - Follow the exact shortest path between (ingress, egress) nodes
 *    - Match the expected VC at each hop
 *    - Begin from the ingress using an assigned VC
 *        - If useStaticVC0 is true then injections are only allowed on VC 0
 *        - Otherwise there is a hashing scheme used to calculate the injection VC
 *
 */

object ShortestPathGeneralizedDatelineRouting {
  // def apply() = (topo: PhysicalTopology) => new RoutingRelation(topo) {
  def apply(
    useStaticVC0: Boolean = false,
    vcHashCoeffs: (Int, Int) = (19, 4)
  ) = (topo: PhysicalTopology) => new RoutingRelation(topo) {
    val result = DatelineAnalyzer.analyze(topo)
    val datelineEdges = result.datelineEdges
    val maxVC = result.vcCount
    val nextHop = result.nextHop

    // Recompute shortest paths and expected VC transitions
    val (shortestPaths, pathVCs): (Map[(Int, Int), List[Int]], Map[(Int, Int), List[Int]]) = {
      val g = new CustomGraph(topo.nNodes, (0 until topo.nNodes).flatMap(u => (0 until topo.nNodes).collect {
        case v if topo.topo(u, v) => (u, v)
      }))
      val sspPaths = g.generateSSPs()
      val pathMap = sspPaths.map { case ((src, dst), edges) => ((src, dst), src +: edges.map(_._2)) }.toMap
      val vcMap = pathMap.map { case ((src, dst), path) =>
        val vcs = if (useStaticVC0) {
          path.sliding(2).scanLeft(0) {
            case (vc, Seq(u, v)) =>
              if (datelineEdges.contains((u, v))) vc + 1 else vc
          }.toList
        } else {
          val (a, b) = vcHashCoeffs
          val baseVC = (src * a + dst * b) % maxVC
          path.sliding(2).scanLeft(baseVC) {
            case (vc, Seq(u, v)) =>
              if (datelineEdges.contains((u, v))) (vc + 1) % maxVC else vc
          }.toList
        }
        ((src, dst), vcs)
      }
      (pathMap, vcMap)
    }

    println(s"[ShortestPathGeneralizedDatelineRouting] Initialized with ${datelineEdges.size} datelines and $maxVC VCs")

    def rel(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Boolean = {

      val flowKey = (flow.ingressNode, flow.egressNode)

      val path = shortestPaths.getOrElse((flow.ingressNode, flow.egressNode), Nil)
      val vcs  = pathVCs.getOrElse((flow.ingressNode, flow.egressNode), Nil)

      if (srcC.src == -1) {
          val validInjectionEdge = path.headOption.contains(nxtC.src) && path.lift(1).contains(nxtC.dst)
          val expectedVC = vcs.lift(1).getOrElse(-1)
          val vcValid = nxtC.vc == expectedVC
          validInjectionEdge && vcValid
      } else if (flow.egressNode == srcC.dst) {
        false
      } else {
        val edge = (srcC.dst, nxtC.dst)
        val pathEdges = path.sliding(2).toList
        val idx = pathEdges.indexWhere { case Seq(a, b) => a == edge._1 && b == edge._2 }

        val validEdge = idx != -1
        val expectedVC = if (idx >= 0 && idx < vcs.length) vcs(idx + 1) else -1
        val vcValid = nxtC.vc == expectedVC

        validEdge && vcValid
      }
    }

    override def getNPrios(src: ChannelRoutingInfo): Int = {
      maxVC
    }

    override def getPrio(srcC: ChannelRoutingInfo, nxtC: ChannelRoutingInfo, flow: FlowRoutingInfo): Int = {
      // Priority is highest for lowest VC to prefer dateline avoidance
      maxVC - nxtC.vc
    }

    override def isEscape(c: ChannelRoutingInfo, vNetId: Int): Boolean = {
      c.isIngress || c.isEgress || c.vc == 0
    }
  }
}

override def getNPrios(src: ChannelRoutingInfo): Int = maxVCs


package constellation.soc

import chisel3._
import chisel3.util._

import org.chipsalliance.cde.config._
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.devices.tilelink._
import freechips.rocketchip.tilelink._
import freechips.rocketchip.subsystem._
import freechips.rocketchip.util._

import constellation.noc.{NoCParams}
import constellation.protocol._

import scala.collection.immutable.{ListMap}

case class ConstellationSystemBusParams(
  sbusParams: SystemBusParams,
  tlNoCParams: TLNoCParams,
  inlineNoC: Boolean
) extends TLBusWrapperInstantiationLike {
  def instantiate(context: HasTileLinkLocations, loc: Location[TLBusWrapper]
  )(implicit p: Parameters): ConstellationSystemBus = {
    val constellation = LazyModule(new ConstellationSystemBus(
      sbusParams, tlNoCParams, loc.name, context, inlineNoC))

    constellation.suggestName(loc.name)
    context.tlBusWrapperLocationMap += (loc -> constellation)
    constellation
  }
}

class ConstellationSystemBus(
  sbus_params: SystemBusParams, noc_params: TLNoCParams, name: String, context: HasTileLinkLocations,
  inlineNoC: Boolean
)(implicit p: Parameters) extends TLBusWrapper(sbus_params, name) {
  private val replicator = sbus_params.replication.map(r => LazyModule(new RegionReplicator(r)))
  val prefixNode = replicator.map { r =>
    r.prefix := addressPrefixNexusNode
    addressPrefixNexusNode
  }

  override def shouldBeInlined = inlineNoC

  private val system_bus_noc = noc_params match {
    case params: GlobalTLNoCParams => context.asInstanceOf[CanHaveGlobalNoC].globalNoCDomain {
      LazyModule(new TLGlobalNoC(params, name))
    }
    case params: SimpleTLNoCParams => LazyModule(new TLNoC(params, name, inlineNoC))
    case params: SplitACDxBETLNoCParams => LazyModule(new TLSplitACDxBENoC(params, name, inlineNoC))
  }

  val inwardNode: TLInwardNode = (system_bus_noc.node :=* TLFIFOFixer(TLFIFOFixer.allVolatile)
    :=* replicator.map(_.node).getOrElse(TLTempNode()))
  val outwardNode: TLOutwardNode = system_bus_noc.node
  def busView: TLEdge = system_bus_noc.node.edges.in.head

  val builtInDevices: BuiltInDevices = BuiltInDevices.attach(sbus_params, outwardNode)
}

case class ConstellationMemoryBusParams(
  mbusParams: MemoryBusParams,
  tlNoCParams: TLNoCParams,
  inlineNoC: Boolean
) extends TLBusWrapperInstantiationLike {
  def instantiate(context: HasTileLinkLocations, loc: Location[TLBusWrapper]
  )(implicit p: Parameters): ConstellationMemoryBus = {
    val constellation = LazyModule(new ConstellationMemoryBus(
      mbusParams, tlNoCParams, loc.name, context, inlineNoC))

    constellation.suggestName(loc.name)
    context.tlBusWrapperLocationMap += (loc -> constellation)
    constellation
  }
}

class ConstellationMemoryBus(mbus_params: MemoryBusParams, noc_params: TLNoCParams, name: String, context: HasTileLinkLocations, inlineNoC: Boolean)
  (implicit p: Parameters) extends TLBusWrapper(mbus_params, name) {
  private val replicator = mbus_params.replication.map(r => LazyModule(new RegionReplicator(r)))
  val prefixNode = replicator.map { r =>
    r.prefix := addressPrefixNexusNode
    addressPrefixNexusNode
  }

  private val memory_bus_noc = noc_params match {
    case params: GlobalTLNoCParams => context.asInstanceOf[CanHaveGlobalNoC].globalNoCDomain {
      LazyModule(new TLGlobalNoC(params, name))
    }
    case params: SimpleTLNoCParams => LazyModule(new TLNoC(params, name, inlineNoC))
    case params: SplitACDxBETLNoCParams => LazyModule(new TLSplitACDxBENoC(params, name, inlineNoC))
  }

  val inwardNode: TLInwardNode =
    replicator.map(memory_bus_noc.node :*=* TLFIFOFixer(TLFIFOFixer.all) :*=* _.node)
        .getOrElse(memory_bus_noc.node :*=* TLFIFOFixer(TLFIFOFixer.all))

  val outwardNode: TLOutwardNode = ProbePicker() :*= memory_bus_noc.node
  def busView: TLEdge = memory_bus_noc.node.edges.in.head

  val builtInDevices: BuiltInDevices = BuiltInDevices.attach(mbus_params, outwardNode)
}

case class ConstellationPeripheryBusParams(
  pbusParams: PeripheryBusParams,
  tlNoCParams: TLNoCParams,
  inlineNoC: Boolean
) extends TLBusWrapperInstantiationLike {
  def instantiate(context: HasTileLinkLocations, loc: Location[TLBusWrapper]
  )(implicit p: Parameters): ConstellationPeripheryBus = {
    val constellation = LazyModule(new ConstellationPeripheryBus(pbusParams, tlNoCParams, loc.name, context, inlineNoC))

    constellation.suggestName(loc.name)
    context.tlBusWrapperLocationMap += (loc -> constellation)
    constellation
  }
}

class ConstellationPeripheryBus(pbus_params: PeripheryBusParams, noc_params: TLNoCParams, name: String, context: HasTileLinkLocations, inlineNoC: Boolean)
  (implicit p: Parameters) extends TLBusWrapper(pbus_params, name) {

  private val replicator = pbus_params.replication.map(r => LazyModule(new RegionReplicator(r)))
  val prefixNode = replicator.map { r =>
    r.prefix := addressPrefixNexusNode
    addressPrefixNexusNode
  }

  def genNoC()(implicit valName: ValName): TLNoCLike = noc_params match {
    case params: GlobalTLNoCParams => context.asInstanceOf[CanHaveGlobalNoC].globalNoCDomain {
      LazyModule(new TLGlobalNoC(params, name))
    }
    case params: SimpleTLNoCParams => LazyModule(new TLNoC(params, name, inlineNoC))
    case params: SplitACDxBETLNoCParams => LazyModule(new TLSplitACDxBENoC(params, name, inlineNoC))
  }

  private val fixer = LazyModule(new TLFIFOFixer(TLFIFOFixer.all))
  private val node: TLNode = pbus_params.atomics.map { pa =>
    val in_xbar = LazyModule(new TLXbar)
    val out_noc = genNoC()
    val fixer_node = replicator.map(fixer.node :*= _.node).getOrElse(fixer.node)
    (out_noc.node
      :*= fixer_node
      :*= TLBuffer(pa.buffer)
      :*= (pa.widenBytes.filter(_ > beatBytes).map { w =>
          TLWidthWidget(w) :*= TLAtomicAutomata(arithmetic = pa.arithmetic)

      :*= in_xbar.node)
  } .getOrElse {
    val noc = genNoC()
    noc.node :*= fixer.node
  }

  def inwardNode: TLInwardNode = node
  def outwardNode: TLOutwardNode = node
  def busView: TLEdge = fixer.node.edges.in.head

  val builtInDevices: BuiltInDevices = BuiltInDevices.attach(pbus_params, outwardNode)
}

} .getOrElse { TLAtomicAutomata(arithmetic = pa.arithmetic) })


package constellation.util

import chisel3._
import chisel3.util._

object WrapInc {
  // "n" is the number of increments, so we wrap at n-1.
  def apply(value: UInt, n: Int): UInt = {
    if (n == 1) {
      0.U
    } else if (isPow2(n)) {
      (value + 1.U)(log2Ceil(n)-1,0)
    } else {

  def apply(value: UInt, n: UInt): UInt = {
    val wrap = value === n - 1.U
    Mux(wrap, 0.U, value + 1.U)
  }
}

      val wrap = (value === (n-1).U)
      Mux(wrap, 0.U, value + 1.U)
    }


package essent

import scala.io.Source
import scala.sys.process._

import logger._

object Driver {
  def main(args: Array[String]): Unit = {

    Logger.setClassLogLevels(Map("essent" -> logger.LogLevel(opt.essentLogLevel)))
    Logger.setClassLogLevels(Map("firrtl" -> logger.LogLevel(opt.firrtlLogLevel)))
    val sourceReader = Source.fromFile(opt.firInputFile)
    val circuit = firrtl.Parser.parse(sourceReader.getLines(), firrtl.Parser.IgnoreInfo)
    sourceReader.close()
    val compiler = new EssentCompiler(opt)
    compiler.compileAndEmit(circuit)
  }

  def compileCPP(dutName: String, buildDir: String): ProcessBuilder = {
    Seq("g++", "-O3", "-std=c++11", s"-I$buildDir", s"-Ifirrtl-sig",
      s"$buildDir/$dutName-harness.cc", "-o", s"$buildDir/$dutName")
  }
}

    (new ArgsParser).getConfig(args.toSeq) match {
      case Some(config) => generate(config)
      case None =>
    }


package essent

import essent.Extract._
import essent.ir._

import firrtl._
import firrtl.annotations._
import firrtl.ir._
import firrtl.Mappers._
import firrtl.PrimOps._
import firrtl.Utils._

import scala.util.Random

object Emitter {
  case class HyperedgeDep(name: String, deps: Seq[String], stmt: Statement)

  // Type Declaration & Initialization
  //----------------------------------------------------------------------------
  def genCppType(tpe: Type) = tpe match {
    case UIntType(IntWidth(w)) => s"UInt<$w>"
    case SIntType(IntWidth(w)) => s"SInt<$w>"
    case AsyncResetType => "UInt<1>"
    case _ => throw new Exception(s"No CPP type implemented for $tpe")
  }

  def initializeVals(topLevel: Boolean)(m: Module, registers: Seq[DefRegister], memories: Seq[DefMemory]) = {
    def initVal(name: String, tpe:Type) = s"$name.rand_init();"
    val regInits = registers map {
      r: DefRegister => initVal(r.name, r.tpe)
    }
    val memInits = memories flatMap { m: DefMemory => {
      if ((m.depth > 1000) && (bitWidth(m.dataType)) <= 64) {
        Seq(s"${m.name}[0].rand_init();",
            s"for (size_t a=0; a < ${m.depth}; a++) ${m.name}[a] = ${m.name}[0].as_single_word() + a;")
      } else
        Seq(s"for (size_t a=0; a < ${m.depth}; a++) ${m.name}[a].rand_init();")
    }}
    val portInits = m.ports flatMap { p => p.tpe match {
      case ClockType => Seq()
      case _ => if (!topLevel) Seq()
                else Seq(initVal(p.name, p.tpe))
    }}
    regInits ++ memInits ++ portInits
  }

  // Prefixing & Replacement
  //----------------------------------------------------------------------------
  def addPrefixToNameStmt(prefix: String)(s: Statement): Statement = {
    val replaced = s match {
      case n: DefNode => n.copy(name = prefix + n.name)
      case r: DefRegister => r.copy(name = prefix + r.name)
      case m: DefMemory => m.copy(name = prefix + m.name)
      case w: DefWire => w.copy(name = prefix + w.name)
      case mw: MemWrite => mw.copy(memName = prefix + mw.memName)
      case _ => s
    }
    replaced map addPrefixToNameStmt(prefix) map addPrefixToNameExpr(prefix)
  }

  def addPrefixToNameExpr(prefix: String)(e: Expression): Expression = {
    val replaced = e match {
      case w: WRef => w.copy(name = prefix + w.name)
      case _ => e
    }
    replaced map addPrefixToNameExpr(prefix)
  }

  def replaceNamesStmt(renames: Map[String, String])(s: Statement): Statement = {
    val nodeReplaced = s match {
      case n: DefNode if (renames.contains(n.name)) => n.copy(name = renames(n.name))
      case cm: CondMux if (renames.contains(cm.name)) => cm.copy(name = renames(cm.name))
      case mw: MemWrite if (renames.contains(mw.memName)) => mw.copy(memName = renames(mw.memName))
      case _ => s
    }
    nodeReplaced map replaceNamesStmt(renames) map replaceNamesExpr(renames)
  }

  def replaceNamesExpr(renames: Map[String, String])(e: Expression): Expression = {
    def findRootKind(e: Expression): Kind = e match {
      case w: WRef => w.kind
      case w: WSubField => findRootKind(w.expr)
    }
    e match {
      case w: WRef => {
        if (renames.contains(w.name)) w.copy(name = renames(w.name))
        else w
      }
      case w: WSubField => {
        val fullName = emitExpr(w)
        // flattens out nested WSubFields
        if (renames.contains(fullName)) WRef(renames(fullName), w.tpe, findRootKind(w), w.flow)
        else w
      }
      case _ => e map replaceNamesExpr(renames)
    }
  }

  // Emission
  //----------------------------------------------------------------------------
  def emitPort(topLevel: Boolean)(p: Port): Seq[String] = p.tpe match {
    case ClockType => if (!topLevel) Seq()
                      else Seq(genCppType(UIntType(IntWidth(1))) + " " + p.name + ";")
      // FUTURE: suppress generation of clock field if not making harness (or used)?
    case _ => if (!topLevel) Seq()
              else Seq(genCppType(p.tpe) + " " + p.name + ";")
  }

  def chunkLitString(litStr: String, chunkWidth:Int = 16): Seq[String] = {
    if (litStr.length % chunkWidth == 0) litStr.grouped(chunkWidth).toSeq
    else Seq(litStr.take(litStr.length % chunkWidth)) ++ chunkLitString(litStr.drop(litStr.length % chunkWidth), chunkWidth)
  }

  // NOTE: assuming no large UIntLiteral is negative
  def splatLargeLiteralIntoRawArray(value: BigInt, width: BigInt): String = {
    val rawHexStr = value.toString(16)
    val isNeg = value < 0
    val asHexStr = if (isNeg) rawHexStr.tail else rawHexStr
    val arrStr = chunkLitString(asHexStr) map { "0x" + _} mkString(",")
    val leadingNegStr = if (isNeg) "(uint64_t) -" else ""
    val numWords = (width + 63) / 64
    s"std::array<uint64_t,$numWords>({$leadingNegStr$arrStr})"
  }

  def emitExpr(e: Expression)(implicit rn: Renamer = null): String = e match {
    case w: WRef => if (rn != null) rn.emit(w.name) else w.name
    case u: UIntLiteral => {
      val maxIn64Bits = (BigInt(1) << 64) - 1
      val width = bitWidth(u.tpe)
      val asHexStr = u.value.toString(16)
      if ((width <= 64) || (u.value <= maxIn64Bits)) s"UInt<$width>(0x$asHexStr)"
      else s"UInt<$width>(${splatLargeLiteralIntoRawArray(u.value, width)})"
    }
    case u: SIntLiteral => {
      val width = bitWidth(u.tpe)
      if (width <= 64) s"SInt<$width>(${u.value.toString(10)})"
      else s"SInt<$width>(${splatLargeLiteralIntoRawArray(u.value, width)})"
    }
    case m: Mux => {
      val condName = emitExprWrap(m.cond)
      val tvalName = emitExprWrap(m.tval)
      val fvalName = emitExprWrap(m.fval)
      s"$condName ? $tvalName : $fvalName"
    }
    case w: WSubField => {
      val result = s"${emitExpr(w.expr)(null)}.${w.name}"
      if (rn != null)
        rn.emit(result)
      else
        result
    }
    case w: WSubAccess => s"${emitExpr(w.expr)}[${emitExprWrap(w.index)}.as_single_word()]"
    case p: DoPrim => p.op match {
      case Add => p.args map emitExprWrap mkString(" + ")

      case Sub => p.args map emitExprWrap mkString(" - ")
      case Subw => s"${emitExprWrap(p.args(0))}.subw(${emitExprWrap(p.args(1))})"
      case Mul => p.args map emitExprWrap mkString(" * ")
      case Div => p.args map emitExprWrap mkString(" / ")
      case Rem => p.args map emitExprWrap mkString(" % ")
      case Lt  => p.args map emitExprWrap mkString(" < ")
      case Leq => p.args map emitExprWrap mkString(" <= ")
      case Gt  => p.args map emitExprWrap mkString(" > ")
      case Geq => p.args map emitExprWrap mkString(" >= ")
      case Eq => p.args map emitExprWrap mkString(" == ")
      case Neq => p.args map emitExprWrap mkString(" != ")
      case Pad => s"${emitExprWrap(p.args.head)}.pad<${bitWidth(p.tpe)}>()"
      case AsUInt => s"${emitExprWrap(p.args.head)}.asUInt()"
      case AsSInt => s"${emitExprWrap(p.args.head)}.asSInt()"
      case AsClock => throw new Exception("AsClock unimplemented!")
      case AsAsyncReset => emitExpr(p.args.head) // TODO: make async
      case Shl => s"${emitExprWrap(p.args.head)}.shl<${p.consts.head.toInt}>()"
      // case Shlw => s"${emitExprWrap(p.args.head)}.shlw<${p.consts.head.toInt}>()"
      case Shr => s"${emitExprWrap(p.args.head)}.shr<${p.consts.head.toInt}>()"
      case Dshl => p.args map emitExprWrap mkString(" << ")
      case Dshlw => s"${emitExprWrap(p.args(0))}.dshlw(${emitExpr(p.args(1))})"
      case Dshr => p.args map emitExprWrap mkString(" >> ")
      case Cvt => s"${emitExprWrap(p.args.head)}.cvt()"
      case Neg => s"-${emitExprWrap(p.args.head)}"
      case Not => s"~${emitExprWrap(p.args.head)}"
      case And => p.args map emitExprWrap mkString(" & ")
      case Or => p.args map emitExprWrap mkString(" | ")
      case Xor => p.args map emitExprWrap mkString(" ^ ")
      case Andr => s"${emitExprWrap(p.args.head)}.andr()"
      case Orr => s"${emitExprWrap(p.args.head)}.orr()"
      case Xorr => s"${emitExprWrap(p.args.head)}.xorr()"
      case Cat => s"${emitExprWrap(p.args(0))}.cat(${emitExpr(p.args(1))})"
      case Bits => s"${emitExprWrap(p.args.head)}.bits<${p.consts(0).toInt},${p.consts(1).toInt}>()"
      case Head => s"${emitExprWrap(p.args.head)}.head<${p.consts.head.toInt}>()"
      case Tail => s"${emitExprWrap(p.args.head)}.tail<${p.consts.head.toInt}>()"
    }
    case _ => throw new Exception(s"Don't yet support $e")
  }

  def emitExprWrap(e: Expression)(implicit rn: Renamer): String = e match {
    case DoPrim(_,_,_,_) | Mux(_,_,_,_) => s"(${emitExpr(e)})"
    case _ => emitExpr(e)
  }

  def emitStmt(s: Statement)(implicit rn: Renamer): Seq[String] = s match {
    case b: Block => b.stmts flatMap emitStmt
    case d: DefNode => {
      val lhs_orig = d.name
      val lhs = rn.emit(lhs_orig)
      val rhs = emitExpr(d.value)
      if (rn.decLocal(lhs_orig)) Seq(s"${genCppType(d.value.tpe)} $lhs = $rhs;")
      else Seq(s"$lhs = $rhs;")
    }
    case c: Connect => {
      val lhs_orig = emitExpr(c.loc)(null)
      val lhs = rn.emit(lhs_orig)
      val rhs = emitExpr(c.expr)
      if (rn.decLocal(lhs_orig)) Seq(s"${genCppType(c.loc.tpe)} $lhs = $rhs;")
      else Seq(s"$lhs = $rhs;")
    }
    case p: Print => {
      val formatters = "(%h)|(%x)|(%d)|(%ld)".r.findAllIn(p.string.serialize).toList
      val argWidths = p.args map {e: Expression => bitWidth(e.tpe)}
      if (!(argWidths forall { _ <= 64 })) throw new Exception(s"Can't print wide signals")
      val replacements = formatters zip argWidths map { case(format, width) =>
        if (format == "%h" || format == "%x") {
          val printWidth = math.ceil(width.toDouble/4).toInt
          (format, s"""%0${printWidth}" PRIx64 """")
        } else {
          val printWidth = math.ceil(math.log10((1L<<width.toInt).toDouble)).toInt
          (format, s"""%${printWidth}" PRIu64 """")
        }
      }
      val formatString = replacements.foldLeft(p.string.serialize){
        case (str, (searchFor, replaceWith)) => str.replaceFirst(searchFor, replaceWith)
      }
      val printfArgs = Seq(s""""$formatString"""") ++
                        (p.args map {arg => s"${emitExprWrap(arg)}.as_single_word()"})
      Seq(s"if (UNLIKELY(done_reset && update_registers && verbose && ${emitExprWrap(p.en)})) printf(${printfArgs mkString(", ")});")
    }
    case st: Stop => {
      Seq(s"if (UNLIKELY(${emitExpr(st.en)})) {assert_triggered = true; assert_exit_code = ${st.ret};}")
    }
    case mw: MemWrite => {
      Seq(s"if (UNLIKELY(update_registers && ${emitExprWrap(mw.wrEn)} && ${emitExprWrap(mw.wrMask)})) ${mw.memName}[${emitExprWrap(mw.wrAddr)}.as_single_word()] = ${emitExpr(mw.wrData)};")
    }
    case ru: RegUpdate => Seq(s"if (update_registers) ${emitExpr(ru.regRef)} = ${emitExpr(ru.expr)};")
    case r: DefRegister => Seq()
    case w: DefWire => Seq()
    case m: DefMemory => Seq()
    case i: WDefInstance => Seq()
    case _ => throw new Exception(s"Don't yet support $s")
  }
}

case Addw => s"${emitExprWrap(p.args(0))}.addw(${emitExprWrap(p.args(1))})"


package essent.ir

import firrtl._
import firrtl.ir._

// ESSENT's additions to the IR for optimization

case class RegUpdate(info: Info, regRef: Expression, expr: Expression) extends Statement {
  def serialize: String =  s"${regRef.serialize} <= ${expr.serialize}" + info.serialize
  def mapStmt(f: Statement => Statement): Statement = this
  def mapExpr(f: Expression => Expression): Statement = this.copy(regRef = f(regRef), expr = f(expr))
  def mapType(f: Type => Type): Statement = this
  def mapString(f: String => String): Statement = this
  def mapInfo(f: Info => Info): Statement = this
  def foreachExpr(f: firrtl.ir.Expression => Unit): Unit = ???
  def foreachInfo(f: firrtl.ir.Info => Unit): Unit = ???
  def foreachStmt(f: firrtl.ir.Statement => Unit): Unit = ???
  def foreachString(f: String => Unit): Unit = ???
  def foreachType(f: firrtl.ir.Type => Unit): Unit = ???
}

case class MemWrite(memName: String,
                    portName: String,
                    wrEn: Expression,
                    wrMask: Expression,
                    wrAddr: Expression,
                    wrData: Expression) extends Statement {

  def mapStmt(f: Statement => Statement): Statement = this
  def mapExpr(f: Expression => Expression): Statement = {
    MemWrite(memName, portName, f(wrEn), f(wrMask), f(wrAddr), f(wrData))
  }
  def mapType(f: Type => Type): Statement = this
  def mapString(f: String => String): Statement = this
  def mapInfo(f: Info => Info): Statement = this
  def nodeName(): String = s"$memName.$portName"
  def foreachExpr(f: firrtl.ir.Expression => Unit): Unit = ???
  def foreachInfo(f: firrtl.ir.Info => Unit): Unit = ???
  def foreachStmt(f: firrtl.ir.Statement => Unit): Unit = ???
  def foreachString(f: String => Unit): Unit = ???
  def foreachType(f: firrtl.ir.Type => Unit): Unit = ???
}

case class CondMux(name: String, mux: Mux, tWay: Seq[Statement], fWay: Seq[Statement]) extends Statement {
  def serialize: String =  "conditional mux"
  def mapStmt(f: Statement => Statement): Statement = this.copy(tWay = tWay map f, fWay = fWay map f)
  def mapExpr(f: Expression => Expression): Statement = this
  def mapType(f: Type => Type): Statement = this
  def mapString(f: String => String): Statement = this
  def mapInfo(f: Info => Info): Statement = this
  def foreachExpr(f: firrtl.ir.Expression => Unit): Unit = ???
  def foreachInfo(f: firrtl.ir.Info => Unit): Unit = ???
  def foreachStmt(f: firrtl.ir.Statement => Unit): Unit = ???
  def foreachString(f: String => Unit): Unit = ???
  def foreachType(f: firrtl.ir.Type => Unit): Unit = ???
}

case class CondPart(
    id: Int,
    alwaysActive: Boolean,
    inputs: Seq[String],
    memberStmts: Seq[Statement],
    outputsToDeclare: Map[String,firrtl.ir.Type]) extends Statement {
  def serialize: String = s"CondPart #$id"
  def mapStmt(f: Statement => Statement): Statement = this.copy(memberStmts = memberStmts map f)
  def mapExpr(f: Expression => Expression): Statement = this
  def mapType(f: Type => Type): Statement = this
  def mapString(f: String => String): Statement = this
  def mapInfo(f: Info => Info): Statement = this
  def foreachExpr(f: firrtl.ir.Expression => Unit): Unit = ???
  def foreachInfo(f: firrtl.ir.Info => Unit): Unit = ???
  def foreachStmt(f: firrtl.ir.Statement => Unit): Unit = ???
  def foreachString(f: String => Unit): Unit = ???
  def foreachType(f: firrtl.ir.Type => Unit): Unit = ???
}

def serialize: String = s"if (${wrEn.serialize} && ${wrMask.serialize}) $memName[${wrAddr.serialize}] = ${wrData.serialize}"


package essent

import Graph._

import collection.mutable.ArrayBuffer

class MFFC(val g: Graph) {
  import MFFC.{Unclaimed,Excluded}

  // numeric vertex ID -> MFFC ID
  val mffc = ArrayBuffer.fill(g.numNodes())(Unclaimed)

  def overrideMFFCs(newAssignments: ArrayBuffer[NodeID]): Unit = {
    mffc.clear()
    mffc ++= newAssignments
  }

  def findMFFCs(): ArrayBuffer[NodeID] = {
    val unvisitedSinks = g.nodeRange() filter {
      id => mffc(id) == Unclaimed && g.outNeigh(id).isEmpty
    }
    val visited = g.nodeRange() filter { id => mffc(id) != Unclaimed }
    val fringe = (visited flatMap(g.inNeigh)).distinct
    val unvisitedFringe = fringe filter { mffc(_) == Unclaimed }
    val newMFFCseeds = unvisitedSinks.toSet ++ unvisitedFringe
    if (newMFFCseeds.isEmpty) {
      mffc
    } else {
      newMFFCseeds foreach { id => mffc(id) = id }
      maximizeFFCs(newMFFCseeds)
      findMFFCs()
    }
  }

  def maximizeFFCs(fringe: Set[NodeID]): Unit = {
    val fringeAncestors = fringe flatMap g.inNeigh filter { mffc(_) == Unclaimed }
    val newMembers = fringeAncestors flatMap { parent => {
      val childrenMFFCs = (g.outNeigh(parent) map mffc).distinct
      if ((childrenMFFCs.size == 1) && (childrenMFFCs.head != Unclaimed)) {
        mffc(parent) = childrenMFFCs.head
        Seq(parent)
      } else Seq()
    }}
    if (newMembers.nonEmpty)
      maximizeFFCs(newMembers)
  }
}

object MFFC {
  val Unclaimed = -1
  val Excluded  = -2

  def apply(g: Graph, excludeSet: Set[NodeID] = Set()): ArrayBuffer[NodeID] = {
    val worker = new MFFC(g)
    excludeSet foreach { id => worker.mffc(id) = Excluded }
    val mffc = worker.findMFFCs()


package essent

import firrtl.ir._

import essent.Emitter._
import essent.Extract._
import essent.ir._

import collection.mutable.{ArrayBuffer, BitSet, HashMap}
import scala.reflect.ClassTag

// Extends Graph to include more attributes per node
//  - Associates a name (String) and Statement with each node
//  - Name must be unique, since can find nodes by name too
//  - Nodes can have an EmptyStatement if no need to emit

class StatementGraph extends Graph {
  // Access companion object's type aliases without prefix
  // TODO: type alias for name type? Hard to imagine other than String?
  import Graph.{NodeID, AdjacencyList}

  // Internal data structures
  //----------------------------------------------------------------------------
  // Vertex name (string of destination variable) -> numeric ID
  val nameToID = HashMap[String,NodeID]()
  // Numeric vertex ID -> name (string destination variable)
  val idToName = ArrayBuffer[String]()
  // Numeric vertex ID -> firrtl statement (Block used for aggregates)
  val idToStmt = ArrayBuffer[Statement]()
  // Numeric vertex ID -> Boolean indicating whether node should be emitted
  val validNodes = BitSet()

  // Graph building
  //----------------------------------------------------------------------------
  def getID(vertexName: String) = {
    if (nameToID contains vertexName) nameToID(vertexName)
    else {
      val newID = nameToID.size
      nameToID(vertexName) = newID
      idToName += vertexName
      idToStmt += EmptyStmt
      growNeighsIfNeeded(newID)
      // TODO: is it better to trust Graph to grow for new ID? (current)
      newID
    }
  }

  def addEdge(sourceName: String, destName: String): Unit = {
    super.addEdge(getID(sourceName), getID(destName))
  }

  def addEdgeIfNew(sourceName: String, destName: String): Unit = {
    super.addEdgeIfNew(getID(sourceName), getID(destName))
  }

  def addStatementNode(resultName: String, depNames: Seq[String],
                       stmt: Statement = EmptyStmt): Unit = {
    val potentiallyNewDestID = getID(resultName)
    depNames foreach {depName : String => addEdge(depName, resultName)}
    if (potentiallyNewDestID >= idToStmt.size) {
      val numElemsToGrow = potentiallyNewDestID - idToStmt.size + 1
      idToStmt.appendAll(ArrayBuffer.fill(numElemsToGrow)(EmptyStmt))
    }
    idToStmt(potentiallyNewDestID) = stmt
    // Don't want to emit state element declarations
    if (!stmt.isInstanceOf[DefRegister] && !stmt.isInstanceOf[DefMemory])
      validNodes += potentiallyNewDestID
  }

  def buildFromBodies(bodies: Seq[Statement]): Unit = {
    val bodyHE = bodies flatMap {
      case b: Block => b.stmts flatMap findDependencesStmt
      case s => findDependencesStmt(s)
    }
    bodyHE foreach { he => addStatementNode(he.name, he.deps, he.stmt) }
  }

  // Traversal / Queries / Extraction
  //----------------------------------------------------------------------------
  def collectValidStmts(ids: Seq[NodeID]): Seq[Statement] = ids filter validNodes map idToStmt

  def containsStmtOfType[T <: Statement]()(implicit tag: ClassTag[T]): Boolean = {
    (idToStmt collectFirst { case s: T => s }).isDefined
  }

  def findIDsOfStmtOfType[T <: Statement]()(implicit tag: ClassTag[T]): Seq[NodeID] = {
    (idToStmt.zipWithIndex collect { case (s: T , id: Int) => id }).toSeq
  }

  def allRegDefs(): Seq[DefRegister] = (idToStmt collect {
    case dr: DefRegister => dr
  }).toSeq

  def stateElemNames(): Seq[String] = (idToStmt collect {
    case dr: DefRegister => dr.name
    case dm: DefMemory => dm.name
  }).toSeq

  def stateElemIDs() = findIDsOfStmtOfType[DefRegister]() ++ findIDsOfStmtOfType[DefMemory]()

  def mergeIsAcyclic(nameA: String, nameB: String): Boolean = {
    val idA = nameToID(nameA)
    val idB = nameToID(nameB)
    super.mergeIsAcyclic(idA, idB)
  }

  def extractSourceIDs(e: Expression): Seq[NodeID] = findDependencesExpr(e) map nameToID

  // Mutation
  //----------------------------------------------------------------------------
  def addOrderingDepsForStateUpdates(): Unit = {
    def addOrderingEdges(writerID: NodeID, readerTargetID: NodeID): Unit = {
      outNeigh(readerTargetID) foreach {
        readerID => if (readerID != writerID) addEdgeIfNew(readerID, writerID)
      }
    }
    idToStmt.zipWithIndex foreach { case(stmt, id) => {
      val readerTargetName = stmt match {
        case ru: RegUpdate => Some(emitExpr(ru.regRef))
        case mw: MemWrite => Some(mw.memName)
        case _ => None
      }
      readerTargetName foreach {
        name => if (nameToID.contains(name)) addOrderingEdges(id, nameToID(name))
      }
    }}
  }

  def mergeStmtsMutably(mergeDest: NodeID, mergeSources: Seq[NodeID], mergeStmt: Statement): Unit = {
    val mergedID = mergeDest
    val idsToRemove = mergeSources
    idsToRemove foreach { id => idToStmt(id) = EmptyStmt }
    // NOTE: keeps mappings of name (idToName & nameToID) for debugging dead nodes
    mergeNodesMutably(mergeDest, mergeSources)
    idToStmt(mergeDest) = mergeStmt
    validNodes(mergeDest) = (mergeSources :+ mergeDest) exists { validNodes }
    validNodes --= idsToRemove
  }

  // Stats
  //----------------------------------------------------------------------------
  def numValidNodes() = validNodes.size

  def numNodeRefs() = idToName.size

  def makeStatsString() =
    s"Graph has ${numNodes()} nodes (${numValidNodes()} valid) and ${numEdges()} edges"
}

object StatementGraph {
  def apply(bodies: Seq[Statement]): StatementGraph = {
    val sg = new StatementGraph
    sg.buildFromBodies(bodies)
    sg.addOrderingDepsForStateUpdates()
    sg
  }

  def apply(circuit: Circuit, removeFlatConnects: Boolean = true): StatementGraph =
    apply(flattenWholeDesign(circuit, removeFlatConnects))
}

def stmtsOrdered(): Seq[Statement] = collectValidStmts(TopologicalSort(this).toSeq)


package disabled

import collection.mutable.HashMap

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import firrtl.passes._
import firrtl.PrimOps._
import firrtl.Utils._

// Replaces a tail following an add with an addw if:
// - result of add is used only in the one tail
// - tail only clips top bit
// - also does sub -> subw

object InferAddw extends Pass {
  def desc = "Replace tail on add results with addw (or sub with subw)"

  val WRefCounts = HashMap[String, Int]()

  // return names of all nodes with add expressions
  def findAddSigs(s: Statement): Seq[(String,DoPrim)] = s match {
    case b: Block => b.stmts flatMap findAddSigs
    case DefNode(_, name: String, primE: DoPrim) => primE.op match {
      case Add => Seq((name, primE copy (tpe = oneNarrower(primE.tpe))))
      // case Sub => Seq((name, primE copy (tpe = oneNarrower(primE.tpe))))
      case _ => Seq()
    }
    case _ => Seq()
  }

  def countWRefsExpr(addSigs: Set[String])(e: Expression): Expression = {
    e match {
      case WRef(name, _, _, _) => {
        if (addSigs.contains(name)) WRefCounts(name) = WRefCounts.getOrElse(name,0) + 1
      }
      case _ => e
    }
    e map countWRefsExpr(addSigs)
  }

  def countWRefsStmt(addSigs: Set[String])(s: Statement): Statement = {
    s map countWRefsStmt(addSigs) map countWRefsExpr(addSigs)
  }

  // return names of all nodes that tail an add expressions
  def findDependentTails(addSigs: Set[String])(s: Statement): Seq[(String,String)] = {
    s match {
      case b: Block => b.stmts flatMap findDependentTails(addSigs)
      case DefNode(_, tName: String, primE: DoPrim) => primE.op match {
        case Tail => {

  }

  def oneNarrower(tpe: Type): Type = tpe match {
    case UIntType(w) => UIntType(w - IntWidth(1))
    case SIntType(w) => SIntType(w - IntWidth(1))
    case _ => throw new Exception("can't narrow a non-ground type!")
  }

  def convertExpToW(primE: DoPrim): DoPrim = {
    val newOp = primE.op match {
      case Add => Addw
      // case Sub => Subw
    }
    primE copy (op = newOp)
  }

  def replaceInferAddwStmt(tailMap: Map[String,DoPrim], addsToRemove: Set[String])(s: Statement): Statement = {
    val replaced = s match {
      case DefNode(info, nodeName: String, primE: DoPrim) => primE.op match {
        case Tail => {
          if (tailMap.contains(nodeName))
            DefNode(info, nodeName, convertExpToW(tailMap(nodeName)))
          else s
        }
        case Add => {
          if (addsToRemove.contains(nodeName)) EmptyStmt
          else s
        }
        // case Sub => {
        //   if (addsToRemove.contains(nodeName)) EmptyStmt
        //   else s
        // }
        case _ => s
      }
      case _ => s
    }
    replaced map replaceInferAddwStmt(tailMap, addsToRemove)
  }

  def InferAddwModule(m: Module): Module = {
    val addMap = findAddSigs(m.body).toMap
    WRefCounts.clear()
    m.body map countWRefsStmt(addMap.keySet)
    val singleUseAddSigs = addMap.filter { case (k, v) => WRefCounts(k) == 1 }
    val tailPairs = findDependentTails(singleUseAddSigs.keySet)(m.body)
    val tailMap = (tailPairs map { case (tailSig: String, addSig: String) =>
      (tailSig, addMap(addSig)) }).toMap
    val addsToRemove = (tailPairs map { _._2 }).toSet
    val newBody = m.body map replaceInferAddwStmt(tailMap, addsToRemove)
    Module(m.info, m.name, m.ports, squashEmpty(newBody))
  }

  def run(c: Circuit): Circuit = {
    val modulesx = c.modules.map {
      case m: ExtModule => m
      case m: Module => InferAddwModule(m)
    }
    Circuit(c.info, modulesx, c.main)
  }
}

          val eName = primE.args.head match { case w: WRef => w.name }
          if ((addSigs.contains(eName)) && (primE.consts.head == 1) && (bitWidth(primE.tpe) == 64))
            Seq((tName, eName))
          else Seq()
        }


package disabled

import essent.Emitter._

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import firrtl.passes._
import firrtl.Utils._

import scala.math.BigInt
import scala.util.Random

object RandInitInvalids extends Pass {
  def desc = "Randomly initialize invalid signals"

  def genRandomLiteral(width: Int): BigInt = BigInt(width, Random)

  def randInitStmt(portNames: Set[String])(s: Statement): Statement = {
    val replaced = s match {
      case i: IsInvalid => {
        val randLit = i.expr.tpe match {
          case UIntType(IntWidth(w)) => UIntLiteral(genRandomLiteral(w.toInt), IntWidth(w))
          case SIntType(IntWidth(w)) => SIntLiteral(genRandomLiteral(w.toInt), IntWidth(w))
        }
        if (portNames.contains(emitExpr(i.expr)))
          Connect(i.info, i.expr, randLit)
        else
          DefNode(i.info, emitExpr(i.expr), randLit)
      }
      case _ => s
    }
    replaced map randInitStmt(portNames)
  }

  def randInitModule(m: Module): Module = {

    Module(m.info, m.name, m.ports, squashEmpty(randInitStmt(portNames)(m.body)))
  }

  def run(c: Circuit): Circuit = {
    val modulesx = c.modules.map {
      case m: ExtModule => m
      case m: Module => randInitModule(m)
    }
    Circuit(c.info, modulesx, c.main)
  }
}

val portNames = (m.ports map { _.name }).toSet


package disabled

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import firrtl.passes._
import firrtl.PrimOps._

object ZeroFromBits extends Pass {
  def desc = "Replaces bit extracts on literal zeros with zeros"

  def simpBitsExpr(e: Expression): Expression = {
    val bigZero = BigInt(0)
    val replaced = e match {
      case p: DoPrim => p.op match {
        case Bits => p.args.head match {
          case UIntLiteral(bigZero, _) => UIntLiteral(bigZero, IntWidth(bitWidth(p.tpe)))
          case _ => p
        }
        case _ => p
      }
      case _ => e
    }
    replaced map simpBitsExpr
  }

  def simpBitsStmt(s: Statement): Statement = {
    s map simpBitsStmt map simpBitsExpr
  }

  def simpBitsModule(m: Module): Module = {
    Module(m.info, m.name, m.ports, simpBitsStmt(m.body))
  }

  def run(c: Circuit): Circuit = {
    val modulesx = c.modules.map {
      case m: ExtModule => m
      case m: Module => simpBitsModule(m)
    }


package essent.passes

import essent.Emitter._

import firrtl._
import firrtl.ir._
import firrtl.Mappers._
import firrtl.options.Dependency
import firrtl.passes._
import firrtl.Utils._

object RegFromMem1 extends Pass {
  def desc = "Replaces single-element mems with a register"

  override def prerequisites = Seq(Dependency(essent.passes.NoClockConnects))
  override def optionalPrerequisites = firrtl.stage.Forms.LowFormOptimized
  override def invalidates(a: Transform) = false

  def memHasRightParams(m: DefMemory) = {
    (m.depth == 1) && (m.writeLatency == 1) && (m.readLatency == 0) &&
      (m.readwriters.isEmpty) && (m.writers.size <= 1)
  }

  def grabMemConnects(s: Statement): Seq[(String, Expression)] = s match {
    case b: Block => b.stmts flatMap grabMemConnects
    case c: Connect if firrtl.Utils.kind(c.loc) == firrtl.MemKind => Seq(emitExpr(c.loc) -> c.expr)
    case _ => Seq()
  }

  def dropMemConnects(memsToReplace: Set[String])(s: Statement): Statement = {
    val noConnects = s match {
      case Connect(_,WSubField(WSubField(WRef(name: String,_,_,_),_,_,_),_,_,_),_) => {
        if (memsToReplace.contains(name)) EmptyStmt
        else s
      }
      case _ => s
    }
    noConnects map dropMemConnects(memsToReplace)
  }

  // replace mem def's and mem reads
  def replaceMemsStmt(memsToTypes: Map[String,Type])(s: Statement): Statement = {
    val memsReplaced = s match {
      case mem: DefMemory if (memsToTypes.contains(mem.name)) => {
        // FUTURE: what is clock for a mem? assuming it is based on surrounding context
        DefRegister(mem.info, mem.name, mem.dataType, WRef("clock", ClockType, PortKind),
                    UIntLiteral(0,IntWidth(1)), UIntLiteral(0,IntWidth(1)))
      }
      case _ => s
    }
    memsReplaced map replaceMemsStmt(memsToTypes) map replaceMemsExpr(memsToTypes)
  }

  def replaceMemsExpr(memsToTypes: Map[String,Type])(e: Expression): Expression = {
    val replaced = e match {
      case WSubField(WSubField(WRef(name: String, _, _, f: Flow),_,_,_),"data",_,_) => {
        if (memsToTypes.contains(name)) WRef(name, memsToTypes(name), firrtl.RegKind, f)
        else e
      }
      case _ => e
    }
    replaced map replaceMemsExpr(memsToTypes)
  }

  // insert reg write muxes
  def generateRegUpdates(memsToReplace: Seq[DefMemory], body: Statement): Seq[Statement] = {
    val memConnects = grabMemConnects(body).toMap

    memsWithWrites flatMap { mem => mem.writers map { writePortName => {
      val selfRef = WRef(mem.name, mem.dataType, firrtl.RegKind, firrtl.DuplexFlow)
      val enSig = memConnects(s"${mem.name}.$writePortName.en")
      val wrDataSig = memConnects(s"${mem.name}.$writePortName.data")
      val wrEnableMux = Mux(enSig, wrDataSig, selfRef, mem.dataType)
      Connect(NoInfo, selfRef, wrEnableMux)
    }}}
  }

  def memReplaceModule(m: Module): Module = {
    val allMems = essent.Extract.findInstancesOf[DefMemory](m.body)
    // FUTURE: need to explicitly handle read enables?
    val singleElementMems = allMems filter memHasRightParams
    if (singleElementMems.isEmpty) m
    else {
      val memNamesToReplace = (singleElementMems map { _.name }).toSet
      val memConnectsRemoved = squashEmpty(dropMemConnects(memNamesToReplace)(m.body))
      val memsToTypes = (singleElementMems map { mem => (mem.name, mem.dataType)}).toMap
      val memsReplaced = replaceMemsStmt(memsToTypes)(memConnectsRemoved)
      val regUpdateStmts = generateRegUpdates(singleElementMems, m.body)
      val newBlock = Block(Seq(memsReplaced) ++ regUpdateStmts)
      m.copy(body = newBlock)
    }
  }

  def run(c: Circuit): Circuit = {
    val modulesx = c.modules.map {
      case m: ExtModule => m
      case m: Module => memReplaceModule(m)
    }
    Circuit(c.info, modulesx, c.main)
  }
}

val memsWithWrites = memsToReplace filter { _.writers.nonEmpty }


package essent
import org.scalatest.flatspec.AnyFlatSpec
import essent.passes.ReplaceRsvdKeywords
import firrtl.CircuitState
import firrtl.options.Dependency

import java.io.{File, FileWriter}
import scala.io.Source

class ReplaceRsvdKeyTest extends AnyFlatSpec{
   "Mypass" should "Replace all reserve keyword" in {
     val sourceReader = Source.fromURL(getClass.getResource("/ReplacedRsvdKey.fir"))
     val circuit = firrtl.Parser.parse(sourceReader.getLines(), firrtl.Parser.IgnoreInfo)
     sourceReader.close()
     val deps = firrtl.stage.Forms.LowFormOptimized ++ Seq(Dependency(ReplaceRsvdKeywords))
     val firrtlCompiler = new firrtl.stage.transforms.Compiler(deps)

     val resultState = firrtlCompiler.execute(CircuitState(circuit, Seq()))
     val CorrectReader = Source.fromURL(getClass.getResource("/ReplacedRsvdKey_correct.fir"))
     val correctString = CorrectReader.getLines().mkString("\n")
     assert(correctString == resultState.circuit.serialize)


package essent

import firrtl.ir._

import org.scalatest.flatspec.AnyFlatSpec

class StatementGraphSpec extends AnyFlatSpec {
  "A NamedGraph" should "grow as necessary for new edges" in {
    val sg = new StatementGraph
    sg.addEdge("alpha", "beta")
    assertResult(0) { sg.numValidNodes() }
    assertResult(2) { sg.numNodeRefs() }
    assertResult(1) { sg.numEdges() }
    sg.addEdge("gamma", "zeta")
    assertResult(0) { sg.numValidNodes() }
    assertResult(4) { sg.numNodeRefs() }
    assertResult(2) { sg.numEdges() }
  }

  it should "not add duplicate edges (if requested)" in {
    val sg = new StatementGraph
    sg.addEdgeIfNew("alpha", "beta")
    assertResult(0) { sg.numValidNodes() }

    assertResult(1) { sg.numEdges() }
    sg.addEdgeIfNew("alpha", "beta")
    assertResult(0) { sg.numValidNodes() }
<TARGET>
    assertResult(1) { sg.numEdges() }
  }

  it should "be buildable from hyperedges" in {
    val sg = new StatementGraph
    sg.addStatementNode("child", Seq("parent0","parent1"))
    assertResult(1) { sg.numValidNodes() }
    assertResult(3) { sg.numNodeRefs() }
    assertResult(2) { sg.numEdges() }
    assert(sg.idToStmt(sg.nameToID("child")) == EmptyStmt)
    sg.addStatementNode("sibling", Seq("parent0","parent1"), Block(Seq()))
    assertResult(2) { sg.numValidNodes() }
    assertResult(4) { sg.numNodeRefs() }
    assertResult(4) { sg.numEdges() }
    assert(sg.idToStmt(sg.nameToID("sibling")) == Block(Seq()))
  }

  it should "fill in idToStmts as necessary with EmptyStmt" in {
    val sg = new StatementGraph
    sg.addEdge("alpha", "beta")
    assertResult(EmptyStmt)(sg.idToStmt(sg.nameToID("alpha")))
    sg.addStatementNode("sibling", Seq("parent0","parent1"), Block(Seq()))
    assertResult(EmptyStmt)(sg.idToStmt(sg.nameToID("parent1")))
  }

  it should "test safety or merges (whether acyclic)" in {
    val sg = new StatementGraph
    sg.addEdge("a","b")
    sg.addEdge("b","c")
    assert( sg.mergeIsAcyclic("a","b"))
    assert( sg.mergeIsAcyclic("b","a"))
    assert(!sg.mergeIsAcyclic("a","c"))
    assert(!sg.mergeIsAcyclic("c","a"))
    sg.addEdge("a","c")
    assert(!sg.mergeIsAcyclic("a","c"))
    sg.addEdge("x","y")
    assert( sg.mergeIsAcyclic("a","x"))
    assert( sg.mergeIsAcyclic("y","a"))
  }

  it should "be able to merge Statements" in {
    val sg = new StatementGraph
    sg.addStatementNode("b", Seq("a"), Attach(NoInfo,Seq()))
    sg.addStatementNode("c", Seq("b","a"), Attach(NoInfo,Seq()))
    sg.mergeStmtsMutably(sg.nameToID("b"), Seq(sg.nameToID("c")), Block(Seq()))
    assertResult(Block(Seq()))(sg.idToStmt(sg.nameToID("b")))
    assertResult(EmptyStmt)(sg.idToStmt(sg.nameToID("c")))
    assert( sg.validNodes(sg.nameToID("b")))
    assert(!sg.validNodes(sg.nameToID("c")))
    // trusts Graph's test for mergeNodesMutably
  }

  it should "be able to detect if Statement type anywhere in graph" in {
    val sg = new StatementGraph
    sg.addStatementNode("child", Seq("parent0","parent1"), Block(Seq()))
    assert( sg.containsStmtOfType[Block]())
    assert(!sg.containsStmtOfType[DefWire]())
  }

  it should "find IDs of requested Statement type" in {
    val sg = new StatementGraph
    sg.addStatementNode("child", Seq("parent0","parent1"), Block(Seq()))
    assertResult(Seq(sg.nameToID("child"))) { sg.findIDsOfStmtOfType[Block]() }
    assertResult(Seq()) { sg.findIDsOfStmtOfType[DefWire]() }
  }

  it should "collect valid statements from subset" in {
    val sg = new StatementGraph
    sg.addStatementNode("b", Seq("a"), Attach(NoInfo,Seq()))
    sg.addStatementNode("c", Seq("b","a"), Block(Seq()))
    val goal = Seq(Attach(NoInfo,Seq()), Block(Seq()))
    val result = sg.collectValidStmts(Seq(sg.nameToID("b"), sg.nameToID("c")))
    assertResult(goal.toSet)(result.toSet)
  }

  it should "be able to handle a 1 node graph with no edges" in {
    val stmt = DefNode(NoInfo,"dummy",UIntLiteral(0,IntWidth(1)))
    val sg = StatementGraph(Seq(stmt))
    assertResult(Seq(stmt)) { sg.stmtsOrdered() }
  }
}


package treecorel2

import chisel3._
import chisel3.util._

trait IOConfig {
  val XLen       = 64
  val InstLen    = 32
  val RegfileLen = 5
  val RegfileNum = 1 << RegfileLen
  val ISALen     = 6
  // mem
  val MaskLen = 8
  val LDSize  = 3
  // branch prediction
  val GHRLen    = 5
  val PHTSize   = 1 << GHRLen
  val BTBIdxLen = 5
  val BTBPcLen  = XLen - BTBIdxLen
  val BTBTgtLen = XLen
  val BTBSize   = 1 << BTBIdxLen
}

trait InstConfig extends IOConfig {
  val SoCEna   = false
  val CacheEna = false

  val FlashStartAddr    = "h0000000030000000".U(XLen.W)
  val SimStartAddr      = "h0000000080000000".U(XLen.W)
  val DiffStartBaseAddr = "h0000000080000000".U(XLen.W)
  val SoCStartBaseAddr  = "h0000000000000000".U(XLen.W)
  val DifftestAddrMask  = "hfffffffffffffff8".U(XLen.W)
  val SoCAddrMask       = "hffffffffffffffff".U(XLen.W)
  val InstSoCRSize      = 2.U
  val InstDiffRSize     = 3.U
  val DiffRWSize        = 3.U

  val NOPInst = 0x13.U
  // inst type
  // nop is equal to [addi x0, x0, 0], so the oper is same as 'addi' inst
  val InstTypeLen = 3
  val nopInstType = 2.U(InstTypeLen.W)
  val rInstType   = 1.U(InstTypeLen.W)
  val iInstType   = 2.U(InstTypeLen.W)
  val sInstType   = 3.U(InstTypeLen.W)
  val bInstType   = 4.U(InstTypeLen.W)

  val jInstType   = 6.U(InstTypeLen.W)
  val wtRegTrue   = true.B
  val wtRegFalse  = false.B
  // inst
  val InstValLen  = 6
  val instADDI    = 0.U(InstValLen.W)
  val instADDIW   = 1.U(InstValLen.W)
  val instSLTI    = 2.U(InstValLen.W)
  val instSLTIU   = 3.U(InstValLen.W)
  val instANDI    = 4.U(InstValLen.W)
  val instORI     = 5.U(InstValLen.W)
  val instXORI    = 6.U(InstValLen.W)
  val instSLLI    = 7.U(InstValLen.W)
  val instSLLIW   = 8.U(InstValLen.W)
  val instSRLI    = 9.U(InstValLen.W)
  val instSRLIW   = 10.U(InstValLen.W)
  val instSRAI    = 11.U(InstValLen.W)
  val instSRAIW   = 12.U(InstValLen.W)
  val instLUI     = 13.U(InstValLen.W)
  val instAUIPC   = 14.U(InstValLen.W)
  val instADD     = 15.U(InstValLen.W)
  val instADDW    = 16.U(InstValLen.W)
  val instSLT     = 17.U(InstValLen.W)
  val instSLTU    = 18.U(InstValLen.W)
  val instAND     = 19.U(InstValLen.W)
  val instOR      = 20.U(InstValLen.W)
  val instXOR     = 21.U(InstValLen.W)
  val instSLL     = 22.U(InstValLen.W)
  val instSLLW    = 23.U(InstValLen.W)
  val instSRL     = 24.U(InstValLen.W)
  val instSRLW    = 25.U(InstValLen.W)
  val instSUB     = 26.U(InstValLen.W)
  val instSUBW    = 27.U(InstValLen.W)
  val instSRA     = 28.U(InstValLen.W)
  val instSRAW    = 29.U(InstValLen.W)
  val instNOP     = 30.U(InstValLen.W)
  val instJAL     = 31.U(InstValLen.W)
  val instJALR    = 32.U(InstValLen.W)
  val instBEQ     = 33.U(InstValLen.W)
  val instBNE     = 34.U(InstValLen.W)
  val instBLT     = 35.U(InstValLen.W)
  val instBLTU    = 36.U(InstValLen.W)
  val instBGE     = 37.U(InstValLen.W)
  val instBGEU    = 38.U(InstValLen.W)
  val instLB      = 39.U(InstValLen.W)
  val instLBU     = 40.U(InstValLen.W)
  val instLH      = 41.U(InstValLen.W)
  val instLHU     = 42.U(InstValLen.W)
  val instLW      = 43.U(InstValLen.W)
  val instLWU     = 44.U(InstValLen.W)
  val instLD      = 45.U(InstValLen.W)
  val instSB      = 46.U(InstValLen.W)
  val instSH      = 47.U(InstValLen.W)
  val instSW      = 48.U(InstValLen.W)
  val instSD      = 49.U(InstValLen.W)
  val instCSRRW   = 50.U(InstValLen.W)
  val instCSRRS   = 51.U(InstValLen.W)
  val instCSRRC   = 52.U(InstValLen.W)
  val instCSRRWI  = 53.U(InstValLen.W)
  val instCSRRSI  = 54.U(InstValLen.W)
  val instCSRRCI  = 55.U(InstValLen.W)
  val instECALL   = 56.U(InstValLen.W)
  val instMRET    = 57.U(InstValLen.W)
  val instFENCE   = 58.U(InstValLen.W)
  val instFENCE_I = 59.U(InstValLen.W)
  val instCUST    = 60.U(InstValLen.W)

  // cache
  val NWay          = 4
  val NBank         = 4
  val NSet          = 32
  val CacheLineSize = XLen * NBank
  val ICacheSize    = NWay * NSet * CacheLineSize
  val DCacheSize    = NWay * NSet * CacheLineSize

  // clint
  val ClintBaseAddr  = 0x02000000.U(XLen.W)
  val ClintBoundAddr = 0x0200bfff.U(XLen.W)
  val MSipOffset     = 0x0.U(XLen.W)
  val MTimeOffset    = 0xbff8.U(XLen.W)
  val MTimeCmpOffset = 0x4000.U(XLen.W)
  // csr addr
  val CSRAddrLen   = 12
  val mhartidAddr  = 0xf14.U(CSRAddrLen.W)
  val mstatusAddr  = 0x300.U(CSRAddrLen.W)
  val mieAddr      = 0x304.U(CSRAddrLen.W)
  val mtvecAddr    = 0x305.U(CSRAddrLen.W)
  val mscratchAddr = 0x340.U(CSRAddrLen.W)
  val mepcAddr     = 0x341.U(CSRAddrLen.W)
  val mcauseAddr   = 0x342.U(CSRAddrLen.W)
  val mipAddr      = 0x344.U(CSRAddrLen.W)
  val mcycleAddr   = 0xb00.U(CSRAddrLen.W)
  val medelegAddr  = 0x302.U(CSRAddrLen.W)
  val timeCause    = "h8000_0000_0000_0007".U(XLen.W)
  val ecallCause   = "h0000_0000_0000_000b".U(XLen.W)

  // special inst
  val customInst = "h0000007b".U(InstLen.W)
  val haltInst   = "h0000006b".U(InstLen.W)
}


package treecorel2

import chisel3._
import chisel3.util._

class ALU extends Module with InstConfig {
  val io = IO(new Bundle {
    val isa  = Input(UInt(InstValLen.W))
    val src1 = Input(UInt(XLen.W))
    val src2 = Input(UInt(XLen.W))
    val imm  = Input(UInt(XLen.W))
    val res  = Output(UInt(XLen.W))
  })

  io.res := MuxLookup(
    io.isa,
    0.U(XLen.W),
    Seq(
      instADDI  -> (io.src1 + io.imm),
      instADD   -> (io.src1 + io.src2),
      instLUI   -> (io.imm),
      instSUB   -> (io.src1 - io.src2),
      instADDIW -> SignExt((io.src1 + io.imm)(31, 0), 64),
      instADDW  -> SignExt((io.src1 + io.src2)(31, 0), 64),
      instSUBW  -> SignExt((io.src1 - io.src2)(31, 0), 64),
      instANDI  -> (io.src1 & io.imm),
      instAND   -> (io.src1 & io.src2),
      instORI   -> (io.src1 | io.imm),
      instOR    -> (io.src1 | io.src2),
      instXORI  -> (io.src1 ^ io.imm),
      instXOR   -> (io.src1 ^ io.src2),
      instSLT   -> Mux(io.src1.asSInt < io.src2.asSInt, 1.U(XLen.W), 0.U(XLen.W)),
      instSLTI  -> Mux(io.src1.asSInt < io.imm.asSInt, 1.U(XLen.W), 0.U(XLen.W)),

      instSRAI  -> (io.src1.asSInt >> io.imm(5, 0)).asUInt,
      instSLLW  -> SignExt((io.src1 << io.src2(4, 0))(31, 0), 64),
      instSRLW  -> SignExt((io.src1(31, 0) >> io.src2(4, 0)), 64),
      instSRAW  -> SignExt((io.src1(31, 0).asSInt >> io.src2(4, 0)).asUInt, 64),
      instSLLIW -> SignExt((io.src1 << io.imm(4, 0))(31, 0), 64),
      instSRLIW -> SignExt((io.src1(31, 0) >> io.imm(4, 0)), 64),
      instSRAIW -> SignExt((io.src1(31, 0).asSInt >> io.imm(4, 0)).asUInt, 64)
    )
  )
}

      instSLTU  -> Mux(io.src1.asUInt < io.src2.asUInt, 1.U(XLen.W), 0.U(XLen.W)),
      instSLTIU -> Mux(io.src1.asUInt < io.imm.asUInt, 1.U(XLen.W), 0.U(XLen.W)),
      instSLL   -> (io.src1 << io.src2(5, 0))(63, 0),
      instSRL   -> (io.src1 >> io.src2(5, 0)),