///|
pub(all) enum EditPosition {
  Before(Int)
  After(Int)
  Edge(source_block~ : Int, successor_index~ : Int)
} derive(Eq, Hash, Debug)

///|
pub(all) struct AllocationEdit {
  value : VirtualReg
  from : Location
  to : Location
  position : EditPosition
} derive(Eq, Debug)

///|
pub struct OperandAssignment {
  instruction : Int
  operand : Int
  location : Location
} derive(Eq, Debug)

///|
pub struct SpillSlotSpec {
  size : Int
  alignment : Int
} derive(Eq, Debug)

///|
pub struct AllocationPlan {
  priv value_locations : Array[Location?]
  priv operand_locations : Array[Array[Location?]]
  priv edits : Array[AllocationEdit]
  priv spill_slots : Array[SpillSlotSpec]
  priv spill_slot_owners : Array[Int]
}

///|
pub impl Debug for AllocationPlan with fn to_repr(self) {
  Repr::record({
    "value_locations": Repr(self.value_locations),
    "operand_assignments": Repr(self.operand_assignments()),
    "edits": Repr(self.edits),
    "spill_slots": Repr(self.spill_slots),
  })
}

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

///|
fn AllocationPlan::new(value_count : Int) -> AllocationPlan {
  {
    value_locations: Array::make(value_count, None),
    operand_locations: [],
    edits: [],
    spill_slots: [],
    spill_slot_owners: [],
  }
}

///|
fn AllocationPlan::assign_value(
  self : AllocationPlan,
  vreg : VirtualReg,
  location : Location,
) -> Unit {
  if vreg.id >= 0 && vreg.id < self.value_locations.length() {
    self.value_locations[vreg.id] = Some(location)
  }
}

///|
fn AllocationPlan::assign_operand(
  self : AllocationPlan,
  instruction : Int,
  operand : Int,
  location : Location,
) -> Unit {
  while self.operand_locations.length() <= instruction {
    self.operand_locations.push([])
  }
  let locations = self.operand_locations[instruction]
  while locations.length() <= operand {
    locations.push(None)
  }
  locations[operand] = Some(location)
}

///|
fn AllocationPlan::add_edit(
  self : AllocationPlan,
  edit : AllocationEdit,
) -> Unit {
  self.edits.push(edit)
}

///|
fn AllocationPlan::create_spill_slot(
  self : AllocationPlan,
  owner : Int,
  size : Int,
  alignment : Int,
) -> Int {
  let slot = self.spill_slots.length()
  self.spill_slots.push({ size, alignment, })
  self.spill_slot_owners.push(owner)
  slot
}

///|
fn[F : FunctionView] AllocationPlan::seal_operand_storage(
  self : AllocationPlan,
  function : F,
) -> Unit {
  for block in 0.. Location? {
  self.value_locations.get(value).bind(location => location)
}

///|
pub fn AllocationPlan::operand_location(
  self : AllocationPlan,
  instruction : Int,
  operand : Int,
) -> Location? {
  self.operand_locations
  .get(instruction)
  .bind(locations => locations.get(operand).bind(location => location))
}

///|
pub fn AllocationPlan::operand_assignments(
  self : AllocationPlan,
) -> Array[OperandAssignment] {
  let assignments : Array[OperandAssignment] = []
  for instruction, locations in self.operand_locations {
    for operand, location in locations {
      if location is Some(location) {
        assignments.push({ instruction, operand, location, })
      }
    }
  }
  assignments
}

///|
pub fn AllocationPlan::operand_instruction_count(self : AllocationPlan) -> Int {
  self.operand_locations.length()
}

///|
pub fn AllocationPlan::instruction_operand_count(
  self : AllocationPlan,
  instruction : Int,
) -> Int {
  match self.operand_locations.get(instruction) {
    Some(locations) => locations.length()
    None => 0
  }
}

///|
pub fn AllocationPlan::edits(self : AllocationPlan) -> Array[AllocationEdit] {
  self.edits.copy()
}

///|
pub fn AllocationPlan::edit_count(self : AllocationPlan) -> Int {
  self.edits.length()
}

///|
pub fn AllocationPlan::edit_at(
  self : AllocationPlan,
  index : Int,
) -> AllocationEdit? {
  self.edits.get(index)
}

///|
pub fn AllocationPlan::spill_count(self : AllocationPlan) -> Int {
  self.spill_slots.length()
}

///|
pub fn AllocationPlan::spill_slot(
  self : AllocationPlan,
  index : Int,
) -> SpillSlotSpec? {
  self.spill_slots.get(index)
}

///|
pub fn AllocationPlan::spill_slot_owner(
  self : AllocationPlan,
  index : Int,
) -> Int? {
  self.spill_slot_owners.get(index)
}

///|
/// Shares the allocator's finalized dense value-location table with an
/// embedding. The plan must not be mutated after this table is adopted.
pub fn AllocationPlan::value_location_storage(
  self : AllocationPlan,
) -> Array[Location?] {
  self.value_locations
}

///|
/// Shares the allocator's finalized dense operand-location tables with an
/// embedding. `seal_operand_storage` establishes the source function's exact
/// instruction and operand shape before allocation returns.
pub fn AllocationPlan::operand_location_storage(
  self : AllocationPlan,
) -> Array[Array[Location?]] {
  self.operand_locations
}