///|
priv struct PlanState {
  values : Map[(Int32, Int32), Array[VirtualReg]]
}

///|
fn PlanState::new() -> PlanState {
  { values: Map([]), }
}

///|
fn PlanState::copy(self : PlanState) -> PlanState {
  let values : Map[(Int32, Int32), Array[VirtualReg]] = Map([])
  for key, value in self.values {
    values[key] = value.copy()
  }
  { values, }
}

///|
fn PlanState::equals(self : PlanState, other : PlanState) -> Bool {
  if self.values.length() != other.values.length() {
    return false
  }
  for key, values in self.values {
    guard other.values.get(key) is Some(other_values) else { return false }
    if values.length() != other_values.length() ||
      !values.all(value => other_values.contains(value)) {
      return false
    }
  }
  true
}

///|
fn PlanState::meet_with(self : PlanState, other : PlanState) -> Unit {
  for key in self.values.keys().to_array() {
    match other.values.get(key) {
      None => self.values.remove(key)
      Some(other_values) => {
        let retained = self.values[key].filter(value => {
          other_values.contains(value)
        })
        if retained.is_empty() {
          self.values.remove(key)
        } else {
          self.values[key] = retained
        }
      }
    }
  }
}

///|
fn PlanState::remove_value(self : PlanState, value : VirtualReg) -> Unit {
  for key in self.values.keys().to_array() {
    self.values[key].retain(other => other != value)
    if self.values[key].is_empty() {
      self.values.remove(key)
    }
  }
}

///|
fn PlanState::remove_register(self : PlanState, register : PhysicalReg) -> Unit {
  self.values.remove(physical_reg_key(register))
}

///|
fn[S : AllocationSink] plan_location_is_legal(
  environment : MachineEnv,
  plan : S,
  value : VirtualReg,
  location : Location,
) -> Bool {
  match location.register() {
    Some(reg) =>
      reg.class == value.class &&
      (
        environment.allocatable_regs.contains(reg) ||
        environment.scratch_regs.contains(reg)
      )
    None => {
      let slot = location.spill_index().unwrap()
      if plan.spill_slot(slot) is Some(spec) {
        spec.size > 0 && spec.alignment > 0
      } else {
        false
      }
    }
  }
}

///|
fn[S : AllocationSink] require_plan_location(
  environment : MachineEnv,
  plan : S,
  value : VirtualReg,
  location : Location,
) -> Unit raise VerifyError {
  if !plan_location_is_legal(environment, plan, value, location) {
    raise InvalidPlan(message="illegal allocation location for v\{value.id}")
  }
}

///|
fn[S : AllocationSink] operand_location_is_legal(
  environment : MachineEnv,
  plan : S,
  operand : Operand,
  location : Location,
) -> Bool {
  if operand.constraint is PreservedHome {
    return plan.value_location(operand.vreg.id) == Some(location)
  }
  if plan_location_is_legal(environment, plan, operand.vreg, location) {
    return true
  }
  match operand.constraint {
    FixedReg(_) =>
      match location.register() {
        Some(reg) =>
          reg.class == operand.vreg.class &&
          environment.fixed_operand_regs.contains(reg)
        None => false
      }
    AnyReg | AnyLocation | PreservedHome => false
  }
}

///|
fn[S : AllocationSink] require_operand_location(
  environment : MachineEnv,
  plan : S,
  operand : Operand,
  location : Location,
) -> Unit raise VerifyError {
  if !operand_location_is_legal(environment, plan, operand, location) {
    raise InvalidPlan(
      message="illegal operand allocation location for v\{operand.vreg.id}",
    )
  }
}

///|
fn require_value_at(
  state : PlanState,
  value : VirtualReg,
  location : Location,
) -> Unit raise VerifyError {
  if !state.values
    .get(location_key(location))
    .map(values => values.contains(value))
    .unwrap_or(false) {
    raise IncorrectValue(vreg=value, location~)
  }
}

///|
fn apply_plan_edit(
  state : PlanState,
  edit : AllocationEdit,
) -> Unit raise VerifyError {
  require_value_at(state, edit.value, edit.from)
  if edit.from.is_spill() && edit.to.is_spill() {
    raise InvalidPlan(message="stack-to-stack allocation edit")
  }
  state.values[location_key(edit.to)] = state.values[location_key(edit.from)].copy()
}

///|
/// Plan edits bucketed by position.
///
/// Verification walks every instruction of every block, and a block may be
/// re-processed when dataflow refines its incoming state. Filtering the
/// function-wide edit array at each step made verification quadratic in
/// practice — a deeply nested function multiplied thousands of instructions
/// by thousands of edits on every visit. One bucketing pass keeps each
/// lookup proportional to the edits that actually apply, and bucket order
/// preserves the plan's edit order.
fn index_plan_edits(
  edits : Array[AllocationEdit],
) -> Map[EditPosition, Array[AllocationEdit]] {
  let index : Map[EditPosition, Array[AllocationEdit]] = Map([])
  for edit in edits {
    match index.get(edit.position) {
      Some(bucket) => bucket.push(edit)
      None => index[edit.position] = [edit]
    }
  }
  index
}

///|
fn apply_position_edits(
  state : PlanState,
  edits_at : Map[EditPosition, Array[AllocationEdit]],
  position : EditPosition,
) -> Unit raise VerifyError {
  guard edits_at.get(position) is Some(bucket) else { return }
  for edit in bucket {
    apply_plan_edit(state, edit)
  }
}

///|
fn[F : FunctionView, S : AllocationSink] process_operand_phase(
  state : PlanState,
  function : F,
  environment : MachineEnv,
  plan : S,
  instruction : Int32,
  timing : OperandTiming,
) -> Unit raise VerifyError {
  for operand_index, operand in function.instruction_operands(instruction) {
    if operand.timing != timing {
      continue
    }
    guard plan.operand_location(instruction, operand_index) is Some(location) else {
      raise InvalidPlan(
        message="missing allocation for instruction \{instruction} operand \{operand_index}",
      )
    }
    require_operand_location(environment, plan, operand, location)
    if operand.constraint is AnyReg && location.is_spill() {
      raise InvalidPlan(
        message="register-constrained operand allocated outside a register",
      )
    }
    if operand.constraint is FixedReg(required) &&
      location != reg_location(required) {
      raise FixedConstraintViolation(
        vreg=operand.vreg,
        required~,
        actual=location,
      )
    }
    if operand.role is Use || operand.role is UseDef {
      require_value_at(state, operand.vreg, location)
    }
    if operand.role is Def || operand.role is UseDef {
      state.remove_value(operand.vreg)
      state.values[location_key(location)] = [operand.vreg]
    }
  }
}

///|
fn[F : FunctionView, S : AllocationSink] process_plan_block(
  state : PlanState,
  function : F,
  environment : MachineEnv,
  plan : S,
  block : Int32,
  edits_at : Map[EditPosition, Array[AllocationEdit]],
) -> PlanState raise VerifyError {
  let current = state.copy()
  for instruction_index in 0.. Unit raise VerifyError {
  for parameter_index in 0.. PlanState raise VerifyError {
  let state = output.copy()
  let source_id = function.block_id_at(block)
  let edge_position = Edge(source_block=source_id, successor_index=successor)
  if edits_at.get(edge_position) is Some(bucket) {
    let original = state.copy()
    for edit in bucket {
      require_value_at(original, edit.value, edit.from)
      state.values[location_key(edit.to)] = original.values[location_key(
        edit.from,
      )].copy()
    }
  }
  let argument_count = function.edge_argument_count(block, successor)
  let parameter_count = function.block_parameter_count(target)
  if argument_count != parameter_count {
    raise InvalidPlan(
      message="edge argument count does not match block parameters",
    )
  }
  for index in 0..
        if !aliases.contains(parameter) {
          aliases.push(parameter)
        }
      None => state.values[home] = [parameter]
    }
  }
  state
}

///|
fn[F : FunctionView, S : AllocationSink] verify_function_allocation(
  function : F,
  environment : MachineEnv,
  plan : S,
) -> Unit raise VerifyError {
  for value in 0.. {
    Array::makei(function.block_successor_count(block), successor => {
      function.block_successor_at(block, successor)
    })
  })
  let has_predecessor = Array::make(block_count, false)
  for source_successors in successors {
    for target in source_successors {
      if target >= 0 && target < block_count {
        has_predecessor[target] = true
      }
    }
  }
  // Process pending blocks in reverse postorder. A plain stack revisits a
  // join-dense region once per refinement, and nested blocks lower to long
  // chains of forward joins, so verification cost grew with nesting depth.
  // In reverse postorder every acyclic region converges in a single pass;
  // loops add one pass per refinement round. Blocks unreachable from any
  // root keep their index order at the end, matching the staged seeding
  // below.
  let priority = Array::make(block_count, block_count)
  let order = reverse_postorder_from_roots(successors, has_predecessor)
  for index, block in order {
    priority[block] = index
  }
  let order = Array::makei(block_count, block => block)
  order.sort_by((left, right) => {
    let by_priority = priority[left] - priority[right]
    if by_priority != 0 {
      by_priority
    } else {
      left - right
    }
  })
  let incoming : Array[PlanState?] = Array::make(block_count, None)
  let pending = Array::make(block_count, false)
  for block in 0.. {
              incoming[target] = Some(next)
              pending[target] = true
            }
            Some(previous) => {
              let merged = previous.copy()
              merged.meet_with(next)
              if !merged.equals(previous) {
                incoming[target] = Some(merged)
                pending[target] = true
              }
            }
          }
        }
      }
    }
    while next_unreachable_root < block_count &&
          incoming[next_unreachable_root] is Some(_) {
      next_unreachable_root = next_unreachable_root + 1
    }
    if next_unreachable_root >= block_count {
      break
    }
    let state = PlanState::new()
    initialize_block_parameters(state, function, plan, next_unreachable_root)
    incoming[next_unreachable_root] = Some(state)
    pending[next_unreachable_root] = true
  }
}