package ifu
import include._
import lib._
import chisel3._
import chisel3.util._

class el2_ifu_bp_ctl extends Module with el2_lib {
  val io = IO (new Bundle {
    val clk = Input(Bool())
    val active_clk = Input(Bool())
    val rst_l = Input(Bool())
    val ic_hit_f = Input(Bool())
    val ifc_fetch_addr_f = Input(UInt(32.W))
    val ifc_fetch_req_f = Input(Bool())           // Fetch request generated by the IFC
    // Decode packet containing information if its a brnach or not
    val dec_tlu_br0_r_pkt = Input(new el2_br_tlu_pkt_t)
    val exu_i0_br_fghr_r = Input(UInt(BHT_GHR_SIZE.W)) // Updated GHR from the exu
    val exu_i0_br_index_r = Input(UInt((BTB_ADDR_HI-BTB_ADDR_LO).W)) // Way from where the btb got a hit
    val dec_tlu_flush_lower_wb = Input(Bool())
    val dec_tlu_flush_leak_one_wb = Input(Bool())
    val dec_tlu_bpred_disable = Input(Bool())
    // Exu misprediction packet
    val exu_mp_pkt = Input(new el2_predict_pkt_t)
    val exu_mp_eghr = Input(UInt(BHT_GHR_SIZE.W))
    val exu_mp_fghr = Input(UInt(BHT_GHR_SIZE.W))
    val exu_mp_index = Input(UInt((BTB_ADDR_HI-BTB_ADDR_LO).W)) // Misprediction index
    val exu_mp_btag = Input(UInt(BTB_BTAG_SIZE.W))
    val exu_flush_final = Input(Bool())

    val ifu_bp_hit_taken_f = Output(Bool())
    val ifu_bp_btb_target_f = Output(UInt(31.W))
    val ifu_bp_inst_mask_f = Output(Bool())
    val ifu_bp_fghr_f = Output(UInt(BHT_GHR_SIZE.W))
    val ifu_bp_way_f = Output(UInt(2.W))
    val ifu_bp_ret_f = Output(UInt(2.W))
    val ifu_bp_hist1_f = Output(UInt(2.W))
    val ifu_bp_hist0_f = Output(UInt(2.W))
    val ifu_bp_pc4_f = Output(UInt(2.W))
    val ifu_bp_valid_f = Output(UInt(2.W))
    val ifu_bp_poffset_f = Output(UInt(12.W))
  })
  io.ifu_bp_hit_taken_f := 0.U
  io.ifu_bp_btb_target_f := 0.U
  io.ifu_bp_inst_mask_f := 0.U
  io.ifu_bp_fghr_f := 0.U
  io.ifu_bp_way_f := 0.U
  io.ifu_bp_ret_f := 0.U
  io.ifu_bp_hist1_f := 0.U
  io.ifu_bp_hist0_f := 0.U
  io.ifu_bp_pc4_f := 0.U
  io.ifu_bp_valid_f := 0.U
  io.ifu_bp_poffset_f := 0.U

  val TAG_START = 16+BTB_BTAG_SIZE
  val PC4 = 4
  val BOFF = 3
  val CALL = 2
  val RET = 1
  val BV = 0
  val LRU_SIZE = BTB_ARRAY_DEPTH
  val NUM_BHT_LOOP = if(BHT_ARRAY_DEPTH > 16) 16 else BHT_ARRAY_DEPTH
  val NUM_BHT_LOOP_INNER_HI = if(BHT_ARRAY_DEPTH > 16) BHT_ADDR_LO+3 else BHT_ADDR_HI
  val NUM_BHT_LOOP_OUTER_LO = if(BHT_ARRAY_DEPTH > 16) BHT_ADDR_LO+4 else BHT_ADDR_LO
  val BHT_NO_ADDR_MATCH = BHT_ARRAY_DEPTH <= 16

  val leak_one_f = WireInit(Bool(), 0.U)
  val exu_mp_valid = io.exu_mp_pkt.misp & ~leak_one_f
  val exu_mp_boffset = io.exu_mp_pkt.boffset
  val exu_mp_pc4 = io.exu_mp_pkt.pc4
  val exu_mp_call = io.exu_mp_pkt.pcall
  val exu_mp_ret = io.exu_mp_pkt.pret
  val exu_mp_ja = io.exu_mp_pkt.pja
  val exu_mp_way = io.exu_mp_pkt.way
  val exu_mp_hist = io.exu_mp_pkt.hist
  val exu_mp_tgt = io.exu_mp_pkt.toffset
  val exu_mp_addr = io.exu_mp_index
  val exu_mp_ataken = io.exu_mp_pkt.ataken

  // LM: Its a commit or update packet
  val dec_tlu_br0_v_wb = io.dec_tlu_br0_r_pkt.valid
  val dec_tlu_br0_hist_wb = io.dec_tlu_br0_r_pkt.hist
  val dec_tlu_br0_addr_wb = io.exu_i0_br_index_r
  val dec_tlu_br0_error_wb = io.dec_tlu_br0_r_pkt.br_error
  val dec_tlu_br0_middle_wb = io.dec_tlu_br0_r_pkt.middle
  val dec_tlu_br0_way_wb = io.dec_tlu_br0_r_pkt.way
  val dec_tlu_br0_start_error_wb = io.dec_tlu_br0_r_pkt.br_start_error
  val exu_i0_br_fghr_wb = io.exu_i0_br_fghr_r

  val fetch_rd_tag_p1_f = WireInit(UInt(BTB_BTAG_SIZE.W), 0.U)
  val fetch_rd_tag_f = WireInit(UInt(BTB_BTAG_SIZE.W), 0.U)
  val bht_dir_f = WireInit(UInt(2.W), 0.U)
  val dec_tlu_error_wb = WireInit(Bool(), 0.U)
  val btb_error_addr_wb = WireInit(UInt((BTB_ADDR_HI-BTB_ADDR_LO).W), 0.U)
  val btb_bank0_rd_data_way0_f = WireInit(UInt((TAG_START+1).W), 0.U)
  val btb_bank0_rd_data_way1_f = WireInit(UInt((TAG_START+1).W), 0.U)
  val btb_bank0_rd_data_way0_p1_f = WireInit(UInt((TAG_START+1).W), 0.U)
  val btb_bank0_rd_data_way1_p1_f = WireInit(UInt((TAG_START+1).W), 0.U)
  val eoc_mask = WireInit(Bool(), 0.U)
  val btb_lru_b0_f = WireInit(UInt(LRU_SIZE.W), init = 0.U)
  val dec_tlu_way_wb = WireInit(Bool(), 0.U)


  // Hash the first PC
  val btb_rd_addr_f = el2_btb_addr_hash(io.ifc_fetch_addr_f)
  // Second pc = pc +4
  val fetch_addr_p1_f = io.ifc_fetch_addr_f + 4.U
  // Hash the second pc
  val btb_rd_addr_p1_f = el2_btb_addr_hash(fetch_addr_p1_f)

  // TODO
  val btb_sel_f = Cat(~bht_dir_f(0),bht_dir_f(0))

  // Checking of the pc is a multiple of 4, if it is fetch-start will be "01"
  val fetch_start_f = Cat(io.ifc_fetch_addr_f(1),~io.ifc_fetch_addr_f(1))

  // If there is an error write-back from the dec check if the current pc is equal to the write-bcak pc
  val branch_error_collision_f = dec_tlu_error_wb & (btb_error_addr_wb === btb_rd_addr_f)
  val branch_error_collision_p1_f = dec_tlu_error_wb & (btb_error_addr_wb === btb_rd_addr_p1_f)

  // If there is an error write back but the address are from different bank
  val branch_error_bank_conflict_f = branch_error_collision_f & dec_tlu_error_wb
  val branch_error_bank_conflict_p1_f = branch_error_collision_p1_f & dec_tlu_error_wb

  // There is a misprediction and the exu is writing back
  val fetch_mp_collision_f = (io.exu_mp_btag === fetch_rd_tag_f) & exu_mp_valid & io.ifc_fetch_req_f & (exu_mp_addr === btb_rd_addr_f)
  val fetch_mp_collision_p1_f = (io.exu_mp_btag === fetch_rd_tag_p1_f) & exu_mp_valid & io.ifc_fetch_req_f & (exu_mp_addr === btb_rd_addr_p1_f)

  val leak_one_f_d1 = RegNext(leak_one_f, init = 0.U)
  val dec_tlu_way_wb_f = RegNext(dec_tlu_way_wb, init = 0.U)
  val exu_mp_way_f = RegNext(exu_mp_way, init = 0.U)
  val exu_flush_final_d1 = RegNext(io.exu_flush_final, init = 0.U)
  // TODO
  leak_one_f := (io.dec_tlu_flush_leak_one_wb & io.dec_tlu_flush_lower_wb) | (leak_one_f_d1 & io.dec_tlu_flush_lower_wb)

  // For a tag to match the branch should be valid tag should match and a fetch request should be generated
  val tag_match_way0_f = btb_bank0_rd_data_way0_f(BV) & (btb_bank0_rd_data_way0_f(TAG_START,17) === fetch_rd_tag_f) &
  ~(dec_tlu_way_wb_f & branch_error_bank_conflict_f) & io.ifc_fetch_req_f & ~leak_one_f

  val tag_match_way1_f = btb_bank0_rd_data_way1_f(BV) & (btb_bank0_rd_data_way1_f(TAG_START,17) === fetch_rd_tag_f) &
  ~(dec_tlu_way_wb_f & branch_error_bank_conflict_f) & io.ifc_fetch_req_f & ~leak_one_f

  val tag_match_way0_p1_f = btb_bank0_rd_data_way0_p1_f(BV) & (btb_bank0_rd_data_way0_p1_f(TAG_START,17) === fetch_rd_tag_p1_f) &
    ~(dec_tlu_way_wb_f & branch_error_bank_conflict_f) & io.ifc_fetch_req_f & ~leak_one_f

  val tag_match_way1_p1_f = btb_bank0_rd_data_way1_p1_f(BV) & (btb_bank0_rd_data_way1_p1_f(TAG_START,17) === fetch_rd_tag_p1_f) &
    ~(dec_tlu_way_wb_f & branch_error_bank_conflict_f) & io.ifc_fetch_req_f & ~leak_one_f

  // Reordering to avoid multiple hit
  val tag_match_way0_expanded_f = Cat(tag_match_way0_f &  (btb_bank0_rd_data_way0_f(BOFF) ^ btb_bank0_rd_data_way0_f(PC4)),
    tag_match_way0_f & ~(btb_bank0_rd_data_way0_f(BOFF) ^ btb_bank0_rd_data_way0_f(PC4)))

  val tag_match_way1_expanded_f = Cat(tag_match_way1_f &  (btb_bank0_rd_data_way1_f(BOFF) ^ btb_bank0_rd_data_way1_f(PC4)),
    tag_match_way1_f & ~(btb_bank0_rd_data_way1_f(BOFF) ^ btb_bank0_rd_data_way1_f(PC4)))


  val tag_match_way0_expanded_p1_f = Cat(tag_match_way0_p1_f &  (btb_bank0_rd_data_way0_p1_f(BOFF) ^ btb_bank0_rd_data_way0_p1_f(PC4)),
    tag_match_way0_p1_f & ~(btb_bank0_rd_data_way0_p1_f(BOFF) ^ btb_bank0_rd_data_way0_p1_f(PC4)))

  val tag_match_way1_expanded_p1_f = Cat(tag_match_way1_p1_f &  (btb_bank0_rd_data_way1_p1_f(BOFF) ^ btb_bank0_rd_data_way1_p1_f(PC4)),
    tag_match_way1_p1_f & ~(btb_bank0_rd_data_way1_p1_f(BOFF) ^ btb_bank0_rd_data_way1_p1_f(PC4)))

  val wayhit_f = tag_match_way0_expanded_f | tag_match_way1_expanded_f

  val wayhit_p1_f = tag_match_way0_expanded_p1_f | tag_match_way1_expanded_p1_f

  // Chopping off the ways that had a hit
  val btb_bank0e_rd_data_f = Mux1H(Seq(tag_match_way0_expanded_f(0).asBool->btb_bank0_rd_data_way0_f,
    tag_match_way1_expanded_f(0).asBool->btb_bank0_rd_data_way1_f))

  val btb_bank0o_rd_data_f = Mux1H(Seq(tag_match_way0_expanded_f(1).asBool->btb_bank0_rd_data_way0_f,
    tag_match_way1_expanded_f(1).asBool->btb_bank0_rd_data_way1_f))

  val btb_bank0e_rd_data_p1_f = Mux1H(Seq(tag_match_way0_expanded_p1_f(0).asBool->btb_bank0_rd_data_way0_p1_f,
    tag_match_way1_expanded_p1_f(1).asBool->btb_bank0_rd_data_way1_p1_f))

  // Making virtual banks, made bit 1 of the pc to check
  val btb_vbank0_rd_data_f = Mux1H(Seq(~io.ifc_fetch_addr_f(1)->btb_bank0e_rd_data_f,
    io.ifc_fetch_addr_f(1)->btb_bank0o_rd_data_f))

  val btb_vbank1_rd_data_f = Mux1H(Seq(~io.ifc_fetch_addr_f(0)->btb_bank0o_rd_data_f,
    io.ifc_fetch_addr_f(0)->btb_bank0e_rd_data_p1_f))

  // Implimenting the LRU for a 2-way BTB
  val mp_wrindex_dec = 1.U(LRU_SIZE) << exu_mp_addr
  val fetch_wrindex_dec = 1.U(LRU_SIZE) << btb_rd_addr_f
  val fetch_wrindex_p1_dec = 1.U(LRU_SIZE) << btb_rd_addr_p1_f
  val mp_wrlru_b0 = mp_wrindex_dec & Fill(LRU_SIZE, exu_mp_valid)

  val vwayhit_f = Mux1H(Seq(~io.ifc_fetch_addr_f(1).asBool->wayhit_f,
    io.ifc_fetch_addr_f(1).asBool->Cat(wayhit_p1_f(0), wayhit_f(1)))) & Cat(eoc_mask, 1.U(1.W))
  val lru_update_valid_f = (vwayhit_f(0) | vwayhit_f(1)) & io.ifc_fetch_req_f & ~leak_one_f

  val fetch_wrlru_b0 = fetch_wrindex_dec & Fill(fetch_wrindex_dec.getWidth, lru_update_valid_f)
  val fetch_wrlru_p1_b0 = fetch_wrindex_p1_dec & Fill(fetch_wrindex_dec.getWidth, lru_update_valid_f)

  val btb_lru_b0_hold = ~mp_wrlru_b0 & ~fetch_wrlru_b0
  val use_mp_way = fetch_mp_collision_f
  val use_mp_way_p1 = fetch_mp_collision_p1_f

  val btb_lru_b0_ns = Mux1H(Seq(~exu_mp_way.asBool->mp_wrlru_b0,
    tag_match_way0_f.asBool->fetch_wrlru_b0,tag_match_way0_p1_f.asBool->fetch_wrlru_p1_b0)) | btb_lru_b0_hold & btb_lru_b0_f

  val btb_lru_rd_f = Mux(use_mp_way.asBool, exu_mp_way_f, (fetch_wrindex_dec & btb_lru_b0_f).orR)
  val btb_lru_rd_p1_f = Mux(use_mp_way_p1.asBool, exu_mp_way_f, (fetch_wrindex_p1_dec & btb_lru_b0_f).orR)

  val btb_vlru_rd_f = Mux1H(Seq(~io.ifc_fetch_addr_f(1).asBool->Cat(btb_lru_rd_f, btb_lru_rd_f),
    io.ifc_fetch_addr_f(1).asBool->Cat(btb_lru_rd_p1_f, btb_lru_rd_f)))

  val tag_match_vway1_expanded_f = Mux1H(Seq(~io.ifc_fetch_addr_f(1).asBool->tag_match_way1_expanded_f,
    io.ifc_fetch_addr_f(1).asBool->Cat(tag_match_way1_expanded_p1_f(0),tag_match_way1_expanded_f(1))))

  val way_raw = tag_match_vway1_expanded_f | (~vwayhit_f & btb_vlru_rd_f)

  //val btb_lru_b0_f = RegNext(btb_lru_b0_ns, init = 0.U)
}


object ifu_bp extends App {
  println((new chisel3.stage.ChiselStage).emitVerilog(new el2_ifu_bp_ctl()))
}