package icenet

import chisel3._
import chisel3.util._
import freechips.rocketchip.unittest.UnitTest
import freechips.rocketchip.util.TwoWayCounter
import scala.util.Random
import IceNetConsts._

class BufferBRAM[T <: Data](n: Int, typ: T) extends Module {
  val addrBits = log2Ceil(n)
  val io = IO(new Bundle {
    // The value in data becomes valid one cycle after enable is asserted.
    // The value is held so long as enable is false
    val read = new Bundle {
      val en = Input(Bool())
      val addr = Input(UInt(addrBits.W))
      val data = Output(typ)
    }
    val write = new Bundle {
      val en = Input(Bool())
      val addr = Input(UInt(addrBits.W))
      val data = Input(typ)
    }
  })

  val ram = Mem(n, typ)

  val wen = RegNext(io.write.en, false.B)
  val waddr = RegNext(io.write.addr)
  val wdata = RegNext(io.write.data)

  when (wen) { ram.write(waddr, wdata) }

  val rread_data = RegEnable(ram.read(io.read.addr), io.read.en)
  val rbypass = RegEnable(io.read.addr === waddr && wen, io.read.en)
  val rbypass_data = RegEnable(wdata, io.read.en)

  io.read.data := Mux(rbypass, rbypass_data, rread_data)
}

/**
 * Buffers incoming packets without backpressuring
 * Drops at packet boundaries if buffer fills us.
 * @bufWords Number of flits held by the buffer
 * @maxBytes Maximum number of bytes in a packet
 * @headerBytes Number of bytes in a header
 * @headerType Bundle type for header
 * @wordBytes Number of bytes in a flit
 * @dropChecks Sequence of functions that can trigger a drop by asserting Bool
 * @dropless If true, dropped packet causes assertion failure
 */
class NetworkPacketBuffer[T <: Data](
    bufWords: Int,
    maxBytes: Int = ETH_JUMBO_MAX_BYTES,
    headerBytes: Int = ETH_HEAD_BYTES,
    headerType: T = new EthernetHeader,
    wordBytes: Int = NET_IF_WIDTH / 8,
    dropChecks: Seq[(T, StreamChannel, Bool) => Bool] = Nil,
    dropless: Boolean = false) extends Module {

  val maxWords = (maxBytes - 1) / wordBytes + 1
  val headerWords = headerBytes / wordBytes
  val wordBits = wordBytes * 8
  val nPackets = (bufWords - 1) / (headerWords + 1) + 1

  require(headerBytes % wordBytes == 0)

  val idxBits = log2Ceil(bufWords)
  val phaseBits = log2Ceil(nPackets + 1)
  val lenBits = log2Ceil(maxBytes + 1)
  val byteOffset = log2Ceil(wordBytes)

  val io = IO(new Bundle {
    val stream = new StreamIO(wordBytes * 8)
    val length = Valid(UInt(lenBits.W))
    val count = Output(UInt(log2Ceil(nPackets+1).W))
    val free = Output(UInt(8.W))
  })

  def discontinuous(bits: UInt, w: Int): Bool =
    (1 until w).map(i => bits(i) && !bits(i-1)).reduce(_ || _)

  assert(!io.stream.in.valid ||
    Mux(io.stream.in.bits.last,
      !discontinuous(io.stream.in.bits.keep, wordBytes),
      io.stream.in.bits.keep.andR),
    "NetworkPacketBuffer does not handle missing data")

  val dataBuffer = Module(new BufferBRAM(bufWords, Bits(wordBits.W)))
  val headerVec = Reg(Vec(headerWords, Bits(wordBits.W)))
  val header = headerVec.asTypeOf(headerType)

  val lenBuffer = Module(new BufferBRAM(nPackets, UInt(lenBits.W)))

  val bufHead = RegInit(0.U(idxBits.W))
  val bufTail = RegInit(0.U(idxBits.W))
  val startHead = Reg(UInt(idxBits.W))
  val revertHead = WireInit(false.B)

  val maybeFull = RegInit(false.B)
  val ptrMatch = bufHead === bufTail
  val bufFull = maybeFull && ptrMatch
  val bufEmpty = !maybeFull && ptrMatch

  val inIdx = RegInit(0.U(idxBits.W))
  val inDrop = RegInit(false.B)
  val inLen = RegInit(0.U(lenBits.W))
  val outIdx = RegInit(0.U(idxBits.W))

  val inPhase = RegInit(0.U(phaseBits.W))
  val outPhase = RegInit(0.U(phaseBits.W))
  val pktCount = RegInit(0.U(phaseBits.W))
  val nextLen = inLen + PopCount(io.stream.in.bits.keep)
  val hasPackets = pktCount > 0.U

  assert(pktCount <= nPackets.U, "Accepted more packets than possible")

  def bytesToWords(nbytes: UInt): UInt =
    nbytes(lenBits-1, byteOffset) + nbytes(byteOffset-1, 0).orR

  def finalKeep(nbytes: UInt): UInt = {
    val remBytes = nbytes(byteOffset-1, 0)
    val finalBytes = Mux(remBytes.orR, remBytes, wordBytes.U)
    (1.U << finalBytes) - 1.U
  }

  val length = lenBuffer.io.read.data
  val numWords = bytesToWords(length)
  // Need one cycle to read out packet length
  val lengthKnown = RegInit(false.B)
  val readLen = !lengthKnown && hasPackets

  val outLast = outIdx === (numWords - 1.U)
  val outValidReg = RegInit(false.B)

  val ren = (io.stream.out.ready || !outValidReg) && lengthKnown && !bufEmpty
  val wen = WireInit(false.B)
  val hwen = wen && inIdx < headerWords.U

  val setLength = WireInit(false.B)
  val clearLength = ren && outLast

  pktCount := pktCount + setLength - clearLength

  val outLastReg = RegEnable(outLast, ren)
  val outIdxReg = RegEnable(outIdx, ren)

  io.stream.out.valid := outValidReg
  io.stream.out.bits.data := dataBuffer.io.read.data
  io.stream.out.bits.last := outLastReg
  io.stream.out.bits.keep := Mux(outLastReg, finalKeep(length), ~0.U(wordBytes.W))
  io.stream.in.ready := true.B
  io.length.bits := lenBuffer.io.read.data
  io.length.valid := lengthKnown
  io.count := pktCount

  def wrapInc(x: UInt, n: Int) = Mux(x === (n - 1).U, 0.U, x + 1.U)

  dataBuffer.io.read.en := ren
  dataBuffer.io.read.addr := bufTail
  dataBuffer.io.write.en := wen
  dataBuffer.io.write.addr := bufHead
  dataBuffer.io.write.data := io.stream.in.bits.data

  lenBuffer.io.read.en := readLen
  lenBuffer.io.read.addr := outPhase
  lenBuffer.io.write.en := setLength
  lenBuffer.io.write.addr := inPhase
  lenBuffer.io.write.data := nextLen

  val headerValid = inIdx >= headerWords.U
  val customDrop = dropChecks.map(check => check(
                                    header,
                                    io.stream.in.bits,
                                    io.stream.in.fire && headerValid))
                             .foldLeft(false.B)(_ || _)
  val startDropping =
    (inPhase === outPhase && hasPackets) ||
    (inIdx === maxWords.U) || bufFull

  val capDrop = startDropping || inDrop
  val nonCapDrop = !headerValid || customDrop
  val anyDrop = capDrop || nonCapDrop
  val dropLastFire = anyDrop && io.stream.in.fire && io.stream.in.bits.last

  // io.free indicates the number of flits being freed up on a cycle
  io.free := ren + Mux(dropLastFire, inIdx + 1.U, 0.U)

  when (io.stream.out.fire) { outValidReg := false.B }

  when (readLen) { lengthKnown := true.B }

  when (ren) {
    outValidReg := true.B
    bufTail := wrapInc(bufTail, bufWords)
    outIdx := outIdx + 1.U

    when (outLast) {
      outIdx := 0.U
      outPhase := wrapInc(outPhase, nPackets)
      lengthKnown := false.B
    }
  }

  when (hwen) { headerVec(inIdx) := io.stream.in.bits.data }

  when (wen) { bufHead := wrapInc(bufHead, bufWords) }

  when (ren =/= wen) { maybeFull := wen }

  when (io.stream.in.fire) {

    when (startDropping) { inDrop := true.B }
    wen := !startDropping && !inDrop
    when (inIdx =/= (bufWords - 1).U) { inIdx := inIdx + 1.U }

    inLen := nextLen

    when (io.stream.in.bits.last) {
      val nextPhase = wrapInc(inPhase, nPackets)
      // Drop packets if there aren't more than headerBytes amount of data
      when (!anyDrop) {
        setLength := true.B
        inPhase := nextPhase
      } .otherwise {
        wen := false.B
        revertHead := inIdx =/= 0.U
        if (dropless) {
          assert(!capDrop,
            "Packet dropped by buffer due to insufficient capacity")
        } else {
          printf("WARNING: dropped packet with %d flits\n", inIdx + 1.U)
        }
      }
      inIdx := 0.U
      inLen := 0.U
      inDrop := false.B
    }
  }

  when (revertHead) {
    bufHead := startHead
    maybeFull := false.B
  }
}

class NetworkPacketBufferTest extends UnitTest(100000) {
  val buffer = Module(new NetworkPacketBuffer(
    bufWords = 12,
    maxBytes = 32,
    headerBytes = 8,
    headerType = UInt(64.W),
    wordBytes = 4))

  val inPackets = Seq(
    (10, false), // drop because too long
    (4,  true),
    (1,  false), // drop because too short
    (5,  false),
    (7,  true),
    (8,  false),
    (6,  false), // drop because buffer full
    (4,  true),
    (3,  false),
    (5,  true),
    (6,  false),
    (4,  true),
    (5,  true))

  val outPackets = Seq(
    (4, true),
    (5, false),
    (7, true),
    (8, false),
    (4, true),
    (3, false),
    (5, true),
    (6, true),
    (4, false),
    (5, true))

  val phaseBits = log2Ceil(inPackets.length)
  val idxBits = 4

  val inPacketLengths = VecInit(inPackets.map {
    case (x, _) => (x - 1).U(idxBits.W)
  })
  val inPacketSwitch = VecInit(inPackets.map(_._2.B))

  val outPacketLengths = VecInit(outPackets.map {
    case (x, _) => (x - 1).U(idxBits.W)
  })
  val outPacketSwitch = VecInit(outPackets.map(_._2.B))

  val inPhase = RegInit(0.U(phaseBits.W))
  val outPhase = RegInit(0.U(phaseBits.W))

  val inIdx = RegInit(0.U(idxBits.W))
  val outIdx = RegInit(0.U(idxBits.W))

  val s_start :: s_input :: s_output :: s_done :: Nil = Enum(4)
  val state = RegInit(s_start)

  buffer.io.stream.in.valid := state === s_input
  buffer.io.stream.in.bits.data := inIdx
  buffer.io.stream.in.bits.keep := ~0.U(32.W)
  buffer.io.stream.in.bits.last := inIdx === inPacketLengths(inPhase)
  buffer.io.stream.out.ready := state === s_output

  when (io.start && state === s_start) {
    state := s_input
  }

  when (buffer.io.stream.in.fire) {
    inIdx := inIdx + 1.U
    when (buffer.io.stream.in.bits.last) {
      inIdx := 0.U
      inPhase := inPhase + 1.U
      when (inPacketSwitch(inPhase)) { state := s_output }
    }
  }

  when (buffer.io.stream.out.fire) {
    outIdx := outIdx + 1.U
    when (buffer.io.stream.out.bits.last) {
      outIdx := 0.U
      outPhase := outPhase + 1.U
      when (outPacketSwitch(outPhase)) { state := s_input }
      when (outPhase === (outPackets.length - 1).U) { state := s_done }
    }
  }

  assert(!buffer.io.stream.out.valid || buffer.io.stream.out.bits.data === outIdx,
    "NetworkPacketBufferTest: got wrong output data")
  assert(!buffer.io.stream.out.valid || !buffer.io.stream.out.bits.last ||
    outIdx === outPacketLengths(outPhase),
    "NetworkPacketBufferTest: got output packet with wrong length")

  io.finished := state === s_done
}

class ReservationBufferAlloc(nXacts: Int, nWords: Int) extends Bundle {
  private val xactIdBits = log2Ceil(nXacts)
  private val countBits = log2Ceil(nWords + 1)

  val id = UInt(xactIdBits.W)
  val count = UInt(countBits.W)

}

class ReservationBufferData(nXacts: Int, dataBits: Int) extends Bundle {
  private val xactIdBits = log2Ceil(nXacts)

  val id = UInt(xactIdBits.W)
  val data = new StreamChannel(dataBits)

}

class ReservationBuffer(nXacts: Int, nWords: Int, dataBits: Int) extends Module {
  private val xactIdBits = log2Ceil(nXacts)
  private val countBits = log2Ceil(nWords + 1)

  require(nXacts <= nWords)

  val io = IO(new Bundle {
    val alloc = Flipped(Decoupled(new ReservationBufferAlloc(nXacts, nWords)))
    val in = Flipped(Decoupled(new ReservationBufferData(nXacts, dataBits)))
    val out = Decoupled(new StreamChannel(dataBits))
  })

  def incWrap(cur: UInt, inc: UInt): UInt = {
    val unwrapped = cur +& inc
    Mux(unwrapped >= nWords.U, unwrapped - nWords.U, unwrapped)
  }

  val buffer = Module(new BufferBRAM(nWords, new StreamChannel(dataBits)))
  val bufValid = RegInit(0.U(nWords.W))

  val head = RegInit(0.U(countBits.W))
  val tail = RegInit(0.U(countBits.W))
  val count = RegInit(0.U(countBits.W))

  val full = (count + io.alloc.bits.count) > nWords.U
  val xactHeads = Reg(Vec(nXacts, UInt(countBits.W)))
  val curXactHead = xactHeads(io.in.bits.id)

  val occupied = RegInit(false.B)
  val ren = (!occupied || io.out.ready) && (bufValid >> tail)(0)

  count := count +
              Mux(io.alloc.fire, io.alloc.bits.count, 0.U) -
              Mux(io.out.fire, 1.U, 0.U)
  bufValid := (bufValid | Mux(io.in.fire, UIntToOH(curXactHead), 0.U)) &
                         ~Mux(ren, UIntToOH(tail), 0.U)

  io.alloc.ready := !full
  io.in.ready := true.B
  io.out.valid := occupied
  io.out.bits := buffer.io.read.data

  buffer.io.write.en := io.in.fire
  buffer.io.write.addr := curXactHead
  buffer.io.write.data := io.in.bits.data
  buffer.io.read.en := ren
  buffer.io.read.addr := tail

  when (io.alloc.fire) {
    xactHeads(io.alloc.bits.id) := head
    head := incWrap(head, io.alloc.bits.count)
  }

  when (io.in.fire) {
    xactHeads(io.in.bits.id) := incWrap(curXactHead, 1.U)
  }

  when (io.out.fire) { occupied := false.B }

  when (ren) {
    occupied := true.B
    tail := incWrap(tail, 1.U)
  }
}

class PacketCollectionBuffer(bufWords: Int) extends Module {
  val io = IO(new StreamIO(NET_IF_WIDTH))

  val headerWords = ETH_HEAD_BYTES / NET_IF_BYTES
  val maxPackets = (bufWords - 1) / (headerWords + 1) + 1

  val queue = Module(new Queue(new StreamChannel(NET_IF_WIDTH), bufWords))
  val pktCount = TwoWayCounter(
    queue.io.enq.fire && queue.io.enq.bits.last,
    queue.io.deq.fire && queue.io.deq.bits.last,
    maxPackets)
  val canDequeue = pktCount > 0.U

  queue.io.enq <> io.in
  io.out.valid := queue.io.deq.valid && canDequeue
  io.out.bits := queue.io.deq.bits
  queue.io.deq.ready := io.out.ready && canDequeue
}

when (inIdx === 0.U) { startHead := bufHead }


package icenet

import chisel3._
import chisel3.util._
import freechips.rocketchip.unittest.UnitTest
import freechips.rocketchip.util.{DecoupledHelper, UIntIsOneOf}
import IceNetConsts._
import NetworkHelpers._

class ChecksumCalcRequest extends Bundle {
  val check = Bool()
  val start = UInt(16.W)
  val init  = UInt(16.W)
}

class ChecksumRewriteRequest extends Bundle {
  val check  = Bool()
  val offset = UInt(16.W)
  val start  = UInt(16.W)
  val init   = UInt(16.W)
}

class ChecksumCalc(dataBits: Int) extends Module {
  val dataBytes = dataBits / 8

  val io = IO(new Bundle {
    val req = Flipped(Decoupled(new ChecksumCalcRequest))
    val stream = new StreamIO(dataBits)
    val result = Decoupled(UInt(16.W))
  })

  val csum = Reg(UInt((dataBits + 16).W))
  val check = Reg(Bool())
  val start  = Reg(UInt(16.W))
  val startPos = Reg(UInt(16.W))
  val nextStartPos = startPos + dataBytes.U
  val sumMask = (0 until dataBytes).map { i =>
    io.stream.in.bits.keep(i) && (startPos + i.U) >= start
  }
  val sumData = io.stream.in.bits.data & FillInterleaved(8, sumMask)

  val s_req :: s_stream :: s_fold :: s_result :: Nil = Enum(4)
  val state = RegInit(s_req)

  io.req.ready := state === s_req
  io.stream.out.valid := state === s_stream && io.stream.in.valid
  io.stream.in.ready := state === s_stream && io.stream.out.ready
  io.stream.out.bits := io.stream.in.bits
  io.result.valid := state === s_result
  io.result.bits := csum(15, 0)

  when (io.req.fire) {
    check := io.req.bits.check
    start := io.req.bits.start
    csum := io.req.bits.init
    startPos := 0.U
    state := s_stream
  }

  when (io.stream.in.fire) {
    when (check) {
      csum := csum + sumData
      startPos := nextStartPos
    }

    when (io.stream.in.bits.last) {
      state := Mux(check, s_fold, s_req)
    }
  }

  when (state === s_fold) {
    val upper = csum(15 + dataBits, 16)
    val lower = csum(15, 0)

    when (upper === 0.U) {
      csum := ~lower
      state := s_result
    } .otherwise {
      csum := upper + lower
    }
  }

  when (io.result.fire) { state := s_req }
}

class ChecksumRewrite(dataBits: Int, nBufFlits: Int) extends Module {
  val dataBytes = dataBits / 8

  val io = IO(new Bundle {
    val req = Flipped(Decoupled(new ChecksumRewriteRequest))
    val stream = new StreamIO(dataBits)
  })

  val reqq = Module(new Queue(new ChecksumRewriteRequest, 2))
  val calc = Module(new ChecksumCalc(dataBits))
  val buffer = Module(new Queue(new StreamChannel(dataBits), nBufFlits))
  val offset = Reg(UInt(16.W))
  val check  = Reg(Bool())
  val csum = Reg(UInt(16.W))

  val reqHelper = DecoupledHelper(
    io.req.valid,
    calc.io.req.ready,
    reqq.io.enq.ready)

  io.req.ready := reqHelper.fire(io.req.valid)
  calc.io.req.valid := reqHelper.fire(calc.io.req.ready)
  calc.io.req.bits.check := io.req.bits.check
  calc.io.req.bits.start := io.req.bits.start
  calc.io.req.bits.init := io.req.bits.init
  reqq.io.enq.valid := reqHelper.fire(reqq.io.enq.ready)
  reqq.io.enq.bits := io.req.bits
  calc.io.stream.in <> io.stream.in
  buffer.io.enq <> calc.io.stream.out

  val byteOffBits = log2Ceil(dataBytes)
  val startPos = Reg(UInt(16.W))
  val nextStartPos = startPos + dataBytes.U
  val baseData = buffer.io.deq.bits.data
  val shiftAmt = Cat((offset - startPos)(byteOffBits-1, 0), 0.U(3.W))
  val dataMask = ~(~0.U(16.W) << shiftAmt)
  val csumShifted = csum << shiftAmt
  val replace = check && (offset >= startPos) && (offset < nextStartPos)
  val outData = Mux(replace, (baseData & dataMask) | csumShifted, baseData)

  val s_req :: s_wait :: s_flush :: Nil = Enum(3)
  val state = RegInit(s_req)

  when (reqq.io.deq.fire) {
    check := reqq.io.deq.bits.check
    offset := reqq.io.deq.bits.offset
    startPos := 0.U
    state := Mux(reqq.io.deq.bits.check, s_wait, s_flush)
  }

  when (calc.io.result.fire) {
    csum := calc.io.result.bits
    state := s_flush
  }

  when (io.stream.out.fire) {
    startPos := nextStartPos
    when (io.stream.out.bits.last) { state := s_req }
  }

  val deqOK = (state === s_flush || nextStartPos <= offset)

  reqq.io.deq.ready := state === s_req
  calc.io.result.ready := state === s_wait

  io.stream.out.valid := buffer.io.deq.valid && deqOK
  buffer.io.deq.ready := io.stream.out.ready && deqOK
  io.stream.out.bits := buffer.io.deq.bits
  io.stream.out.bits.data := outData
}

class ChecksumTest extends UnitTest {
  val offset = 6
  val init = 0x4315
  val start = 2
  val data = Seq(0xdead, 0xbeef, 0x7432, 0x0000, 0xf00d, 0x3163, 0x9821, 0x1543)
  val take = Seq(true,   true,   true,   true,   true,   true,   true,   false)

  var csum = init + data.zip(take).drop(start/2)
                        .filter(_._2).map(_._1).reduce(_ + _)
  while (csum > 0xffff) {
    csum = (csum >> 16) + (csum & 0xffff)
  }
  csum = ~csum & 0xffff

  val expected = data.take(offset/2) ++
    Seq(csum) ++ data.drop(offset/2+1)

  def seqToVec(seq: Seq[Int], step: Int, nbits: Int) = {
    VecInit((0 until seq.length by step).map { i =>
      Cat((i until (i + step)).map(seq(_).U(nbits.W)).reverse)
    })
  }

  val dataBits = 32
  val dataBytes = dataBits / 8
  val shortsPerFlit = dataBits / 16
  val dataVec = seqToVec(data, shortsPerFlit, 16)
  val expectedVec = seqToVec(expected, shortsPerFlit, 16)
  val keepVec = seqToVec(keep, shortsPerFlit, 2)

  val s_start :: s_req :: s_input :: s_output :: s_done :: Nil = Enum(5)
  val state = RegInit(s_start)

  val rewriter = Module(new ChecksumRewrite(
    dataBits, data.length/shortsPerFlit))

  val (inIdx, inDone) = Counter(rewriter.io.stream.in.fire, dataVec.length)
  val (outIdx, outDone) = Counter(rewriter.io.stream.out.fire, expectedVec.length)

  rewriter.io.req.valid := state === s_req
  rewriter.io.req.bits.check  := true.B
  rewriter.io.req.bits.start  := start.U
  rewriter.io.req.bits.init   := init.U
  rewriter.io.req.bits.offset := offset.U
  rewriter.io.stream.in.valid := state === s_input
  rewriter.io.stream.in.bits.data := dataVec(inIdx)
  rewriter.io.stream.in.bits.keep := keepVec(inIdx)
  rewriter.io.stream.in.bits.last := inIdx === (dataVec.length-1).U
  rewriter.io.stream.out.ready := state === s_output
  io.finished := state === s_done

  when (state === s_start && io.start) { state := s_req }
  when (rewriter.io.req.fire) { state := s_input }
  when (inDone) { state := s_output }
  when (outDone) { state := s_done }

  assert(!rewriter.io.stream.out.valid ||
    rewriter.io.stream.out.bits.data === expectedVec(outIdx),
    "ChecksumTest: got wrong data")

  assert(!rewriter.io.stream.out.valid ||
    rewriter.io.stream.out.bits.keep === keepVec(outIdx),
    "ChecksumTest: got wrong keep")
}

class TCPChecksumOffloadResult extends Bundle {
  val correct = Bool()
  val checked = Bool()
}

class TCPChecksumOffload(dataBits: Int) extends Module {
  val io = IO(new Bundle {
    val in = Flipped(Decoupled(new StreamChannel(dataBits)))
    val result = Decoupled(new TCPChecksumOffloadResult)
  })

  class FullHeader extends Bundle {
    val tcp = new TCPHeader
    val ipv4 = new IPv4Header
    val eth = new EthernetHeader
  }
  class PseudoHeader extends Bundle {
    val tcp = new TCPHeader
    val length = UInt(16.W)
    val protocol = UInt(8.W)
    val zeros = UInt(8.W)
    val dest_ip = UInt(32.W)
    val source_ip = UInt(32.W)
  }

  val dataBytes = dataBits/8

  val headerBytes = ETH_HEAD_BYTES + IPV4_HEAD_BYTES + TCP_HEAD_BYTES
  val headerWords = headerBytes / dataBytes
  val headerVec = Reg(Vec(headerWords, UInt(dataBits.W)))
  val header = headerVec.asTypeOf(new FullHeader)
  val headerIdx = RegInit(0.U(log2Ceil(headerWords).W))

  val pseudoHeaderBytes = 12 + TCP_HEAD_BYTES
  val pseudoHeaderWords = pseudoHeaderBytes / dataBytes
  val pseudoHeader = Wire(new PseudoHeader)
  val pseudoHeaderVec = pseudoHeader.asTypeOf(
    Vec(pseudoHeaderWords, UInt(dataBits.W)))

  pseudoHeader.tcp := header.tcp
  pseudoHeader.length := htons(ntohs(header.ipv4.total_length) - IPV4_HEAD_BYTES.U)
  pseudoHeader.protocol := header.ipv4.protocol
  pseudoHeader.zeros := 0.U
  pseudoHeader.dest_ip := header.ipv4.dest_ip
  pseudoHeader.source_ip := header.ipv4.source_ip

  require(dataBits >= 16)
  require(headerBytes % dataBytes == 0)
  require(pseudoHeaderBytes % dataBytes == 0)

  val (s_header_in :: s_csum_req ::
       s_header_csum :: s_body_csum ::
       s_passthru :: s_result :: Nil) = Enum(6)
  val state = RegInit(s_header_in)

  val headerOK =
    header.eth.ethType === IPV4_ETHTYPE.U &&
    header.ipv4.protocol === TCP_PROTOCOL.U &&
    header.ipv4.ihl === 5.U
  val headerChannel = Wire(new StreamChannel(dataBits))
  headerChannel.data := pseudoHeaderVec(headerIdx)
  headerChannel.keep := ~0.U(dataBits.W)
  headerChannel.last := false.B

  val resultExpected = RegInit(false.B)

  val csum = Module(new ChecksumCalc(dataBits))

  csum.io.req.valid := state === s_csum_req
  csum.io.req.bits.check := true.B
  csum.io.req.bits.start := 0.U
  csum.io.req.bits.init  := 0.U

  csum.io.stream.in.valid :=
    (state === s_header_csum) || (state === s_body_csum && io.in.valid)
  csum.io.stream.in.bits := Mux(
    state === s_header_csum, headerChannel, io.in.bits)
  csum.io.stream.out.ready := true.B // just ignore output

  io.in.ready := state.isOneOf(s_header_in, s_passthru) ||
                (state === s_body_csum && csum.io.stream.in.ready)

  io.result.valid := state === s_result && (!resultExpected || csum.io.result.valid)
  io.result.bits.correct := csum.io.result.bits === 0.U
  io.result.bits.checked := resultExpected
  csum.io.result.ready := state === s_result && resultExpected && io.result.ready

  when (io.in.fire) {
    when (io.in.bits.last) {
      state := s_result
    } .elsewhen (state === s_header_in) {
      headerVec(headerIdx) := io.in.bits.data
      headerIdx := headerIdx + 1.U
      when (headerIdx === (headerWords-1).U) {
        resultExpected := headerOK
        state := Mux(headerOK, s_csum_req, s_passthru)
      }
    }
  }

  when (csum.io.req.fire) {
    headerIdx := 0.U
    state := s_header_csum
  }

  when (state === s_header_csum && csum.io.stream.in.ready) {
    headerIdx := headerIdx + 1.U
    when (headerIdx === (pseudoHeaderWords-1).U) {
      state := s_body_csum
    }
  }

  when (io.result.fire) {
    headerIdx := 0.U
    resultExpected := false.B
    state := s_header_in
  }
}

class ChecksumTCPVerify extends UnitTest {
  val packets = Seq("""00 00 00 12 6d 00 00 03 00 12 6d 00 00 02 08 00
                      |45 00 00 3c 45 9b 40 00 40 06 9c fb ac 10 00 02
                      |ac 10 00 03 c2 a8 14 51 22 ad 9e d6 00 00 00 00
                      |a0 02 72 10 bf 07 00 00 02 04 05 b4 04 02 08 0a
                      |43 ec e2 58 00 00 00 00 01 03 03 07 00 00 00 00""",
                    """00 00 00 12 6d 00 00 02 00 12 6d 00 00 03 08 06
                      |00 01 08 00 06 04 00 02 00 12 6d 00 00 03 ac 10
                      |00 03 00 12 6d 00 00 02 ac 10 00 02 00 02 00 00""")
  val dataBytes = packets.map(packetStr =>
      packetStr.stripMargin.replaceAll("\n", " ")
               .split(" ").map(BigInt(_, 16)))
  val dataWords = VecInit(dataBytes.flatMap(
    bytes => (0 until bytes.length by 8).map(
      i => Cat(bytes.slice(i, i + 8).map(_.U(8.W)).reverse))))
  val dataLast = VecInit(dataBytes.map(_.length / 8).flatMap(
    nWords => Seq.fill(nWords-1)(false.B) :+ true.B))

  val expectedResults = VecInit(3.U, 2.U)

  val started = RegInit(false.B)
  val inputValid = RegInit(false.B)
  val outputReady = RegInit(false.B)
  val resultReady = RegInit(false.B)

  val offload = Module(new TCPChecksumOffload(NET_IF_WIDTH))

  val (inIdx, inDone) = Counter(offload.io.in.fire, dataWords.length)
  val (resultIdx, resultDone) = Counter(offload.io.result.fire, expectedResults.length)

  offload.io.in.valid := inputValid
  offload.io.in.bits.data := dataWords(inIdx)
  offload.io.in.bits.keep := NET_FULL_KEEP
  offload.io.in.bits.last := dataLast(inIdx)
  offload.io.result.ready := resultReady
  io.finished := started && !inputValid && !resultReady

  when (!started && io.start) {
    started := true.B
    inputValid := true.B
    resultReady := true.B
  }
  when (inDone) { inputValid := false.B }
  when (resultDone) { resultReady := false.B }

  assert(!offload.io.result.valid ||
    offload.io.result.bits.asUInt === expectedResults(resultIdx),
    "ChecksumTCPVerify: checksum was not correct")
}

val keep = take.map(if (_) 0x3 else 0x0)


package icenet

import chisel3._
import freechips.rocketchip.subsystem.BaseSubsystemConfig
import org.chipsalliance.cde.config.{Parameters, Config}
import freechips.rocketchip.unittest.UnitTests

class WithIceNetUnitTests extends Config((site, here, up) => {
  case NICKey => Some(NICConfig())
  case UnitTests => (p: Parameters) => {
    Seq(
      Module(new NetworkPacketBufferTest),
      Module(new PauserTest),
      Module(new NetworkTapTest),
      Module(new RateLimiterTest),
      Module(new AlignerTest),
      Module(new ChecksumTest),
      Module(new ChecksumTCPVerify),
      Module(new IceNicSendTestWrapper()(p)),
      Module(new IceNicRecvTestWrapper()(p)),
      Module(new IceNicTestWrapper()(p)),
      Module(new MisalignedTestWrapper()(p)))
  }
})

class WithIceNIC(inBufFlits: Int = 1800, usePauser: Boolean = false, ctrlQueueDepth: Int = 64)
    extends Config((site, here, up) => {
  case NICKey => Some(NICConfig(
    inBufFlits = inBufFlits,
    ctrlQueueDepth = ctrlQueueDepth,
    usePauser = usePauser,
    checksumOffload = true))
})

class WithNoIceNIC extends Config((site, here, up) => {
  case NICKey => None
})

class IceNetUnitTestConfig extends Config(
  new WithIceNetUnitTests ++ new BaseSubsystemConfig)


package icenet

import chisel3._

object IceNetConsts {
  val NET_IF_WIDTH = 64
  val NET_IF_BYTES = NET_IF_WIDTH/8
  val NET_LEN_BITS = 16

  val NET_IP_ALIGN = 2
  val ETH_HEAD_BYTES = 16

  val ETH_STANDARD_MTU = 1500
  val ETH_STANDARD_MAX_BYTES = ETH_STANDARD_MTU + ETH_HEAD_BYTES
  val ETH_STANDARD_MAX_FLITS = (ETH_STANDARD_MAX_BYTES - 1) / NET_IF_BYTES + 1

  val ETH_JUMBO_MTU = 9000
  val ETH_JUMBO_MAX_BYTES = ETH_JUMBO_MTU + ETH_HEAD_BYTES
  val ETH_JUMBO_MAX_FLITS = (ETH_JUMBO_MAX_BYTES - 1) / NET_IF_BYTES + 1

  val ETH_SUPER_JUMBO_MTU = 64 << 10
  val ETH_SUPER_JUMBO_MAX_BYTES = ETH_SUPER_JUMBO_MTU + ETH_HEAD_BYTES
  val ETH_SUPER_JUMBO_MAX_FLITS = (ETH_SUPER_JUMBO_MAX_BYTES - 1) / NET_IF_BYTES + 1

  val ETH_MAC_BITS = 48
  val ETH_TYPE_BITS = 16
  val ETH_PAD_BITS = 16

  val IPV4_HEAD_BYTES = 20
  val UDP_HEAD_BYTES = 8
  val TCP_HEAD_BYTES = 20

  val RLIMIT_MAX_INC = 256
  val RLIMIT_MAX_PERIOD = 256
  val RLIMIT_MAX_SIZE = 256

  def ETH_BCAST_MAC = ~0.U(ETH_MAC_BITS.W)

  val IPV4_ETHTYPE = 0x0008
  val TCP_PROTOCOL = 0x06
}

def NET_FULL_KEEP = ~0.U(NET_IF_BYTES.W)


package icenet

import chisel3._
import chisel3.util._
import org.chipsalliance.cde.config.Parameters
import freechips.rocketchip.diplomacy.{LazyModule, LazyModuleImp, IdRange}
import freechips.rocketchip.util.DecoupledHelper
import freechips.rocketchip.tilelink._

class StreamReadRequest extends Bundle {
  val address = UInt(48.W)
  val length = UInt(15.W)
  val partial = Bool()
}

class StreamReader(nXacts: Int, outFlits: Int, maxBytes: Int)
    (implicit p: Parameters) extends LazyModule {

  val core = LazyModule(new StreamReaderCore(nXacts, outFlits, maxBytes))
  val node = core.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val dataBits = core.module.dataBits

    val io = IO(new Bundle {
      val req = Flipped(Decoupled(new StreamReadRequest))
      val resp = Decoupled(Bool())
      val out = Decoupled(new StreamChannel(dataBits))
    })

    core.module.io.req <> io.req
    io.resp <> core.module.io.resp

    val buffer = Module(new ReservationBuffer(nXacts, outFlits, dataBits))
    buffer.io.alloc <> core.module.io.alloc
    buffer.io.in <> core.module.io.out

    val aligner = Module(new Aligner(dataBits))
    aligner.io.in <> buffer.io.out
    io.out <> aligner.io.out
  }
}

class StreamReaderCore(nXacts: Int, outFlits: Int, maxBytes: Int)
    (implicit p: Parameters) extends LazyModule {
  val node = TLClientNode(Seq(TLMasterPortParameters.v1(Seq(TLClientParameters(
    name = "stream-reader", sourceId = IdRange(0, nXacts))))))

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val (tl, edge) = node.out(0)
    val dataBits = tl.params.dataBits
    val beatBytes = dataBits / 8
    val byteAddrBits = log2Ceil(beatBytes)
    val addrBits = tl.params.addressBits
    val lenBits = 15

    require (edge.manager.minLatency > 0)

    val io = IO(new Bundle {
      val req = Flipped(Decoupled(new StreamReadRequest))
      val resp = Decoupled(Bool())
      val alloc = Decoupled(new ReservationBufferAlloc(nXacts, outFlits))
      val out = Decoupled(new ReservationBufferData(nXacts, dataBits))
    })

    val s_idle :: s_read :: s_resp :: Nil = Enum(3)
    val state = RegInit(s_idle)

    // Physical (word) address in memory
    val sendaddr = Reg(UInt(addrBits.W))
    // Number of words to send
    val sendlen  = Reg(UInt(lenBits.W))
    // 0 if last packet in sequence, 1 otherwise
    val sendpart = Reg(Bool())

    val xactBusy = RegInit(0.U(nXacts.W))
    val xactOnehot = PriorityEncoderOH(~xactBusy)
    val xactId = OHToUInt(xactOnehot)
    val xactLast = Reg(UInt(nXacts.W))
    val xactLeftKeep = Reg(Vec(nXacts, UInt(beatBytes.W)))
    val xactRightKeep = Reg(Vec(nXacts, UInt(beatBytes.W)))

    val reqSize = MuxCase(byteAddrBits.U,
      (log2Ceil(maxBytes) until byteAddrBits by -1).map(lgSize =>
          // Use the largest size (beatBytes <= size <= maxBytes)
          // s.t. sendaddr % size == 0 and sendlen > size
          (sendaddr(lgSize-1,0) === 0.U &&
            (sendlen >> lgSize.U) =/= 0.U) -> lgSize.U))
    val isLast = (xactLast >> tl.d.bits.source)(0) && edge.last(tl.d)
    val canSend = state === s_read && !xactBusy.andR

    val fullKeep = ~0.U(beatBytes.W)
    val loffset = Reg(UInt(byteAddrBits.W))
    val roffset = Reg(UInt(byteAddrBits.W))
    val lkeep = fullKeep << loffset
    val rkeep = fullKeep >> roffset
    val first = Reg(Bool())

    xactBusy := (xactBusy | Mux(tl.a.fire, xactOnehot, 0.U)) &
                    ~Mux(tl.d.fire && edge.last(tl.d),
                          UIntToOH(tl.d.bits.source), 0.U)

    val helper = DecoupledHelper(tl.a.ready, io.alloc.ready)

    io.req.ready := state === s_idle
    io.alloc.valid := helper.fire(io.alloc.ready, canSend)
    io.alloc.bits.id := xactId
    io.alloc.bits.count := (1.U << (reqSize - byteAddrBits.U))
    tl.a.valid := helper.fire(tl.a.ready, canSend)
    tl.a.bits := edge.Get(
      fromSource = xactId,
      toAddress = sendaddr,
      lgSize = reqSize)._2

    val outLeftKeep = xactLeftKeep(tl.d.bits.source)
    val outRightKeep = xactRightKeep(tl.d.bits.source)

    io.out.valid := tl.d.valid
    io.out.bits.id := tl.d.bits.source
    io.out.bits.data.data := tl.d.bits.data
    io.out.bits.data.keep := MuxCase(fullKeep, Seq(
      (edge.first(tl.d) && edge.last(tl.d)) -> (outLeftKeep & outRightKeep),
      edge.first(tl.d) -> outLeftKeep,
      edge.last(tl.d)  -> outRightKeep))
    io.out.bits.data.last := isLast
    tl.d.ready := io.out.ready
    io.resp.valid := state === s_resp
    io.resp.bits := true.B

    when (io.req.fire) {
      val req = io.req.bits
      val lastaddr = req.address + req.length
      val startword = req.address(addrBits-1, byteAddrBits)
      val endword = lastaddr(addrBits-1, byteAddrBits) +
                      Mux(lastaddr(byteAddrBits-1, 0) === 0.U, 0.U, 1.U)

      loffset := req.address(byteAddrBits-1, 0)
      roffset := Cat(endword, 0.U(byteAddrBits.W)) - lastaddr
      first := true.B

      sendaddr := Cat(startword, 0.U(byteAddrBits.W))
      sendlen  := Cat(endword - startword, 0.U(byteAddrBits.W))
      sendpart := req.partial
      state := s_read

      assert(req.length > 0.U, "request length must be >0")
    }

    when (tl.a.fire) {
      val reqBytes = 1.U << reqSize
      sendaddr := sendaddr + reqBytes
      sendlen  := sendlen - reqBytes
      when (sendlen === reqBytes) {
        xactLast := (xactLast & ~xactOnehot) | Mux(sendpart, 0.U, xactOnehot)
        xactRightKeep(xactId) := rkeep
        state := s_resp
      } .otherwise {
        xactLast := xactLast & ~xactOnehot
        xactRightKeep(xactId) := fullKeep
      }
      when (first) {
        first := false.B
        xactLeftKeep(xactId) := lkeep
      } .otherwise {
        xactLeftKeep(xactId) := fullKeep
      }
    }

    when (io.resp.fire) {
      state := s_idle
    }
  }
}

class StreamWriteRequest extends Bundle {
  val address = UInt(48.W)
  val length = UInt(16.W)
}

class StreamWriter(nXacts: Int, maxBytes: Int)
    (implicit p: Parameters) extends LazyModule {
  val node = TLClientNode(Seq(TLMasterPortParameters.v1(Seq(TLClientParameters(
    name = "stream-writer", sourceId = IdRange(0, nXacts))))))

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val (tl, edge) = node.out(0)
    val dataBits = tl.params.dataBits
    val beatBytes = dataBits / 8
    val byteAddrBits = log2Ceil(beatBytes)
    val addrBits = tl.params.addressBits
    val lenBits = 16

    require (edge.manager.minLatency > 0)

    val io = IO(new Bundle {
      val req = Flipped(Decoupled(new StreamWriteRequest))
      val resp = Decoupled(UInt(lenBits.W))
      val in = Flipped(Decoupled(new StreamChannel(dataBits)))
    })

    val s_idle :: s_data :: s_resp :: Nil = Enum(3)
    val state = RegInit(s_idle)

    val length = Reg(UInt(lenBits.W))
    val baseAddr = Reg(UInt(addrBits.W))
    val offset = Reg(UInt(addrBits.W))
    val addrMerged = baseAddr + offset
    val bytesToSend = length - offset
    val baseByteOff = baseAddr(byteAddrBits-1, 0)
    val byteOff = addrMerged(byteAddrBits-1, 0)
    val extraBytes = Mux(baseByteOff === 0.U, 0.U, beatBytes.U - baseByteOff)

    val xactBusy = RegInit(0.U(nXacts.W))
    val xactOnehot = PriorityEncoderOH(~xactBusy)
    val xactId = OHToUInt(xactOnehot)

    val maxBeats = maxBytes / beatBytes
    val beatIdBits = log2Ceil(maxBeats)

    val beatsLeft = Reg(UInt(beatIdBits.W))
    val headAddr = Reg(UInt(addrBits.W))
    val headXact = Reg(UInt(log2Ceil(nXacts).W))
    val headSize = Reg(UInt(log2Ceil(maxBytes + 1).W))

    val newBlock = beatsLeft === 0.U
    val canSend = !xactBusy.andR || !newBlock

    val reqSize = MuxCase(0.U,
      (log2Ceil(maxBytes) until 0 by -1).map(lgSize =>
          (addrMerged(lgSize-1,0) === 0.U &&
            (bytesToSend >> lgSize.U) =/= 0.U) -> lgSize.U))

    xactBusy := (xactBusy | Mux(tl.a.fire && newBlock, xactOnehot, 0.U)) &
                    ~Mux(tl.d.fire, UIntToOH(tl.d.bits.source), 0.U)

    val overhang = RegInit(0.U(dataBits.W))
    val sendTrail = bytesToSend <= extraBytes
    val fulldata = (overhang | (io.in.bits.data << Cat(baseByteOff, 0.U(3.W))))

    val fromSource = Mux(newBlock, xactId, headXact)
    val toAddress = Mux(newBlock, addrMerged, headAddr)
    val lgSize = Mux(newBlock, reqSize, headSize)
    val wdata = fulldata(dataBits-1, 0)
    val wmask = Cat((0 until beatBytes).map(
      i => (i.U >= byteOff) && (i.U < bytesToSend)).reverse)
    val wpartial = !wmask.andR

    val putPartial = edge.Put(
      fromSource = xactId,
      toAddress = addrMerged & ~(beatBytes-1).U(addrBits.W),
      lgSize = log2Ceil(beatBytes).U,
      data = Mux(sendTrail, overhang, wdata),
      mask = wmask)._2

    val putFull = edge.Put(
      fromSource = fromSource,
      toAddress = toAddress,
      lgSize = lgSize,
      data = wdata)._2

    io.req.ready := state === s_idle
    tl.a.valid := (state === s_data) && (io.in.valid || sendTrail) && canSend
    tl.a.bits := Mux(wpartial, putPartial, putFull)
    tl.d.ready := xactBusy.orR
    io.in.ready := state === s_data && canSend && !sendTrail && tl.a.ready
    io.resp.valid := state === s_resp && !xactBusy.orR
    io.resp.bits := length

    when (io.req.fire) {
      offset := 0.U
      baseAddr := io.req.bits.address
      length := io.req.bits.length
      beatsLeft := 0.U
      state := s_data
    }

    when (tl.a.fire) {
      when (!newBlock) {
        beatsLeft := beatsLeft - 1.U
      } .elsewhen (reqSize > byteAddrBits.U) {
        val nBeats = 1.U << (reqSize - byteAddrBits.U)
        beatsLeft := nBeats - 1.U
        headAddr := addrMerged
        headXact := xactId
        headSize := reqSize
      }

      val bytesSent = PopCount(wmask)
      offset := offset + bytesSent
      overhang := fulldata >> dataBits.U

      when (bytesSent === bytesToSend) { state := s_resp }
    }


package icenet

import chisel3._
import chisel3.util._
import IceNetConsts._

object NetworkHelpers {
  def reverse_bytes(a: UInt, n: Int) = {
    val bytes = (0 until n).map(i => a((i + 1) * 8 - 1, i * 8))
    Cat(bytes)
  }

  def ntohl(a: UInt) = reverse_bytes(a, 4)
  def htons(a: UInt) = reverse_bytes(a, 2)
  def ntohs(a: UInt) = reverse_bytes(a, 2)
}

class EthernetHeader extends Bundle {
  val ethType = UInt(ETH_TYPE_BITS.W)
  val srcmac  = UInt(ETH_MAC_BITS.W)
  val dstmac  = UInt(ETH_MAC_BITS.W)
  val padding = UInt(ETH_PAD_BITS.W)

  def toWords(w: Int = NET_IF_WIDTH) =
    this.asTypeOf(Vec(ETH_HEAD_BYTES * 8 / w, UInt(w.W)))
}

object EthernetHeader {
  def apply(dstmac: UInt, srcmac: UInt, ethType: UInt) = {
    val header = Wire(new EthernetHeader)
    header.dstmac := dstmac
    header.srcmac := srcmac
    header.ethType := ethType
    header.padding := DontCare
    header
  }

  def apply(words: Seq[UInt], w: Int = NET_IF_WIDTH) = {
    val headerWords = ETH_HEAD_BYTES * 8 / w
    Cat(words.take(headerWords).reverse).asTypeOf(new EthernetHeader)
  }
}

class IPv4Header extends Bundle {
  val dest_ip = UInt(32.W)
  val source_ip = UInt(32.W)
  val header_checksum = UInt(16.W)
  val protocol = UInt(8.W)
  val ttl = UInt(8.W)
  val frag_off2 = UInt(8.W)
  val flags = UInt(3.W)
  val frag_off1 = UInt(5.W)
  val ident = UInt(16.W)
  val total_length = UInt(16.W)
  val dcsp = UInt(6.W)
  val ecn = UInt(2.W)
  val version = UInt(4.W)
  val ihl = UInt(4.W)
}

class UDPHeader extends Bundle {
  val checksum = UInt(16.W)
  val length = UInt(16.W)
  val dest_port = UInt(16.W)
  val source_port = UInt(16.W)
}

class TCPHeader extends Bundle {
  val urgent_pointer = UInt(16.W)
  val checksum = UInt(16.W)
  val window_size = UInt(16.W)
  val cwr = Bool()
  val ece = Bool()
  val urg = Bool()
  val ack = Bool()
  val psh = Bool()
  val rst = Bool()
  val syn = Bool()
  val fin = Bool()
  val data_offset = UInt(4.W)
  val reserved = UInt(3.W)
  val ns = Bool()
  val acknum = UInt(32.W)
  val seqnum = UInt(32.W)
  val dest_port = UInt(16.W)
  val source_port = UInt(16.W)
}

def htonl(a: UInt) = reverse_bytes(a, 4)


package icenet

import chisel3._
import chisel3.util._
import freechips.rocketchip.unittest.UnitTest
import IceNetConsts._

class RateLimiterSettings extends Bundle {
  val incBits = log2Ceil(RLIMIT_MAX_INC)
  val periodBits = log2Ceil(RLIMIT_MAX_PERIOD)
  val sizeBits = log2Ceil(RLIMIT_MAX_SIZE)

  /*
   * Given a clock frequency of X, you can achieve an output bandwidth
   * of Y = X * (N / D), where N <= D, by setting inc to N and period to (D - 1).
   * The field size should be set to the number of consecutive beats that
   * can be sent before rate-limiting kicks in.
   */
  val inc = UInt(incBits.W)
  val period = UInt(periodBits.W)
  val size = UInt(sizeBits.W)
}

class RateLimiter[T <: Data](typ: T) extends Module {
  val incBits = log2Ceil(RLIMIT_MAX_INC)
  val periodBits = log2Ceil(RLIMIT_MAX_PERIOD)
  val sizeBits = log2Ceil(RLIMIT_MAX_SIZE)

  val io = IO(new Bundle {
    val in = Flipped(Decoupled(typ))
    val out = Decoupled(typ)
    val settings = Input(new RateLimiterSettings)
  })

  val tokens = RegInit(0.U(sizeBits.W))
  val counter = RegInit(0.U(periodBits.W))

  when (counter === 0.U) {
    counter := io.settings.period
  } .otherwise {
    counter := counter - 1.U
  }

  val inc_trigger = counter === 0.U
  val ok_to_send = tokens > 0.U

  io.in.ready := io.out.ready && ok_to_send
  io.out.valid := io.in.valid && ok_to_send
  io.out.bits := io.in.bits

  def uint_min(a: UInt, b: UInt) = Mux(a < b, a, b)

  when (inc_trigger && io.out.fire) {
    val next_tokens = tokens +& (io.settings.inc - 1.U)
    tokens := uint_min(next_tokens, io.settings.size)
  } .elsewhen (inc_trigger) {
    val next_tokens = tokens +& io.settings.inc
    tokens := uint_min(next_tokens, io.settings.size)
  } .elsewhen (io.out.fire) {
    tokens := tokens - 1.U
  }
}

object RateLimiter {
  def apply[T <: Data](in: DecoupledIO[T], inc: Int, period: Int, size: Int) = {
    if (period > 1) {
      val limiter = Module(new RateLimiter(in.bits.cloneType))
      limiter.io.in <> in
      limiter.io.settings.inc := inc.U
      limiter.io.settings.period := (period - 1).U
      limiter.io.settings.size := size.U
      limiter.io.out

  }
}

class RateLimiterTest extends UnitTest {
  val nFlits = 48
  val started = RegInit(false.B)
  val sending = RegInit(false.B)
  val receiving = RegInit(false.B)

  val limiter = Module(new RateLimiter(Bool()))
  limiter.io.in.valid := sending
  limiter.io.in.bits := true.B
  limiter.io.out.ready := receiving
  limiter.io.settings.inc := 1.U
  limiter.io.settings.period := 3.U
  limiter.io.settings.size := 2.U

  val (sendCount, sendDone) = Counter(limiter.io.in.fire, nFlits)
  val (recvCount, recvDone) = Counter(limiter.io.out.fire, nFlits)

  when (io.start && !started) {
    started := true.B
    sending := true.B
    receiving := true.B
  }

  when (sendDone) { sending := false.B }
  when (recvDone) { receiving := false.B }

  io.finished := started && !sending && !receiving
}


package icenet

import chisel3._
import chisel3.util._
import chisel3.reflect.DataMirror
import freechips.rocketchip.subsystem.{BaseSubsystem, TLBusWrapperLocation, PBUS, FBUS}
import org.chipsalliance.cde.config.{Field, Parameters}
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.prci._
import freechips.rocketchip.regmapper._
import freechips.rocketchip.interrupts._
import freechips.rocketchip.tilelink._
import freechips.rocketchip.util._
import IceNetConsts._

// This is copied from testchipip to avoid dependencies
class ClockedIO[T <: Data](private val gen: T) extends Bundle {
  val clock = Output(Clock())
  val bits = DataMirror.internal.chiselTypeClone[T](gen)
}

/**
 * @inBufFlits How many flits in the input buffer(s)
 * @outBufFlits Number of flits in the output buffer
 * @nMemXacts Maximum number of transactions that the send/receive path can send to memory
 * @maxAcquireBytes Cache block size
 * @ctrlQueueDepth Depth of the MMIO control queues
 * @usePauser Hardware support for Ethernet pause frames
 * @checksumOffload TCP checksum offload engine
 * @packetMaxBytes Maximum number of bytes in a packet (header size + MTU)
 */
case class NICConfig(
  inBufFlits: Int  = 2 * ETH_STANDARD_MAX_BYTES / NET_IF_BYTES,
  outBufFlits: Int = 2 * ETH_STANDARD_MAX_BYTES / NET_IF_BYTES,
  nMemXacts: Int = 8,
  maxAcquireBytes: Int = 64,
  ctrlQueueDepth: Int = 10,
  usePauser: Boolean = false,
  checksumOffload: Boolean = false,
  packetMaxBytes: Int = ETH_STANDARD_MAX_BYTES)

case class NICAttachParams(
  masterWhere: TLBusWrapperLocation = FBUS,
  slaveWhere: TLBusWrapperLocation = PBUS
)

case object NICKey extends Field[Option[NICConfig]](None)
case object NICAttachKey extends Field[NICAttachParams](NICAttachParams())

trait HasNICParameters {
  implicit val p: Parameters
  val nicExternal = p(NICKey).get
  val inBufFlits = nicExternal.inBufFlits
  val outBufFlits = nicExternal.outBufFlits
  val nMemXacts = nicExternal.nMemXacts
  val maxAcquireBytes = nicExternal.maxAcquireBytes
  val ctrlQueueDepth = nicExternal.ctrlQueueDepth
  val usePauser = nicExternal.usePauser
  val checksumOffload = nicExternal.checksumOffload
  val packetMaxBytes = nicExternal.packetMaxBytes
}

abstract class NICLazyModule(implicit p: Parameters)
  extends LazyModule with HasNICParameters

abstract class NICModule(implicit val p: Parameters)
  extends Module with HasNICParameters

abstract class NICBundle(implicit val p: Parameters)
  extends Bundle with HasNICParameters

class PacketArbiter(arbN: Int, rr: Boolean = false)
  extends HellaPeekingArbiter(
    new StreamChannel(NET_IF_WIDTH), arbN,
    (ch: StreamChannel) => ch.last, rr = rr)

class IceNicSendIO extends Bundle {
  val req = Decoupled(UInt(NET_IF_WIDTH.W))
  val comp = Flipped(Decoupled(Bool()))
}

class IceNicRecvIO extends Bundle {
  val req = Decoupled(UInt(NET_IF_WIDTH.W))
  val comp = Flipped(Decoupled(UInt(NET_LEN_BITS.W)))
}

trait IceNicControllerBundle extends Bundle {
  val send = new IceNicSendIO
  val recv = new IceNicRecvIO
  val macAddr = Input(UInt(ETH_MAC_BITS.W))
  val txcsumReq = Decoupled(new ChecksumRewriteRequest)
  val rxcsumRes = Flipped(Decoupled(new TCPChecksumOffloadResult))
  val csumEnable = Output(Bool())
}

case class IceNicControllerParams(address: BigInt, beatBytes: Int)

/*
 * Take commands from the CPU over TL2, expose as Queues
 */
class IceNicController(c: IceNicControllerParams)(implicit p: Parameters)
    extends RegisterRouter(RegisterRouterParams("ice-nic", Seq("ucb-bar,ice-nic"),
      c.address, beatBytes=c.beatBytes))
    with HasTLControlRegMap
    with HasInterruptSources
    with HasNICParameters {
  override def nInterrupts = 2
  def tlRegmap(mapping: RegField.Map*): Unit = regmap(mapping:_*)
  override lazy val module = new IceNiCControllerModuleImp(this)
}

class IceNiCControllerModuleImp(outer: IceNicController)(implicit p: Parameters) extends LazyModuleImp(outer) with HasNICParameters {
  val io = IO(new Bundle with IceNicControllerBundle)

  val sendCompDown = WireInit(false.B)

  val qDepth = ctrlQueueDepth
  require(qDepth < (1 << 8))

  def queueCount[T <: Data](qio: QueueIO[T], depth: Int): UInt =
    TwoWayCounter(qio.enq.fire, qio.deq.fire, depth)

  // hold (len, addr) of packets that we need to send out
  val sendReqQueue = Module(new HellaQueue(qDepth)(UInt(NET_IF_WIDTH.W)))
  val sendReqCount = queueCount(sendReqQueue.io, qDepth)
  // hold addr of buffers we can write received packets into
  val recvReqQueue = Module(new HellaQueue(qDepth)(UInt(NET_IF_WIDTH.W)))
  val recvReqCount = queueCount(recvReqQueue.io, qDepth)
  // count number of sends completed
  val sendCompCount = TwoWayCounter(io.send.comp.fire, sendCompDown, qDepth)
  // hold length of received packets
  val recvCompQueue = Module(new HellaQueue(qDepth)(UInt(NET_LEN_BITS.W)))
  val recvCompCount = queueCount(recvCompQueue.io, qDepth)

  val sendCompValid = sendCompCount > 0.U
  val intMask = RegInit(0.U(2.W))

  io.send.req <> sendReqQueue.io.deq
  io.recv.req <> recvReqQueue.io.deq
  io.send.comp.ready := sendCompCount < qDepth.U
  recvCompQueue.io.enq <> io.recv.comp

  outer.interrupts(0) := sendCompValid && intMask(0)
  outer.interrupts(1) := recvCompQueue.io.deq.valid && intMask(1)

  val sendReqSpace = (qDepth.U - sendReqCount)
  val recvReqSpace = (qDepth.U - recvReqCount)

  def sendCompRead = (ready: Bool) => {
    sendCompDown := sendCompValid && ready
    (sendCompValid, true.B)
  }

  val txcsumReqQueue = Module(new HellaQueue(qDepth)(UInt(49.W)))
  val rxcsumResQueue = Module(new HellaQueue(qDepth)(UInt(2.W)))
  val csumEnable = RegInit(false.B)

  io.txcsumReq.valid := txcsumReqQueue.io.deq.valid
  io.txcsumReq.bits := txcsumReqQueue.io.deq.bits.asTypeOf(new ChecksumRewriteRequest)
  txcsumReqQueue.io.deq.ready := io.txcsumReq.ready

  rxcsumResQueue.io.enq.valid := io.rxcsumRes.valid
  rxcsumResQueue.io.enq.bits := io.rxcsumRes.bits.asUInt
  io.rxcsumRes.ready := rxcsumResQueue.io.enq.ready

  io.csumEnable := csumEnable

  outer.tlRegmap(
    0x00 -> Seq(RegField.w(NET_IF_WIDTH, sendReqQueue.io.enq)),
    0x08 -> Seq(RegField.w(NET_IF_WIDTH, recvReqQueue.io.enq)),
    0x10 -> Seq(RegField.r(1, sendCompRead)),
    0x12 -> Seq(RegField.r(NET_LEN_BITS, recvCompQueue.io.deq)),
    0x14 -> Seq(
      RegField.r(8, sendReqSpace),
      RegField.r(8, recvReqSpace),
      RegField.r(8, sendCompCount),
      RegField.r(8, recvCompCount)),
    0x18 -> Seq(RegField.r(ETH_MAC_BITS, io.macAddr)),
    0x20 -> Seq(RegField(2, intMask)),
    0x28 -> Seq(RegField.w(49, txcsumReqQueue.io.enq)),
    0x30 -> Seq(RegField.r(2, rxcsumResQueue.io.deq)),
    0x31 -> Seq(RegField(1, csumEnable)))
}

class IceNicSendPath(nInputTaps: Int = 0)(implicit p: Parameters)
    extends NICLazyModule {
  val reader = LazyModule(new StreamReader(
    nMemXacts, outBufFlits, maxAcquireBytes))
  val node = reader.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle {
      val send = Flipped(new IceNicSendIO)
      val tap = Flipped(Vec(nInputTaps, Decoupled(new StreamChannel(NET_IF_WIDTH))))
      val out = Decoupled(new StreamChannel(NET_IF_WIDTH))
      val rlimit = Input(new RateLimiterSettings)
      val csum = checksumOffload.option(new Bundle {
        val req = Flipped(Decoupled(new ChecksumRewriteRequest))
        val enable = Input(Bool())
      })
    })

    val readreq = reader.module.io.req
    io.send.req.ready := readreq.ready
    readreq.valid := io.send.req.valid
    readreq.bits.address := io.send.req.bits(47, 0)
    readreq.bits.length  := io.send.req.bits(62, 48)
    readreq.bits.partial := io.send.req.bits(63)
    io.send.comp <> reader.module.io.resp

    val preArbOut = if (checksumOffload) {
      val readerOut = reader.module.io.out
      val arb = Module(new PacketArbiter(2))
      val bufFlits = (packetMaxBytes - 1) / NET_IF_BYTES + 1
      val rewriter = Module(new ChecksumRewrite(NET_IF_WIDTH, bufFlits))
      val enable = io.csum.get.enable

      rewriter.io.req <> io.csum.get.req

      arb.io.in(0) <> rewriter.io.stream.out
      arb.io.in(1).valid := !enable && readerOut.valid
      arb.io.in(1).bits  := readerOut.bits
      rewriter.io.stream.in.valid := enable && readerOut.valid
      rewriter.io.stream.in.bits := readerOut.bits
      readerOut.ready := Mux(enable,
        rewriter.io.stream.in.ready, arb.io.in(1).ready)

      arb.io.out
    } else { reader.module.io.out }

    val unlimitedOut = if (nInputTaps > 0) {
      val bufWords = (packetMaxBytes - 1) / NET_IF_BYTES + 1
      val inputs = (preArbOut +: io.tap).map { in =>
        // The packet collection buffer doesn't allow sending the first flit
        // of a packet until the last flit is received.
        // This ensures that we don't lock the arbiter while waiting for data
        // to arrive, which could cause deadocks.
        val buffer = Module(new PacketCollectionBuffer(bufWords))
        buffer.io.in <> in
        buffer.io.out
      }
      val arb = Module(new PacketArbiter(inputs.size, rr = true))
      arb.io.in <> inputs
      arb.io.out
    } else { preArbOut }

    val limiter = Module(new RateLimiter(new StreamChannel(NET_IF_WIDTH)))
    limiter.io.in <> unlimitedOut
    limiter.io.settings := io.rlimit
    io.out <> limiter.io.out
  }
}

class IceNicWriter(implicit p: Parameters) extends NICLazyModule {
  val writer = LazyModule(new StreamWriter(nMemXacts, maxAcquireBytes))
  val node = writer.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle {
      val recv = Flipped(new IceNicRecvIO)
      val in = Flipped(Decoupled(new StreamChannel(NET_IF_WIDTH)))
      val length = Flipped(Valid(UInt(NET_LEN_BITS.W)))
    })

    val streaming = RegInit(false.B)
    val byteAddrBits = log2Ceil(NET_IF_BYTES)
    val helper = DecoupledHelper(
      io.recv.req.valid,
      writer.module.io.req.ready,
      io.length.valid, !streaming)

    writer.module.io.req.valid := helper.fire(writer.module.io.req.ready)
    writer.module.io.req.bits.address := io.recv.req.bits
    writer.module.io.req.bits.length := io.length.bits
    io.recv.req.ready := helper.fire(io.recv.req.valid)

    writer.module.io.in.valid := io.in.valid && streaming
    writer.module.io.in.bits := io.in.bits
    io.in.ready := writer.module.io.in.ready && streaming

    io.recv.comp <> writer.module.io.resp

    when (io.recv.req.fire) { streaming := true.B }
    when (io.in.fire && io.in.bits.last) { streaming := false.B }
  }
}

/*
 * Recv frames
 */
class IceNicRecvPath(val tapFuncs: Seq[EthernetHeader => Bool] = Nil)
    (implicit p: Parameters) extends LazyModule {
  val writer = LazyModule(new IceNicWriter)
  val node = TLIdentityNode()
  node := writer.node
  lazy val module = new IceNicRecvPathModule(this)
}

class IceNicRecvPathModule(val outer: IceNicRecvPath)
    extends LazyModuleImp(outer) with HasNICParameters {
  val io = IO(new Bundle {
    val recv = Flipped(new IceNicRecvIO)
    val in = Flipped(Decoupled(new StreamChannel(NET_IF_WIDTH))) // input stream
    val tap = Vec(outer.tapFuncs.length, Decoupled(new StreamChannel(NET_IF_WIDTH)))
    val csum = checksumOffload.option(new Bundle {
      val res = Decoupled(new TCPChecksumOffloadResult)
      val enable = Input(Bool())
    })
    val buf_free = Output(Vec(1 + outer.tapFuncs.length, UInt(8.W)))
  })

  def tapOutToDropCheck(tapOut: EthernetHeader => Bool) = {
    (header: EthernetHeader, ch: StreamChannel, update: Bool) => {
      val first = RegInit(true.B)
      val drop = tapOut(header) && first
      val dropReg = RegInit(false.B)

      when (update && first) { first := false.B; dropReg := drop }
      when (update && ch.last) { first := true.B; dropReg := false.B }

      drop || dropReg
    }
  }

  def duplicateStream(in: DecoupledIO[StreamChannel], outs: Seq[DecoupledIO[StreamChannel]]) = {
    outs.foreach { out =>
      out.valid := in.valid
      out.bits := in.bits
    }
    in.ready := outs.head.ready
    val outReadys = Cat(outs.map(_.ready))
    assert(outReadys.andR || !outReadys.orR,
      "Duplicated streams must all be ready simultaneously")
    outs
  }

  def invertCheck(check: (EthernetHeader, StreamChannel, Bool) => Bool) =
    (eth: EthernetHeader, ch: StreamChannel, up: Bool) => !check(eth, ch, up)

  val tapDropChecks = outer.tapFuncs.map(func => tapOutToDropCheck(func))

  val allDropChecks =
    // Drop checks for the primary buffer
    // Drop if the packet should be tapped out or is a pause frame
    Seq(tapDropChecks ++ pauseDropCheck.toSeq) ++
    // Drop checks for the tap buffers
    // For each tap, drop if the packet doesn't match the tap function or is a pause frame
    tapDropChecks.map(check => invertCheck(check) +: pauseDropCheck.toSeq)

  val buffers = allDropChecks.map(dropChecks =>
    Module(new NetworkPacketBuffer(
      inBufFlits,
      maxBytes = packetMaxBytes,
      dropChecks = dropChecks, dropless = usePauser)))
  duplicateStream(io.in, buffers.map(_.io.stream.in))

  io.buf_free := buffers.map(_.io.free)

  io.tap <> buffers.tail.map(_.io.stream.out)
  val bufout = buffers.head.io.stream.out
  val buflen = buffers.head.io.length

  val (csumout, recvreq) = (if (checksumOffload) {
    val offload = Module(new TCPChecksumOffload(NET_IF_WIDTH))
    val offloadReady = offload.io.in.ready || !io.csum.get.enable

    val out = Wire(Decoupled(new StreamChannel(NET_IF_WIDTH)))
    val recvreq = Wire(Decoupled(UInt(NET_IF_WIDTH.W)))
    val reqq = Module(new Queue(UInt(NET_IF_WIDTH.W), 1))

    val enqHelper = DecoupledHelper(
      io.recv.req.valid, reqq.io.enq.ready, recvreq.ready)
    val deqHelper = DecoupledHelper(
      bufout.valid, offloadReady, out.ready, reqq.io.deq.valid)

    reqq.io.enq.valid := enqHelper.fire(reqq.io.enq.ready)
    reqq.io.enq.bits := io.recv.req.bits
    io.recv.req.ready := enqHelper.fire(io.recv.req.valid)
    recvreq.valid := enqHelper.fire(recvreq.ready)
    recvreq.bits := io.recv.req.bits

    out.valid := deqHelper.fire(out.ready)
    out.bits  := bufout.bits
    offload.io.in.valid := deqHelper.fire(offloadReady, io.csum.get.enable)
    offload.io.in.bits := bufout.bits
    bufout.ready := deqHelper.fire(bufout.valid)
    reqq.io.deq.ready := deqHelper.fire(reqq.io.deq.valid, bufout.bits.last)

    io.csum.get.res <> offload.io.result

    (out, recvreq)
  } else { (bufout, io.recv.req) })

  val writer = outer.writer.module
  writer.io.recv.req <> Queue(recvreq, 1)
  io.recv.comp <> writer.io.recv.comp
  writer.io.in <> csumout
  writer.io.length.valid := buflen.valid
  writer.io.length.bits  := buflen.bits
}

class NICIO extends StreamIO(NET_IF_WIDTH) {
  val macAddr = Input(UInt(ETH_MAC_BITS.W))
  val rlimit = Input(new RateLimiterSettings)
  val pauser = Input(new PauserSettings)

}

/*
 * A simple NIC
 *
 * Expects ethernet frames (see below), but uses a custom transport
 * (see ExtBundle)
 *
 * Ethernet Frame format:
 *   2 bytes |  6 bytes  |  6 bytes    | 2 bytes  | 46-1500B
 *   Padding | Dest Addr | Source Addr | Type/Len | Data
 *
 * @address Starting address of MMIO control registers
 * @beatBytes Width of memory interface (in bytes)
 * @tapOutFuncs Sequence of functions for each output tap.
 *              Each function takes the header of an Ethernet frame
 *              and returns Bool that is true if matching and false if not.
 * @nInputTaps Number of input taps
 *
 */
class IceNIC(address: BigInt, beatBytes: Int = 8,
    tapOutFuncs: Seq[EthernetHeader => Bool] = Nil,
    nInputTaps: Int = 0)
    (implicit p: Parameters) extends NICLazyModule {

  val control = LazyModule(new IceNicController(
    IceNicControllerParams(address, beatBytes)))
  val sendPath = LazyModule(new IceNicSendPath(nInputTaps))
  val recvPath = LazyModule(new IceNicRecvPath(tapOutFuncs))

  val mmionode = TLIdentityNode()
  val dmanode = TLIdentityNode()
  val intnode = control.intXing(NoCrossing)

  control.node := TLAtomicAutomata() := mmionode
  dmanode := TLWidthWidget(NET_IF_BYTES) := sendPath.node
  dmanode := TLWidthWidget(NET_IF_BYTES) := recvPath.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle {
      val ext = new NICIO
      val tapOut = Vec(tapOutFuncs.length, Decoupled(new StreamChannel(NET_IF_WIDTH)))
      val tapIn = Flipped(Vec(nInputTaps, Decoupled(new StreamChannel(NET_IF_WIDTH))))
    })

    sendPath.module.io.send <> control.module.io.send
    recvPath.module.io.recv <> control.module.io.recv

    // connect externally
    if (usePauser) {
      val pauser = Module(new Pauser(inBufFlits, 1 + tapOutFuncs.length))
      pauser.io.int.out <> sendPath.module.io.out
      recvPath.module.io.in <> pauser.io.int.in
      io.ext.out <> pauser.io.ext.out
      pauser.io.ext.in <> io.ext.in
      pauser.io.in_free := recvPath.module.io.buf_free
      pauser.io.macAddr := io.ext.macAddr
      pauser.io.settings := io.ext.pauser
    } else {
      recvPath.module.io.in <> io.ext.in
      io.ext.out <> sendPath.module.io.out
    }

    control.module.io.macAddr := io.ext.macAddr
    sendPath.module.io.rlimit := io.ext.rlimit

    io.tapOut <> recvPath.module.io.tap
    sendPath.module.io.tap <> io.tapIn

    if (checksumOffload) {
      sendPath.module.io.csum.get.req <> control.module.io.txcsumReq
      sendPath.module.io.csum.get.enable := control.module.io.csumEnable
      control.module.io.rxcsumRes <> recvPath.module.io.csum.get.res
      recvPath.module.io.csum.get.enable := control.module.io.csumEnable
    } else {
      control.module.io.txcsumReq.ready := false.B
      control.module.io.rxcsumRes.valid := false.B
      control.module.io.rxcsumRes.bits := DontCare
    }
  }
}

class SimNetwork extends BlackBox with HasBlackBoxResource {
  val io = IO(new Bundle {
    val clock = Input(Clock())
    val reset = Input(Bool())
    val net = Flipped(new NICIOvonly)
  })
  addResource("/vsrc/SimNetwork.v")
  addResource("/csrc/SimNetwork.cc")
  addResource("/csrc/device.h")
  addResource("/csrc/device.cc")
  addResource("/csrc/switch.h")
  addResource("/csrc/switch.cc")
  addResource("/csrc/packet.h")
}

class NICIOvonly extends Bundle {
  val in = Flipped(Valid(new StreamChannel(NET_IF_WIDTH)))
  val out = Valid(new StreamChannel(NET_IF_WIDTH))
  val macAddr = Input(UInt(ETH_MAC_BITS.W))
  val rlimit = Input(new RateLimiterSettings)
  val pauser = Input(new PauserSettings)

}

object NICIOvonly {
  def apply(nicio: NICIO): NICIOvonly = {
    val vonly = Wire(new NICIOvonly)
    vonly.out.valid := nicio.out.valid
    vonly.out.bits  := nicio.out.bits
    nicio.out.ready := true.B
    nicio.in.valid  := vonly.in.valid
    nicio.in.bits   := vonly.in.bits
    assert(!vonly.in.valid || nicio.in.ready, "NIC input not ready for valid")
    nicio.macAddr := vonly.macAddr
    nicio.rlimit  := vonly.rlimit
    nicio.pauser  := vonly.pauser
    vonly
  }
}

object NICIO {
  def apply(vonly: NICIOvonly): NICIO = {
    val nicio = Wire(new NICIO)
    assert(!vonly.out.valid || nicio.out.ready)
    nicio.out.valid := vonly.out.valid
    nicio.out.bits  := vonly.out.bits
    vonly.in.valid  := nicio.in.valid
    vonly.in.bits   := nicio.in.bits
    nicio.in.ready  := true.B
    vonly.macAddr   := nicio.macAddr
    vonly.rlimit    := nicio.rlimit
    vonly.pauser    := nicio.pauser
    nicio
  }

}
trait CanHavePeripheryIceNIC  { this: BaseSubsystem =>
  private val address = BigInt(0x10016000)
  private val portName = "Ice-NIC"

  val icenicOpt = p(NICKey).map { params =>
    val manager = locateTLBusWrapper(p(NICAttachKey).slaveWhere)
    val client = locateTLBusWrapper(p(NICAttachKey).masterWhere)
    // TODO: currently the controller is in the clock domain of the bus which masters it
    // we assume this is same as the clock domain of the bus the controller masters
    val domain = manager.generateSynchronousDomain.suggestName("icenic_domain")

    val icenic = domain { LazyModule(new IceNIC(address, manager.beatBytes)) }

    manager.coupleTo(portName) { icenic.mmionode := TLFragmenter(manager.beatBytes, manager.blockBytes) := _ }
    client.coupleFrom(portName) { _ :=* icenic.dmanode }
    ibus.fromSync := icenic.intnode

    val inner_io = domain { InModuleBody {
      val inner_io = IO(new NICIOvonly).suggestName("nic")
      inner_io <> NICIOvonly(icenic.module.io.ext)
      inner_io
    } }

    val outer_io = InModuleBody {
      val outer_io = IO(new ClockedIO(new NICIOvonly)).suggestName("nic")
      outer_io.bits <> inner_io
      outer_io.clock := domain.module.clock
      outer_io
    }
    outer_io
  }
}

object NicLoopback {
  def connect(net: Option[NICIOvonly], nicConf: Option[NICConfig], qDepth: Int, latency: Int = 10): Unit = {
    net.foreach { netio =>
      import PauseConsts.BT_PER_QUANTA
      val packetWords = nicConf.get.packetMaxBytes / NET_IF_BYTES
      val packetQuanta = (nicConf.get.packetMaxBytes * 8) / BT_PER_QUANTA
      netio.macAddr := PlusArg("macaddr")
      netio.rlimit.inc := PlusArg("rlimit-inc", 1)
      netio.rlimit.period := PlusArg("rlimit-period", 1)
      netio.rlimit.size := PlusArg("rlimit-size", 8)
      netio.pauser.threshold := PlusArg("pauser-threshold", 2 * packetWords + latency)
      netio.pauser.quanta := PlusArg("pauser-quanta", 2 * packetQuanta)
      netio.pauser.refresh := PlusArg("pauser-refresh", packetWords)

      if (nicConf.get.usePauser) {
        val pauser = Module(new PauserComplex(qDepth))
        pauser.io.ext.flipConnect(NetDelay(NICIO(netio), latency))
        pauser.io.int.out <> pauser.io.int.in
        pauser.io.macAddr := netio.macAddr + (1 << 40).U
        pauser.io.settings := netio.pauser
      } else {

        netio.in := Pipe(netio.out, latency)
      }
      netio.in.bits.keep := NET_FULL_KEEP
    }
  }

  def connect(net: NICIOvonly, nicConf: NICConfig): Unit = {
    val packetWords = nicConf.packetMaxBytes / NET_IF_BYTES
    NicLoopback.connect(Some(net), Some(nicConf), 4 * packetWords)
  }
}

object SimNetwork {
  def connect(net: Option[NICIOvonly], clock: Clock, reset: Bool) {
    net.foreach { netio =>
      val sim = Module(new SimNetwork)
      sim.io.clock := clock
      sim.io.reset := reset
      sim.io.net <> netio
    }
  }
}

val pauseDropCheck = if (usePauser) Some(PauseDropCheck(_, _, _)) else None


package icenet

import chisel3._
import chisel3.util._

import org.chipsalliance.cde.config.Parameters
import freechips.rocketchip.devices.tilelink.TLROM
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.tilelink._
import freechips.rocketchip.unittest.{UnitTest, UnitTestIO}
import freechips.rocketchip.util.{LatencyPipe, TwoWayCounter, UIntIsOneOf}
import scala.math.max
import IceNetConsts._

class IceNicTestSendDriver(
    sendReqs: Seq[(Int, Int, Boolean)],
    sendData: Seq[BigInt])(implicit p: Parameters) extends LazyModule {
  val node = TLClientNode(Seq(TLMasterPortParameters.v1(Seq(TLClientParameters(
    name = "test-send-driver", sourceId = IdRange(0, 1))))))

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO {
      val send = new IceNicSendIO
    })

    val (tl, edge) = node.out(0)
    val dataBits = tl.params.dataBits
    val beatBytes = dataBits / 8
    val byteAddrBits = log2Ceil(beatBytes)
    val lenBits = NET_LEN_BITS - 1
    val addrBits = NET_IF_WIDTH - NET_LEN_BITS

    val (s_start :: s_write_req :: s_write_resp ::
         s_send :: s_done :: Nil) = Enum(5)
    val state = RegInit(s_start)

    val sendReqVec = VecInit(sendReqs.map {
      case (addr, len, part) => Cat(part.B, len.U(lenBits.W), addr.U(addrBits.W))})

    val sendReqAddrVec = VecInit(sendReqs.map{case (addr, _, _) => addr.U(addrBits.W)})
    val sendReqCounts = sendReqs.map { case (_, len, _) => len / beatBytes }
    val sendReqCountVec = VecInit(sendReqCounts.map(_.U))
    val sendReqBase = VecInit((0 +: (1 until sendReqs.size).map(
      i => sendReqCounts.slice(0, i).reduce(_ + _))).map(_.U(addrBits.W)))
    val maxMemCount = sendReqCounts.reduce(max(_, _))
    val totalMemCount = sendReqCounts.reduce(_ + _)

    require(totalMemCount == sendData.size)

    val sendDataVec = VecInit(sendData.map(_.U(dataBits.W)))

    val reqIdx = Reg(UInt(log2Ceil(sendReqs.size).W))
    val memIdx = Reg(UInt(log2Ceil(totalMemCount).W))

    val outSend = TwoWayCounter(
      io.send.req.fire, io.send.comp.fire, sendReqs.size)

    val writeAddr = sendReqAddrVec(reqIdx) + (memIdx << byteAddrBits.U)
    val writeData = sendDataVec(sendReqBase(reqIdx) + memIdx)

    tl.a.valid := state === s_write_req
    tl.a.bits := edge.Put(
      fromSource = 0.U,
      toAddress = writeAddr,
      lgSize = byteAddrBits.U,
      data = writeData)._2
    tl.d.ready := state === s_write_resp

    io.send.req.valid := state === s_send
    io.send.req.bits := sendReqVec(reqIdx)
    io.send.comp.ready := outSend =/= 0.U

    io.finished := state === s_done && outSend === 0.U

    when (state === s_start && io.start) {
      reqIdx := 0.U
      memIdx := 0.U
      state := s_write_req
    }

    when (tl.a.fire) { state := s_write_resp }

    when (tl.d.fire) {
      memIdx := memIdx + 1.U
      state := s_write_req

      when (memIdx === (sendReqCountVec(reqIdx) - 1.U)) {
        memIdx := 0.U
        reqIdx := reqIdx + 1.U
        when (reqIdx === (sendReqs.size - 1).U) {
          reqIdx := 0.U
          state := s_send
        }
      }
    }

    when (io.send.req.fire) {
      reqIdx := reqIdx + 1.U
      when (reqIdx === (sendReqs.size - 1).U) {
        reqIdx := 0.U
        state := s_done
      }
    }
  }
}

class IceNicTestRecvDriver(recvReqs: Seq[Int], recvData: Seq[BigInt])
    (implicit p: Parameters) extends LazyModule {
  val node = TLClientNode(Seq(TLMasterPortParameters.v1(Seq(TLClientParameters(
    name = "test-recv-driver", sourceId = IdRange(0, 1))))))

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO {
      val recv = new IceNicRecvIO
    })

    val (tl, edge) = node.out(0)
    val dataBits = tl.params.dataBits
    val beatBytes = dataBits / 8
    val byteAddrBits = log2Ceil(beatBytes)

    val (s_start :: s_recv :: s_wait ::
         s_check_req :: s_check_resp :: s_done :: Nil) = Enum(6)
    val state = RegInit(s_start)

    val recvReqVec  = VecInit(recvReqs.map(_.U(NET_IF_WIDTH.W)))
    val recvDataVec = VecInit(recvData.map(_.U(dataBits.W)))

    val reqIdx = Reg(UInt(log2Ceil(recvReqs.size).W))
    val memIdx = Reg(UInt(log2Ceil(recvData.size).W))

    val outRecv = TwoWayCounter(
      io.recv.req.fire, io.recv.comp.fire, recvReqVec.size)

    tl.a.valid := state === s_check_req
    tl.a.bits := edge.Get(
      fromSource = 0.U,
      toAddress = recvReqVec.head + (memIdx << byteAddrBits.U),
      lgSize = byteAddrBits.U)._2
    tl.d.ready := state === s_check_resp

    io.recv.req.valid := state === s_recv
    io.recv.req.bits := recvReqVec(reqIdx)
    io.recv.comp.ready := outRecv =/= 0.U

    io.finished := state === s_done

    when (state === s_start && io.start) {
      reqIdx := 0.U
      memIdx := 0.U
      state := s_recv
    }

    when (io.recv.req.fire) {
      reqIdx := reqIdx + 1.U
      when (reqIdx === (recvReqVec.size - 1).U) {
        reqIdx := 0.U
        state := s_wait
      }
    }

    when (state === s_wait && outRecv === 0.U) {
      state := s_check_req
    }

    when (state === s_check_req && tl.a.ready) {
      state := s_check_resp
    }

    when (state === s_check_resp && tl.d.valid) {
      memIdx := memIdx + 1.U
      state := s_check_req
      when (memIdx === (recvData.size - 1).U) {
        memIdx := 0.U
        state := s_done
      }
    }

    assert(!tl.d.valid || tl.d.bits.data === recvDataVec(memIdx),
      "IceNicTest: Received wrong data")
  }
}

class IceNicRecvTest(implicit p: Parameters) extends NICLazyModule {
  val recvReqs = Seq(0, 1440, 1464)
  // The 90-flit packet should be dropped
  val recvLens = Seq(180, 3, 90, 7)
  val testData = Seq.tabulate(280) { i => BigInt(i << 4) }
  val recvData = testData.take(183) ++ testData.drop(273)

  val nicParams = p.alterPartial({
    case NICKey => Some(p(NICKey).get.copy(inBufFlits = 200))
  })
  val recvDriver = LazyModule(new IceNicTestRecvDriver(recvReqs, recvData))
  val recvPath = LazyModule(new IceNicRecvPath()(nicParams))
  val xbar = LazyModule(new TLXbar)
  val mem = LazyModule(new TLRAM(
    AddressSet(0, 0x7ff), beatBytes = NET_IF_BYTES))

  val MEM_LATENCY = 32
  val RLIMIT_INC = 1
  val RLIMIT_PERIOD = 4
  val RLIMIT_SIZE = 8

  xbar.node := recvDriver.node
  xbar.node := recvPath.node
  mem.node := TLFragmenter(NET_IF_BYTES, maxAcquireBytes) := TLBuffer.chainNode(MEM_LATENCY) := xbar.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO)

    val gen = Module(new PacketGen(recvLens, testData))
    gen.io.start := io.start
    recvDriver.module.io.start := gen.io.finished
    recvPath.module.io.recv <> recvDriver.module.io.recv
    recvPath.module.io.in <> RateLimiter(
      gen.io.out, RLIMIT_INC, RLIMIT_PERIOD, RLIMIT_SIZE)
    io.finished := recvDriver.module.io.finished
  }
}

class IceNicRecvTestWrapper(implicit p: Parameters) extends UnitTest(20000) {
  val test = Module(LazyModule(new IceNicRecvTest).module)
  test.io.start := io.start
  io.finished := test.io.finished
}

class IceNicSendTest(implicit p: Parameters) extends LazyModule {
  val sendReqs = Seq(
    (2, 10, true),
    (17, 6, false),
    (24, 12, false))
  val sendData = Seq(
    BigInt("7766554433221100", 16),
    BigInt("FFEEDDCCBBAA9988", 16),
    BigInt("0123456789ABCDEF", 16),
    BigInt("FEDCBA9876543210", 16),
    BigInt("76543210FDECBA98", 16))

  val recvData = Seq(
    BigInt("9988776655443322", 16),
    BigInt("23456789ABCDBBAA", 16),
    BigInt("FEDCBA9876543210", 16),
    BigInt("00000000FDECBA98", 16))
  val recvKeep = Seq(0xFF, 0xFF, 0xFF, 0x0F)
  val recvLast = Seq(false, true, false, true)

  val sendPath = LazyModule(new IceNicSendPath)
  val rom = LazyModule(new TLROM(0, 64,
    sendData.flatMap(
      data => (0 until 8).map(i => ((data >> (i * 8)) & 0xff).toByte)),
    beatBytes = 8))

  rom.node := TLFragmenter(NET_IF_BYTES, 16) := TLBuffer() := sendPath.node

  val RLIMIT_INC = 1
  val RLIMIT_PERIOD = 0
  val RLIMIT_SIZE = 8

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO)

    val sendPathIO = sendPath.module.io
    val sendReqVec = VecInit(sendReqs.map { case (start, len, part) =>
      Cat(part.B, len.U(15.W), start.U(48.W))
    })
    val (sendReqIdx, sendReqDone) = Counter(sendPathIO.send.req.fire, sendReqs.size)
    val (sendCompIdx, sendCompDone) = Counter(sendPathIO.send.comp.fire, sendReqs.size)

    val started = RegInit(false.B)
    val requesting = RegInit(false.B)
    val completing = RegInit(false.B)

    sendPathIO.send.req.valid := requesting
    sendPathIO.send.req.bits := sendReqVec(sendReqIdx)
    sendPathIO.send.comp.ready := completing

    when (!started && io.start) {
      requesting := true.B
      completing := true.B
    }
    when (sendReqDone)  { requesting := false.B }

    sendPathIO.rlimit.inc := RLIMIT_INC.U
    sendPathIO.rlimit.period := RLIMIT_PERIOD.U
    sendPathIO.rlimit.size := RLIMIT_SIZE.U

    val check = Module(new PacketCheck(recvData, recvKeep, recvLast))
    check.io.in <> sendPathIO.out
    io.finished := check.io.finished && !completing && !requesting
  }
}

class IceNicSendTestWrapper(implicit p: Parameters) extends UnitTest {
  val test = Module(LazyModule(new IceNicSendTest).module)
  test.io.start := io.start
  io.finished := test.io.finished
}

class IceNicTest(implicit p: Parameters) extends NICLazyModule {
  val sendReqs = Seq(
    (0, 128, true),
    (144, 160, false),
    (320, 64, false))
  val recvReqs = Seq(256, 544)
  val testData = Seq.tabulate(44)(i => BigInt(i << 4))

  val sendDriver = LazyModule(new IceNicTestSendDriver(sendReqs, testData))
  val recvDriver = LazyModule(new IceNicTestRecvDriver(recvReqs, testData))
  val sendPath = LazyModule(new IceNicSendPath)
  val recvPath = LazyModule(new IceNicRecvPath)
  val xbar = LazyModule(new TLXbar)
  val mem = LazyModule(new TLRAM(
    AddressSet(0, 0x1ff), beatBytes = NET_IF_BYTES))

  val NET_LATENCY = 64
  val MEM_LATENCY = 32
  val RLIMIT_INC = 1
  val RLIMIT_PERIOD = 0
  val RLIMIT_SIZE = 8

  xbar.node := sendDriver.node
  xbar.node := recvDriver.node
  xbar.node := sendPath.node
  xbar.node := recvPath.node
  mem.node := TLFragmenter(NET_IF_BYTES, maxAcquireBytes) := TLBuffer.chainNode(MEM_LATENCY) := xbar.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO)

    sendPath.module.io.send <> sendDriver.module.io.send
    recvPath.module.io.recv <> recvDriver.module.io.recv

    sendPath.module.io.rlimit.inc := RLIMIT_INC.U
    sendPath.module.io.rlimit.period := RLIMIT_PERIOD.U
    sendPath.module.io.rlimit.size := RLIMIT_SIZE.U

    recvPath.module.io.in <> LatencyPipe(sendPath.module.io.out, NET_LATENCY)

    sendDriver.module.io.start := io.start
    recvDriver.module.io.start := io.start
    io.finished := sendDriver.module.io.finished && recvDriver.module.io.finished

    val count_start :: count_up :: count_print :: count_done :: Nil = Enum(4)
    val count_state = RegInit(count_start)
    val cycle_count = Reg(UInt(64.W))
    val recv_count = Reg(UInt(1.W))

    when (count_state === count_start && sendPath.module.io.send.req.fire) {
      count_state := count_up
      cycle_count := 0.U
      recv_count := 1.U
    }
    when (count_state === count_up) {
      cycle_count := cycle_count + 1.U
      when (recvPath.module.io.recv.comp.fire) {
        recv_count := recv_count - 1.U
        when (recv_count === 0.U) { count_state := count_print }
      }
    }
    when (count_state === count_print) {
      printf("NIC test completed in %d cycles\n", cycle_count)
      count_state := count_done
    }
  }
}

class IceNicTestWrapper(implicit p: Parameters) extends UnitTest(50000) {
  val test = Module(LazyModule(new IceNicTest).module)
  test.io.start := io.start
  io.finished := test.io.finished
}

class MisalignedTestDriver(implicit p: Parameters) extends Module {
  val io = IO(new Bundle with UnitTestIO {
    val net = new StreamIO(NET_IF_WIDTH)
    val send = new IceNicSendIO
    val recv = new IceNicRecvIO
  })

  val fullData = Seq(
    "hDEADBEEF", "h01234567", "h555999ff", "h09341423",
    "h1384AEDF", "hABCD1234", "h1093567A", "hBADD00D1",
    "hA325B246", "h49230923", "h11113333", "hAEDC1445").map(_.U(32.W))

  val outData1 = VecInit(
    Cat(fullData(8),  fullData(7)),
    Cat(fullData(10), fullData(9)),
    Cat(0.U(32.W),    fullData(11)))
  val outKeep1 = VecInit(Seq.fill(2)(~0.U(8.W)) :+ "h0f".U)

  val outData2 = VecInit(
    Cat(fullData(1), fullData(0)),
    Cat(fullData(3), fullData(2)),
    Cat(fullData(5), fullData(4)),
    Cat(0.U(32.W),   fullData(6)))
  val outKeep2 = VecInit(Seq.fill(3)(~0.U(8.W)) :+ "h0f".U)

  val expData = VecInit(
    Cat(fullData(6),  fullData(5)),
    Cat(fullData(8),  fullData(7)),
    Cat(fullData(10), fullData(9)))

  val (s_start :: s_sendreq :: s_sendcomp :: s_recvreq :: s_recvcomp ::
       s_outdata1 :: s_outdata2 :: s_indata :: s_done :: Nil) = Enum(9)
  val state = RegInit(s_start)

  val (recvReqIdx, recvReqDone) = Counter(io.recv.req.fire, 2)
  val (recvCompIdx, recvCompDone) = Counter(io.recv.comp.fire, 2)

  val (outIdx1, outDone1) = Counter(
    state === s_outdata1 && io.net.out.ready, outData1.size)
  val (outIdx2, outDone2) = Counter(
    state === s_outdata2 && io.net.out.ready, outData2.size)
  val (inIdx, inDone) = Counter(io.net.in.fire, expData.size)

  val outBits1 = Wire(new StreamChannel(NET_IF_WIDTH))
  outBits1.data := outData1(outIdx1)
  outBits1.keep := outKeep1(outIdx1)
  outBits1.last := outIdx1 === (outData1.size - 1).U

  val outBits2 = Wire(new StreamChannel(NET_IF_WIDTH))
  outBits2.data := outData2(outIdx2)
  outBits2.keep := outKeep2(outIdx2)
  outBits2.last := outIdx2 === (outData2.size - 1).U

  io.send.req.valid := state === s_sendreq
  io.send.req.bits := Cat((expData.size * 8).U(16.W), 20.U(48.W))
  io.send.comp.ready := state === s_sendcomp

  io.recv.req.valid := state === s_recvreq
  io.recv.req.bits  := Mux(recvReqIdx === 0.U, 28.U, 0.U)
  io.recv.comp.ready := state === s_recvcomp

  io.net.out.valid := state.isOneOf(s_outdata1, s_outdata2)
  io.net.out.bits := Mux(state === s_outdata1, outBits1, outBits2)
  io.net.in.ready := state === s_indata

  io.finished := state === s_done

  when (state === s_start && io.start) { state := s_recvreq }

  when (recvReqDone) { state := s_outdata1 }

  when (outDone1) { state := s_outdata2 }
  when (outDone2) { state := s_recvcomp }

  when (recvCompDone) { state := s_sendreq }

  when (io.send.req.fire) { state := s_indata }

  when (inDone) { state := s_done }

  assert(!io.net.in.valid || io.net.in.bits.data === expData(inIdx),
    "MisalignedTest: input data does not match expected")
  assert(!io.net.in.valid || io.net.in.bits.keep.andR,
    "MisalignedTest: input keep does not match expected")
  assert(!io.net.in.valid || io.net.in.bits.last === (inIdx === (expData.size-1).U),
    "MisalignedTest: input last does not match expected")

}

class MisalignedTest(implicit p: Parameters) extends NICLazyModule {
  val sendpath = LazyModule(new IceNicSendPath)
  val recvpath = LazyModule(new IceNicRecvPath)

  val xbar = LazyModule(new TLXbar)
  val mem = LazyModule(new TLRAM(
    AddressSet(0, 0x7ff), beatBytes = NET_IF_BYTES))

  xbar.node := sendpath.node
  xbar.node := recvpath.node
  mem.node := TLFragmenter(NET_IF_BYTES, maxAcquireBytes) :=
              TLBuffer() := xbar.node

  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    val io = IO(new Bundle with UnitTestIO)

    val driver = Module(new MisalignedTestDriver)

    driver.io.start := io.start
    io.finished := driver.io.finished

    driver.io.net.in <> sendpath.module.io.out
    recvpath.module.io.in <> driver.io.net.out

    sendpath.module.io.send.req <> Queue(driver.io.send.req, 1)
    recvpath.module.io.recv.req <> Queue(driver.io.recv.req, 2)
    driver.io.send.comp <> Queue(sendpath.module.io.send.comp, 1)
    driver.io.recv.comp <> Queue(recvpath.module.io.recv.comp, 2)

    val rlimit = sendpath.module.io.rlimit
    rlimit.inc := 1.U
    rlimit.period := 0.U
    rlimit.size := 8.U
  }
}

class MisalignedTestWrapper(implicit p: Parameters) extends UnitTest {
  val test = Module(LazyModule(new MisalignedTest).module)
  test.io.start := io.start
  io.finished := test.io.finished
}

when (sendCompDone) { completing := false.B }


package icenet

import chisel3._
import chisel3.util._
import freechips.rocketchip.util.{HellaPeekingArbiter}
import freechips.rocketchip.unittest.UnitTest
import IceNetConsts._

object PauseConsts {
  val MAC_ETHTYPE = 0x8808
  val PAUSE_CTRL = 0x0001
  val BT_PER_QUANTA = 512
  val CYCLES_PER_QUANTA = BT_PER_QUANTA / NET_IF_WIDTH
  val MULTICAST_MACADDR = 0x010000C28001L
}
import PauseConsts._

trait NetworkEndianHelpers {
  def ntohs(x: UInt) = Cat(x(7, 0), x(15, 8))
  def htons(x: UInt) = ntohs(x)
  def ntohl(x: UInt) = Cat(ntohs(x(15, 0)), ntohs(x(31, 16)))
  def htonl(x: UInt) = ntohl(x)
}

class PauserSettings extends Bundle {
  val threshold = UInt(16.W)
  val quanta    = UInt(16.W)
  val refresh   = UInt(16.W)
}

object PauseDropCheck extends NetworkEndianHelpers {
  def apply(header: EthernetHeader, ch: StreamChannel, update: Bool): Bool = {
    val first = RegInit(true.B)
    val isPause = ntohs(header.ethType) === MAC_ETHTYPE.U &&
                  ntohs(ch.data(15, 0)) === PAUSE_CTRL.U && first
    val isPauseReg = RegInit(false.B)

    when (update && first)  { first := false.B; isPauseReg := isPause }
    when (update && ch.last) { first := true.B; isPauseReg := false.B }

    isPause || isPauseReg
  }
}

/**
 * Flow control unit using Ethernet pause frames
 * See https://en.wikipedia.org/wiki/Ethernet_flow_control#Pause_frame
 * @creditInit Size of each buffer being tracked
 * @nBuckets Number of buffers being tracked
 */
class Pauser(creditInit: Int, nBuckets: Int) extends Module
    with NetworkEndianHelpers {
  val timerBits = 16 + log2Ceil(CYCLES_PER_QUANTA)
  val creditBits = log2Ceil(creditInit + 1)

  val io = IO(new Bundle {
    val ext = new StreamIO(NET_IF_WIDTH)
    val int = Flipped(new StreamIO(NET_IF_WIDTH))
    val in_free = Input(Vec(nBuckets, UInt(8.W)))
    val macAddr = Input(UInt(48.W))
    val settings = Input(new PauserSettings)
  })

  val credits = RegInit(VecInit(Seq.fill(nBuckets)(creditInit.U(16.W))))
  val outPauseTimer = RegInit(0.U(timerBits.W))
  val outPaused = outPauseTimer > 0.U

  io.int.in <> io.ext.in

  val headerBeats = ETH_HEAD_BYTES / NET_IF_BYTES
  val inHeaderVec = Reg(Vec(headerBeats, UInt(NET_IF_WIDTH.W)))
  val inHeader = inHeaderVec.asTypeOf(new EthernetHeader)

  val s_head :: s_mac :: s_tail :: Nil = Enum(3)
  val state = RegInit(s_head)

  val inHeadIdx = RegInit(0.U(log2Ceil(headerBeats).W))

  when (outPaused) { outPauseTimer := outPauseTimer - 1.U }

  when (io.int.in.fire) {
    val data = io.int.in.bits.data

    switch (state) {
      is (s_head) {
        inHeaderVec(inHeadIdx) := data
        inHeadIdx := inHeadIdx + 1.U
        when (inHeadIdx === (headerBeats-1).U) { state := s_mac }
      }
      is (s_mac) {
        val isMac = ntohs(inHeader.ethType) === MAC_ETHTYPE.U
        val isPause = ntohs(data(15, 0)) === PAUSE_CTRL.U
        val quanta = ntohs(data(31, 16))
        val cycles = quanta << log2Ceil(CYCLES_PER_QUANTA).U
        when (isMac && isPause) { outPauseTimer := cycles }
        state := s_tail
      }
    }

    when (io.int.in.bits.last) {
      state := s_head
      inHeadIdx := 0.U
    }
  }

  for (i <- 0 until nBuckets) {
    credits(i) := credits(i) - io.int.in.fire + io.in_free(i)
  }

  val arb = Module(new PacketArbiter(2))

  val outHeader = EthernetHeader(
    MULTICAST_MACADDR.U,
    io.macAddr,
    htons(MAC_ETHTYPE.U(16.W)))
  val outVec = VecInit(outHeader.toWords() :+ Cat(
    htons(io.settings.quanta), htons(PAUSE_CTRL.U(16.W))))
  val sendPause = RegInit(false.B)
  val (outIdx, outDone) = Counter(arb.io.in(0).fire, outVec.size)

  arb.io.in(0).valid := sendPause
  arb.io.in(0).bits.data := outVec(outIdx)
  arb.io.in(0).bits.keep := NET_FULL_KEEP
  arb.io.in(0).bits.last := outIdx === (outVec.size-1).U

  val inPauseTimer = RegInit(0.U(16.W))
  val inPaused = inPauseTimer > 0.U

  val creditsLow = credits.map(_ < io.settings.threshold).reduce(_ || _)

  when (outDone) {
    inPauseTimer := io.settings.refresh
    sendPause := false.B
  }
  when (inPaused) { inPauseTimer := inPauseTimer - 1.U }
  when (creditsLow && !sendPause && !inPaused) {
    sendPause := true.B
  }

  val outInProgress = RegInit(false.B)
  val canForward = !outPaused || outInProgress

  arb.io.in(1).valid := canForward && io.int.out.valid
  arb.io.in(1).bits := io.int.out.bits
  io.int.out.ready := canForward && arb.io.in(1).ready

  when (arb.io.in(1).fire) {
    when (!outInProgress) { outInProgress := true.B }

  }

  io.ext.out <> arb.io.out
}

class PauserComplex(nFlits: Int) extends Module {
  val creditBits = log2Ceil(nFlits + 1)
  val io = IO(new Bundle {
    val ext = new StreamIO(NET_IF_WIDTH)
    val int = Flipped(new StreamIO(NET_IF_WIDTH))
    val macAddr = Input(UInt(48.W))
    val settings = Input(new PauserSettings)
  })

  val pauser = Module(new Pauser(nFlits, 1))
  val buffer = Module(new NetworkPacketBuffer(
    nFlits, dropChecks = Seq(PauseDropCheck(_, _, _)), dropless = true))

  io.ext <> pauser.io.ext
  pauser.io.macAddr := io.macAddr
  pauser.io.settings := io.settings

  pauser.io.int.out <> io.int.out
  buffer.io.stream.in <> pauser.io.int.in
  pauser.io.in_free(0) := buffer.io.free
  io.int.in <> buffer.io.stream.out
}

class PauserTest extends UnitTest {
  val nFlits = 400
  val latency = 128
  val packetWords = ETH_STANDARD_MAX_BYTES / NET_IF_BYTES
  val threshold = latency + packetWords
  val pauseQuanta = threshold / CYCLES_PER_QUANTA
  val pauseRefresh = packetWords

  val lcomplex = Module(new PauserComplex(nFlits))
  val rcomplex = Module(new PauserComplex(nFlits))

  rcomplex.io.ext.flipConnect(NetDelay(lcomplex.io.ext, latency))
  rcomplex.io.int.out <> rcomplex.io.int.in

  val lmacaddr = (0x2L << 40).U(48.W)
  val rmacaddr = (0x3L << 40).U(48.W)

  lcomplex.io.macAddr := lmacaddr
  rcomplex.io.macAddr := rmacaddr
  lcomplex.io.settings.threshold := threshold.U
  rcomplex.io.settings.threshold := threshold.U
  lcomplex.io.settings.quanta := pauseQuanta.U
  rcomplex.io.settings.quanta := pauseQuanta.U
  lcomplex.io.settings.refresh := pauseRefresh.U
  rcomplex.io.settings.refresh := pauseRefresh.U

  val ethHeader = EthernetHeader(rmacaddr, lmacaddr, 0x0000.U)
  val pktData = VecInit(
    ethHeader.toWords() ++ Seq.tabulate(130)(_.U(NET_IF_WIDTH.W)))
  val nPackets = 8

  val started = RegInit(false.B)

  val sending = RegInit(false.B)
  val (sendIdx, sendPktDone) = Counter(lcomplex.io.int.out.fire, pktData.size)
  val (sendPhase, sendDone) = Counter(sendPktDone, nPackets)

  val receiving = RegInit(false.B)
  val (recvIdx, recvPktDone) = Counter(lcomplex.io.int.in.fire, pktData.size)
  val (recvPhase, recvDone) = Counter(recvPktDone, nPackets)

  lcomplex.io.int.out.valid := sending
  lcomplex.io.int.out.bits.data := pktData(sendIdx)
  lcomplex.io.int.out.bits.keep := NET_FULL_KEEP
  lcomplex.io.int.out.bits.last := sendIdx === (pktData.size-1).U
  lcomplex.io.int.in.ready := receiving

  when (!started && io.start) {
    started := true.B
    sending := true.B
    receiving := true.B
  }
  when (sendDone) { sending := false.B }
  when (recvDone) { receiving := false.B }

  val recv = lcomplex.io.int.in
  assert(!recv.valid ||
    (recv.bits.data === pktData(recvIdx) &&
     recv.bits.keep === NET_FULL_KEEP &&
     recv.bits.last === (recvIdx === (pktData.size-1).U)),
    "PauserTest: received incorrect data, keep, or last")

  io.finished := started && !sending && !receiving
}

when (arb.io.in(1).bits.last) { outInProgress := false.B }


package icenet

import chisel3._
import chisel3.util._

// These classes originated in testchipip, but are copied
// here to reduce dependencies
class StreamChannel(val w: Int) extends Bundle {
  val data = UInt(w.W)
  val keep = UInt((w/8).W)
  val last = Bool()
}

class StreamIO(val w: Int) extends Bundle {
  val in = Flipped(Decoupled(new StreamChannel(w)))
  val out = Decoupled(new StreamChannel(w))

  def flipConnect(other: StreamIO) {
    in <> other.out
    other.in <> out
  }


package icenet

import chisel3._
import chisel3.util._
import org.chipsalliance.cde.config.Parameters
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.tilelink._

class TCAMMatchIO(n: Int, dataBits: Int) extends Bundle {
  val addrBits = log2Ceil(n)
  val data = Output(UInt(dataBits.W))
  val addr = Input(UInt(addrBits.W))
  val found = Input(Bool())

}

class TCAM(address: BigInt, val n: Int, val dataBits: Int, val nPorts: Int)
    (implicit p: Parameters) extends LazyModule {
  val addrBits = log2Ceil(n)
  val byteAddrBits = log2Ceil(dataBits / 8)
  val beatBytes = 1 << byteAddrBits
  val addrMask = (1 << (1 + addrBits + byteAddrBits)) - 1

  val node = TLManagerNode(Seq(TLSlavePortParameters.v1(Seq(TLSlaveParameters.v1(
    address = Seq(AddressSet(address, addrMask)),
    regionType = RegionType.PUT_EFFECTS,
    supportsGet = TransferSizes(1, beatBytes),
    supportsPutFull = TransferSizes(1, beatBytes))), beatBytes)))

  lazy val module = new TCAMModule(this)
}

class TCAMModule(outer: TCAM) extends LazyModuleImp(outer) {
  val io = IO(new Bundle {
    val tcam = Flipped(
      Vec(outer.nPorts, new TCAMMatchIO(outer.n, outer.dataBits)))
  })

  val (tl, edge) = outer.node.in(0)

  val dataArr = Reg(Vec(outer.n, UInt(outer.dataBits.W)))
  val maskArr = Reg(Vec(outer.n, UInt(outer.dataBits.W)))

  val acq = Queue(tl.a, 1)
  val regsel = acq.bits.address(outer.addrBits + outer.byteAddrBits)
  val wordaddr = acq.bits.address(
    outer.addrBits + outer.byteAddrBits - 1, outer.byteAddrBits)

  tl.d.valid := acq.valid
  acq.ready := tl.d.ready
  tl.d.bits := edge.AccessAck(acq.bits, 0.U)
  tl.d.bits.opcode := Mux(edge.hasData(acq.bits),
    TLMessages.AccessAck, TLMessages.AccessAckData)
  tl.d.bits.data := Mux(regsel, maskArr(wordaddr), dataArr(wordaddr))

  when (acq.fire && edge.hasData(acq.bits)) {
    when (regsel) {

  io.tcam.foreach { m =>
    val matches = dataArr.zip(maskArr).map { case (data, mask) =>
      !((m.data ^ data) & mask).orR
    }
    m.addr := PriorityEncoder(matches)
    m.found := matches.reduce(_ || _)
  }
}

      maskArr(wordaddr) := acq.bits.data(outer.dataBits - 1, 0)
    } .otherwise {
      dataArr(wordaddr) := acq.bits.data(outer.dataBits - 1, 0)


package icenet

import chisel3._
import chisel3.util._
import freechips.rocketchip.unittest.UnitTest
import freechips.rocketchip.util.UIntIsOneOf
import scala.math.max
import IceNetConsts._

class NetworkTap[T <: Data](
    selectFuncs: Seq[T => Bool],
    headerTyp: T = new EthernetHeader,
    headerBytes: Int = ETH_HEAD_BYTES,
    wordBytes: Int = NET_IF_WIDTH / 8) extends Module {

  val wordBits = wordBytes * 8
  val headerWords = headerBytes / wordBytes
  val n = selectFuncs.length

  val io = IO(new Bundle {
    val inflow = Flipped(Decoupled(new StreamChannel(wordBits)))
    val passthru = Decoupled(new StreamChannel(wordBits))
    val tapout = Vec(n, Decoupled(new StreamChannel(wordBits)))
  })

  assert(n > 0, "NetworkTap must have at least one output tap")

  val headerVec = Reg(Vec(headerWords, UInt(wordBits.W)))
  val header = headerVec.asTypeOf(headerTyp)

  val idxBits = log2Ceil(headerWords)
  val headerIdx = RegInit(0.U(idxBits.W))
  val headerLen = Reg(UInt(idxBits.W))
  val bodyLess = Reg(Bool())

  val (s_collect_header :: s_check_header ::
       s_output_header :: s_forward_body :: Nil) = Enum(4)
  val state = RegInit(s_collect_header)
  val route = Reg(Vec(n, Bool()))

  val selectedReady = MuxCase(io.passthru.ready,
    route.zip(io.tapout.map(_.ready)))

  io.inflow.ready := MuxLookup(state, false.B)(Seq(
    s_collect_header -> true.B,
    s_forward_body -> selectedReady))

  val hasTapout = route.reduce(_ || _)

  io.passthru.valid := MuxLookup(Cat(state, hasTapout), false.B)(Seq(
    Cat(s_output_header, false.B) -> true.B,
    Cat(s_forward_body,  false.B) -> io.inflow.valid))
  io.passthru.bits.data := MuxLookup(Cat(state, hasTapout), 0.U)(Seq(
    Cat(s_output_header, false.B) -> headerVec(headerIdx),
    Cat(s_forward_body,  false.B) -> io.inflow.bits.data))
  io.passthru.bits.last := MuxLookup(Cat(state, hasTapout), false.B)(Seq(
    Cat(s_output_header, false.B) -> (bodyLess && headerIdx === headerLen),
    Cat(s_forward_body,  false.B) -> io.inflow.bits.last))
  io.passthru.bits.keep := NET_FULL_KEEP

  io.tapout.zip(route).foreach { case (tapout, sel) =>
    tapout.valid := MuxLookup(Cat(state, sel), false.B)(Seq(
      Cat(s_output_header, true.B) -> true.B,
      Cat(s_forward_body,  true.B) -> io.inflow.valid))
    tapout.bits.data := MuxLookup(Cat(state, sel), 0.U)(Seq(
      Cat(s_output_header, true.B) -> headerVec(headerIdx),
      Cat(s_forward_body,  true.B) -> io.inflow.bits.data))
    tapout.bits.last := MuxLookup(Cat(state, sel), false.B)(Seq(
      Cat(s_output_header, true.B) -> (bodyLess && headerIdx === headerLen),
      Cat(s_forward_body,  true.B) -> io.inflow.bits.last))
    tapout.bits.keep := NET_FULL_KEEP
  }

  when (state === s_collect_header && io.inflow.valid) {
    headerIdx := headerIdx + 1.U
    headerVec(headerIdx) := io.inflow.bits.data

    val headerLast = headerIdx === (headerWords-1).U

    when (io.inflow.bits.last || headerLast) {
      headerLen := headerIdx
      headerIdx := 0.U
      bodyLess := io.inflow.bits.last

      when (headerLast) {
        state := s_check_header
      } .otherwise {
        route.foreach(_ := false.B)
        state := s_output_header
      }
    }
  }

  when (state === s_check_header) {
    route := selectFuncs.map(_(header))
    state := s_output_header
  }

  val headerFire = state === s_output_header && selectedReady
  val bodyFire = state === s_forward_body && io.inflow.fire

  when (headerFire) {
    headerIdx := headerIdx + 1.U
    when (headerIdx === headerLen) {
      headerIdx := 0.U
      state := Mux(bodyLess, s_collect_header, s_forward_body)
    }
  }

  when (bodyFire && io.inflow.bits.last) { state := s_collect_header }
}

class NetworkTapTest extends UnitTest {
  val sendPayloads = Seq(
    Seq(0L, 0x800L << 48, 23L, 13L, 56L, 12L),
    Seq(0L, 0x800L << 48),
    Seq(0L, 0x801L << 48, 22, 16),
    Seq(0L))
  val sendLengths = sendPayloads.map(_.size)
  val sendData = sendPayloads.flatten.map(l => BigInt(l))

  val genIn = Module(new PacketGen(sendLengths, sendData))
  genIn.io.start := io.start

  val tapPayloads = sendPayloads.take(2)
  val tapData = tapPayloads.flatten.map(l => BigInt(l))
  val tapKeep = Seq.fill(tapData.size) { 0xff }
  val tapLast = tapPayloads.flatMap {
<TARGET>
  }
  val checkTap = Module(new PacketCheck(tapData, tapKeep, tapLast))

  val passPayloads = sendPayloads.drop(2)
  val passData = passPayloads.flatten.map(l => BigInt(l))
  val passKeep = Seq.fill(passData.size) { 0xff }
  val passLast = passPayloads.flatMap {

  }
  val checkPass = Module(new PacketCheck(passData, passKeep, passLast))

  val tap = Module(new NetworkTap(
    Seq((header: EthernetHeader) => header.ethType === 0x800.U)))
  tap.io.inflow <> genIn.io.out
  checkTap.io.in <> tap.io.tapout(0)
  checkPass.io.in <> tap.io.passthru

  io.finished := checkTap.io.finished || checkPass.io.finished
}

pl => (Seq.fill(pl.size-1) { false } ++ Seq(true))