// Licensing scope: see ../../LICENSE and this directory's LICENSE and NOTICE.
// Copyright 2023 WebAssembly Community Group participants.
// Retained historical upstream notice; current derivation remains unconfirmed.
// Conservative attribution from pinned Binaryen version_130 counterpart:
// Copyright 2016 WebAssembly Community Group participants
// Counterpart: src/passes/MergeLocals.cpp (5d704ad5); origin remains unconfirmed.
// Potentially incorporated Binaryen expression retains Apache-2.0 terms.
// Original Starshine contributions use the scoped MIT grant in ../../LICENSE.
// Modified for Starshine through 2026: implementation and validation use
// Starshine's representations and APIs; this is not an unchanged upstream file.
// Conservative notice retention does not establish copying or independence.

///|
pub fn merge_locals_descriptor() -> @ir.HotPassDescriptor {
  @ir.HotPassDescriptor::new(
    "merge-locals",
    requires=[@ir.HotAnalysis::cfg()],
    invalidates=[
      @ir.HotAnalysis::cfg(),
      @ir.HotAnalysis::dominance(),
      @ir.HotAnalysis::liveness(),
      @ir.HotAnalysis::use_def(),
      @ir.HotAnalysis::effects(),
      @ir.HotAnalysis::loop_info(),
      @ir.HotAnalysis::ssa(),
    ],
  )
}

///|
pub fn merge_locals_summary() -> String {
  "Balance copy-shaped local live ranges with CFG-backed bidirectional retargeting and rollback."
}

///|
priv struct MergeLocalsCopy {
  copy_id : @ir.NodeId
  trivial_id : @ir.NodeId
  source_local : @ir.LocalId
  target_local : @ir.LocalId
}

///|
priv enum MergeLocalsOrientation {
  ToCopy
  ToTrivial
}

///|
priv struct MergeLocalsOptimizedCopy {
  copy : MergeLocalsCopy
  orientation : MergeLocalsOrientation
  influenced_gets : Array[@ir.NodeId]
}

///|
fn merge_locals_get_has_only_source(
  graph : @ir.HotLocalGraph,
  get_id : @ir.NodeId,
  write_id : @ir.NodeId,
) -> Bool {
  let sources = @ir.local_graph_reaching_sources_for_get(graph, get_id)
  if sources.length() != 1 {
    return false
  }
  match sources[0] {
    @ir.HotLocalGraphSource::SetSource(found) => found == write_id
    _ => false
  }
}

///|
fn merge_locals_get_value_type(
  func : @ir.HotFunc,
  get_id : @ir.NodeId,
) -> @lib.ValType? {
  if get_id < 0 || !@ir.hot_node_is_live(func, get_id) {
    return None
  }
  let node = @ir.hot_node_get(func, get_id)
  match node.op {
    @ir.HotOp::LocalGet =>
      match @ir.hot_type_get(func, node.ty) {
        @ir.HotTypeInfo::Value(ty) => Some(ty)
        _ => None
      }
    _ => None
  }
}

///|
fn merge_locals_gets_match_target_type(
  func : @ir.HotFunc,
  gets : Array[@ir.NodeId],
  target_local : @ir.LocalId,
) -> Bool {
  let target_type = match @ir.hot_local_type(func, target_local) {
    Some(ty) => ty
    None => return false
  }
  for get_id in gets {
    match merge_locals_get_value_type(func, get_id) {
      Some(ty) => if ty != target_type { return false }
      None => return false
    }
  }
  true
}

///|
fn merge_locals_gets_are_single_source(
  graph : @ir.HotLocalGraph,
  gets : Array[@ir.NodeId],
  write_id : @ir.NodeId,
) -> Bool {
  for get_id in gets {
    if !merge_locals_get_has_only_source(graph, get_id, write_id) {
      return false
    }
  }
  true
}

///|
fn merge_locals_gets_follow_copy(
  gets : Array[@ir.NodeId],
  copy_id : @ir.NodeId,
) -> Bool {
  for get_id in gets {
    if get_id <= copy_id {
      // A source-older get can remain stack-carried until after this write.
      // Retargeting it from the graph's region order would change the value
      // already evaluated before the copy.
      return false
    }
  }
  true
}

///|
fn merge_locals_retarget_get(
  func : @ir.HotFunc,
  get_id : @ir.NodeId,
  local_id : @ir.LocalId,
) -> Unit {
  match @ir.hot_node_get(func, get_id).op {
    @ir.HotOp::LocalGet => @ir.hot_local_node_id_set(func, get_id, local_id)
    _ => abort("merge-locals expected an influenced local.get")
  }
}

///|
fn merge_locals_retarget_gets(
  func : @ir.HotFunc,
  gets : Array[@ir.NodeId],
  local_id : @ir.LocalId,
) -> Unit {
  for get_id in gets {
    merge_locals_retarget_get(func, get_id, local_id)
  }
}

///|
fn merge_locals_instrument_copies(func : @ir.HotFunc) -> Array[MergeLocalsCopy] {
  let copies : Array[MergeLocalsCopy] = []
  let original_node_count = @ir.hot_node_count(func)
  for copy_id = 0; copy_id < original_node_count; copy_id = copy_id + 1 {
    if !@ir.hot_node_is_live(func, copy_id) {
      continue
    }
    let copy = @ir.hot_node_get(func, copy_id)
    match copy.op {
      @ir.HotOp::LocalSet | @ir.HotOp::LocalTee => {
        if @ir.hot_child_count(func, copy_id) != 1 {
          continue
        }
        let value_id = @ir.hot_child_get(func, copy_id, 0)
        if value_id < 0 || !@ir.hot_node_is_live(func, value_id) {
          continue
        }
        let value = @ir.hot_node_get(func, value_id)
        match value.op {
          @ir.HotOp::LocalGet if value.imm0 != copy.imm0 => {
            let trivial_id = @ir.hot_build_local_tee(func, value.imm0, value_id)
            @ir.hot_child_set(func, copy_id, 0, trivial_id)
            copies.push({
              copy_id,
              trivial_id,
              source_local: value.imm0,
              target_local: copy.imm0,
            })
          }
          _ => ()
        }
      }
      _ => ()
    }
  }
  copies
}

///|
fn merge_locals_optimize_copies(
  func : @ir.HotFunc,
  graph : @ir.HotLocalGraph,
  copies : Array[MergeLocalsCopy],
) -> Array[MergeLocalsOptimizedCopy] {
  let optimized : Array[MergeLocalsOptimizedCopy] = []
  for copy in copies {
    let trivial_influences = @ir.local_graph_influenced_gets_for_write(
      graph,
      copy.trivial_id,
    )
    let can_optimize_to_copy = trivial_influences.length() > 0 &&
      merge_locals_gets_follow_copy(trivial_influences, copy.copy_id) &&
      merge_locals_gets_are_single_source(
        graph,
        trivial_influences,
        copy.trivial_id,
      ) &&
      merge_locals_gets_match_target_type(
        func,
        trivial_influences,
        copy.target_local,
      )
    if can_optimize_to_copy {
      merge_locals_retarget_gets(func, trivial_influences, copy.target_local)
      optimized.push({
        copy,
        orientation: MergeLocalsOrientation::ToCopy,
        influenced_gets: trivial_influences,
      })
      continue
    }

    if trivial_influences.length() == 0 {
      continue
    }
    let copy_influences = @ir.local_graph_influenced_gets_for_write(
      graph,
      copy.copy_id,
    )
    let can_optimize_to_trivial = copy_influences.length() > 0 &&
      merge_locals_gets_follow_copy(copy_influences, copy.copy_id) &&
      merge_locals_gets_are_single_source(graph, copy_influences, copy.copy_id) &&
      merge_locals_gets_match_target_type(
        func,
        copy_influences,
        copy.source_local,
      )
    if can_optimize_to_trivial {
      merge_locals_retarget_gets(func, copy_influences, copy.source_local)
      optimized.push({
        copy,
        orientation: MergeLocalsOrientation::ToTrivial,
        influenced_gets: copy_influences,
      })
    }
  }
  optimized
}

///|
fn merge_locals_verify_and_rollback(
  func : @ir.HotFunc,
  post_graph : @ir.HotLocalGraph,
  optimized : Array[MergeLocalsOptimizedCopy],
) -> Int32 {
  let mut kept = 0
  for rewrite in optimized {
    let expected_write = match rewrite.orientation {
      MergeLocalsOrientation::ToCopy => rewrite.copy.copy_id
      MergeLocalsOrientation::ToTrivial => rewrite.copy.trivial_id
    }
    let mut valid = true
    for get_id in rewrite.influenced_gets {
      if !merge_locals_get_has_only_source(post_graph, get_id, expected_write) {
        valid = false
        break
      }
    }
    if valid {
      kept += 1
    } else {
      let original_local = match rewrite.orientation {
        MergeLocalsOrientation::ToCopy => rewrite.copy.source_local
        MergeLocalsOrientation::ToTrivial => rewrite.copy.target_local
      }
      merge_locals_retarget_gets(func, rewrite.influenced_gets, original_local)
    }
  }
  kept
}

///|
fn merge_locals_remove_instrumentation(
  func : @ir.HotFunc,
  copies : Array[MergeLocalsCopy],
) -> Unit {
  for copy in copies {
    if !@ir.hot_node_is_live(func, copy.copy_id) ||
      !@ir.hot_node_is_live(func, copy.trivial_id) {
      continue
    }
    let original_get = @ir.hot_child_get(func, copy.trivial_id, 0)
    @ir.hot_child_set(func, copy.copy_id, 0, original_get)
    @ir.hot_delete_node(func, copy.trivial_id)
  }
}

///|
priv struct MergeLocalsRawCandidate {
  set_idx : Int32
  source_local : Int32
  target_local : Int32
  virtual_source : Int32
  copy_source : Int32
}

///|
priv struct MergeLocalsRawRewrite {
  supported : Bool
  func : @lib.Func
  changed : Bool
}

///|
fn merge_locals_raw_idx(idx : @lib.LocalIdx) -> Int32 {
  let @lib.LocalIdx(raw) = idx
  raw.reinterpret_as_int()
}

///|
fn merge_locals_raw_local_idx(idx : Int32) -> @lib.LocalIdx {
  @lib.LocalIdx::new(idx.reinterpret_as_uint())
}

///|
fn merge_locals_raw_local_types(
  params : Array[@lib.ValType],
  locals : @lib.Locals,
) -> Array[@lib.ValType] {
  let out = params.copy()
  for run in locals.runs() {
    for _ in 0.. Bool {
  for instr in instrs {
    match instr {
      @lib.Block(_, _)
      | @lib.Loop(_, _)
      | @lib.If(_, _, _)
      | @lib.Try(_, _, _, _)
      | @lib.TryTable(_, _, _)
      | @lib.Br(_)
      | @lib.BrIf(_)
      | @lib.BrTable(_, _)
      | @lib.BrOnNull(_)
      | @lib.BrOnNonNull(_)
      | @lib.BrOnCast(_, _, _, _)
      | @lib.BrOnCastFail(_, _, _, _)
      | @lib.Return
      | @lib.ReturnCall(_)
      | @lib.ReturnCallIndirect(_, _)
      | @lib.ReturnCallRef(_)
      | @lib.Throw(_)
      | @lib.ThrowRef => return false
      _ => ()
    }
  }
  true
}

///|
fn merge_locals_raw_candidates(
  instrs : Array[@lib.Instruction],
  local_count : Int32,
) -> Array[MergeLocalsRawCandidate] {
  let out : Array[MergeLocalsRawCandidate] = []
  for idx = 1; idx < instrs.length(); idx = idx + 1 {
    match (instrs[idx - 1], instrs[idx]) {
      (
        @lib.LocalGet(source_idx),
        @lib.LocalSet(target_idx)
        | @lib.LocalTee(target_idx),
      ) => {
        let source = merge_locals_raw_idx(source_idx)
        let target = merge_locals_raw_idx(target_idx)
        if source != target {
          out.push({
            set_idx: idx,
            source_local: source,
            target_local: target,
            virtual_source: local_count + idx * 2,
            copy_source: local_count + idx * 2 + 1,
          })
        }
      }
      _ => ()
    }
  }
  out
}

///|
fn merge_locals_raw_candidate_map(
  candidates : Array[MergeLocalsRawCandidate],
  instr_count : Int32,
) -> Array[Int32] {
  let map = Array::make(instr_count, -1)
  for idx = 0; idx < candidates.length(); idx = idx + 1 {
    let set_idx = candidates[idx].set_idx
    if set_idx >= 0 && set_idx < map.length() {
      map[set_idx] = idx
    }
  }
  map
}

///|
fn merge_locals_raw_build_influences(
  instrs : Array[@lib.Instruction],
  local_count : Int32,
  candidates : Array[MergeLocalsRawCandidate],
) -> Array[Array[Int32]] {
  let source_count = local_count + instrs.length() * 2 + 2
  // Empty buckets share a read-only sentinel. Allocate a private bucket on its
  // first get, so instruction-only sources do not each allocate an empty array.
  let influences : Array[Array[Int32]] = Array::make(source_count, [])
  let current : Array[Int32] = []
  for local_id = 0; local_id < local_count; local_id = local_id + 1 {
    current.push(local_id)
  }
  let candidate_map = merge_locals_raw_candidate_map(
    candidates,
    instrs.length(),
  )
  for idx = 0; idx < instrs.length(); idx = idx + 1 {
    let candidate_idx = candidate_map[idx]
    if candidate_idx >= 0 {
      let candidate = candidates[candidate_idx]
      current[candidate.source_local] = candidate.virtual_source
    }
    match instrs[idx] {
      @lib.LocalGet(local_idx) => {
        let local_id = merge_locals_raw_idx(local_idx)
        if local_id >= 0 && local_id < current.length() {
          let source = current[local_id]
          if influences[source].is_empty() {
            influences[source] = [idx]
          } else {
            influences[source].push(idx)
          }
        }
      }
      @lib.LocalSet(local_idx) | @lib.LocalTee(local_idx) => {
        let local_id = merge_locals_raw_idx(local_idx)
        if local_id >= 0 && local_id < current.length() {
          current[local_id] = local_count + idx * 2 + 1
        }
      }
      _ => ()
    }
  }
  influences
}

///|
fn merge_locals_raw_retarget_gets(
  instrs : Array[@lib.Instruction],
  get_indices : Array[Int32],
  local_id : Int32,
) -> Unit {
  for idx in get_indices {
    instrs[idx] = @lib.Instruction::local_get(
      merge_locals_raw_local_idx(local_id),
    )
  }
}

///|
fn merge_locals_raw_build_get_sources(
  instrs : Array[@lib.Instruction],
  local_count : Int32,
  candidates : Array[MergeLocalsRawCandidate],
) -> Array[Int32] {
  let get_sources = Array::make(instrs.length(), -1)
  let current : Array[Int32] = []
  for local_id = 0; local_id < local_count; local_id = local_id + 1 {
    current.push(local_id)
  }
  let candidate_map = merge_locals_raw_candidate_map(
    candidates,
    instrs.length(),
  )
  for idx = 0; idx < instrs.length(); idx = idx + 1 {
    let candidate_idx = candidate_map[idx]
    if candidate_idx >= 0 {
      let candidate = candidates[candidate_idx]
      current[candidate.source_local] = candidate.virtual_source
    }
    match instrs[idx] {
      @lib.LocalGet(local_idx) => {
        let local_id = merge_locals_raw_idx(local_idx)
        if local_id >= 0 && local_id < current.length() {
          get_sources[idx] = current[local_id]
        }
      }
      @lib.LocalSet(local_idx) | @lib.LocalTee(local_idx) => {
        let local_id = merge_locals_raw_idx(local_idx)
        if local_id >= 0 && local_id < current.length() {
          current[local_id] = local_count + idx * 2 + 1
        }
      }
      _ => ()
    }
  }
  get_sources
}

///|
fn merge_locals_run_raw_straight_line_func(
  func : @lib.Func,
  params : Array[@lib.ValType],
) -> MergeLocalsRawRewrite {
  let @lib.Func(locals, @lib.Expr(original)) = func
  if !merge_locals_raw_is_straight_line(original) {
    return { supported: false, func, changed: false, }
  }
  let local_types = merge_locals_raw_local_types(params, locals)
  let candidates = merge_locals_raw_candidates(original, local_types.length())
  if candidates.is_empty() {
    return { supported: true, func, changed: false, }
  }
  let influences = merge_locals_raw_build_influences(
    original,
    local_types.length(),
    candidates,
  )
  let rewritten = original.copy()
  let orientations = Array::make(candidates.length(), 0)
  let rewrite_gets : Array[Array[Int32]] = []
  for candidate_idx = 0
      candidate_idx < candidates.length()
      candidate_idx = candidate_idx + 1 {
    let candidate = candidates[candidate_idx]
    let trivial_gets = influences[candidate.virtual_source]
    let same_type = candidate.source_local >= 0 &&
      candidate.target_local >= 0 &&
      candidate.source_local < local_types.length() &&
      candidate.target_local < local_types.length() &&
      local_types[candidate.source_local] == local_types[candidate.target_local]
    if same_type && !trivial_gets.is_empty() {
      merge_locals_raw_retarget_gets(
        rewritten,
        trivial_gets,
        candidate.target_local,
      )
      orientations[candidate_idx] = 1
      rewrite_gets.push(trivial_gets)
    } else {
      rewrite_gets.push([])
    }
  }
  let post_sources = merge_locals_raw_build_get_sources(
    rewritten,
    local_types.length(),
    candidates,
  )
  let mut kept = 0
  for candidate_idx = 0
      candidate_idx < candidates.length()
      candidate_idx = candidate_idx + 1 {
    let orientation = orientations[candidate_idx]
    if orientation == 0 {
      continue
    }
    let candidate = candidates[candidate_idx]
    let expected = candidate.copy_source
    let mut valid = true
    for get_idx in rewrite_gets[candidate_idx] {
      if get_idx < 0 ||
        get_idx >= post_sources.length() ||
        post_sources[get_idx] != expected {
        valid = false
        break
      }
    }
    if valid {
      kept += 1
    } else {
      merge_locals_raw_retarget_gets(
        rewritten,
        rewrite_gets[candidate_idx],
        candidate.source_local,
      )
    }
  }
  if kept == 0 {
    { supported: true, func, changed: false, }
  } else {
    {
      supported: true,
      func: @lib.Func::new(locals, @lib.Expr::new(rewritten)),
      changed: true,
    }
  }
}

///|
fn merge_locals_run_raw_region_expr(
  expr : @lib.Expr,
  locals : @lib.Locals,
  params : Array[@lib.ValType],
) -> (@lib.Expr, Bool) {
  let @lib.Expr(original) = expr
  let rewritten = original.copy()
  let mut changed = false
  for idx = 0; idx < rewritten.length(); idx = idx + 1 {
    match rewritten[idx] {
      @lib.Block(block_type, body) => {
        let (next_body, body_changed) = merge_locals_run_raw_region_expr(
          body, locals, params,
        )
        if body_changed {
          rewritten[idx] = @lib.Instruction::block(block_type, next_body)
          changed = true
        }
      }
      @lib.Loop(block_type, body) => {
        let (next_body, body_changed) = merge_locals_run_raw_region_expr(
          body, locals, params,
        )
        if body_changed {
          rewritten[idx] = @lib.Instruction::loop_(block_type, next_body)
          changed = true
        }
      }
      @lib.If(block_type, then_body, else_body) => {
        let (next_then, then_changed) = merge_locals_run_raw_region_expr(
          @lib.Expr::new(then_body),
          locals,
          params,
        )
        let (next_else, else_changed) = match else_body {
          Some(body) => {
            let (next, did_change) = merge_locals_run_raw_region_expr(
              @lib.Expr::new(body),
              locals,
              params,
            )
            let @lib.Expr(next_instrs) = next
            (Some(next_instrs), did_change)
          }
          None => (None, false)
        }
        if then_changed || else_changed {
          let @lib.Expr(next_then_instrs) = next_then
          rewritten[idx] = @lib.Instruction::if_(
            block_type, next_then_instrs, next_else,
          )
          changed = true
        }
      }
      @lib.Try(block_type, body, catches, delegate) => {
        let (next_body, body_changed) = merge_locals_run_raw_region_expr(
          body, locals, params,
        )
        let next_catches : Array[@lib.LegacyCatch] = []
        let mut catches_changed = false
        for catch_ in catches {
          match catch_ {
            @lib.LegacyCatch(tag_idx, catch_body) => {
              let (next_catch, catch_changed) = merge_locals_run_raw_region_expr(
                catch_body, locals, params,
              )
              next_catches.push(@lib.LegacyCatch::tagged(tag_idx, next_catch))
              catches_changed = catches_changed || catch_changed
            }
            @lib.LegacyCatchAll(catch_body) => {
              let (next_catch, catch_changed) = merge_locals_run_raw_region_expr(
                catch_body, locals, params,
              )
              next_catches.push(@lib.LegacyCatch::all(next_catch))
              catches_changed = catches_changed || catch_changed
            }
          }
        }
        if body_changed || catches_changed {
          rewritten[idx] = @lib.Instruction::try_(
            block_type,
            next_body,
            next_catches,
            delegate~,
          )
          changed = true
        }
      }
      @lib.TryTable(block_type, catches, body) => {
        let (next_body, body_changed) = merge_locals_run_raw_region_expr(
          body, locals, params,
        )
        if body_changed {
          rewritten[idx] = @lib.Instruction::try_table(
            block_type, catches, next_body,
          )
          changed = true
        }
      }
      _ => ()
    }
  }
  let leaf = merge_locals_run_raw_straight_line_func(
    @lib.Func::new(locals, @lib.Expr::new(rewritten)),
    params,
  )
  if leaf.supported && leaf.changed {
    let @lib.Func(_, next_expr) = leaf.func
    (next_expr, true)
  } else {
    (@lib.Expr::new(rewritten), changed)
  }
}

///|
fn merge_locals_run_raw_legacy_func(
  func : @lib.Func,
  params : Array[@lib.ValType],
) -> MergeLocalsRawRewrite {
  let @lib.Func(locals, body) = func
  let (rewritten, changed) = merge_locals_run_raw_region_expr(
    body, locals, params,
  )
  {
    supported: true,
    func: if changed {
      @lib.Func::new(locals, rewritten)
    } else {
      func
    },
    changed,
  }
}

///|
fn merge_locals_collect_subtree_min_local_reads(
  func : @ir.HotFunc,
  node_id : Int32,
  minima : @hashmap.HashMap[Int32, Int32],
  second_minima : @hashmap.HashMap[Int32, Int32],
  pending : Array[Int32],
  visit_epochs : Array[Int32],
  epoch : Int32,
  region_worklist : Array[@ir.HotRegionRef],
  queued_controls : Array[Bool],
  visits~ : Array[Int32]? = None,
) -> Unit {
  pending.clear()
  pending.push(node_id)
  while !pending.is_empty() {
    let current = pending.pop().unwrap()
    if current < 0 ||
      current >= visit_epochs.length() ||
      !@ir.hot_node_is_live(func, current) ||
      visit_epochs[current] == epoch {
      continue
    }
    visit_epochs[current] = epoch
    match visits {
      Some(count) => count[0] += 1
      None => ()
    }
    let node = @ir.hot_node_get(func, current)
    if !queued_controls[current] {
      match node.op {
        @ir.HotOp::Block | @ir.HotOp::Loop => {
          queued_controls[current] = true
          region_worklist.push(@ir.hot_region_body_ref(func, current))
        }
        @ir.HotOp::If => {
          queued_controls[current] = true
          region_worklist.push(@ir.hot_region_then_ref(func, current))
          match @ir.hot_region_else_ref(func, current) {
            Some(else_ref) => region_worklist.push(else_ref)
            None => ()
          }
        }
        @ir.HotOp::Try => {
          queued_controls[current] = true
          region_worklist.push(@ir.hot_region_body_ref(func, current))
          match @ir.hot_region_catch_ref(func, current) {
            Some(catch_ref) => region_worklist.push(catch_ref)
            None => ()
          }
        }
        @ir.HotOp::TryTable => {
          queued_controls[current] = true
          region_worklist.push(@ir.hot_region_body_ref(func, current))
          region_worklist.push(@ir.hot_region_catch_list_ref(func, current))
        }
        _ => ()
      }
    }
    match node.op {
      @ir.HotOp::LocalGet =>
        match minima.get(node.imm0) {
          Some(previous) =>
            if current != previous && current < previous {
              second_minima.set(node.imm0, previous)
              minima.set(node.imm0, current)
            } else if current != previous {
              match second_minima.get(node.imm0) {
                Some(second) =>
                  if current < second {
                    second_minima.set(node.imm0, current)
                  }
                None => second_minima.set(node.imm0, current)
              }
            }
          None => minima.set(node.imm0, current)
        }
      _ => ()
    }
    for slot = 0; slot < @ir.hot_child_count(func, current); slot = slot + 1 {
      pending.push(@ir.hot_child_get(func, current, slot))
    }
  }
}

///|
fn merge_locals_region_has_stack_carried_copy(
  func : @ir.HotFunc,
  region_ref : @ir.HotRegionRef,
  visits~ : Array[Int32]? = None,
) -> Bool {
  let region_worklist = [region_ref]
  let node_count = @ir.hot_node_count(func)
  let visit_epochs = Array::make(node_count, 0)
  let queued_controls = Array::make(node_count, false)
  let pending_nodes : Array[Int32] = []
  let epoch = 1
  let later_min_reads : @hashmap.HashMap[Int32, Int32] = @hashmap.new()
  let later_second_min_reads : @hashmap.HashMap[Int32, Int32] = @hashmap.new()
  while !region_worklist.is_empty() {
    let current_region = region_worklist.pop().unwrap()
    let count = @ir.hot_region_root_count(func, current_region)
    for idx = count - 1; idx >= 0; idx = idx - 1 {
      let root_id = @ir.hot_region_get(func, current_region, idx)
      if root_id < 0 || !@ir.hot_node_is_live(func, root_id) {
        continue
      }
      let root = @ir.hot_node_get(func, root_id)
      match root.op {
        @ir.HotOp::LocalSet | @ir.HotOp::LocalTee if root.child_count == 1 => {
          let value_id = @ir.hot_child_get(func, root_id, 0)
          if value_id >= 0 && @ir.hot_node_is_live(func, value_id) {
            let value = @ir.hot_node_get(func, value_id)
            match value.op {
              @ir.HotOp::LocalGet if value.imm0 != root.imm0 =>
                match later_min_reads.get(value.imm0) {
                  Some(read_id) => {
                    let distinct_read_id = if read_id != value_id {
                      read_id
                    } else {
                      match later_second_min_reads.get(value.imm0) {
                        Some(second) => second
                        None => -1
                      }
                    }
                    if distinct_read_id >= 0 && distinct_read_id < root_id {
                      return true
                    }
                  }
                  _ => ()
                }
              _ => ()
            }
          }
        }
        _ => ()
      }
      merge_locals_collect_subtree_min_local_reads(
        func,
        root_id,
        later_min_reads,
        later_second_min_reads,
        pending_nodes,
        visit_epochs,
        epoch,
        region_worklist,
        queued_controls,
        visits~,
      )
    }
  }
  false
}

///|
fn merge_locals_run(ctx : HotPassContext, func : @ir.HotFunc) -> HotPassResult {
  if pass_has_repeated_entry_evaluation(func) {
    return HotPassResult::unchanged()
  }
  if merge_locals_region_has_stack_carried_copy(func, @ir.HotRegionRef::root()) {
    return HotPassResult::unchanged()
  }
  let copies = merge_locals_instrument_copies(func)
  if copies.is_empty() {
    return HotPassResult::unchanged()
  }

  // Carried control must execute once at its source position in both graphs.
  let pre_graph = @ir.local_graph_build_full_flow(
    func,
    @ir.cfg_build(func, expand_operand_control=true),
  )
  let optimized = merge_locals_optimize_copies(func, pre_graph, copies)
  let kept = if optimized.is_empty() {
    0
  } else {
    let post_graph = @ir.local_graph_build_full_flow(
      func,
      @ir.cfg_build(func, expand_operand_control=true),
    )
    merge_locals_verify_and_rollback(func, post_graph, optimized)
  }
  merge_locals_remove_instrumentation(func, copies)
  pass_mark_mutated(ctx, func)
  if kept > 0 {
    HotPassResult::changed()
  } else {
    HotPassResult::unchanged()
  }
}