package nagicore

import chisel3._
import chisel3.util._
import _root_.circt.stage._

object Main extends App {
    val target = args(0)
    val build_dir = "./build"
    println(target)
    def exportVerilog(core: () => chisel3.RawModule): Unit = {
        println("Export Verilog Started")
        val chiselStageOption = Seq(
            chisel3.stage.ChiselGeneratorAnnotation(() => core()),
            CIRCTTargetAnnotation(CIRCTTarget.Verilog)
        )
        val firtoolOptions = Seq(
            // FirtoolOption("--lowering-options=disallowLocalVariables,locationInfoStyle=wrapInAtSquareBracket,noAlwaysComb"),
            FirtoolOption("--lowering-options=disallowLocalVariables,locationInfoStyle=wrapInAtSquareBracket,noAlwaysComb"),
//            FirtoolOption("--lowering-options=disallowLocalVariables,disallowPackedArrays,locationInfoStyle=wrapInAtSquareBracket,noAlwaysComb"),

            FirtoolOption("--split-verilog"),
            FirtoolOption("-o=" + build_dir),
            FirtoolOption("--disable-all-randomization"),
            FirtoolOption("--preserve-aggregate=none"),
        )
        val executeOptions = chiselStageOption ++ firtoolOptions
        val executeArgs = Array("-td", build_dir)
        (new ChiselStage).execute(executeArgs, executeOptions)
    }
    target match {
        case "NSCSCC" => {
            GlobalConfg.SIM = false
            exportVerilog(() => new nagicore.loongarch.nscscc2024.CoreNSCSCC)
        }
        // case "TEST" => {
        //     exportVerilog(() => new Module{
        //         val io = IO(new Bundle {
        //             val clk = Input(Clock())
        //         })
        //         val a = "h123".U
        //         val xbar = Module(new nagicore.unit.ip.axi_corssbar.AXI4XBar(32, 32, List((0, nagicore.unit.ip.axi_corssbar.Axi4RW.RW)), List(("0x80000000", "0x807FFFFF"))))
        //         xbar.io.masters <> DontCare
        //         xbar.io.slaves <> DontCare
        //         xbar.io.slaves(0).ar.addr := 2.U(32.W)
        //     })
        // }

object GlobalConfg{
    var SIM = true
}

        case _ => {
            exportVerilog(() => new nagicore.loongarch.nscscc2024.Core)
        }


package nagicore.loongarch.nscscc2024

import chisel3._
import chisel3.util._

trait Config{
    def XLEN = 32
    def GPR_NUM = 32
    def GPR_LEN = log2Up(GPR_NUM)

    def PC_START = "h80000000".U(XLEN.W)

    def ICACHE_WAYS = 2

    def ICACHE_WORDS = 4

    def WBUFF_LEN = 4

    def BTB_ENTRYS = 8

    def AXI4IDBITS = 4
}


package nagicore.loongarch.nscscc2024.stages

import chisel3._
import chisel3.util._
import nagicore.bus.AXI4IO
//import nagicore.unit.{InstrsBuff, InstrsBuffCacheBundle}
import nagicore.unit.cache.Cache
import nagicore.GlobalConfg
import nagicore.unit.cache.CacheReplaceType
import nagicore.unit.BTBPredOutIO
import nagicore.loongarch.nscscc2024.{Config, CtrlFlags}
import nagicore.bus.RamType

class if2idBits extends Bundle with Config{
    val pc          = UInt(XLEN.W)
    val bpu_out     = new BTBPredOutIO(BTB_ENTRYS, XLEN)
    val instr       = UInt(XLEN.W)

    val valid       = Bool()
}

class if2idIO extends Bundle{
    val bits = Output(new if2idBits)
    val stall = Input(Bool())
}

class IF extends Module with Config{
    val io = IO(new Bundle {
        val preif2if = Flipped(new preif2ifIO)
        val if2id = new if2idIO
        val isram = new AXI4IO(XLEN, XLEN)
    })
    // 2-stages 1cyc cache
    val icache = Module(new Cache(XLEN, XLEN, ICACHE_WAYS, ICACHE_LINES, ICACHE_WORDS, () => new preif2ifBits(), CacheReplaceType.LRU,
        dataRamType = RamType.RAM_1CYC,
        tagVRamType = RamType.RAM_1CYC,
        debug_id = 0))
    icache.io.axi <> io.isram

    icache.io.master.front.bits.addr := io.preif2if.bits.pc
    icache.io.master.front.bits.size := 2.U
    icache.io.master.front.bits.uncache := false.B
    icache.io.master.front.bits.wmask := 0.U
    icache.io.master.front.bits.valid := io.preif2if.bits.valid
    icache.io.master.front.bits.wdata := DontCare
    icache.io.master.front.bits.pipedata := io.preif2if.bits
    icache.io.master.back.stall := io.if2id.stall

    io.if2id.bits.instr := icache.io.master.back.bits.rdata
    io.if2id.bits.valid := icache.io.master.back.bits.valid
    io.if2id.bits.pc := icache.io.master.back.bits.pipedata.pc
    io.if2id.bits.bpu_out := icache.io.master.back.bits.pipedata.bpu_out

    io.preif2if.stall := icache.io.master.front.stall

    if(GlobalConfg.SIM){
        import nagicore.unit.DPIC_PERF_PIPE
        val perf_pipe_if = Module(new DPIC_PERF_PIPE())
        perf_pipe_if.io.clk := clock
        perf_pipe_if.io.rst := reset


package nagicore.loongarch.nscscc2024.stages

import chisel3._
import chisel3.util._
import nagicore.bus.AXI4IO
import nagicore.unit.cache.CacheMini
import nagicore.utils.Flags
import nagicore.GlobalConfg
import nagicore.unit.cache.CacheReplaceType
import nagicore.loongarch.nscscc2024.{Config, CtrlFlags}
import nagicore.unit.cache.UnCache

class mem2idIO extends Bundle with Config{
    // effective signal
    val bypass_rc   = Output(UInt(GPR_LEN.W))
    val bypass_val  = Output(UInt(XLEN.W))
    val bypass_en   = Output(Bool())

    val gpr_wid     = Output(UInt(GPR_LEN.W))
    val gpr_wdata   = Output(UInt(XLEN.W))
    val gpr_wen     = Output(Bool())
}

class MEM extends Module with Config{
    val io = IO(new Bundle {
        val ex2mem = Flipped(new ex2memIO())
        val mem2id = new mem2idIO()
        val dmem = new AXI4IO(XLEN, XLEN)
        val stall_all = Input(Bool())
    })

    class dcachePipeT extends Bundle {
        val instr       = UInt(XLEN.W)
        val alu_out     = UInt(XLEN.W)
        val rc          = UInt(GPR_LEN.W)
        val ld_type     = CtrlFlags.ldType()
        val pc          = UInt(XLEN.W)
        val no_ldst     = Bool()

        val valid       = Bool()
    }

    // val dcache = Module(new CacheMini(XLEN, XLEN, 8, 8, 1))
    val dcache = Module(new UnCache(XLEN, XLEN, WBUFF_LEN, 1))

    // pipeline registers
    val preg = RegEnable(io.ex2mem.bits, !dcache.io.out.busy && !io.stall_all)
    io.ex2mem.stall := dcache.io.out.busy || io.stall_all

    dcache.io.axi <> io.dmem

    val addr = preg.alu_out

    dcache.io.in.bits.addr := addr
    // dcache.io.in.bits.uncache := addr(31, 28) === "hb".U
    dcache.io.in.bits.we := !Flags.OHis(preg.st_type, CtrlFlags.stType.x)
    dcache.io.in.bits.wdata := Flags.onehotMux(preg.st_type, Seq(
        CtrlFlags.stType.x  -> 0.U,
        CtrlFlags.stType.b  -> Fill(XLEN/8, preg.rb_val(7, 0)),
        CtrlFlags.stType.h  -> Fill(XLEN/16, preg.rb_val(15, 0)),
        CtrlFlags.stType.w  -> preg.rb_val(31, 0),
    ))
    dcache.io.in.bits.size := Flags.onehotMux(preg.st_type, Seq(
        CtrlFlags.stType.x  -> 0.U,
        CtrlFlags.stType.b  -> 0.U,
        CtrlFlags.stType.h  -> 1.U,
        CtrlFlags.stType.w  -> 2.U,
    )) | Flags.onehotMux(preg.ld_type, Seq(
        CtrlFlags.ldType.x  -> 0.U,
        CtrlFlags.ldType.b  -> 0.U,
        CtrlFlags.ldType.bu -> 0.U,
        CtrlFlags.ldType.h  -> 1.U,
        CtrlFlags.ldType.hu -> 1.U,
        CtrlFlags.ldType.w  -> 2.U,
    ))
    dcache.io.in.bits.wmask := Flags.onehotMux(preg.st_type, Seq(
        CtrlFlags.stType.x  -> 0.U,
        CtrlFlags.stType.b  -> ("b1".U<<addr(1, 0)),
        CtrlFlags.stType.h  -> ("b11".U<<(addr(1)##0.U(1.W))),
        CtrlFlags.stType.w  -> "b1111".U,
    ))
    // 不走Cache的指令
    val nolr = Flags.OHis(preg.ld_type, CtrlFlags.ldType.x) && Flags.OHis(preg.st_type, CtrlFlags.stType.x)
    dcache.io.in.req := preg.valid && !nolr && RegNext(!dcache.io.out.busy) && !io.stall_all

    val rdata_raw = dcache.io.out.rdata

    val byteData = Mux(addr(0), halfData(15, 8), halfData(7, 0))

    val rdata_mem = Flags.onehotMux(preg.ld_type, Seq(
        CtrlFlags.ldType.x  -> (0.U).zext,
        CtrlFlags.ldType.b  -> byteData.asSInt,
        CtrlFlags.ldType.bu -> byteData.zext,
        CtrlFlags.ldType.h  -> halfData.asSInt,
        CtrlFlags.ldType.hu -> halfData.zext,
        CtrlFlags.ldType.w  -> wordData.zext,
    )).asUInt

    val mem_valid = preg.valid && !dcache.io.out.busy

    io.mem2id.bypass_rc := preg.rc
    io.mem2id.bypass_en := mem_valid
    val wb_data = Mux(Flags.OHis(preg.ld_type, CtrlFlags.ldType.x), preg.alu_out, rdata_mem)
    io.mem2id.bypass_val := wb_data

    // when(nolr){
    //     io.mem2id.bypass_rc := Mux(preg.valid, preg.rc, 0.U)
    //     io.mem2id.bypass_val := preg.alu_out
    // }.otherwise{
    //     io.mem2id.bypass_rc := Mux(preg.valid && preg.ld_type === Flags.bp(CtrlFlags.ldType.w), preg.rc, 0.U)
    //     io.mem2id.bypass_val := dcache.io.out.rdata
    // }

    io.mem2id.gpr_wid := preg.rc
    io.mem2id.gpr_wdata := wb_data
    io.mem2id.gpr_wen := mem_valid

    if(GlobalConfg.SIM){
        import nagicore.unit.DPIC_TRACE_MEM
        val dpic_trace_mem_w = Module(new DPIC_TRACE_MEM(XLEN, XLEN))
        dpic_trace_mem_w.io.clk := clock
        dpic_trace_mem_w.io.rst := reset
        dpic_trace_mem_w.io.valid := dcache.io.in.req && dcache.io.in.bits.wmask.orR
        dpic_trace_mem_w.io.addr := dcache.io.in.bits.addr
        dpic_trace_mem_w.io.size := dcache.io.in.bits.size
        dpic_trace_mem_w.io.data := dcache.io.in.bits.wdata
        dpic_trace_mem_w.io.wmask := dcache.io.in.bits.wmask

        import nagicore.unit.DPIC_PERF_PIPE
        val perf_pipe_dcache = Module(new DPIC_PERF_PIPE())
        perf_pipe_dcache.io.clk := clock
        perf_pipe_dcache.io.rst := reset
        perf_pipe_dcache.io.id := 2.U
        perf_pipe_dcache.io.invalid := !mem_valid
        perf_pipe_dcache.io.stall := io.ex2mem.stall

        import nagicore.unit.DPIC_UPDATE_PC
        val dpic_update_pc = Module(new DPIC_UPDATE_PC(XLEN))
        dpic_update_pc.io.clk := clock
        dpic_update_pc.io.rst := reset
        dpic_update_pc.io.pc := preg.pc
        dpic_update_pc.io.wen := mem_valid

        import nagicore.unit.DPIC_TRACE_MEM
        val dpic_trace_mem_r = Module(new DPIC_TRACE_MEM(XLEN, XLEN))
        dpic_trace_mem_r.io.clk := clock
        dpic_trace_mem_r.io.rst := reset
        dpic_trace_mem_r.io.valid := mem_valid && preg.ld_type =/= Flags.bp(CtrlFlags.ldType.x)
        dpic_trace_mem_r.io.addr := preg.alu_out
        dpic_trace_mem_r.io.size := Flags.onehotMux(preg.ld_type, Seq(
            CtrlFlags.ldType.x  -> 0.U,
            CtrlFlags.ldType.b  -> 0.U,
            CtrlFlags.ldType.bu -> 0.U,
            CtrlFlags.ldType.h  -> 1.U,
            CtrlFlags.ldType.hu -> 1.U,
            CtrlFlags.ldType.w  -> 2.U,
        ))
        dpic_trace_mem_r.io.data := rdata_mem
        dpic_trace_mem_r.io.wmask := 0.U

    }
}

    val wordData = if(XLEN == 64) Mux(addr(2), rdata_raw(63, 32), rdata_raw(31, 0))
                    else rdata_raw(31, 0)


package nagicore.loongarch.nscscc2024Dual

import chisel3._
import chisel3.util._

trait Config{
    def XLEN = 32
    def GPR_NUM = 32
    def GPR_LEN = log2Up(GPR_NUM)

    def PC_START = "h80000000".U(XLEN.W)

    def ICACHE_WAYS = 2

    def ICACHE_WORDS = 4

    def WBUFF_LEN = 8

    def BTB_ENTRYS = 8

    def AXI4IDBITS = 4
}


package nagicore.loongarch.nscscc2024Dual.stages

import chisel3._
import chisel3.util._
import nagicore.unit.GPR
import nagicore.unit.ALU_OP
import nagicore.unit.BR_TYPE
import nagicore.unit.BTBPredOutIO
import nagicore.loongarch.nscscc2024Dual.{Config, CtrlFlags, Decoder}
import nagicore.GlobalConfg
import nagicore.unit.RingBuff
import nagicore.utils.Flags

class is2exBits extends Bundle with Config{
    val instr1      = UInt(XLEN.W)
    val pc1         = UInt(XLEN.W)
    val ra1_val     = UInt(XLEN.W)
    val alu1A_sel   = CtrlFlags.aluASel()
    val rb1_val     = UInt(XLEN.W)
    val alu1B_sel   = CtrlFlags.aluBSel()
    val alu1_op     = ALU_OP()
    val rc1         = UInt(GPR_LEN.W)
    val imm1        = UInt(XLEN.W)
    val pc_refill1  = Bool()

    val br_type     = BR_TYPE()
    val brpcAdd_sel = CtrlFlags.brpcAddSel()
    val ld_type     = CtrlFlags.ldType()
    val st_type     = CtrlFlags.stType()
    val bpu_out     = new BTBPredOutIO(BTB_ENTRYS, XLEN)

    val valid1       = Bool()

    val instr2      = UInt(XLEN.W)
    val pc2         = UInt(XLEN.W)
    val ra2_val     = UInt(XLEN.W)
    val alu2A_sel   = CtrlFlags.aluASel()
    val rb2_val     = UInt(XLEN.W)
    val alu2B_sel   = CtrlFlags.aluBSel()
    val alu2_op     = ALU_OP()
    val rc2         = UInt(GPR_LEN.W)
    val imm2        = UInt(XLEN.W)
    val pc_refill2  = Bool()

    val valid2       = Bool()
}

class is2exIO extends Bundle{
    val bits = Output(new is2exBits)
    val stall = Input(Bool())
}

class IS extends Module with Config{
    val io = IO(new Bundle{
        val id2is = Flipped(new id2isIO)
        val is2ex = new is2exIO

        val ex2is = Flipped(new ex2isIO)
        val mem2is = Flipped(new mem2isIO)
    })

    val issue_buffer = Module(new RingBuff(()=>new id2isBits, 8, rchannel=2, debug_id=0))

    issue_buffer.io.push := io.id2is.bits.valid
    issue_buffer.io.wdata := io.id2is.bits
    issue_buffer.io.clear := io.ex2is.clear_is

    val is1 = issue_buffer.io.rdatas(0)

    issue_buffer.io.popN := issue_double

//    io.id2is.stall := io.is2ex.stall
    io.id2is.stall := issue_buffer.io.full

    val gpr     = Module(new GPR(XLEN, GPR_NUM, 4, 2))
    gpr.io.wen(0) := io.mem2is.gpr_wen1 && (!io.mem2is.gpr_wen2 || io.mem2is.gpr_wid2 =/= io.mem2is.gpr_wid1)
    gpr.io.waddr(0) := io.mem2is.gpr_wid1
    gpr.io.wdata(0) := io.mem2is.gpr_wdata1

    gpr.io.wen(1) := io.mem2is.gpr_wen2
    gpr.io.waddr(1) := io.mem2is.gpr_wid2
    gpr.io.wdata(1) := io.mem2is.gpr_wdata2

    if(GlobalConfg.SIM){
        import nagicore.unit.DPIC_UPDATE_GPR2
        val dpic_update_gpr = Module(new DPIC_UPDATE_GPR2(XLEN, GPR_NUM))
        dpic_update_gpr.io.clk := clock
        dpic_update_gpr.io.rst := reset

        dpic_update_gpr.io.id1 := gpr.io.waddr(0)
        dpic_update_gpr.io.wen1 := gpr.io.wen(0)
        dpic_update_gpr.io.wdata1 := gpr.io.wdata(0)

        dpic_update_gpr.io.id2 := gpr.io.waddr(1)
        dpic_update_gpr.io.wen2 := gpr.io.wen(1)
        dpic_update_gpr.io.wdata2 := gpr.io.wdata(1)
    }

    def bypass_unit(rx: UInt, gpr_rdata: UInt):UInt = {
        Mux(rx === 0.U, 0.U,
            Mux(io.ex2is.bypass_rc2 === rx && io.ex2is.bypass_en2, io.ex2is.bypass_val2,
                Mux(io.ex2is.bypass_rc1 === rx && io.ex2is.bypass_en1, io.ex2is.bypass_val1,
                    Mux(io.mem2is.bypass_rc2 === rx && io.mem2is.bypass_en2, io.mem2is.bypass_val2,
                        Mux(io.mem2is.bypass_rc1 === rx && io.mem2is.bypass_en1, io.mem2is.bypass_val1,
                            gpr_rdata
                        )
                    )
                )
            )
        )
    }

    gpr.io.raddr(0) := is1.ra
    io.is2ex.bits.ra1_val := bypass_unit(is1.ra, gpr.io.rdata(0))
    io.is2ex.bits.alu1A_sel := is1.aluA_sel
    gpr.io.raddr(1) := is1.rb
    io.is2ex.bits.rb1_val := bypass_unit(is1.rb, gpr.io.rdata(1))
    io.is2ex.bits.alu1B_sel := is1.aluB_sel

    gpr.io.raddr(2) := is2.ra
    io.is2ex.bits.ra2_val := bypass_unit(is2.ra, gpr.io.rdata(2))
    io.is2ex.bits.alu2A_sel := is2.aluA_sel
    gpr.io.raddr(3) := is2.rb
    io.is2ex.bits.rb2_val := bypass_unit(is2.rb, gpr.io.rdata(3))
    io.is2ex.bits.alu2B_sel := is2.aluB_sel

    io.is2ex.bits.instr1 := is1.instr
    io.is2ex.bits.instr2 := is2.instr

    io.is2ex.bits.pc1 := is1.pc
    io.is2ex.bits.pc2 := is2.pc

    io.is2ex.bits.alu1_op := is1.alu_op
    io.is2ex.bits.alu2_op := is2.alu_op

    io.is2ex.bits.rc1 := is1.rc
    io.is2ex.bits.rc2 := is2.rc

    io.is2ex.bits.imm1 := is1.imm
    io.is2ex.bits.imm2 := is2.imm

    io.is2ex.bits.br_type := is1.br_type
    io.is2ex.bits.brpcAdd_sel := is1.brpcAdd_sel
    io.is2ex.bits.ld_type := is1.ld_type
    io.is2ex.bits.st_type := is1.st_type
    io.is2ex.bits.bpu_out := is1.bpu_out

    io.is2ex.bits.pc_refill1 := is1.pc_refill
    io.is2ex.bits.pc_refill2 := is2.pc_refill

    io.is2ex.bits.valid1 := issue_buffer.io.rvalids(0)
    io.is2ex.bits.valid2 := issue_buffer.io.rvalids(1) && issue_double
}

    val is2 = issue_buffer.io.rdatas(1)
    val data_hazard = (is1.rc === is2.ra || is1.rc === is2.rb) && is1.rc =/= 0.U
    // 只双发is2是ALU类，且无数据冒险的指令
    val issue_double =
        Flags.is(is2.instr_type, CtrlFlags.InstrType.alu) &&


package nagicore.unit

import chisel3._
import chisel3.util._
import nagicore.utils.Flags

object DIVU_IMP extends Enumeration {
    type DIVU_IMP = Value
    val none, radix2 = Value
}

class DIVU(dataBits: Int, imp_way: DIVU_IMP.DIVU_IMP = DIVU_IMP.radix2) extends Module{
    val io = IO(new Bundle{
        val a       = Input(UInt(dataBits.W))
        val b       = Input(UInt(dataBits.W))
        val signed  = Input(Bool())
        val quo     = Output(UInt(dataBits.W))
        val rem     = Output(UInt(dataBits.W))
        val valid   = Input(Bool())
        val busy    = Output(Bool())
    })

    imp_way match {

            val sign_r = Mux(io.signed, io.a(dataBits-1), false.B)
            val src1 = Mux(io.signed && io.a(dataBits-1), ~io.a + 1.U, io.a)
            val src2 = Mux(io.signed && io.b(dataBits-1), ~io.b + 1.U, io.b)

            // get highest 1 in src1
            // TODO use log2
            val high_rev = PriorityEncoder(Reverse(src1))

            val cnt             = RegInit(0.U(6.W))
            val stage1_fire     = cnt === 0.U

            val src1_reg1       = RegEnable(src1, stage1_fire)
            val src2_reg1       = RegEnable(src2, stage1_fire)
            val signed_reg1     = RegEnable(io.signed, stage1_fire)
            val sign_s_reg1     = RegEnable(sign_s, stage1_fire)
            val sign_r_reg1     = RegEnable(sign_r, stage1_fire)
            val en_reg1         = RegEnable(io.valid, stage1_fire)
            val high_rev_reg1   = RegEnable(high_rev, stage1_fire)

            /* stage2+: div */
            val stage2_init     = en_reg1 && cnt === 0.U

            val src2_reg2       = RegEnable(src2_reg1, stage2_init)
            val signed_reg2     = RegEnable(signed_reg1, stage2_init)
            val sign_s_reg2     = RegEnable(sign_s_reg1, stage2_init)
            val sign_r_reg2     = RegEnable(sign_r_reg1, stage2_init)

            when(cnt =/= 0.U){
                cnt := cnt - 1.U
            }.elsewhen(en_reg1){
                cnt := (dataBits+1).U - high_rev_reg1
            }

            val quo_rem_reg = RegInit(0.U((dataBits*2+1).W))
            val quo = quo_rem_reg(dataBits-1, 0)
            val rem = quo_rem_reg(dataBits*2-1, dataBits)
            when(cnt =/= 0.U){
                val mins = rem - src2_reg2
                when(rem >= src2_reg2){
                    quo_rem_reg := mins(dataBits-1, 0) ## quo ## 1.U(1.W)
                }.otherwise{
                    quo_rem_reg := quo_rem_reg(dataBits*2-1, 0) ## 0.U(1.W)
                }
            }.elsewhen(en_reg1){
                quo_rem_reg := (0.U((dataBits+1).W) ## src1_reg1) << high_rev_reg1
            }

            io.busy := cnt =/= 0.U || en_reg1

            io.quo := Mux(signed_reg2,
                Mux(sign_s_reg2, ~quo + 1.U, quo),
                quo
            )

            val rem_res = quo_rem_reg(dataBits*2, dataBits+1)
            io.rem := Mux(signed_reg2,
                Mux(sign_r_reg2, ~rem_res + 1.U, rem_res),
                rem_res
            )
        }
        case DIVU_IMP.none => {
            io.busy := false.B
            io.quo := DontCare
            io.rem := DontCare
        }
    }
}

/*
class DIVU(dataBits: Int) extends Module{
    val io = IO(new Bundle{
        val a       = Input(UInt(dataBits.W))
        val b       = Input(UInt(dataBits.W))
        val signed  = Input(Bool())
        val quo     = Output(UInt(dataBits.W))
        val rem     = Output(UInt(dataBits.W))
        val valid   = Input(Bool())
        val busy   = Output(Bool())
    })

    /* ref: https://github.com/MaZirui2001/LAdataBitsR-pipeline-scala */

    /* stage1: solve sign */
    val sign_s = Mux(io.signed, io.a(dataBits-1) ^ io.b(dataBits-1), false.B)
    val sign_r = Mux(io.signed, io.a(dataBits-1), false.B)
    val src1 = Mux(io.signed && io.a(dataBits-1), ~io.a + 1.U, io.a)
    val src2 = Mux(io.signed && io.b(dataBits-1), ~io.b + 1.U, io.b)

    // get highest 1 in src1
    val high_rev = PriorityEncoder(Reverse(src1))

    val src1_reg1       = ShiftRegister(src1, 1, !io.busy)
    val src2_reg1       = ShiftRegister(src2, 1, !io.busy)
    val signed_reg1     = ShiftRegister(io.signed, 1, !io.busy)
    val sign_s_reg1     = ShiftRegister(sign_s, 1, !io.busy)
    val sign_r_reg1     = ShiftRegister(sign_r, 1, !io.busy)
    val en_reg1         = ShiftRegister(io.valid, 1, !io.busy)
    val high_rev_reg1   = ShiftRegister(high_rev, 1, !io.busy)

    /* stage2+: div */
    val cnt             = RegInit(0.U(6.W))
    val stage2_init     = en_reg1 && cnt === 0.U

    val src2_reg2       = RegEnable(src2_reg1, stage2_init)
    val signed_reg2     = RegEnable(signed_reg1, stage2_init)
    val sign_s_reg2     = RegEnable(sign_s_reg1, stage2_init)
    val sign_r_reg2     = RegEnable(sign_r_reg1, stage2_init)

    when(cnt =/= 0.U){
        cnt := cnt - 1.U
    }.elsewhen(en_reg1){
        cnt := (dataBits+1).U - high_rev_reg1
    }

    val quo_rem_reg = RegInit(0.U((dataBits*2+1).W))
    val quo = quo_rem_reg(dataBits-1, 0)
    val rem = quo_rem_reg(dataBits*2-1, dataBits)
    when(cnt =/= 0.U){
        val mins = rem - src2_reg2
        when(rem >= src2_reg2){
            quo_rem_reg := mins ## quo ## 1.U(1.W)
        }.otherwise{
            quo_rem_reg := quo_rem_reg(dataBits*2-1, 0) ## 0.U(1.W)
        }
    }.elsewhen(en_reg1){
        quo_rem_reg := (0.U((dataBits+1).W) ## src1_reg1) << high_rev_reg1
    }

    !io.busy := cnt === 0.U

    io.quo := Mux(signed_reg2,
        Mux(sign_s_reg2, ~quo + 1.U, quo),
        quo
    )

    io.rem := Mux(signed_reg2,
        Mux(sign_r_reg2, ~quo_rem_reg(dataBits*2, dataBits+1) + 1.U, quo_rem_reg(dataBits*2, dataBits+1)),
        quo_rem_reg(dataBits*2, dataBits+1)
    )
}
 */

        case DIVU_IMP.radix2 => {
            /* ref: https://github.com/MaZirui2001/LAdataBitsR-pipeline-scala */


package nagicore.unit

import chisel3._
import chisel3.util._

class GPRIO(dataBits: Int, addrBits: Int, rchannel: Int, wchannel: Int) extends Bundle {
    val raddr   = Input(Vec(rchannel, UInt(addrBits.W)))
    val rdata   = Output(Vec(rchannel, UInt(dataBits.W)))

}

class GPR(dataBits: Int, regNum: Int, rchannel: Int, wchannel: Int) extends Module {
    val io = IO(new GPRIO(dataBits, log2Up(regNum), rchannel, wchannel))
    val regs = Reg(Vec(regNum, UInt(dataBits.W)))
    // val regs = Reg(VecInit.fill(regNum)(0.U(dataBits.W)))
    // val regs = Mem(regNum, UInt(dataBits.W))
    for(i <- 0 until rchannel){
        io.rdata(i) := Mux(io.raddr(i) =/= 0.U, regs(io.raddr(i)), 0.U)
    }
    for(i <- 0 until wchannel){
        when(io.wen(i) && io.waddr(i) =/= 0.U){
            regs(io.waddr(i)) := io.wdata(i)
        }
    }
}

    val wen     = Input(Vec(wchannel, Bool()))
    val waddr   = Input(Vec(wchannel, UInt(addrBits.W)))


package nagicore.unit.cache

import chisel3._
import chisel3.util._

import nagicore.bus._
import chisel3.util.random.LFSR
import nagicore.utils.isPowerOf2
import nagicore.GlobalConfg

class CachePipedIOFrontBits[T <: Bundle](addrBits: Int, dataBits: Int, pipedataT: () => T) extends Bundle{
    val valid     = Bool()
    val addr      = UInt(addrBits.W)
    val wmask     = UInt((dataBits/8).W)
    val size      = UInt(2.W)
    val wdata     = UInt(dataBits.W)
    val uncache   = Bool()
    val pipedata  = pipedataT()
}

class CachePipedIOFront[T <: Bundle](addrBits: Int, dataBits: Int, pipedataT: () => T) extends Bundle{
    val bits        = Input(new CachePipedIOFrontBits[T](addrBits, dataBits, pipedataT))
    val stall       = Output(Bool())
}

class CachePipedIOBackBits[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val valid       = Bool()
    val rdata       = UInt(dataBits.W)
    // 整个缓存块，注意uncahce访问时，结果无意义
    val rline       = Vec(blockWords, UInt(dataBits.W))
    val pipedata_s1 = pipedataT()
    val pipedata_s2 = pipedataT()
    // for debug
    val addr        = UInt(addrBits.W)
    val size        = UInt(2.W)
    val wdata       = UInt(dataBits.W)
    val wmask       = UInt((dataBits/8).W)
}

class CachePipedIOBack[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val bits        = Output(new CachePipedIOBackBits[T](addrBits, dataBits, blockWords, pipedataT))
    val stall       = Input(Bool())
}

class CachePipedIO[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val front = new CachePipedIOFront[T](addrBits, dataBits, pipedataT)
    val back = new CachePipedIOBack[T](addrBits, dataBits, blockWords, pipedataT)
}

/**
  * 阻塞式两级流水线Cache（已弃用）
  * @pipeline
  *                   stage1             stage2
  * EX                -> MEM1            -> MEM2           -> WB
  * 向SRAM发出地址      读SRAM和TLB         比较Tag，选择数据   字节选择，符号扩展
  *                       |preg1|------->|preg2|
  *         |-------------------|
  * ADDR--->| Cache RAM(2 cyc)  |------->|rtag, rdata ...|
  *         |-------------------|
  *
  * @param addrBits
  * @param dataBits
  * @param ways         相连数，必须为2的幂次
  * @param sets         行数
  * @param blockWords   块字个数
  * @param pipedataT    流水线其他信号
  */
class CachePiped[T <: Bundle](addrBits: Int, dataBits: Int, ways: Int, sets: Int, blockWords: Int, pipedataT: () => T, id: Int=0, replaceT: CacheReplaceType.CacheReplaceType=CacheReplaceType.Random) extends Module{
    require(isPowerOf2(ways))
    require(isPowerOf2(dataBits))

    val io = IO(new Bundle{
        val axi = new AXI4IO(addrBits, dataBits)
        val master = new CachePipedIO(addrBits, dataBits, blockWords, pipedataT)
    })

    // Block address
    // [ Tag | Index | Word | Byte ]
    //    |      |      |      |
    //   line   set
    val len_byte = log2Ceil(dataBits/8)
    val num_word = blockWords
    val len_word = log2Ceil(num_word)
    val len_idx = log2Ceil(sets)
    val len_tag = addrBits - len_idx - len_word - len_byte
    val beats_len = num_word

    val axi_w_agent = Module(new AXI4WriteAgent(addrBits, dataBits, beats_len))
    axi_w_agent.io.axi.aw <> io.axi.aw
    axi_w_agent.io.axi.w <> io.axi.w
    axi_w_agent.io.axi.b <> io.axi.b
    axi_w_agent.io.cmd.in.req := false.B
    axi_w_agent.io.cmd.in.addr := 0.U
    axi_w_agent.io.cmd.in.wdata := VecInit.fill(blockWords)(0.U)
    axi_w_agent.io.cmd.in.wmask := VecInit.fill(blockWords)(0.U)
    axi_w_agent.io.cmd.in.size := 0.U
    axi_w_agent.io.cmd.in.len := 0.U

    val axi_r_agent = Module(new AXI4ReadAgent(addrBits, dataBits, beats_len))
    axi_r_agent.io.axi.ar <> io.axi.ar
    axi_r_agent.io.axi.r <> io.axi.r
    axi_r_agent.io.cmd.in.req := false.B
    axi_r_agent.io.cmd.in.addr := 0.U
    axi_r_agent.io.cmd.in.len := 0.U
    axi_r_agent.io.cmd.in.size := 0.U

    val random_way = LFSR(16)(log2Up(ways)-1, 0)
    val lru = Reg(Vec(sets, UInt(log2Up(ways).W)))

    val active_way = Reg(UInt(log2Up(ways).W))

    val data_bank = Seq.fill(ways)(Seq.fill(num_word)(Module(new CacheMem(dataBits, sets))))
    val data_bank_io = VecInit(data_bank.map(t => VecInit(t.map(_.io))))
    val tag_v = Seq.fill(ways)(Module(new CacheMem(len_tag+1, sets)))
    val tag_v_io = VecInit(tag_v.map(_.io))
    val dirty = Seq.fill(ways)(Module(new CacheMem(1, sets)))
    val dirty_io = VecInit(dirty.map(_.io))
    val addr_idx = io.master.front.bits.addr(len_idx+len_word+len_byte-1, len_word+len_byte)

    // 流水使能信号
    val pipego = Wire(Bool())

    tag_v_io.map(io => {
        io.addr := addr_idx
        io.en := pipego
        io.re := true.B
        io.we := false.B
        io.din := 0.U
        io.wmask := Fill(len_tag+1, true.B)
    })
    data_bank_io.map(_.map(io => {
        io.addr := addr_idx
        io.en := pipego
        io.re := true.B
        io.we := false.B
        io.din := 0.U
        io.wmask := Fill(dataBits, true.B)
    }))
    dirty_io.map(io => {
        io.addr := addr_idx
        io.en := pipego
        io.re := true.B
        io.we := false.B
        io.din := 0.U
        io.wmask := Fill(1, true.B)
    })

    // stage1 pipeline registers
    val preg1 = RegEnable(io.master.front.bits, pipego)

    // stage2 pipeline registers
    val preg2 = RegEnable(preg1, pipego)
    object Stage2State extends ChiselEnum {
        //  0       1          2        3           4            5            6
        val lookup, writeback, replace, replaceEnd, uncacheWait, uncacheRead, uncacheEnd = Value
    }
    // stage2 state
    val state_s2 = RegInit(Stage2State.lookup)
    val hit = Wire(Bool())

    val pipego_reg = pipego

            tag_v_io(i).dout(0).asBool
        }, pipego_reg)
    val rdirty = RegEnable(VecInit.tabulate(ways){ i =>
            dirty_io(i).dout.asBool
        }, pipego_reg)

    pipego :=
        (
        // 连续命中
        (state_s2 === Stage2State.lookup && hit) ||
        // 无效命令
        !preg2.valid ||
        // 替换完成
        state_s2 === Stage2State.replaceEnd ||
        // uncache访问
        state_s2 === Stage2State.uncacheEnd
        ) &&
        // 下一级无阻塞请求
        !io.master.back.stall

    val addr_s2 = preg2.addr
    // TODO: VIPT, let TLB pass in real tag
    val addr_tag_s2 = addr_s2(addrBits-1, len_idx+len_word+len_byte)
    val addr_idx_s2 = addr_s2(len_idx+len_word+len_byte-1, len_word+len_byte)
    val addr_word_s2 = addr_s2(len_word+len_byte-1, len_byte)

    val rdatas_hit = WireDefault(VecInit.fill(blockWords)(0.U(dataBits.W)))
    val rdatas_replace = RegInit(VecInit.fill(blockWords)(0.U(dataBits.W)))
    val rdata_uncache = RegInit(0.U(dataBits.W))

    io.master.front.stall := !pipego
    io.master.back.bits.valid := pipego && preg2.valid
//    io.master.back.bits.rdata := Mux(state_s2 === Stage2State.replaceEnd, rdata_replace, rdata_hit)
    io.master.back.bits.rdata := MuxCase(rdatas_hit(addr_word_s2), Seq(
        (state_s2 === Stage2State.replaceEnd) -> rdatas_replace(addr_word_s2),
        (state_s2 === Stage2State.uncacheEnd) -> rdata_uncache
    ))
    io.master.back.bits.rline := Mux(state_s2 === Stage2State.replaceEnd, rdatas_replace,
        rdatas_hit
    )
    io.master.back.bits.pipedata_s1 := preg1.pipedata
    io.master.back.bits.pipedata_s2 := preg2.pipedata

    io.master.back.bits.addr := preg2.addr
    io.master.back.bits.size := preg2.size
    io.master.back.bits.wdata := preg2.wdata
    io.master.back.bits.wmask := preg2.wmask

    val bypass_addr = Reg(UInt((len_tag+len_idx).W))
    val bypass_val = Reg(Vec(blockWords, UInt(dataBits.W)))
    val bypass_valid = Reg(Vec(blockWords, Bool()))
    for(i <- 0 until blockWords){
        bypass_valid(i) := 0.U
    }

    val hits = VecInit.tabulate(ways){ i =>
        rtags(i) === addr_tag_s2 && rvalid(i)
    }
    hit := hits.reduceTree(_||_) && !preg2.uncache

    if(GlobalConfg.SIM){
        import  nagicore.unit.DPIC_PERF_CACHE
        val dpic_perf_cache = Module(new DPIC_PERF_CACHE)
        dpic_perf_cache.io.clk := clock
        dpic_perf_cache.io.rst := reset
        dpic_perf_cache.io.valid := preg2.valid && state_s2 === Stage2State.lookup && !io.master.back.stall
        dpic_perf_cache.io.id := id.U
        dpic_perf_cache.io.access_type := Cat(preg2.uncache, hit)
    }

    // 请求读缺失行，并开始替换Cache行
    def load_miss_lines() = {
        axi_r_agent.io.cmd.in.req := true.B
        axi_r_agent.io.cmd.in.addr := addr_tag_s2 ## addr_idx_s2 ## 0.U((len_word+len_byte).W)
        axi_r_agent.io.cmd.in.len := (blockWords-1).U
        axi_r_agent.io.cmd.in.size := log2Up(dataBits).U
        state_s2 := Stage2State.replace
    }

    def writeSyncRam(io: RamIO, addr: UInt, data: UInt) = {
        io.addr := addr
        io.en := true.B
        io.re := false.B
        io.we := true.B
        io.din := data
    }

    def wait_uncache_ready() = {
        when(preg2.wmask.orR){
            when(axi_w_agent.io.cmd.out.ready){
                axi_w_agent.io.cmd.in.addr := preg2.addr
                axi_w_agent.io.cmd.in.req := true.B
                axi_w_agent.io.cmd.in.wdata(0) := preg2.wdata
                axi_w_agent.io.cmd.in.wmask(0) := preg2.wmask
                axi_w_agent.io.cmd.in.size := preg2.size
                axi_w_agent.io.cmd.in.len := 0.U
                state_s2 := Stage2State.uncacheEnd
            }
        }.otherwise{
            when(axi_w_agent.io.cmd.out.ready&&axi_r_agent.io.cmd.out.ready){
                axi_r_agent.io.cmd.in.req := true.B
                axi_r_agent.io.cmd.in.addr := preg2.addr
                axi_r_agent.io.cmd.in.size := preg2.size
                axi_r_agent.io.cmd.in.len := 0.U
                state_s2 := Stage2State.uncacheRead
            }
        }
    }

    switch(state_s2){
        is(Stage2State.lookup){
            when(preg2.valid){
                when(hit){
                    val hit_way = PriorityEncoder(hits)
                    if(replaceT == CacheReplaceType.LRU){
                        lru(addr_idx_s2) := ~hit_way
                    }
                    // 当连续hit wr/ww同一Cache行时，需要前递
                    val hit_rw_bypass_need = bypass_addr === Cat(addr_tag_s2, addr_idx_s2)
                    val rdatas_real = VecInit.tabulate(blockWords){
                        i => Mux(hit_rw_bypass_need && bypass_valid(i),
                            bypass_val(i),
                            rdatas(hit_way)(i)
                        )
                    }
                    when(preg2.wmask.orR){
                        // 命中写
                        val wdata = Cat((0 until (dataBits/8)).reverse.map(i =>
                            Mux(preg2.wmask(i),
                                preg2.wdata(i*8+7, i*8),
                                rdatas_real(addr_word_s2)(i*8+7, i*8)
                            )))
                        writeSyncRam(data_bank_io(hit_way)(addr_word_s2), addr_idx_s2, wdata)
                        bypass_addr := Cat(addr_tag_s2, addr_idx_s2)
                        bypass_valid(addr_word_s2) := true.B
                        bypass_val(addr_word_s2) := wdata
                        writeSyncRam(dirty_io(hit_way), addr_idx_s2, 1.U(1.W))
                    }otherwise{
                        // 命中读
                        rdatas_hit := rdatas_real
                    }
                    state_s2 := Stage2State.lookup
                }.otherwise{
                    when(preg2.uncache){
                        wait_uncache_ready()
                        // or just state_s2 := Stage2State.uncacheWait
                    }.otherwise{
                        val active_way_wire = Wire(chiselTypeOf(active_way))
                        if(replaceT == CacheReplaceType.LRU){
                            active_way_wire := lru(addr_idx_s2)
                        }else{
                            active_way_wire := random_way
                        }
                        active_way := active_way_wire
                        when(rdirty(active_way_wire) && rvalid(active_way_wire)) {
                            state_s2 := Stage2State.writeback
                        }.otherwise {
                            load_miss_lines()
                        }
                    }
                }
            }
        }
        is(Stage2State.writeback){
            // 等待代理空闲
            when(axi_w_agent.io.cmd.out.ready){
                // 交给代理写入将要被替换的脏行
                val addr = rtags(active_way) ## addr_idx_s2 ## 0.U((len_word+len_byte).W)
                axi_w_agent.io.cmd.in.addr := addr
                axi_w_agent.io.cmd.in.req := true.B
                axi_w_agent.io.cmd.in.wdata := rdatas(active_way)
                axi_w_agent.io.cmd.in.wmask := VecInit.fill(blockWords)(1.U)
                axi_w_agent.io.cmd.in.size := log2Up(dataBits).U
                axi_w_agent.io.cmd.in.len := (blockWords-1).U

                load_miss_lines()
            }.otherwise{
                state_s2 := Stage2State.writeback
            }
        }
        is(Stage2State.replace){
            when(axi_r_agent.io.cmd.out.ready){
                val word_i = axi_r_agent.io.cmd.out.order
                val wdata_active = Mux(word_i === addr_word_s2,
                    Cat((0 until (dataBits/8)).reverse.map(i =>
                        Mux(preg2.wmask(i),
                            preg2.wdata(i*8+7, i*8),
                            axi_r_agent.io.cmd.out.rdata(i*8+7, i*8)
                        )
                    )),
                    axi_r_agent.io.cmd.out.rdata
                )

                writeSyncRam(data_bank_io(active_way)(word_i), addr_idx_s2, wdata_active)
                bypass_val(word_i) := wdata_active

                rdatas_replace(word_i) := wdata_active

                when(axi_r_agent.io.cmd.out.last){
                    state_s2 := Stage2State.replaceEnd
                    writeSyncRam(tag_v_io(active_way), addr_idx_s2, addr_tag_s2 ## 1.U(1.W))
                    writeSyncRam(dirty_io(active_way), addr_idx_s2, 0.U(1.W))
                    bypass_addr := Cat(addr_tag_s2, addr_idx_s2)
                }
            }
        }
        // 等待下一级准备好接受
        is(Stage2State.replaceEnd){
            when(!io.master.back.stall){
                state_s2 := Stage2State.lookup
                for(i <- 0 until blockWords){
                    bypass_valid(i) := true.B
                }
            }
        }
        is(Stage2State.uncacheWait){
            wait_uncache_ready()
        }
        is(Stage2State.uncacheRead){
            // 等待代理空闲
            when(axi_r_agent.io.cmd.out.ready){
                rdata_uncache := axi_r_agent.io.cmd.out.rdata
                state_s2 := Stage2State.uncacheEnd
            }
        }
        is(Stage2State.uncacheEnd){
            // 等待下一级准备好接受
            when(!io.master.back.stall){
                state_s2 := Stage2State.lookup
            }
        }

    }
}

    val rdatas = RegEnable(VecInit.tabulate(ways){ i =>
            VecInit.tabulate(num_word){ j =>
                data_bank_io(i)(j).dout
            }
        }, pipego_reg)


package nagicore.unit.cache

import chisel3.VecInit._
import chisel3._
import chisel3.util._
import nagicore.bus._
import chisel3.util.random.LFSR
import nagicore.utils.isPowerOf2
import nagicore.GlobalConfg
import nagicore.unit.RingBuff

class CacheWTIOFrontBits[T <: Bundle](addrBits: Int, dataBits: Int, pipedataT: () => T) extends Bundle{
    val valid     = Bool()
    val addr      = UInt(addrBits.W)
    val we        = Bool()
    val wmask     = UInt((dataBits/8).W)
    val size      = UInt(2.W)
    val wdata     = UInt(dataBits.W)
    val uncache   = Bool()
    val pipedata  = pipedataT()
}

class CacheWTIOFront[T <: Bundle](addrBits: Int, dataBits: Int, pipedataT: () => T) extends Bundle{
    val bits        = Input(new CacheWTIOFrontBits[T](addrBits, dataBits, pipedataT))
    val stall       = Output(Bool())
}

class CacheWTIOBackBits[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val valid       = Bool()
    val rdata       = UInt(dataBits.W)
    // 整个缓存块，注意uncahce访问时，结果无意义
    val rline       = Vec(blockWords, UInt(dataBits.W))
    val pipedata    = pipedataT()
    // for debug
    val addr        = UInt(addrBits.W)
    val size        = UInt(2.W)
    val wdata       = UInt(dataBits.W)
    val wmask       = UInt((dataBits/8).W)
}

class CacheWTIOBack[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val bits        = Output(new CacheWTIOBackBits[T](addrBits, dataBits, blockWords, pipedataT))
    val stall       = Input(Bool())
}

class CacheWTIO[T <: Bundle](addrBits: Int, dataBits: Int, blockWords: Int, pipedataT: () => T) extends Bundle{
    val front = new CacheWTIOFront[T](addrBits, dataBits, pipedataT)
    val back = new CacheWTIOBack[T](addrBits, dataBits, blockWords, pipedataT)
}

/**
  * 单周期写通型Cache，一拍写入到WriteBuffer中
  * @pipeline
  *                   stage0             stage1
                                         比较Tag，选择数据   字节选择，符号扩展
  * CONTROL SIGNAL---------------------->|preg|
  *                   |-------------|
  * ADDR------------> |  Cache RAM  |--->|rtag, rdata, ...|
  *                   |-------------|
  *
  * @param addrBits
  * @param dataBits
  * @param ways         相连数，必须为2的幂次
  * @param sets         行数
  * @param blockWords   块字个数
  * @param pipedataT    流水线其他信号
  */
class CacheWT[T <: Bundle](addrBits: Int, dataBits: Int, ways: Int, sets: Int, blockWords: Int, writeBuffLen: Int,
                        pipedataT: () => T,
                        replaceT: CacheReplaceType.CacheReplaceType=CacheReplaceType.Random,
                        debug_id: Int=0) extends Module{
    require(isPowerOf2(ways))
    require(isPowerOf2(dataBits))

    val io = IO(new Bundle{
        val axi = new AXI4IO(addrBits, dataBits)
        val cmd = new CacheWTIO(addrBits, dataBits, blockWords, pipedataT)
    })

    // Block address
    // [ Tag | Index | Word | Byte ]
    //    |      |      |      |
    //   line   set
    val len_byte = log2Ceil(dataBits/8)
    val num_word = blockWords
    val len_word = log2Ceil(num_word)
    val len_idx = log2Ceil(sets)
    val len_tag = addrBits - len_idx - len_word - len_byte
    val beats_len = num_word

    val axi_w_agent = Module(new AXI4WriteAgent(addrBits, dataBits, beats_len))
    axi_w_agent.io.axi.aw <> io.axi.aw
    axi_w_agent.io.axi.w <> io.axi.w
    axi_w_agent.io.axi.b <> io.axi.b
    axi_w_agent.io.cmd.in := DontCare
    axi_w_agent.io.cmd.in.req := false.B
    val axi_r_agent = Module(new AXI4ReadAgent(addrBits, dataBits, beats_len))
    axi_r_agent.io.axi.ar <> io.axi.ar
    axi_r_agent.io.axi.r <> io.axi.r
    axi_r_agent.io.cmd.in := DontCare
    axi_r_agent.io.cmd.in.req := false.B

    class WriteInfo extends Bundle{
        val addr    = UInt(addrBits.W)
        val size    = UInt(2.W)
        val wmask   = UInt((dataBits/8).W)
        val wdata   = UInt(dataBits.W)
    }
    val write_buff = Module(new RingBuff(()=>new WriteInfo, writeBuffLen, 1, debug_id))
    write_buff.io.push := false.B
    write_buff.io.pop := false.B
    write_buff.io.wdata := DontCare
    write_buff.io.clear := false.B

    val random_way = LFSR(16)(log2Up(ways)-1, 0)
    val lru = Reg(Vec(sets, UInt(log2Up(ways).W)))

    val active_way = Reg(UInt(log2Up(ways).W))

    val data_bank = Seq.fill(ways)(Seq.fill(num_word)(Module(new Ram(dataBits, sets, RamType.RAM_1CYC))))
    val data_bank_io = VecInit(data_bank.map(t => VecInit(t.map(_.io))))
    val tag_v = Seq.fill(ways)(Module(new Ram(len_tag+1, sets, RamType.RAM_1CYC)))
    val tag_v_io = VecInit(tag_v.map(_.io))
    val addr_idx = io.cmd.front.bits.addr(len_idx+len_word+len_byte-1, len_word+len_byte)

    // 流水使能信号
    val pipego = Wire(Bool())

    tag_v_io.foreach(io => {
        io.addr := addr_idx
        io.en := pipego
        io.re := true.B
        io.we := false.B
        io.din := 0.U
        io.wmask := Fill(len_tag+1, true.B)
    })
    data_bank_io.foreach(_.foreach(io => {
        io.addr := addr_idx
        io.en := pipego
        io.re := true.B
        io.we := false.B
        io.din := 0.U
        io.wmask := Fill(dataBits, true.B)
    }))

    // stage1 pipeline registers
    val preg = RegEnable(io.cmd.front.bits, pipego)

    // // stage2 pipeline registers
    // val preg2 = RegEnable(preg, pipego)
    object StageState extends ChiselEnum {
        //  0       1               2           3               4           5                   6
        val lookup, replaceWait,    replace,    replaceWrite,   replaceEnd, uncacheReadWait,    uncacheRead,
        //  7           8
            uncacheEnd, waitWriteBuff = Value
    }
    // stage2 state
    val state = RegInit(StageState.lookup)
    val hit = Wire(Bool())

    val pipego_reg = pipego

    val rdatas = RegEnable(tabulate(ways){ i =>
            tabulate(num_word){ j =>
                data_bank_io(i)(j).dout
            }
        }, pipego_reg)
    val rtags = RegEnable(tabulate(ways){ i =>
            tag_v_io(i).dout(len_tag, 1)
        }, pipego_reg)
    val rvalid = RegEnable(tabulate(ways){ i =>
            tag_v_io(i).dout(0).asBool
        }, pipego_reg)

    val data_ok = (
        // 当hit read/write 或者 uncache write且write_buffer未满 的时候，可以连续命中，单周期返回
        (state === StageState.lookup && ((!preg.we && hit) || (preg.we && hit && !write_buff.io.full) || (preg.we && preg.uncache && !write_buff.io.full)) ) ||
        // 替换完成
        (state === StageState.replaceEnd) ||
        // uncache访问
        (state === StageState.uncacheEnd) ||
        // 成功写入WBUFF
        (state === StageState.waitWriteBuff && !write_buff.io.full)
    )
    pipego := (data_ok ||
        // 无效命令
        !preg.valid) &&
        // 下一级无阻塞请求
        !io.cmd.back.stall

    val addr_reg = preg.addr
    // TODO: VIPT, let TLB pass in real tag
    val addr_tag_reg = addr_reg(addrBits-1, len_idx+len_word+len_byte)
    val addr_idx_reg = addr_reg(len_idx+len_word+len_byte-1, len_word+len_byte)

    val rdatas_hit = WireDefault(fill(blockWords)(0.U(dataBits.W)))
    val rdatas_replace = RegInit(fill(blockWords)(0.U(dataBits.W)))
    val rdata_uncache = RegInit(0.U(dataBits.W))

    io.cmd.front.stall := !pipego
    io.cmd.back.bits.valid := data_ok && preg.valid
//    io.master.back.bits.rdata := Mux(state_s2 === Stage2State.replaceEnd, rdata_replace, rdata_hit)
    io.cmd.back.bits.rdata := MuxCase(rdatas_hit(addr_word_reg), Seq(
        (state === StageState.replaceEnd) -> rdatas_replace(addr_word_reg),
        (state === StageState.uncacheEnd) -> rdata_uncache
    ))
    io.cmd.back.bits.rline := Mux(state === StageState.replaceEnd, rdatas_replace,
        rdatas_hit
    )
    io.cmd.back.bits.pipedata := preg.pipedata

    io.cmd.back.bits.addr := addr_reg
    io.cmd.back.bits.size := preg.size
    io.cmd.back.bits.wdata := preg.wdata
    io.cmd.back.bits.wmask := preg.wmask

    val bypass_addr = Reg(UInt((len_tag+len_idx).W))
    val bypass_val = Reg(Vec(blockWords, UInt(dataBits.W)))
    val bypass_valid = Reg(Vec(blockWords, Bool()))
    for(i <- 0 until blockWords){
        bypass_valid(i) := 0.U
    }

    val hits = tabulate(ways){ i =>
        rtags(i) === addr_tag_reg && rvalid(i)
    }
    hit := hits.reduceTree(_||_) && !preg.uncache

    if(GlobalConfg.SIM){
        import  nagicore.unit.DPIC_PERF_CACHE
        val dpic_perf_cache = Module(new DPIC_PERF_CACHE)
        dpic_perf_cache.io.clk := clock
        dpic_perf_cache.io.rst := reset
        dpic_perf_cache.io.valid := preg.valid && state === StageState.lookup && !io.cmd.back.stall
        dpic_perf_cache.io.id := debug_id.U
        dpic_perf_cache.io.access_type := Cat(preg.uncache, hit)
    }

    def writeRAM(io: RamIO, addr: UInt, data: UInt) = {
        io.addr := addr
        io.en := true.B
        io.re := false.B
        io.we := true.B
        io.din := data
    }

    // 当WriteBuff和AXIW都写完了之后，才能保证读到的是正确的
    val axi_ok_to_read = write_buff.io.empty && axi_w_agent.io.cmd.out.ready && axi_r_agent.io.cmd.out.ready

    def launch_uncahe_read() = {
        axi_r_agent.io.cmd.in.req := true.B
        axi_r_agent.io.cmd.in.addr := preg.addr
        axi_r_agent.io.cmd.in.size := preg.size
        axi_r_agent.io.cmd.in.len := 0.U
        state := StageState.uncacheRead
    }

    def lanuch_replace_read() = {
        axi_r_agent.io.cmd.in.req := true.B
        axi_r_agent.io.cmd.in.addr := addr_tag_reg ## addr_idx_reg ## 0.U((len_word+len_byte).W)
        axi_r_agent.io.cmd.in.len := (blockWords-1).U
        axi_r_agent.io.cmd.in.size := log2Up(dataBits).U
    }

    val hit_way = PriorityEncoder(hits)
    // 当连续hit wr/ww同一Cache行时，需要前递
    val hit_rw_bypass_need = bypass_addr === Cat(addr_tag_reg, addr_idx_reg)
    val rdatas_real = tabulate(blockWords){
        i => Mux(hit_rw_bypass_need && bypass_valid(i),
            bypass_val(i),
            rdatas(hit_way)(i)
        )
    }
    val wdata_real = Cat((0 until (dataBits/8)).reverse.map(i =>
        Mux(preg.wmask(i),
            preg.wdata(i*8+7, i*8),
            rdatas_real(addr_word_reg)(i*8+7, i*8)
        )))

    def write_write_buff() = {
        write_buff.io.push := true.B
        write_buff.io.wdata.addr := addr_reg
        write_buff.io.wdata.size := preg.size
        write_buff.io.wdata.wdata := wdata_real
        write_buff.io.wdata.wmask := Fill(dataBits/8, 1.U)
    }

    switch(state){
        is(StageState.lookup){
            when(preg.valid){
                when(hit){
                    if(replaceT == CacheReplaceType.LRU){
                        lru(addr_idx_reg) := ~hit_way
                    }
                    when(preg.we){
                        // 命中写
                        writeRAM(data_bank_io(hit_way)(addr_word_reg), addr_idx_reg, wdata_real)
                        bypass_addr := Cat(addr_tag_reg, addr_idx_reg)
                        bypass_valid(addr_word_reg) := true.B
                        bypass_val(addr_word_reg) := wdata_real
                        when(!write_buff.io.full){
                            when(!io.cmd.back.stall){ // 防止后级阻塞时，造成多余写入
                                write_write_buff()
                            }
                        }.otherwise{
                            state := StageState.waitWriteBuff
                        }
                    }otherwise{
                        // 命中读
                        rdatas_hit := rdatas_real
                    }
                    state := StageState.lookup
                }.otherwise{
                    when(preg.uncache){
                        when(preg.we){
                            when(!write_buff.io.full){
                                when(!io.cmd.back.stall){ // 防止后级阻塞时，造成多余写入
                                    write_write_buff()
                                }
                            }.otherwise{
                                state := StageState.waitWriteBuff
                            }
                        }.otherwise{
                            when(axi_ok_to_read){
                                launch_uncahe_read()
                            }.otherwise{
                                state := StageState.uncacheReadWait
                            }
                        }
                        // or just state_s2 := Stage2State.uncacheWait
                    }.otherwise{
                        val active_way_wire = Wire(chiselTypeOf(active_way))
                        if(replaceT == CacheReplaceType.LRU){
                            active_way_wire := lru(addr_idx_reg)
                        }else{
                            active_way_wire := random_way
                        }
                        active_way := active_way_wire
                        when(axi_ok_to_read){
                            lanuch_replace_read()
                            state := StageState.replace
                        }.otherwise{
                            state := StageState.replaceWait
                        }
                    }
                }
            }
        }
        is(StageState.replaceWait){
            when(axi_ok_to_read){
                lanuch_replace_read()
                state := StageState.replace
            }
        }
        is(StageState.replace){
            when(axi_r_agent.io.cmd.out.ready){
                val word_i = axi_r_agent.io.cmd.out.order
                val wdata_active = Mux(word_i === addr_word_reg,
                    Cat((0 until (dataBits/8)).reverse.map(i =>
                        Mux(preg.wmask(i),
                            preg.wdata(i*8+7, i*8),
                            axi_r_agent.io.cmd.out.rdata(i*8+7, i*8)
                        )
                    )),
                    axi_r_agent.io.cmd.out.rdata
                )

                writeRAM(data_bank_io(active_way)(word_i), addr_idx_reg, wdata_active)
                bypass_val(word_i) := wdata_active

                rdatas_replace(word_i) := wdata_active

                when(axi_r_agent.io.cmd.out.last){
                    state := StageState.replaceWrite
                    writeRAM(tag_v_io(active_way), addr_idx_reg, addr_tag_reg ## 1.U(1.W))
                    bypass_addr := Cat(addr_tag_reg, addr_idx_reg)
                }
            }
        }
        is(StageState.replaceWrite){
            when(!write_buff.io.full){
                write_buff.io.push := true.B
                write_buff.io.wdata.addr := preg.addr
                write_buff.io.wdata.size := preg.size
                write_buff.io.wdata.wdata := bypass_val(addr_word_reg)
                write_buff.io.wdata.wmask := Fill(dataBits/8, 1.U)
                state := StageState.replaceEnd
            }
        }
        // 等待下一级准备好接受并且成功写入到WriteBuff
        is(StageState.replaceEnd){
            when(!io.cmd.back.stall){
                state := StageState.lookup

                for(i <- 0 until blockWords){
                    bypass_valid(i) := true.B
                }
            }
        }
        is(StageState.uncacheReadWait){
            when(axi_ok_to_read){
                launch_uncahe_read()
            }
        }
        is(StageState.uncacheRead){
            // 等待代理空闲并且WriteBuff为空
            when(axi_r_agent.io.cmd.out.ready){
                rdata_uncache := axi_r_agent.io.cmd.out.rdata
                state := StageState.uncacheEnd
            }
        }
        is(StageState.uncacheEnd){
            // 等待下一级准备好接受
            when(!io.cmd.back.stall){
                state := StageState.lookup
            }
        }

        is(StageState.waitWriteBuff){
            when(!write_buff.io.full){
                write_write_buff()

                state := StageState.lookup
            }
        }
    }

    when(!write_buff.io.empty){
        when(axi_w_agent.io.cmd.out.ready){
            axi_w_agent.io.cmd.in.req := true.B
            axi_w_agent.io.cmd.in.addr := write_buff.io.rdatas(0).addr
            axi_w_agent.io.cmd.in.len := 0.U
            axi_w_agent.io.cmd.in.size := write_buff.io.rdatas(0).size
            axi_w_agent.io.cmd.in.wdata(0) := write_buff.io.rdatas(0).wdata
            axi_w_agent.io.cmd.in.wmask(0) := write_buff.io.rdatas(0).wmask

            write_buff.io.pop := true.B
        }
    }
}

val addr_word_reg = if(len_word!=0) addr_reg(len_word+len_byte-1, len_byte) else 0.U


package constellation.channel

import chisel3._
import chisel3.experimental.SourceInfo
import org.chipsalliance.cde.config.{Parameters, Field}
import freechips.rocketchip.diplomacy._

case class ChannelEdgeParams(cp: ChannelParams, p: Parameters)

object ChannelImp extends SimpleNodeImp[EmptyParams, ChannelParams, ChannelEdgeParams, Channel] {
  def edge(pd: EmptyParams, pu: ChannelParams, p: Parameters, sourceInfo: SourceInfo) = {
    ChannelEdgeParams(pu, p)
  }
  def bundle(e: ChannelEdgeParams) = new Channel(e.cp)(e.p)
  def render(e: ChannelEdgeParams) = if (e.cp.possibleFlows.size == 0) {
    RenderedEdge(colour = "ffffff", label = "X")
  } else {
    RenderedEdge(colour = "#0000ff", label = e.cp.payloadBits.toString)
  }

  override def monitor(bundle: Channel, edge: ChannelEdgeParams): Unit = {
    val monitor = Module(new NoCMonitor(edge.cp)(edge.p))
    monitor.io.in := bundle
  }
  // TODO: Add nodepath stuff? override def mixO, override def mixI
}

case class ChannelSourceNode(val destId: Int)(implicit valName: ValName) extends SourceNode(ChannelImp)(Seq(EmptyParams()))
case class ChannelDestNode(val destParams: ChannelParams)(implicit valName: ValName) extends SinkNode(ChannelImp)(Seq(destParams))
case class ChannelAdapterNode(
  slaveFn:  ChannelParams => ChannelParams = { d => d })(
  implicit valName: ValName) extends AdapterNode(ChannelImp)((e: EmptyParams) => e, slaveFn)
case class ChannelIdentityNode()(implicit valName: ValName) extends IdentityNode(ChannelImp)()
case class ChannelEphemeralNode()(implicit valName: ValName) extends EphemeralNode(ChannelImp)()

case class IngressChannelEdgeParams(cp: IngressChannelParams, p: Parameters)
case class EgressChannelEdgeParams(cp: EgressChannelParams, p: Parameters)

object IngressChannelImp extends SimpleNodeImp[EmptyParams, IngressChannelParams, IngressChannelEdgeParams, IngressChannel] {
  def edge(pd: EmptyParams, pu: IngressChannelParams, p: Parameters, sourceInfo: SourceInfo) = {
    IngressChannelEdgeParams(pu, p)
  }
  def bundle(e: IngressChannelEdgeParams) = new IngressChannel(e.cp)(e.p)
  def render(e: IngressChannelEdgeParams) = if (e.cp.possibleFlows.size == 0) {
    RenderedEdge(colour = "ffffff", label = "X")
  } else {
    RenderedEdge(colour = "#00ff00", label = e.cp.payloadBits.toString)
  }
}

object EgressChannelImp extends SimpleNodeImp[EmptyParams, EgressChannelParams, EgressChannelEdgeParams, EgressChannel] {
  def edge(pd: EmptyParams, pu: EgressChannelParams, p: Parameters, sourceInfo: SourceInfo) = {
    EgressChannelEdgeParams(pu, p)
  }
  def bundle(e: EgressChannelEdgeParams) = new EgressChannel(e.cp)(e.p)
  def render(e: EgressChannelEdgeParams) = if (e.cp.possibleFlows.size == 0) {
    RenderedEdge(colour = "ffffff", label = "X")
  } else {
    RenderedEdge(colour = "#ff0000", label = e.cp.payloadBits.toString)
  }
}

case class IngressChannelSourceNode(val destId: Int)(implicit valName: ValName) extends SourceNode(IngressChannelImp)(Seq(EmptyParams()))
case class IngressChannelDestNode(val destParams: IngressChannelParams)(implicit valName: ValName) extends SinkNode(IngressChannelImp)(Seq(destParams))
case class EgressChannelSourceNode(val egressId: Int)(implicit valName: ValName) extends SourceNode(EgressChannelImp)(Seq(EmptyParams()))
case class EgressChannelDestNode(val destParams: EgressChannelParams)(implicit valName: ValName) extends SinkNode(EgressChannelImp)(Seq(destParams))

case class IngressChannelAdapterNode(
  slaveFn:  IngressChannelParams => IngressChannelParams = { d => d })(
  implicit valName: ValName) extends AdapterNode(IngressChannelImp)(m => m, slaveFn)

case class EgressChannelAdapterNode(
  slaveFn:  EgressChannelParams => EgressChannelParams = { d => d })(
  implicit valName: ValName) extends AdapterNode(EgressChannelImp)(m => m, slaveFn)

case class IngressChannelIdentityNode()(implicit valName: ValName) extends IdentityNode(IngressChannelImp)()
case class EgressChannelIdentityNode()(implicit valName: ValName) extends IdentityNode(EgressChannelImp)()

case class IngressChannelEphemeralNode()(implicit valName: ValName) extends EphemeralNode(IngressChannelImp)()
case class EgressChannelEphemeralNode()(implicit valName: ValName) extends EphemeralNode(EgressChannelImp)()


package constellation.channel

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

object WidthWidget {
  def split[T <: BaseFlit](in: IrrevocableIO[T], out: IrrevocableIO[T]) = {
    val inBits = in.bits.payload.getWidth
    val outBits = out.bits.payload.getWidth
    require(inBits > outBits && inBits % outBits == 0)
    val ratio = inBits / outBits
    val count = RegInit(0.U(log2Ceil(ratio).W))
    val first = count === 0.U
    val last = count === (ratio - 1).U
    val stored = Reg(UInt((inBits - outBits).W))

    out.valid := in.valid
    out.bits := in.bits
    out.bits.head := in.bits.head && first
    out.bits.tail := in.bits.tail && last
    out.bits.payload := Mux(first, in.bits.payload, stored)
    in.ready := last && out.ready
    when (out.fire) {
      count := Mux(last, 0.U, count + 1.U)
      stored := Mux(first, in.bits.payload, stored) >> outBits
    }
  }

  def merge[T <: BaseFlit](in: IrrevocableIO[T], out: IrrevocableIO[T]) = {
    val inBits = in.bits.payload.getWidth
    val outBits = out.bits.payload.getWidth
    require(outBits > inBits && outBits % inBits == 0)
    val ratio = outBits / inBits
    val count = RegInit(0.U(log2Ceil(ratio).W))
    val first = count === 0.U
    val last = count === (ratio - 1).U
    val flit = Reg(out.bits.cloneType)
    val payload = Reg(Vec(ratio-1, UInt(inBits.W)))

    out.valid := in.valid && last
    out.bits := flit
    out.bits.tail := last && in.bits.tail
    out.bits.payload := Cat(in.bits.payload, payload.asUInt)
    in.ready := !last || out.ready
    when (in.fire) {
      count := Mux(last, 0.U, count + 1.U)
      payload(count) := in.bits.payload
      when (first) {
        flit := in.bits
      }
    }
  }

  def apply[T <: BaseFlit](in: IrrevocableIO[T], out: IrrevocableIO[T]) = {
    val inBits = in.bits.payload.getWidth
    val outBits = out.bits.payload.getWidth

    if (inBits == outBits) {
      out <> in
    } else if (inBits < outBits) {
      merge(in, out)
    } else {
      split(in, out)
    }
  }
}

class IngressWidthWidget(srcBits: Int)(implicit p: Parameters) extends LazyModule {
  val node = new IngressChannelAdapterNode(
    slaveFn = { s => s.copy(payloadBits=srcBits) }
  )
  lazy val module = new LazyModuleImp(this) {
    (node.in zip node.out) foreach { case ((in, edgeIn), (out, edgeOut)) =>
      WidthWidget(in.flit, out.flit)
    }
  }
}

object IngressWidthWidget {
  def apply(destBits: Int, srcBits: Int)(implicit p: Parameters) = {
    if (destBits == srcBits) {
      val node = IngressChannelEphemeralNode()
      node
    } else {
      val ingress_width_widget = LazyModule(new IngressWidthWidget(srcBits))
      ingress_width_widget.node
    }
  }
}

class EgressWidthWidget(srcBits: Int)(implicit p: Parameters) extends LazyModule {
  val node = new EgressChannelAdapterNode(

  )
  lazy val module = new LazyModuleImp(this) {
    (node.in zip node.out) foreach { case ((in, edgeIn), (out, edgeOut)) =>
      WidthWidget(in.flit, out.flit)
    }
  }
}

object EgressWidthWidget {
  def apply(destBits: Int, srcBits: Int)(implicit p: Parameters) = {
    if (destBits == srcBits) {
      val node = EgressChannelEphemeralNode()
      node
    } else {
      val egress_width_widget = LazyModule(new EgressWidthWidget(srcBits))
      egress_width_widget.node
    }
  }
}

class ChannelWidthWidget(srcBits: Int)(implicit p: Parameters) extends LazyModule {
  val node = new ChannelAdapterNode(
    slaveFn = { s =>
      val destBits = s.payloadBits
      if (srcBits > destBits) {
        val ratio = srcBits / destBits
        val virtualChannelParams = s.virtualChannelParams.map { vP =>
          require(vP.bufferSize >= ratio)
          vP.copy(bufferSize = vP.bufferSize / ratio)
        }
        s.copy(payloadBits=srcBits, virtualChannelParams=virtualChannelParams)
      } else {
        val ratio = destBits / srcBits
        val virtualChannelParams = s.virtualChannelParams.map { vP =>
          vP.copy(bufferSize = vP.bufferSize * ratio)
        }
        s.copy(payloadBits=srcBits, virtualChannelParams=virtualChannelParams)
      }
    }
  )
  lazy val module = new LazyModuleImp(this) {
    (node.in zip node.out) foreach { case ((in, edgeIn), (out, edgeOut)) =>
      val destBits = out.cParam.payloadBits
      in.vc_free := out.vc_free
      if (srcBits == destBits) {
        in.credit_return := out.credit_return
        out.flit <> in.flit
      } else if (srcBits > destBits) {
        require(srcBits % destBits == 0, s"$srcBits, $destBits")

        (in.flit zip out.flit) foreach { case (iF, oF) =>
          val in_q = Module(new Queue(new Flit(in.cParam.payloadBits),
            in.cParam.virtualChannelParams.map(_.bufferSize).sum, pipe=true, flow=true))
          in_q.io.enq.valid := iF.valid
          in_q.io.enq.bits := iF.bits
          assert(!(in_q.io.enq.valid && !in_q.io.enq.ready))

          val in_flit = Wire(Irrevocable(new Flit(in.cParam.payloadBits)))
          in_flit.valid := in_q.io.deq.valid
          in_flit.bits := in_q.io.deq.bits
          in_q.io.deq.ready := in_flit.ready

          val out_flit = Wire(Irrevocable(new Flit(out.cParam.payloadBits)))
          oF.valid := out_flit.valid
          oF.bits := out_flit.bits
          out_flit.ready := true.B
          WidthWidget(in_flit, out_flit)
        }

        val ratio = srcBits / destBits
        val counts = RegInit(VecInit(Seq.fill(in.nVirtualChannels) { 0.U(log2Ceil(ratio).W) }))
        val in_credit_return = Wire(Vec(in.nVirtualChannels, Bool()))
        in.credit_return := in_credit_return.asUInt
        for (i <- 0 until in.nVirtualChannels) {
          in_credit_return(i) := false.B
          when (out.credit_return(i)) {
            val last = counts(i) === (ratio-1).U
            counts(i) := Mux(last, 0.U, counts(i) + 1.U)
            when (last) {
              in_credit_return(i) := true.B
            }
          }
        }
      } else {
        require(destBits % srcBits == 0)

        val in_flits = Wire(Vec(in.nVirtualChannels, Irrevocable(new Flit(in.cParam.payloadBits))))
        val out_flits = Wire(Vec(in.nVirtualChannels, Irrevocable(new Flit(out.cParam.payloadBits))))
        for (i <- 0 until in.nVirtualChannels) {
          val sel = in.flit.map(f => f.valid && f.bits.virt_channel_id === i.U)
          in_flits(i).valid := sel.orR
          in_flits(i).bits := Mux1H(sel, in.flit.map(_.bits))

          out_flits(i).ready := sel.orR
          WidthWidget(in_flits(i), out_flits(i))
        }
        for (i <- 0 until in.flit.size) {
          val sel = UIntToOH(in.flit(i).bits.virt_channel_id)
          out.flit(i).valid := Mux1H(sel, out_flits.map(_.valid))
          out.flit(i).bits := Mux1H(sel, out_flits.map(_.bits))
        }

        val ratio = destBits / srcBits
        val credits = RegInit(VecInit((0 until in.nVirtualChannels).map { i =>
          0.U(log2Ceil(ratio*in.cParam.virtualChannelParams(i).bufferSize).W)
        }))
        val in_credit_return = Wire(Vec(in.nVirtualChannels, Bool()))
        in.credit_return := in_credit_return.asUInt
        for (i <- 0 until in.nVirtualChannels) {
          when (out.credit_return(i)) {
            in_credit_return(i) := true.B
            credits(i) := credits(i) + (ratio - 1).U
          } .otherwise {
            val empty = credits(i) === 0.U
            in_credit_return(i) := !empty
            credits(i) := Mux(empty, 0.U, credits(i) - 1.U)
          }
        }
      }
    }
  }
}

object ChannelWidthWidget {
  def apply(destBits: Int, srcBits: Int)(implicit p: Parameters) = {
    if (destBits == srcBits) {
      val node = ChannelEphemeralNode()
      node
    } else {
      val channel_width_widget = LazyModule(new ChannelWidthWidget(srcBits))
      channel_width_widget.node
    }
  }
}


package constellation.noc

import chisel3._
import chisel3.util._

import org.chipsalliance.cde.config.{Field, Parameters}
import freechips.rocketchip.diplomacy.{LazyModule, LazyModuleImp, BundleBridgeSink, InModuleBody}
import freechips.rocketchip.util.ElaborationArtefacts
import freechips.rocketchip.prci._
import constellation.router._
import constellation.channel._
import constellation.routing.{RoutingRelation, ChannelRoutingInfo}
import constellation.topology.{PhysicalTopology, UnidirectionalLine}

class NoCTerminalIO(
  val ingressParams: Seq[IngressChannelParams],
  val egressParams: Seq[EgressChannelParams])(implicit val p: Parameters) extends Bundle {
  val ingress = MixedVec(ingressParams.map { u => Flipped(new IngressChannel(u)) })
  val egress = MixedVec(egressParams.map { u => new EgressChannel(u) })
}

class NoC(nocParams: NoCParams)(implicit p: Parameters) extends LazyModule {

  override def shouldBeInlined = nocParams.inlineNoC

  val internalParams = InternalNoCParams(nocParams)
  val allChannelParams = internalParams.channelParams
  val allIngressParams = internalParams.ingressParams
  val allEgressParams = internalParams.egressParams
  val allRouterParams = internalParams.routerParams

  val iP = p.alterPartial({ case InternalNoCKey => internalParams })
  val nNodes = nocParams.topology.nNodes
  val nocName = nocParams.nocName
  val skipValidationChecks = nocParams.skipValidationChecks

  val clockSourceNodes = Seq.tabulate(nNodes) { i => ClockSourceNode(Seq(ClockSourceParameters())) }
  val router_sink_domains = Seq.tabulate(nNodes) { i =>
    val router_sink_domain = LazyModule(new ClockSinkDomain(ClockSinkParameters(
      name = Some(s"${nocName}_router_$i")
    )))
    router_sink_domain.clockNode := clockSourceNodes(i)
    router_sink_domain
  }

  val routers = Seq.tabulate(nNodes) { i => router_sink_domains(i) {
    val inParams = allChannelParams.filter(_.destId == i).map(
      _.copy(payloadBits=allRouterParams(i).user.payloadBits)
    )
    val outParams = allChannelParams.filter(_.srcId == i).map(
      _.copy(payloadBits=allRouterParams(i).user.payloadBits)
    )
    val ingressParams = allIngressParams.filter(_.destId == i).map(
      _.copy(payloadBits=allRouterParams(i).user.payloadBits)
    )
    val egressParams = allEgressParams.filter(_.srcId == i).map(
      _.copy(payloadBits=allRouterParams(i).user.payloadBits)
    )
    val noIn = inParams.size + ingressParams.size == 0
    val noOut = outParams.size + egressParams.size == 0
    if (noIn || noOut) {
      println(s"Constellation WARNING: $nocName router $i seems to be unused, it will not be generated")
      None
    } else {
      Some(LazyModule(new Router(
        routerParams = allRouterParams(i),
        preDiplomaticInParams = inParams,
        preDiplomaticIngressParams = ingressParams,
        outDests = outParams.map(_.destId),
        egressIds = egressParams.map(_.egressId)
      )(iP)))
    }
  }}.flatten

  val ingressNodes = allIngressParams.map { u => IngressChannelSourceNode(u.destId) }
  val egressNodes = allEgressParams.map { u => EgressChannelDestNode(u) }

  // Generate channels between routers diplomatically
  Seq.tabulate(nNodes, nNodes) { case (i, j) => if (i != j) {
    val routerI = routers.find(_.nodeId == i)
    val routerJ = routers.find(_.nodeId == j)
    if (routerI.isDefined && routerJ.isDefined) {
      val sourceNodes: Seq[ChannelSourceNode] = routerI.get.sourceNodes.filter(_.destId == j)
      val destNodes: Seq[ChannelDestNode] = routerJ.get.destNodes.filter(_.destParams.srcId == i)
      require (sourceNodes.size == destNodes.size)
      (sourceNodes zip destNodes).foreach { case (src, dst) =>
        val channelParam = allChannelParams.find(c => c.srcId == i && c.destId == j).get
        router_sink_domains(j) {
          implicit val p: Parameters = iP
          (dst
            := ChannelWidthWidget(routerJ.get.payloadBits, routerI.get.payloadBits)
            := channelParam.channelGen(p)(src)
          )
        }
      }
    }
  }}

  // Generate terminal channels diplomatically
  routers.foreach { dst => router_sink_domains(dst.nodeId) {
    implicit val p: Parameters = iP
    dst.ingressNodes.foreach(n => {
      val ingressId = n.destParams.ingressId
      require(dst.payloadBits <= allIngressParams(ingressId).payloadBits)
      (n
        := IngressWidthWidget(dst.payloadBits, allIngressParams(ingressId).payloadBits)
        := ingressNodes(ingressId)
      )
    })
    dst.egressNodes.foreach(n => {
      val egressId = n.egressId
      require(dst.payloadBits <= allEgressParams(egressId).payloadBits)
      (egressNodes(egressId)
        := EgressWidthWidget(allEgressParams(egressId).payloadBits, dst.payloadBits)
        := n
      )
    })
  }}

  val debugNodes = routers.map { r =>
    val sink = BundleBridgeSink[DebugBundle]()
    sink := r.debugNode
    sink
  }
  val ctrlNodes = if (nocParams.hasCtrl) {
    (0 until nNodes).map { i =>
      routers.find(_.nodeId == i).map { r =>
        val sink = BundleBridgeSink[RouterCtrlBundle]()
        sink := r.ctrlNode.get
        sink
      }
    }
  } else {
    Nil
  }

  println(s"Constellation: $nocName Finished parameter validation")
  lazy val module = new Impl
  class Impl extends LazyModuleImp(this) {
    println(s"Constellation: $nocName Starting NoC RTL generation")
    val io = IO(new NoCTerminalIO(allIngressParams, allEgressParams)(iP) {
      val router_clocks = Vec(nNodes, Input(new ClockBundle(ClockBundleParameters())))
      val router_ctrl = if (nocParams.hasCtrl) Vec(nNodes, new RouterCtrlBundle) else Nil
    })

    (io.ingress zip ingressNodes.map(_.out(0)._1)).foreach { case (l,r) => r <> l }
    (io.egress  zip egressNodes .map(_.in (0)._1)).foreach { case (l,r) => l <> r }
    (io.router_clocks zip clockSourceNodes.map(_.out(0)._1)).foreach { case (l,r) => l <> r }

    if (nocParams.hasCtrl) {
      ctrlNodes.zipWithIndex.map { case (c,i) =>
        if (c.isDefined) {
          io.router_ctrl(i) <> c.get.in(0)._1
        } else {
          io.router_ctrl(i) <> DontCare
        }
      }
    }

    // TODO: These assume a single clock-domain across the entire noc
    val debug_va_stall_ctr = RegInit(0.U(64.W))
    val debug_sa_stall_ctr = RegInit(0.U(64.W))
    val debug_any_stall_ctr = debug_va_stall_ctr + debug_sa_stall_ctr
    debug_va_stall_ctr := debug_va_stall_ctr + debugNodes.map(_.in(0)._1.va_stall.reduce(_+_)).reduce(_+_)
    debug_sa_stall_ctr := debug_sa_stall_ctr + debugNodes.map(_.in(0)._1.sa_stall.reduce(_+_)).reduce(_+_)

    dontTouch(debug_va_stall_ctr)
    dontTouch(debug_sa_stall_ctr)
    dontTouch(debug_any_stall_ctr)

    def prepend(s: String) = Seq(nocName, s).mkString(".")
    ElaborationArtefacts.add(prepend("noc.graphml"), graphML)

    val adjList = routers.map { r =>
      val outs = r.outParams.map(o => s"${o.destId}").mkString(" ")
      val egresses = r.egressParams.map(e => s"e${e.egressId}").mkString(" ")
      val ingresses = r.ingressParams.map(i => s"i${i.ingressId} ${r.nodeId}")

    }.mkString("\n")
    ElaborationArtefacts.add(prepend("noc.adjlist"), adjList)

    val xys = routers.map(r => {
      val n = r.nodeId
      val ids = (Seq(r.nodeId.toString)
        ++ r.egressParams.map(e => s"e${e.egressId}")
        ++ r.ingressParams.map(i => s"i${i.ingressId}")
      )
      val plotter = nocParams.topology.plotter
      val coords = (Seq(plotter.node(r.nodeId))
        ++ Seq.tabulate(r.egressParams.size ) { i => plotter. egress(i, r. egressParams.size, r.nodeId) }
        ++ Seq.tabulate(r.ingressParams.size) { i => plotter.ingress(i, r.ingressParams.size, r.nodeId) }
      )

      (ids zip coords).map { case (i, (x, y)) => s"$i $x $y" }.mkString("\n")
    }).mkString("\n")
    ElaborationArtefacts.add(prepend("noc.xy"), xys)

    val edgeProps = routers.map { r =>
      val outs = r.outParams.map { o =>
        (Seq(s"${r.nodeId} ${o.destId}") ++ (if (o.possibleFlows.size == 0) Some("unused") else None))
          .mkString(" ")
      }
      val egresses = r.egressParams.map { e =>
        (Seq(s"${r.nodeId} e${e.egressId}") ++ (if (e.possibleFlows.size == 0) Some("unused") else None))
          .mkString(" ")
      }
      val ingresses = r.ingressParams.map { i =>
        (Seq(s"i${i.ingressId} ${r.nodeId}") ++ (if (i.possibleFlows.size == 0) Some("unused") else None))
          .mkString(" ")
      }
      (outs ++ egresses ++ ingresses).mkString("\n")
    }.mkString("\n")
    ElaborationArtefacts.add(prepend("noc.edgeprops"), edgeProps)

    println(s"Constellation: $nocName Finished NoC RTL generation")
  }
}

(Seq(s"${r.nodeId} $outs $egresses") ++ ingresses).mkString("\n")


package constellation.protocol

import chisel3._
import chisel3.util._

import constellation.channel._
import constellation.noc._

import org.chipsalliance.cde.config._
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.util._
import freechips.rocketchip.amba.axi4._

trait AXI4FieldHelper {
  val wideBundle: AXI4BundleParameters
}
class AXI4MasterToNoC(
  edgeIn: AXI4EdgeParameters, edgesOut: Seq[AXI4EdgeParameters],
  sourceStart: Int, sourceSize: Int,
  wideBundle: AXI4BundleParameters,
  slaveToEgressOffset: Int => Int,
  flitWidth: Int,
  awQueueDepth: Int
)(implicit p: Parameters) extends Module {
  // Can only keep 1 in-flight per ID because NoC can't guarantee
  // FIFO between bursts with the same ID
  val maxFlightPerId = 1

  val io = IO(new Bundle {
    val axi4 = Flipped(new AXI4Bundle(wideBundle))
    val flits = new Bundle {
      val aw = Decoupled(new IngressFlit(flitWidth))
      val w = Decoupled(new IngressFlit(flitWidth))
      val ar = Decoupled(new IngressFlit(flitWidth))
      val r = Flipped(Decoupled(new EgressFlit(flitWidth)))
      val b = Flipped(Decoupled(new EgressFlit(flitWidth)))
    }
  })

  val port_addrs = edgesOut.map(_.slave.slaves.map(_.address).flatten)
  val routingMask = AddressDecoder(port_addrs)
  val route_addrs = port_addrs.map(seq => AddressSet.unify(seq.map(_.widen(~routingMask)).distinct))
  val outputPorts = route_addrs.map(seq => (addr: UInt) => seq.map(_.contains(addr)).reduce(_||_))

  // To route W need to record where AW went
  class AWInBundle extends Bundle {
    val out = UInt(edgesOut.size.W)
    val id = UInt(wideBundle.idBits.W)
  }
  val awIn = Module(new Queue(new AWInBundle, awQueueDepth, flow=true))

  val requestARIO = VecInit(outputPorts.map { o => o(io.axi4.ar.bits.addr) })
  val requestAWIO = VecInit(outputPorts.map { o => o(io.axi4.aw.bits.addr) })

  // W follows path dicated by AW Q
  awIn.io.enq.bits.out := requestAWIO.asUInt
  awIn.io.enq.bits.id  := io.axi4.aw.bits.id | sourceStart.U
  val requestWIO = awIn.io.deq.bits.out.asBools

  val in = Wire(new AXI4Bundle(wideBundle))
  in <> io.axi4

  // Handle size = 1 gracefully (Chisel3 empty range is broken)
  def trim(id: UInt, size: Int) = if (size <= 1) 0.U else id(log2Ceil(size)-1, 0)

  // Manipulate the AXI IDs to differentiate masters
  in.aw.bits.id := io.axi4.aw.bits.id | sourceStart.U
  in.ar.bits.id := io.axi4.ar.bits.id | sourceStart.U
  io.axi4.r.bits.id := trim(in.r.bits.id, sourceSize)
  io.axi4.b.bits.id := trim(in.b.bits.id, sourceSize)

  // Block A[RW] if we switch ports, to ensure responses stay ordered (also: beware the dining philosophers)
  val endId = edgeIn.master.endId

  val awFIFOMap = WireInit(VecInit(Seq.fill(endId) { true.B }))
  val arSel = UIntToOH(io.axi4.ar.bits.id, endId)
  val awSel = UIntToOH(io.axi4.aw.bits.id, endId)
  val rSel  = UIntToOH(io.axi4.r .bits.id, endId)
  val bSel  = UIntToOH(io.axi4.b .bits.id, endId)
  val arTag = OHToUInt(requestARIO.asUInt, edgesOut.size)
  val awTag = OHToUInt(requestAWIO.asUInt, edgesOut.size)

  for (master <- edgeIn.master.masters) {
    def idTracker(port: UInt, req_fire: Bool, resp_fire: Bool) = {
      if (master.maxFlight == Some(0)) {
        true.B
      } else {
        val legalFlight = master.maxFlight.getOrElse(maxFlightPerId+1)
        val flight = legalFlight min maxFlightPerId
        val canOverflow = legalFlight > flight
        val count = RegInit(0.U(log2Ceil(flight+1).W))
        val last = Reg(UInt(log2Ceil(edgesOut.size).W))
        count := count + req_fire.asUInt - resp_fire.asUInt
        assert (!resp_fire || count =/= 0.U)
        assert (!req_fire  || count =/= flight.U)
        when (req_fire) { last := port }
        // No need to track where it went if we cap it at 1 request
        val portMatch = if (flight == 1) { true.B } else { last === port }
          (count === 0.U || portMatch) && (!canOverflow.B || count =/= flight.U)
      }
    }

    for (id <- master.id.start until master.id.end) {
      arFIFOMap(id) := idTracker(
        arTag,
        arSel(id) && io.axi4.ar.fire,
        rSel(id) && io.axi4.r.fire && io.axi4.r.bits.last)
      awFIFOMap(id) := idTracker(
        awTag,
        awSel(id) && io.axi4.aw.fire,
        bSel(id) && io.axi4.b.fire)
    }
  }

  val allowAR = arFIFOMap(io.axi4.ar.bits.id)
  in.ar.valid := io.axi4.ar.valid && allowAR
  io.axi4.ar.ready := in.ar.ready && allowAR

  // Block AW if we cannot record the W destination
  val allowAW = awFIFOMap(io.axi4.aw.bits.id)
  in.aw.valid := io.axi4.aw.valid && awIn.io.enq.ready && allowAW
  io.axi4.aw.ready := in.aw.ready && awIn.io.enq.ready && allowAW
  awIn.io.enq.valid := io.axi4.aw.valid && in.aw.ready && allowAW

  // Block W if we do not have an AW destination
  in.w.valid := io.axi4.w.valid && awIn.io.deq.valid
  io.axi4.w.ready := in.w.ready && awIn.io.deq.valid
  awIn.io.deq.ready := io.axi4.w.valid && io.axi4.w.bits.last && in.w.ready

  io.flits.aw.valid := in.aw.valid
  in.aw.ready := io.flits.aw.ready
  io.flits.aw.bits.head := true.B
  io.flits.aw.bits.tail := true.B
  io.flits.aw.bits.payload := in.aw.bits.asUInt
  io.flits.aw.bits.egress_id := Mux1H(requestAWIO.zipWithIndex.map { case (r, i) =>
    r -> (slaveToEgressOffset(i) + 0).U
  })

  val w_first = RegInit(true.B)
  io.flits.w.valid := in.w.valid
  in.w.ready := io.flits.w.ready
  io.flits.w.bits.head := w_first
  io.flits.w.bits.tail := in.w.bits.last
  io.flits.w.bits.payload := Cat(awIn.io.deq.bits.id, in.w.bits.asUInt)
  io.flits.w.bits.egress_id := Mux1H(requestWIO.zipWithIndex.map { case (r, i) =>
    r -> (slaveToEgressOffset(i) + 1).U
  })
  when (io.flits.w.fire) { w_first := in.w.bits.last }

  io.flits.ar.valid := in.ar.valid
  in.ar.ready := io.flits.ar.ready
  io.flits.ar.bits.head := true.B
  io.flits.ar.bits.tail := true.B
  io.flits.ar.bits.payload := in.ar.bits.asUInt
  io.flits.ar.bits.egress_id := Mux1H(requestARIO.zipWithIndex.map { case (r, i) =>
    r -> (slaveToEgressOffset(i) + 2).U
  })

  in.b.valid := io.flits.b.valid
  io.flits.b.ready := in.b.ready
  in.b.bits := io.flits.b.bits.payload.asTypeOf(in.b.bits)

  in.r.valid := io.flits.r.valid
  io.flits.r.ready := in.r.ready
  in.r.bits := io.flits.r.bits.payload.asTypeOf(in.r.bits)
}

class AXI4SlaveToNoC(
  edgeOut: AXI4EdgeParameters, edgesIn: Seq[AXI4EdgeParameters],
  wideBundle: AXI4BundleParameters,
  masterToEgressOffset: Int => Int,
  flitWidth: Int,
  awQueueDepth: Int
)(implicit p: Parameters) extends Module {
  val io = IO(new Bundle {
    val axi4 = new AXI4Bundle(wideBundle)
    val flits = new Bundle {
      val aw = Flipped(Decoupled(new EgressFlit(flitWidth)))
      val w = Flipped(Decoupled(new EgressFlit(flitWidth)))
      val ar = Flipped(Decoupled(new EgressFlit(flitWidth)))
      val r = Decoupled(new IngressFlit(flitWidth))
      val b = Decoupled(new IngressFlit(flitWidth))
    }
  })
  val out = Wire(new AXI4Bundle(wideBundle))
  val out_w_in_head = Wire(Bool())
  val out_w_in_id = Wire(UInt((1 << wideBundle.idBits).W))
  io.axi4 <> out

  // Grab the port ID mapping
  val inputIdRanges = AXI4Xbar.mapInputIds(edgesIn.map(_.master))

  val requestROI  = inputIdRanges.map { i => i.contains(io.axi4.r.bits.id) }
  val requestBOI  = inputIdRanges.map { i => i.contains(io.axi4.b.bits.id) }

  // Store the incoming AW messages in a buffer, 1 entry per source ID
  // Order the AW/W based on the arriving order of W bursts
  val aw_val = RegInit(VecInit(0.U((1 << wideBundle.idBits).W).asBools))
  val aw_buf = Reg(Vec(1 << wideBundle.idBits, new AXI4BundleAW(wideBundle)))
  val w_fire = RegInit(false.B)
  val aw_q   = Module(new Queue(new AXI4BundleAW(wideBundle), awQueueDepth, pipe=true))

  io.axi4.aw <> aw_q.io.deq

  out.aw.ready       := !aw_val(out.aw.bits.id)
  when (out.aw.fire) {
    aw_val(out.aw.bits.id) := true.B
    aw_buf(out.aw.bits.id) := out.aw.bits
  }

  aw_q.io.enq.valid := out.w.valid && Mux1H(out_w_in_id, aw_val) && out_w_in_head
  aw_q.io.enq.bits  :=                Mux1H(out_w_in_id, aw_buf)
  when (aw_q.io.enq.fire) {
    aw_val(OHToUInt(out_w_in_id)) := false.B
    w_fire := true.B
  }

  out.w.ready        := io.axi4.w.ready && (
    w_fire || (out_w_in_head && aw_q.io.enq.ready && Mux1H(out_w_in_id, aw_val))
  )
  io.axi4.w.bits      := out.w.bits
  io.axi4.w.valid     := out.w.valid && (
    w_fire || (out_w_in_head && aw_q.io.enq.ready && Mux1H(out_w_in_id, aw_val))
  )

  when (out.w.fire && out.w.bits.last) { w_fire := false.B }

  val r_first = RegInit(true.B)
  io.flits.r.valid := out.r.valid
  out.r.ready := io.flits.r.ready
  io.flits.r.bits.head := r_first
  io.flits.r.bits.tail := out.r.bits.last
  io.flits.r.bits.payload := out.r.bits.asUInt
  io.flits.r.bits.egress_id := Mux1H(requestROI.zipWithIndex.map { case (r, i) =>
    r -> (masterToEgressOffset(i) + 0).U
  })
  when (io.flits.r.fire) { r_first := out.r.bits.last }

  io.flits.b.valid := out.b.valid
  out.b.ready := io.flits.b.ready
  io.flits.b.bits.head := true.B
  io.flits.b.bits.tail := true.B
  io.flits.b.bits.payload := out.b.bits.asUInt
  io.flits.b.bits.egress_id := Mux1H(requestBOI.zipWithIndex.map { case (r, i) =>
    r -> (masterToEgressOffset(i) + 1).U
  })

  out_w_in_head := io.flits.w.bits.head
  out_w_in_id := UIntToOH(io.flits.w.bits.payload.asUInt >> out.w.bits.getWidth)

  out.aw.valid := io.flits.aw.valid
  io.flits.aw.ready := out.aw.ready
  out.aw.bits := io.flits.aw.bits.payload.asTypeOf(out.aw.bits)

  out.w.valid := io.flits.w.valid
  io.flits.w.ready := out.w.ready
  out.w.bits := io.flits.w.bits.payload.asTypeOf(out.w.bits)

  out.ar.valid := io.flits.ar.valid
  io.flits.ar.ready := out.ar.ready
  out.ar.bits := io.flits.ar.bits.payload.asTypeOf(out.ar.bits)
}

class AXI4InterconnectInterface(edgesIn: Seq[AXI4EdgeParameters], edgesOut: Seq[AXI4EdgeParameters])(implicit val p: Parameters) extends Bundle {
  val in = MixedVec(edgesIn.map { e => Flipped(new AXI4Bundle(e.bundle)) })
  val out = MixedVec(edgesOut.map { e => new AXI4Bundle(e.bundle) })
}

// BEGIN: AXI4ProtocolParams
case class AXI4ProtocolParams(
  edgesIn: Seq[AXI4EdgeParameters],
  edgesOut: Seq[AXI4EdgeParameters],
  edgeInNodes: Seq[Int],
  edgeOutNodes: Seq[Int],
  awQueueDepth: Int
) extends ProtocolParams {
  // END: AXI4ProtocolParams
  val wideBundle = AXI4BundleParameters.union(edgesIn.map(_.bundle) ++ edgesOut.map(_.bundle))
  val inputIdRanges = AXI4Xbar.mapInputIds(edgesIn.map(_.master))
  val minPayloadWidth = {
    val b = new AXI4Bundle(wideBundle)
    Seq(b.aw, b.ar, b.w, b.b, b.r).map { t => t.bits match {
      case b: AXI4BundleAW => b.getWidth
      case b: AXI4BundleW => b.getWidth + wideBundle.idBits
      case b: AXI4BundleAR => b.getWidth
      case b: AXI4BundleR => b.getWidth
      case b: AXI4BundleB => b.getWidth
    }}.max
  }
  val ingressNodes = (edgeInNodes.map(u => Seq.fill(3) (u)) ++ edgeOutNodes.map(u => Seq.fill(2) (u))).flatten
  val egressNodes = (edgeInNodes.map(u => Seq.fill(2) (u)) ++ edgeOutNodes.map(u => Seq.fill(3) (u))).flatten
  val nVirtualNetworks = 5
  val vNetBlocking = (blocker: Int, blockee: Int) => blocker < blockee
  val flows = (0 until edgeInNodes.size).map { i => (0 until edgeOutNodes.size).map { o =>
    Seq(FlowParams(i * 3 + 0                    , o * 3 + 0 + edgesIn.size * 2, 4), // AW
        FlowParams(i * 3 + 1                    , o * 3 + 1 + edgesIn.size * 2, 3), // W
        FlowParams(i * 3 + 2                    , o * 3 + 2 + edgesIn.size * 2, 2), // AR
        FlowParams(o * 2 + 0 + edgesIn.size * 3 , i * 2 + 0                   , 1), // R
        FlowParams(o * 2 + 1 + edgesIn.size * 3 , i * 2 + 1                   , 0)  // B
    )
  }}.flatten.flatten

  def genIO()(implicit p: Parameters): Data = new AXI4InterconnectInterface(edgesIn, edgesOut)
  def interface(terminals: NoCTerminalIO,
    ingressOffset: Int, egressOffset: Int, protocol: Data)(implicit p: Parameters) = {
    val ingresses = terminals.ingress
    val egresses = terminals.egress
    protocol match {
      case protocol: AXI4InterconnectInterface => {
        edgesIn.zipWithIndex.map { case (e,i) =>
          val nif_master = Module(new AXI4MasterToNoC(
            e, edgesOut,
            inputIdRanges(i).start, inputIdRanges(i).size,
            wideBundle,
            (s) => s * 3 + edgesIn.size * 2 + egressOffset,
            minPayloadWidth,
            awQueueDepth
          ))
          nif_master.io.axi4 <> protocol.in(i)
          ingresses(i * 3 + 0).flit <> nif_master.io.flits.aw
          ingresses(i * 3 + 1).flit <> nif_master.io.flits.w
          ingresses(i * 3 + 2).flit <> nif_master.io.flits.ar
          nif_master.io.flits.r <> egresses(i * 2 + 0).flit
          nif_master.io.flits.b <> egresses(i * 2 + 1).flit
        }
        edgesOut.zipWithIndex.map { case (e,i) =>
          val nif_slave = Module(new AXI4SlaveToNoC(
            e, edgesIn,
            wideBundle,
            (s) => s * 2 + egressOffset,
            minPayloadWidth,
            awQueueDepth
          ))
          protocol.out(i) <> nif_slave.io.axi4
          ingresses(i * 2 + 0 + edgesIn.size * 3).flit <> nif_slave.io.flits.r
          ingresses(i * 2 + 1 + edgesIn.size * 3).flit <> nif_slave.io.flits.b
          nif_slave.io.flits.aw <> egresses(i * 3 + 0 + edgesIn.size * 2).flit
          nif_slave.io.flits.w  <> egresses(i * 3 + 1 + edgesIn.size * 2).flit
          nif_slave.io.flits.ar <> egresses(i * 3 + 2 + edgesIn.size * 2).flit
        }
      }
    }
  }
}

class AXI4NoCNode(implicit valName: ValName) extends AXI4NexusNode(
  masterFn  = { seq =>
    seq(0).copy(
      echoFields    = BundleField.union(seq.flatMap(_.echoFields)),
      requestFields = BundleField.union(seq.flatMap(_.requestFields)),
      responseKeys  = seq.flatMap(_.responseKeys).distinct,
      masters = (AXI4Xbar.mapInputIds(seq) zip seq) flatMap { case (range, port) =>
        port.masters map { master => master.copy(id = master.id.shift(range.start)) }
      }
    )
  },
  slaveFn = { seq =>
    seq(0).copy(
      responseFields = BundleField.union(seq.flatMap(_.responseFields)),
      requestKeys    = seq.flatMap(_.requestKeys).distinct,
      minLatency = seq.map(_.minLatency).min,
      slaves = seq.flatMap { port =>
        require (port.beatBytes == seq(0).beatBytes,
          s"Xbar data widths don't match: ${port.slaves.map(_.name)} has ${port.beatBytes}B vs ${seq(0).slaves.map(_.name)} has ${seq(0).beatBytes}B")
        port.slaves
      }
    )
  }
)

abstract class AXI4NoCModuleImp(outer: LazyModule) extends LazyModuleImp(outer) {
  val edgesIn: Seq[AXI4EdgeParameters]
  val edgesOut: Seq[AXI4EdgeParameters]
  val nodeMapping: DiplomaticNetworkNodeMapping
  val nocName: String
  val awQueueDepth: Int
  lazy val inNames  = nodeMapping.genUniqueName(edgesIn.map(_.master.masters.map(_.name)))
  lazy val outNames = nodeMapping.genUniqueName(edgesOut.map(_.slave.slaves.map(_.name)))
  lazy val edgeInNodes = nodeMapping.getNodesIn(inNames)
  lazy val edgeOutNodes = nodeMapping.getNodesOut(outNames)
  lazy val protocolParams = AXI4ProtocolParams(
    edgesIn = edgesIn,
    edgesOut = edgesOut,
    edgeInNodes = edgeInNodes.flatten,
    edgeOutNodes = edgeOutNodes.flatten,
    awQueueDepth = awQueueDepth
  )

  def printNodeMappings() {
    println(s"Constellation: AXI4NoC $nocName inwards mapping:")
    for ((n, i) <- inNames zip edgeInNodes) {
      val node = i.map(_.toString).getOrElse("X")
      println(s"  $node <- $n")
    }

    println(s"Constellation: AXI4NoC $nocName outwards mapping:")
    for ((n, i) <- outNames zip edgeOutNodes) {
      val node = i.map(_.toString).getOrElse("X")
      println(s"  $node <- $n")
    }
  }
}

// private AXI4 NoC
// BEGIN: AXI4NoCParams
case class AXI4NoCParams(
  nodeMappings: DiplomaticNetworkNodeMapping,
  nocParams: NoCParams,
  awQueueDepth: Int = 2
)
class AXI4NoC(params: AXI4NoCParams, name: String = "test")(implicit p: Parameters) extends LazyModule {
  // END: AXI4NoCParams
  val node = new AXI4NoCNode
  lazy val module = new AXI4NoCModuleImp(this) {
    val (io_in, edgesIn) = node.in.unzip
    val (io_out, edgesOut) = node.out.unzip
    val nodeMapping = params.nodeMappings
    val nocName = name
    val awQueueDepth = params.awQueueDepth

    printNodeMappings()

    val noc = Module(new ProtocolNoC(ProtocolNoCParams(
      params.nocParams.copy(hasCtrl = false, nocName = name),
      Seq(protocolParams)
    )))

    noc.io.protocol(0) match {
      case protocol: AXI4InterconnectInterface => {
        (protocol.in zip io_in).foreach { case (l,r) => l <> r }
        (io_out zip protocol.out).foreach { case (l,r) => l <> r }
      }
    }
  }
}

val arFIFOMap = WireInit(VecInit(Seq.fill(endId) { true.B }))


package constellation.protocol

import chisel3._
import chisel3.util._

import constellation.channel._
import constellation.noc._
import constellation.router.{RouterCtrlBundle}

import org.chipsalliance.cde.config._
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.subsystem._

import scala.collection.immutable.{ListMap}

// BEGIN: NodeMapping
case class DiplomaticNetworkNodeMapping(
  inNodeMapping: ListMap[String, Int] = ListMap[String, Int](),
  outNodeMapping: ListMap[String, Int] = ListMap[String, Int]()
) {
  // END: NodeMapping
  def genUniqueName(all: Seq[Seq[String]]) = {
    all.zipWithIndex.map { case (strs, i) =>
      val matches = all.take(i).map(_.mkString).count(_ == strs.mkString)
      strs.map(s => s"${s}[${matches}]").mkString(",") + "|"
    }
  }
  def getNode(l: String, nodeMapping: ListMap[String, Int]): Option[Int] = {
    val keys = nodeMapping.keys.toSeq
    val matches = keys.map(k => l.contains(k))
    if (matches.filter(i => i).size == 1) {
      val index = matches.indexWhere(i => i)
      Some(nodeMapping.values.toSeq(index))
    } else {
      None
    }
  }
  def getNodes(ls: Seq[String], mapping: ListMap[String, Int]): Seq[Option[Int]] = {
    ls.map(l => getNode(l, mapping))
  }

  def getNodesIn(ls: Seq[String]): Seq[Option[Int]] = getNodes(ls, inNodeMapping)

}

// BEGIN: ProtocolParams
trait ProtocolParams {
  val minPayloadWidth: Int
  val ingressNodes: Seq[Int]
  val egressNodes: Seq[Int]
  val nVirtualNetworks: Int
  val vNetBlocking: (Int, Int) => Boolean
  val flows: Seq[FlowParams]
  def genIO()(implicit p: Parameters): Data
  def interface(
    terminals: NoCTerminalIO,
    ingressOffset: Int,
    egressOffset: Int,
    protocol: Data)(implicit p: Parameters)
}
// END: ProtocolParams

// BEGIN: ProtocolNoC
case class ProtocolNoCParams(
  nocParams: NoCParams,
  protocolParams: Seq[ProtocolParams],
  widthDivision: Int = 1,
  inlineNoC: Boolean = false
)
class ProtocolNoC(params: ProtocolNoCParams)(implicit p: Parameters) extends Module {
  val io = IO(new Bundle {
    val ctrl = if (params.nocParams.hasCtrl) Vec(params.nocParams.topology.nNodes, new RouterCtrlBundle) else Nil
    val protocol = MixedVec(params.protocolParams.map { u => u.genIO() })
  })
  // END: ProtocolNoC

  if (params.inlineNoC) {
    chisel3.experimental.annotate(this)(Seq(firrtl.passes.InlineAnnotation(toNamed)))
  }

  val protocolParams       = params.protocolParams
  val minPayloadWidth      = protocolParams.map(_.minPayloadWidth).max
  val nocPayloadWidth      = math.ceil(minPayloadWidth.toDouble / params.widthDivision).toInt
  val terminalPayloadWidth = nocPayloadWidth * params.widthDivision
  val ingressOffsets  = protocolParams.map(_.ingressNodes.size).scanLeft(0)(_+_)
  val egressOffsets   = protocolParams.map(_.egressNodes.size).scanLeft(0)(_+_)
  val vNetOffsets     = protocolParams.map(_.nVirtualNetworks).scanLeft(0)(_+_)

  val nocParams = params.nocParams.copy(
    ingresses = protocolParams.map(_.ingressNodes).flatten.map(i =>
      UserIngressParams(i, payloadBits=terminalPayloadWidth)),
    egresses = protocolParams.map(_.egressNodes).flatten.map(i =>
      UserEgressParams(i, payloadBits=terminalPayloadWidth)),
    routerParams = (i) => params.nocParams.routerParams(i).copy(payloadBits=nocPayloadWidth),
    vNetBlocking = (blocker, blockee) => {
      def protocolId(i: Int) = vNetOffsets.drop(1).indexWhere(_ > i)
      if (protocolId(blocker) == protocolId(blockee)) {
        protocolParams(protocolId(blocker)).vNetBlocking(
          blocker - vNetOffsets(protocolId(blocker)),
          blockee - vNetOffsets(protocolId(blockee))
        )
      } else {
        true
      }
    },
    flows = protocolParams.zipWithIndex.map { case (u,i) =>
      u.flows.map(f => f.copy(
        ingressId = f.ingressId + ingressOffsets(i),
        egressId = f.egressId + egressOffsets(i),
        vNetId = f.vNetId + vNetOffsets(i)
      ))
    }.flatten
  )
  val noc = Module(LazyModule(new NoC(nocParams)).module)
  noc.io.router_clocks.foreach(_.clock := clock)
  noc.io.router_clocks.foreach(_.reset := reset)
  (noc.io.router_ctrl zip io.ctrl).foreach { case (l, r) => l <> r }
  (protocolParams zip io.protocol).zipWithIndex.foreach { case ((u, io), x) =>
    val terminals = Wire(new NoCTerminalIO(
      noc.io.ingressParams.drop(ingressOffsets(x)).take(u.ingressNodes.size),
      noc.io.egressParams .drop(egressOffsets(x)) .take(u.egressNodes.size)
    ))
    (terminals.ingress zip noc.io.ingress.drop(ingressOffsets(x))).map { case (l,r) => l <> r }
    (terminals.egress  zip  noc.io.egress.drop (egressOffsets(x))).map { case (l,r) => l <> r }
    u.interface(
      terminals,
      ingressOffsets(x),
      egressOffsets(x),
      io)
  }
}

def getNodesOut(ls: Seq[String]): Seq[Option[Int]] = getNodes(ls, outNodeMapping)