// Cranelift-style regalloc output consumed directly by the emitter.
//
// Design note:
// - Keep MachV instructions unchanged after regalloc.
// - Regalloc produces per-operand allocations plus a stream of edits (moves
//   between regs and spillslots) that the emitter interleaves with original
//   instructions, mirroring Cranelift machinst + regalloc2::Output.

///|
/// A spill slot index (8-byte slot).
pub type SpillSlot = Int

///|
/// A location for a value at a program point: either in a register or in a spill slot.
pub(all) enum Loc {
  Reg(@abi.PReg)
  Spill(SpillSlot)
}

///|
fn Loc::to_string(self : Loc) -> String {
  match self {
    Reg(r) => "\{r}"
    Spill(slot) => "spill(\{slot})"
  }
}

///|
pub impl Show for Loc with fn output(self, logger) {
  logger.write_string(self.to_string())
}

///|
/// An edit produced by regalloc to be inserted at a program point.
///
/// Equivalent to regalloc2::Edit::Move in Cranelift: move between two locations.
pub(all) enum Edit {
  Move(Loc, Loc, @abi.RegClass)
}

///|
fn Edit::to_string(self : Edit) -> String {
  match self {
    Move(from, to, _class) => "move \{from} -> \{to}"
  }
}

///|
pub impl Show for Edit with fn output(self, logger) {
  logger.write_string(self.to_string())
}

///|
/// Regalloc output.
///
/// - `allocs` is a flat array of operand allocations, aligned with `operand_ranges`.
/// - `edits` is a list of edits keyed by a program point (block/inst/pos).
/// - `num_spillslots` is used by stack-frame layout.
pub struct Output {
  // Locations of function parameters after prologue, aligned with `func.params`.
  param_locs : Array[Loc]
  // Flat operand allocations.
  allocs : Array[Loc]
  // (block_id, inst_idx, is_terminator, range_start, def_count, use_count)
  // This mirrors Cranelift’s operand_ranges; we keep block-local inst indexing
  // and treat the terminator as inst_idx==block.insts.length().
  operand_ranges : Array[(Int, Int, Bool, Int, Int, Int)]
  // Dense quick lookups for operand range indices.
  inst_operand_range_index : Array[Array[Int]]
  term_operand_range_index : Array[Int]
  // Edits to insert at program points.
  edits : Array[((Int, Int, ProgPos), Edit)]
  // Fast path for emission: grouped edits per program point (dense by block/inst).
  edits_before_dense : Array[Array[Array[Edit]]]
  edits_after_dense : Array[Array[Array[Edit]]]
  // Any-use preg bitmap (including params/allocs/edits) for fast membership checks.
  used_int_pregs_any : Array[Bool]
  used_fp_pregs_any : Array[Bool]
  // Total number of spill slots reserved.
  mut num_spillslots : Int
}

///|
pub fn Output::Output() -> Output {
  {
    param_locs: [],
    allocs: [],
    operand_ranges: [],
    inst_operand_range_index: [],
    term_operand_range_index: [],
    edits: [],
    edits_before_dense: [],
    edits_after_dense: [],
    used_int_pregs_any: Array::make(128, false),
    used_fp_pregs_any: Array::make(128, false),
    num_spillslots: 0,
  }
}

///|
/// A compact human-readable summary for debugging (used by the CLI explore command).
pub fn Output::summary(self : Output) -> String {
  let mut s = "Regalloc output (Cranelift-style):\n"

  // Edits breakdown: this approximates stack traffic induced by spills/reloads.
  let mut reloads = 0
  let mut spills = 0
  let mut reg_moves = 0
  let mut spill_to_spill = 0
  let slot_traffic : Map[Int, (Int, Int)] = Map([]) // slot -> (reloads, spills)
  for _, entry in self.edits {
    let (_point, edit) = entry
    match edit {
      Move(from, to, _) =>
        match (from, to) {
          (Spill(slot), Reg(_)) => {
            reloads += 1
            let (r, w) = match slot_traffic.get(slot) {
              Some(v) => v
              None => (0, 0)
            }
            slot_traffic.set(slot, (r + 1, w))
          }
          (Reg(_), Spill(slot)) => {
            spills += 1
            let (r, w) = match slot_traffic.get(slot) {
              Some(v) => v
              None => (0, 0)
            }
            slot_traffic.set(slot, (r, w + 1))
          }
          (Reg(_), Reg(_)) => reg_moves += 1
          (Spill(_), Spill(_)) => spill_to_spill += 1
        }
    }
  }

  // Count distinct spill slots referenced (may be less than reserved slots due to reuse).
  let used_slots : @hashset.HashSet[Int] = HashSet([])
  fn record_slot(used : @hashset.HashSet[Int], loc : Loc) -> Unit {
    if loc is Spill(slot) {
      used.add(slot) |> ignore
    }
  }

  for loc in self.param_locs {
    record_slot(used_slots, loc)
  }
  for loc in self.allocs {
    record_slot(used_slots, loc)
  }
  for _, entry in self.edits {
    let (_point, edit) = entry
    match edit {
      Move(from, to, _) => {
        record_slot(used_slots, from)
        record_slot(used_slots, to)
      }
    }
  }

  // Operand allocation breakdown and unique preg usage.
  let mut alloc_reg = 0
  let mut alloc_spill = 0
  let used_int_pregs : @hashset.HashSet[Int] = HashSet([])
  let used_fp_pregs : @hashset.HashSet[Int] = HashSet([])
  let used_vec_pregs : @hashset.HashSet[Int] = HashSet([])
  fn record_preg(
    used_int : @hashset.HashSet[Int],
    used_fp : @hashset.HashSet[Int],
    used_vec : @hashset.HashSet[Int],
    preg : @abi.PReg,
  ) -> Unit {
    match preg.class {
      Int => used_int.add(preg.index) |> ignore
      Float32 | Float64 => used_fp.add(preg.index) |> ignore
      Vector => used_vec.add(preg.index) |> ignore
    }
  }

  for loc in self.param_locs {
    if loc is Reg(p) {
      record_preg(used_int_pregs, used_fp_pregs, used_vec_pregs, p)
    }
  }
  for loc in self.allocs {
    match loc {
      Reg(p) => {
        alloc_reg += 1
        record_preg(used_int_pregs, used_fp_pregs, used_vec_pregs, p)
      }
      Spill(_) => alloc_spill += 1
    }
  }

  // Top spill slots by traffic (loads+stores).
  let hot_slots : Array[(Int, Int, Int, Int)] = [] // (total, slot, reloads, spills)
  for slot, counts in slot_traffic {
    let (r, w) = counts
    hot_slots.push((r + w, slot, r, w))
  }
  hot_slots.sort_by(fn(a, b) {
    let (at, a_slot, _, _) = a
    let (bt, b_slot, _, _) = b
    if at > bt {
      -1
    } else if at < bt {
      1
    } else if a_slot < b_slot {
      -1
    } else if a_slot > b_slot {
      1
    } else {
      0
    }
  })

  // Top edit points by edit count.
  let hot_points : Array[(Int, Int, ProgPos, Int)] = [] // (block, inst, pos, count)
  fn collect_hot_points(
    table : Array[Array[Array[Edit]]],
    pos : ProgPos,
    hot_points : Array[(Int, Int, ProgPos, Int)],
  ) -> Unit {
    for block_id, inst_rows in table {
      for inst_idx, edits in inst_rows {
        if edits.length() > 0 {
          hot_points.push((block_id, inst_idx, pos, edits.length()))
        }
      }
    }
  }
  collect_hot_points(self.edits_before_dense, Before, hot_points)
  collect_hot_points(self.edits_after_dense, After, hot_points)
  hot_points.sort_by(fn(a, b) {
    let (_, _, _, ac) = a
    let (_, _, _, bc) = b
    if ac > bc {
      -1
    } else if ac < bc {
      1
    } else {
      0
    }
  })
  s = s +
    "  spillslots: \{self.num_spillslots} (\{self.num_spillslots * 8} bytes)\n"
  s = s + "  spillslots_used: \{used_slots.length()}\n"
  s = s +
    "  edits: \{self.edits.length()} (reloads=\{reloads}, spills=\{spills}, reg_moves=\{reg_moves}, spill_to_spill=\{spill_to_spill})\n"
  s = s +
    "  operand_allocs: \{self.allocs.length()} (reg=\{alloc_reg}, spill=\{alloc_spill})\n"
  s = s +
    "  regs_used: int=\{used_int_pregs.length()}, fp=\{used_fp_pregs.length()}, vec=\{used_vec_pregs.length()}\n"
  if used_int_pregs.length() > 16 {
    let regs : Array[Int] = []
    for r in used_int_pregs {
      regs.push(r)
    }
    regs.sort()
    s = s + "  regs_used_int_indices: \{to_repr(regs)}\n"
  }
  s = s + "  operand_ranges: \{self.operand_ranges.length()}\n"
  if !hot_slots.is_empty() {
    s = s + "  hot_spillslots: "
    let mut shown = 0
    for entry in hot_slots {
      let (total, slot, r, w) = entry
      if shown >= 8 {
        break
      }
      if shown > 0 {
        s = s + ", "
      }
      s = s + "\{slot}:\{total}(\{r}L+\{w}S)"
      shown += 1
    }
    s = s + "\n"
  }
  if !hot_points.is_empty() {
    fn pos_str(pos : ProgPos) -> String {
      match pos {
        Before => "b"
        After => "a"
      }
    }

    s = s + "  hot_edit_points: "
    let mut shown = 0
    for p in hot_points {
      let (b, i, pos, c) = p
      if shown >= 6 {
        break
      }
      if shown > 0 {
        s = s + ", "
      }
      s = s + "\{b}:\{i}\{pos_str(pos)}=\{c}"
      shown += 1
    }
    s = s + "\n"
  }
  s = s + "  params: "
  for i in 0.. 0 {
      s = s + ", "
    }
    s = s + self.param_locs[i].to_string()
  }
  s + "\n"
}

///|
/// Return stack-traffic counts inferred from regalloc edits:
/// (spills, reloads, reg_moves, spill_to_spill).
pub fn Output::spill_reload_stats(self : Output) -> (Int, Int, Int, Int) {
  let mut reloads = 0
  let mut spills = 0
  let mut reg_moves = 0
  let mut spill_to_spill = 0
  for _, entry in self.edits {
    let (_point, edit) = entry
    match edit {
      Move(from, to, _) =>
        match (from, to) {
          (Spill(_), Reg(_)) => reloads += 1
          (Reg(_), Spill(_)) => spills += 1
          (Reg(_), Reg(_)) => reg_moves += 1
          (Spill(_), Spill(_)) => spill_to_spill += 1
        }
    }
  }
  (spills, reloads, reg_moves, spill_to_spill)
}

///|
/// Validate that the regalloc output only uses legal physical registers for the
/// selected ISA.
///
/// This is a fail-fast guard: silent truncation in x86 encoders can clobber rsp/rbp
/// if an out-of-range preg leaks into emission (e.g. x20 -> rsp).
pub fn Output::validate_for_isa(self : Output, isa : @isa.ISA) -> Unit {
  if !(isa is AMD64) {
    return
  }
  fn alloc_context(self : Output, alloc_idx : Int) -> String {
    for range in self.operand_ranges {
      let (block_id, inst_idx, is_terminator, range_start, def_count, use_count) = range
      let total = def_count + use_count
      if alloc_idx >= range_start && alloc_idx < range_start + total {
        let rel = alloc_idx - range_start
        let kind = if rel < def_count {
          "def[\{rel}]"
        } else {
          "use[\{rel - def_count}]"
        }
        let t = if is_terminator { "t" } else { "i" }
        return "block[\{block_id}] \{t}[\{inst_idx}] \{kind}"
      }
    }
    "unknown"
  }

  fn validate_preg(p : @abi.PReg, ctx_desc_fn : () -> String) -> Unit {
    match p.class {
      Int => {
        guard p.index >= 0 && p.index < 16 else {
          abort(
            "amd64 regalloc output contains invalid Int preg \{p} at \{ctx_desc_fn()}",
          )
        }
        // rsp/rbp are never allocatable as value registers.
        guard p.index != 4 && p.index != 5 else {
          abort(
            "amd64 regalloc output contains reserved Int preg \{p} at \{ctx_desc_fn()}",
          )
        }
      }
      Float32 | Float64 | Vector => {
        guard p.index >= 0 && p.index < 16 else {
          abort(
            "amd64 regalloc output contains invalid FP/Vec preg \{p} at \{ctx_desc_fn()}",
          )
        }
      }
    }
  }

  fn validate_loc(loc : Loc, ctx_desc_fn : () -> String) -> Unit {
    if loc is Reg(p) {
      validate_preg(p, ctx_desc_fn)
    }
  }

  for i in 0.. String {
      match pos {
        Before => "b"
        After => "a"
      }
    }

    let ctx_desc = "edit[\{block_id}:\{inst_idx}:\{pos_str(pos)}]"
    match edit {
      Move(from, to, _class) => {
        validate_loc(from, fn() { ctx_desc + ".from" })
        validate_loc(to, fn() { ctx_desc + ".to" })
      }
    }
  }
}

///|
pub fn Output::get_num_spillslots(self : Output) -> Int {
  self.num_spillslots
}

///|
pub fn Output::get_param_loc(self : Output, idx : Int) -> Loc {
  self.param_locs[idx]
}

///|
pub fn Output::get_num_params(self : Output) -> Int {
  self.param_locs.length()
}

///|
pub fn Output::iter_allocs(self : Output) -> Array[Loc] {
  self.allocs
}

///|
fn Output::mark_used_preg_any(self : Output, loc : Loc) -> Unit {
  if loc is Reg(preg) {
    if preg.index < 0 {
      return
    }
    match preg.class {
      Int =>
        if preg.index < self.used_int_pregs_any.length() {
          self.used_int_pregs_any[preg.index] = true
        }
      Float32 | Float64 | Vector =>
        if preg.index < self.used_fp_pregs_any.length() {
          self.used_fp_pregs_any[preg.index] = true
        }
    }
  }
}

///|
pub fn Output::uses_preg_index_any(
  self : Output,
  preg_idx : Int,
  is_int_class : Bool,
) -> Bool {
  if preg_idx < 0 {
    return false
  }
  if is_int_class {
    preg_idx < self.used_int_pregs_any.length() &&
    self.used_int_pregs_any[preg_idx]
  } else {
    preg_idx < self.used_fp_pregs_any.length() &&
    self.used_fp_pregs_any[preg_idx]
  }
}

///|
pub fn Output::push_param_loc(self : Output, loc : Loc) -> Unit {
  self.param_locs.push(loc)
  self.mark_used_preg_any(loc)
}

///|
pub fn Output::edits_at(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  pos : ProgPos,
) -> Array[Edit]? {
  let table = match pos {
    Before => self.edits_before_dense
    After => self.edits_after_dense
  }
  if block_id < 0 || block_id >= table.length() {
    return None
  }
  let inst_rows = table[block_id]
  if inst_idx < 0 || inst_idx >= inst_rows.length() {
    return None
  }
  let edits = inst_rows[inst_idx]
  if edits.is_empty() {
    None
  } else {
    Some(edits)
  }
}

///|
pub fn Output::push_edit(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  pos : ProgPos,
  edit : Edit,
) -> Unit {
  // Cranelift/regalloc2 does not emit no-op moves; drop them early to reduce
  // code size and make `explore` diagnostics less noisy.
  fn loc_eq(a : Loc, b : Loc) -> Bool {
    match (a, b) {
      (Reg(ra), Reg(rb)) => ra.class == rb.class && ra.index == rb.index
      (Spill(sa), Spill(sb)) => sa == sb
      _ => false
    }
  }

  match edit {
    Move(from, to, _cls) => if loc_eq(from, to) { return }
  }
  self.edits.push(((block_id, inst_idx, pos), edit))
  match edit {
    Move(from, to, _class) => {
      self.mark_used_preg_any(from)
      self.mark_used_preg_any(to)
    }
  }
}

///|
pub fn Output::rebuild_edits_index(self : Output) -> Unit {
  self.edits_before_dense.clear()
  self.edits_after_dense.clear()
  fn push_dense(
    table : Array[Array[Array[Edit]]],
    block_id : Int,
    inst_idx : Int,
    edit : Edit,
  ) -> Unit {
    if block_id < 0 || inst_idx < 0 {
      return
    }
    while table.length() <= block_id {
      table.push([])
    }
    let inst_rows = table[block_id]
    while inst_rows.length() <= inst_idx {
      inst_rows.push([])
    }
    inst_rows[inst_idx].push(edit)
  }
  for entry in self.edits {
    let ((block_id, inst_idx, pos), edit) = entry
    match pos {
      Before => push_dense(self.edits_before_dense, block_id, inst_idx, edit)
      After => push_dense(self.edits_after_dense, block_id, inst_idx, edit)
    }
  }
}

///|
pub fn Output::sort_edits(self : Output) -> Unit {
  fn pos_ord(pos : ProgPos) -> Int {
    match pos {
      Before => 0
      After => 1
    }
  }
  self.edits.sort_by(fn(lhs, rhs) {
    let ((lb, li, lp), _le) = lhs
    let ((rb, ri, rp), _re) = rhs
    if lb != rb {
      lb - rb
    } else if li != ri {
      li - ri
    } else {
      pos_ord(lp) - pos_ord(rp)
    }
  })
}

///|
pub fn Output::finalize_edits(self : Output) -> Unit {
  self.rebuild_edits_index()
}

///|
pub fn Output::iter_edits(self : Output) -> Array[((Int, Int, ProgPos), Edit)] {
  self.edits
}

///|
/// Start recording operand allocations for one instruction/terminator.
/// Returns `(start, total)` where `start` is the first index in `allocs`.
pub fn Output::begin_inst_allocs(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  is_terminator : Bool,
  def_count : Int,
  use_count : Int,
) -> (Int, Int) {
  let start = self.allocs.length()
  let total = def_count + use_count
  let range_idx = self.operand_ranges.length()
  self.operand_ranges.push(
    (block_id, inst_idx, is_terminator, start, def_count, use_count),
  )
  while self.inst_operand_range_index.length() <= block_id {
    self.inst_operand_range_index.push([])
  }
  while self.term_operand_range_index.length() <= block_id {
    self.term_operand_range_index.push(-1)
  }
  if is_terminator {
    self.term_operand_range_index[block_id] = range_idx
  } else {
    let inst_rows = self.inst_operand_range_index[block_id]
    while inst_rows.length() <= inst_idx {
      inst_rows.push(-1)
    }
    inst_rows[inst_idx] = range_idx
  }
  (start, total)
}

///|
pub fn Output::push_inst_alloc_loc(self : Output, loc : Loc) -> Unit {
  self.allocs.push(loc)
  self.mark_used_preg_any(loc)
}

///|
pub fn Output::end_inst_allocs(self : Output, start : Int, total : Int) -> Unit {
  guard self.allocs.length() == start + total else {
    abort("bad inst alloc count")
  }
}

///|
/// Record operand allocations for one instruction/terminator.
/// `locs` must have length `def_count + use_count`, in that order.
pub fn Output::push_inst_allocs(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  is_terminator : Bool,
  def_count : Int,
  use_count : Int,
  locs : Array[Loc],
) -> Unit {
  let (start, total) = self.begin_inst_allocs(
    block_id, inst_idx, is_terminator, def_count, use_count,
  )
  for loc in locs {
    self.push_inst_alloc_loc(loc)
  }
  self.end_inst_allocs(start, total)
}

///|
fn Output::range_info(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  is_terminator : Bool,
) -> (Int, Int, Int) {
  let idx = if is_terminator {
    guard block_id >= 0 && block_id < self.term_operand_range_index.length() else {
      abort("bad block_id for terminator alloc lookup")
    }
    self.term_operand_range_index[block_id]
  } else {
    guard block_id >= 0 && block_id < self.inst_operand_range_index.length() else {
      abort("bad block_id for instruction alloc lookup")
    }
    let inst_rows = self.inst_operand_range_index[block_id]
    guard inst_idx >= 0 && inst_idx < inst_rows.length() else {
      abort("bad inst_idx for instruction alloc lookup")
    }
    inst_rows[inst_idx]
  }
  guard idx >= 0 && idx < self.operand_ranges.length() else {
    abort(
      "missing alloc range for block=\{block_id} inst=\{inst_idx} term=\{is_terminator}",
    )
  }
  let (_b, _i, _t, start, defs, uses) = self.operand_ranges[idx]
  (start, defs, uses)
}

///|
pub fn Output::inst_def_loc(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  is_terminator : Bool,
  def_idx : Int,
) -> Loc {
  let (start, defs, _uses) = self.range_info(block_id, inst_idx, is_terminator)
  guard def_idx >= 0 && def_idx < defs else { abort("bad def_idx") }
  self.allocs[start + def_idx]
}

///|
pub fn Output::inst_use_loc(
  self : Output,
  block_id : Int,
  inst_idx : Int,
  is_terminator : Bool,
  use_idx : Int,
) -> Loc {
  let (start, defs, uses) = self.range_info(block_id, inst_idx, is_terminator)
  guard use_idx >= 0 && use_idx < uses else { abort("bad use_idx") }
  self.allocs[start + defs + use_idx]
}