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

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

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

///|
fn PlanState::equals(self : PlanState, other : PlanState) -> Bool {
  let mut count = 0
  for key, value in self.values {
    count = count + 1
    if other.values.get(key) != Some(value) {
      return false
    }
  }
  let mut other_count = 0
  for _ in other.values {
    other_count = other_count + 1
  }
  count == other_count
}

///|
fn PlanState::meet_with(self : PlanState, other : PlanState) -> Unit {
  let removed : Array[(Int, Int)] = []
  for key, value in self.values {
    if other.values.get(key) != Some(value) {
      removed.push(key)
    }
  }
  for key in removed {
    self.values.remove(key)
  }
}

///|
fn PlanState::remove_value(self : PlanState, value : VirtualReg) -> Unit {
  let removed : Array[(Int, Int)] = []
  for key, existing in self.values {
    if existing == value {
      removed.push(key)
    }
  }
  for key in removed {
    self.values.remove(key)
  }
}

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

///|
fn plan_location_is_legal(
  environment : MachineEnv,
  plan : AllocationPlan,
  value : VirtualReg,
  location : Location,
) -> Bool {
  match location {
    Reg(reg) =>
      reg.class == value.class &&
      (
        environment.allocatable_regs.contains(reg) ||
        environment.scratch_regs.contains(reg)
      )
    Spill(slot) =>
      if plan.spill_slot(slot) is Some(spec) {
        spec.size > 0 && spec.alignment > 0
      } else {
        false
      }
  }
}

///|
fn require_plan_location(
  environment : MachineEnv,
  plan : AllocationPlan,
  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 operand_location_is_legal(
  environment : MachineEnv,
  plan : AllocationPlan,
  operand : Operand,
  location : Location,
) -> Bool {
  if plan_location_is_legal(environment, plan, operand.vreg, location) {
    return true
  }
  match (operand.constraint, location) {
    (FixedReg(_), Reg(reg)) =>
      reg.class == operand.vreg.class &&
      environment.fixed_operand_regs.contains(reg)
    _ => false
  }
}

///|
fn require_operand_location(
  environment : MachineEnv,
  plan : AllocationPlan,
  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)) != Some(value) {
    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)] = edit.value
}

///|
/// 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] process_operand_phase(
  state : PlanState,
  function : F,
  environment : MachineEnv,
  plan : AllocationPlan,
  instruction : Int,
  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(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] process_plan_block(
  state : PlanState,
  function : F,
  environment : MachineEnv,
  plan : AllocationPlan,
  block : Int,
  edits_at : Map[EditPosition, Array[AllocationEdit]],
) -> PlanState raise VerifyError {
  let current = state.copy()
  for instruction in function.block_instructions(block) {
    apply_position_edits(current, edits_at, Before(instruction))
    process_operand_phase(
      current,
      function,
      environment,
      plan,
      instruction,
      Early,
    )
    for clobber in function.instruction_clobbers(instruction) {
      current.remove_register(clobber)
    }
    process_operand_phase(
      current,
      function,
      environment,
      plan,
      instruction,
      Late,
    )
    apply_position_edits(current, edits_at, After(instruction))
  }
  current
}

///|
fn[F : FunctionView] initialize_block_parameters(
  state : PlanState,
  function : F,
  plan : AllocationPlan,
  block : Int,
) -> Unit raise VerifyError {
  for parameter in function.block_parameters(block) {
    guard plan.value_location(parameter.id) is Some(location) else {
      raise Unassigned(vreg=parameter)
    }
    state.remove_value(parameter)
    state.values[location_key(location)] = parameter
  }
}

///|
fn[F : FunctionView] prepare_successor_state(
  output : PlanState,
  function : F,
  plan : AllocationPlan,
  block : Int,
  successor : Int,
  target : Int,
  edits_at : Map[EditPosition, Array[AllocationEdit]],
) -> 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)] = edit.value
    }
  }
  let arguments = function.edge_arguments(block, successor)
  let parameters = function.block_parameters(target)
  if arguments.length() != parameters.length() {
    raise InvalidPlan(
      message="edge argument count does not match block parameters",
    )
  }
  for index, argument in arguments {
    let parameter = parameters[index]
    guard plan.value_location(parameter.id) is Some(parameter_home) else {
      raise Unassigned(vreg=parameter)
    }
    require_value_at(state, argument, parameter_home)
    state.remove_value(parameter)
    state.values[location_key(parameter_home)] = parameter
  }
  state
}

///|
fn[F : FunctionView] verify_function_allocation(
  function : F,
  environment : MachineEnv,
  plan : AllocationPlan,
) -> Unit raise VerifyError {
  for value in 0.. {
    function.block_successors(block)
  })
  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
  }
}