///|
/// Translate an if-else construct
///
/// On-demand loading:
/// - Support block params (multi-value extension)
/// - Pass params to both then and else blocks
/// - Pass the locals either arm can redefine through the continuation as SSA
///   phi nodes; see `region_local_params` for why the rest need none
fn Translator::translate_if(
  self : Translator,
  block_type : @types.BlockType,
  then_body : Array[@types.Instruction],
  else_body : Array[@types.Instruction],
) -> Unit {
  // Save the unreachable state from outer context
  let outer_is_unreachable = self.is_unreachable
  let result_types = get_block_result_types(block_type, self.func_types)
  let param_types = get_block_param_types(block_type, self.func_types)

  // Pop condition first (it's on top of the stack)
  let cond = if outer_is_unreachable {
    // Create a dummy value when unreachable
    self.builder.iconst_i32(0)
  } else {
    self.pop()
  }

  // Pop param values from stack (Note: params are below condition)
  let param_values : Array[Value] = []
  if !outer_is_unreachable {
    for _ in 0.. ignore
    self.builder.add_block_param(else_block, ty) |> ignore
  }

  // Add continuation parameters for explicit results
  for ty in result_types {
    self.builder.add_block_param(continuation, ty) |> ignore
  }

  // Add continuation parameters for the locals either arm can redefine
  let local_params = self.region_local_params([then_body, else_body])
  let local_param_start = result_types.length()
  self.add_local_block_params(continuation, local_params)

  // Branch with param values (only if not unreachable)
  // Note: pass params to both then and else blocks
  if !outer_is_unreachable {
    if param_values.is_empty() {
      // No params, use simple branch
      self.builder.brnz(cond, then_block, else_block)
    } else {
      // Has params, use trampoline blocks to pass arguments
      let then_trampoline = self.builder.create_block()
      let else_trampoline = self.builder.create_block()
      self.builder.brnz(cond, then_trampoline, else_trampoline)
      // Then trampoline jumps to then_block with args
      self.builder.switch_to_block(then_trampoline)
      self.builder.jump(then_block, param_values)
      // Else trampoline jumps to else_block with args
      self.builder.switch_to_block(else_trampoline)
      self.builder.jump(else_block, param_values)
    }
  }

  // Record stack height AFTER popping params (for stack restoration)
  let stack_height_after_params = self.value_stack.length()

  // Push frame for then branch (else block as the "else" continuation)
  let frame : BlockFrame = {
    block: continuation,
    result_types,
    local_params,
    stack_height: stack_height_after_params,
    has_predecessor: false,
  }
  let frame_idx = self.block_stack.length()
  self.block_stack.push(frame)

  // Save locals at entry for else branch
  let saved_locals = self.locals.copy()

  // Translate then body
  self.builder.switch_to_block(then_block)
  self.is_unreachable = outer_is_unreachable // Reset for then body

  // Push block params onto value stack (they become available in the then body)
  for i, _ty in param_types {
    self.push(then_block.params[i].0)
  }
  for instr in then_body {
    self.translate_instruction(instr)
  }

  // Jump to continuation from then block with locals (only if not unreachable)
  let block = self.builder.current_block()
  if !self.is_unreachable && block.terminator is None {
    let args : Array[Value] = []
    for _ in 0.. frame.stack_height {
    self.pop() |> ignore
  }

  // Restore locals to entry state for else branch
  for i, loc in saved_locals {
    self.locals[i] = loc
  }

  // Translate else body
  self.builder.switch_to_block(else_block)
  self.is_unreachable = outer_is_unreachable // Reset for else body

  // Push block params onto value stack (they become available in the else body)
  for i, _ty in param_types {
    self.push(else_block.params[i].0)
  }
  for instr in else_body {
    self.translate_instruction(instr)
  }

  // Jump to continuation from else block with locals (only if not unreachable)
  let block = self.builder.current_block()
  if !self.is_unreachable && block.terminator is None {
    let args : Array[Value] = []
    for _ in 0.. ignore

  // Switch to continuation
  self.builder.switch_to_block(continuation)

  self.is_unreachable = outer_is_unreachable || !continuation_reachable
  if !continuation_reachable {
    self.builder.trap("unreachable if continuation")
  }

  // Push results onto stack
  if continuation_reachable {
    for i, _ty in result_types {
      self.push(continuation.params[i].0)
    }

    // Update the carried locals to use the continuation's phi values
    self.adopt_local_params(continuation, local_params, local_param_start)
  }
}

///|
/// Translate a br instruction
/// We need to pass both explicit params AND current local values
fn Translator::translate_br(self : Translator, depth : Int) -> Unit {
  let idx = self.block_stack.length() - 1 - depth
  if idx >= 0 && idx < self.block_stack.length() {
    // Jump to a block within the function
    let frame = self.block_stack[idx]
    let args : Array[Value] = []
    // Pop explicit block params
    for _ in 0.. Unit {
  let cond = self.pop()
  let idx = self.block_stack.length() - 1 - depth

  // For br_if, we need a fallthrough block and a taken block (critical edge split)
  let fallthrough = self.builder.create_block()
  let taken = self.builder.create_block()
  if idx >= 0 && idx < self.block_stack.length() {
    // Jump to a block within the function
    let frame = self.block_stack[idx]

    // Collect values that would be passed on branch
    let args : Array[Value] = []
    for i in 0.. Unit {
  let ref_val = self.pop()
  let idx = self.block_stack.length() - 1 - depth

  // Check if reference is null.
  let null_sentinel = self.builder.iconst(ref_val.ty, @wasm_milkir.NULL_REF)
  let is_null = self.builder.icmp_eq(ref_val, null_sentinel)

  // For br_on_null, we need a fallthrough block and a taken block (critical edge split)
  let fallthrough = self.builder.create_block()
  let taken = self.builder.create_block()
  if idx >= 0 && idx < self.block_stack.length() {
    let frame = self.block_stack[idx]

    // Collect values that would be passed on branch (not including the ref since it's null)
    let args : Array[Value] = []
    for i in 0.. Unit {
  let ref_val = self.pop()
  let idx = self.block_stack.length() - 1 - depth

  // Check if reference is null.
  let null_sentinel = self.builder.iconst(ref_val.ty, @wasm_milkir.NULL_REF)
  let is_null = self.builder.icmp_eq(ref_val, null_sentinel)

  // For br_on_non_null, we need a fallthrough block and a taken block
  let fallthrough = self.builder.create_block()
  let taken = self.builder.create_block()
  if idx >= 0 && idx < self.block_stack.length() {
    let frame = self.block_stack[idx]

    // Collect values that would be passed on branch
    // For br_on_non_null, we need to collect args BEFORE the ref was popped
    // The block expects result_types.length() values, and the ref is one of them
    // Since we already popped ref_val, we need to include it manually
    let args : Array[Value] = []
    // If the block expects N results, and the ref is one of them (the last one),
    // we need N-1 values from the stack plus the ref
    let stack_values_needed = frame.result_types.length() - 1
    for i in 0.. Unit {
  let ref_val = self.pop()
  let idx = self.block_stack.length() - 1 - depth
  // Get the type index for the target type
  let type_idx = self.extract_type_idx(target_type)
  let nullable = target_type.is_nullable()
  // Use ref_test to check if the reference matches the target type
  let matches = @wasm_milkir.ref_test(self.builder, type_idx, nullable, ref_val)
  // Create fallthrough and taken blocks
  let fallthrough = self.builder.create_block()
  let taken = self.builder.create_block()
  if idx >= 0 && idx < self.block_stack.length() {
    let frame = self.block_stack[idx]
    // Collect values that would be passed on branch (excluding the ref)
    let stack_values_needed = frame.result_types.length() - 1
    let stack_args : Array[Value] = []
    for i in 0.. Unit {
  let ref_val = self.pop()
  let idx = self.block_stack.length() - 1 - depth
  // Get the type index for the target type
  let type_idx = self.extract_type_idx(target_type)
  let nullable = target_type.is_nullable()
  // Use ref_test to check if the reference matches the target type
  let matches = @wasm_milkir.ref_test(self.builder, type_idx, nullable, ref_val)
  // Create fallthrough and taken blocks
  let fallthrough = self.builder.create_block()
  let taken = self.builder.create_block()
  if idx >= 0 && idx < self.block_stack.length() {
    let frame = self.block_stack[idx]
    // Collect values that would be passed on branch (excluding the ref)
    let stack_values_needed = frame.result_types.length() - 1
    let stack_args : Array[Value] = []
    for i in 0.. Unit {
  let index = self.pop()

  // Helper to get a target's block, result types, and carried locals for a
  // given depth. The return continuation carries no locals.
  // Returns (target_block, result_types, local_params)
  fn get_target(
    self : Translator,
    depth : Int,
  ) -> (Block, Array[Type], Array[Int]) {
    let idx = self.block_stack.length() - 1 - depth
    if idx >= 0 && idx < self.block_stack.length() {
      let frame = self.block_stack[idx]
      self.mark_block_frame_reachable(idx)
      (frame.block, frame.result_types, frame.local_params)
    } else {
      // Function level
      let ret_cont = self.get_or_create_return_continuation()
      (ret_cont, self.func_result_types, [])
    }
  }

  // Get default target info
  let (default_block, default_result_types, default_local_params) = get_target(
    self, default_,
  )

  // The explicit results are the same for every target -- all targets share the
  // default's result type. The locals are not: each target carries only what
  // its own region can redefine, so every target builds its own argument list
  // off this shared prefix.
  let result_args : Array[Value] = []
  let num_results = default_result_types.length()
  // Only collect results if we have enough values on the stack
  if self.value_stack.length() >= num_results {
    for i in 0.. Array[Value] {
    let args = result_args.copy()
    self.push_local_args(args, local_params)
    args
  }

  // Phase 1: Create intermediate blocks for unique target depths.
  //
  // br_table payloads can be very large with many repeated labels. Emitting one
  // unique intermediate block per target depth keeps IR size bounded and avoids
  // O(table_size) extra blocks in translator hot paths.
  let intermediate_blocks : Array[Block] = []
  let unique_depths : Array[Int] = []
  let depth_to_intermediate : Map[Int, Block] = Map([])
  let default_intermediate = self.builder.create_block()
  for depth in labels {
    if depth == default_ {
      intermediate_blocks.push(default_intermediate)
      continue
    }
    if depth_to_intermediate.get(depth) is Some(existing) {
      intermediate_blocks.push(existing)
      continue
    }
    let intermediate = self.builder.create_block()
    depth_to_intermediate.set(depth, intermediate)
    unique_depths.push(depth)
    intermediate_blocks.push(intermediate)
  }

  // Phase 2: Emit br_table in original block
  self.builder.br_table(index, intermediate_blocks, default_intermediate)

  // Phase 3: Fill in unique intermediate blocks with jumps to real targets.
  for depth in unique_depths {
    guard depth_to_intermediate.get(depth) is Some(intermediate) else {
      continue
    }
    let (target_block, _, local_params) = get_target(self, depth)
    self.builder.switch_to_block(intermediate)
    self.builder.jump(target_block, target_args(self, local_params))
  }

  // Fill in default intermediate
  self.builder.switch_to_block(default_intermediate)
  self.builder.jump(default_block, target_args(self, default_local_params))

  // br_table is a terminator, code after it is unreachable
  // Don't switch back to original - leave current_block as the last intermediate
  self.is_unreachable = true
}

///|
/// Remap a local function index to global function index for cross-module calls
/// - For imports: use import_remap if available
/// - For local functions: add func_base
fn Translator::remap_func_idx(self : Translator, func_idx : Int) -> Int {
  if func_idx < self.num_imports {
    // Import function - use remap if available
    if self.import_remap.length() > 0 && func_idx < self.import_remap.length() {
      self.import_remap[func_idx]
    } else {
      func_idx
    }
  } else {
    // Local function - add base offset
    self.func_base + func_idx
  }
}

///|
/// Translate a direct call
fn Translator::translate_call(self : Translator, func_idx : Int) -> Unit {
  // Get function type
  let type_idx = if func_idx < self.num_imports {
    // Look up import function type
    if func_idx < self.import_func_type_indices.length() {
      self.import_func_type_indices[func_idx]
    } else {
      0 // Fallback
    }
  } else {
    let local_idx = func_idx - self.num_imports
    if local_idx < self.func_type_indices.length() {
      self.func_type_indices[local_idx]
    } else {
      0
    }
  }
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Pop arguments
    let args : Array[Value] = []
    for _ in 0.. Unit {
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Use original type_idx - the JIT type check (is_subtype_cached) handles
    // canonical indices internally to support both subtyping and structural equivalence
    // Pop callee (element index within the table)
    let elem_idx = self.pop()

    // Bounds check: elem_idx must be < table_size[table_idx]
    // Generate: if (elem_idx >= table_size) trap "out of bounds table access"
    if table_idx < self.table_sizes.length() {
      let table_size = self.table_sizes[table_idx]
      let size_const = self.builder.iconst(elem_idx.ty, table_size.to_int64())
      let out_of_bounds = self.builder.icmp_uge(elem_idx, size_const)

      // Create trap block and continuation block
      let trap_block = self.builder.create_block()
      let continue_block = self.builder.create_block()

      // Branch: if out_of_bounds goto trap_block else continue_block
      self.builder.brnz(out_of_bounds, trap_block, continue_block)

      // Trap block: emit trap instruction
      self.builder.switch_to_block(trap_block)
      self.builder.trap("out of bounds table access")

      // Continue block: proceed with call
      self.builder.switch_to_block(continue_block)
    }

    // Multi-table support: use elem_idx directly (no flattening needed)
    // The lowering phase handles table access via indirect_tables[table_idx]
    // Pop arguments
    let args : Array[Value] = []
    for _ in 0.. Unit {
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Pop the function reference
    let func_ref = self.pop()

    // Null check: lowered null reference is NULL_REF (0)
    let null_sentinel = self.builder.iconst(func_ref.ty, @wasm_milkir.NULL_REF)
    let is_null = self.builder.icmp_eq(func_ref, null_sentinel)

    // Create trap block and continuation block
    let trap_block = self.builder.create_block()
    let continue_block = self.builder.create_block()

    // Branch: if is_null goto trap_block else continue_block
    self.builder.brnz(is_null, trap_block, continue_block)

    // Trap block: emit trap for null reference
    self.builder.switch_to_block(trap_block)
    self.builder.trap("null function reference")

    // Continue block: proceed with call
    self.builder.switch_to_block(continue_block)

    // Pop arguments
    let args : Array[Value] = []
    for _ in 0.. Unit {
  // Get function type
  let type_idx = if func_idx < self.num_imports {
    // Look up import function type
    if func_idx < self.import_func_type_indices.length() {
      self.import_func_type_indices[func_idx]
    } else {
      0 // Fallback
    }
  } else {
    let local_idx = func_idx - self.num_imports
    if local_idx < self.func_type_indices.length() {
      self.func_type_indices[local_idx]
    } else {
      0
    }
  }
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Pop arguments
    let args : Array[Value] = []
    for _ in 0.. Unit {
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Use original type_idx - the JIT type check (is_subtype_cached) handles
    // canonical indices internally to support both subtyping and structural equivalence
    // Pop callee (element index within the table)
    let elem_idx = self.pop()

    // Pop arguments
    let args : Array[Value] = []
    for _ in 0.. Unit {
  if type_idx < self.func_types.length() {
    let func_type = self.func_types[type_idx]
    // Pop the function reference
    let func_ref = self.pop()

    // Null check: lowered null reference is NULL_REF (0)
    let null_sentinel = self.builder.iconst(func_ref.ty, @wasm_milkir.NULL_REF)
    let is_null = self.builder.icmp_eq(func_ref, null_sentinel)

    // Create trap block and continuation block
    let trap_block = self.builder.create_block()
    let continue_block = self.builder.create_block()

    // Branch: if is_null goto trap_block else continue_block
    self.builder.brnz(is_null, trap_block, continue_block)

    // Trap block: emit trap for null reference
    self.builder.switch_to_block(trap_block)
    self.builder.trap("null function reference")

    // Continue block: proceed with call
    self.builder.switch_to_block(continue_block)

    // Pop arguments
    let args : Array[Value] = []
    for _ in 0..