// Licensing scope: see ../../LICENSE and this directory's LICENSE and NOTICE.
// 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.

///|
fn tail_call_label_exits_function(
  label : @lib.LabelIdx,
  tail_labels : Array[Bool],
) -> Bool {
  let @lib.LabelIdx(raw) = label
  let index = tail_labels.length() - 1 - raw.to_int()
  index >= 0 && index < tail_labels.length() && tail_labels[index]
}

///|
fn tail_call_pure_selector(instr : @lib.Instruction) -> Bool {
  match instr {
    @lib.LocalGet(_) | @lib.I32Const(_) | @lib.I64Const(_) => true
    _ => false
  }
}

///|
fn tail_call_transfer_after(
  instrs : Array[@lib.Instruction],
  index : Int32,
  is_tail : Bool,
  tail_labels : Array[Bool],
) -> Bool {
  if index + 1 >= instrs.length() {
    return is_tail
  }
  match instrs[index + 1] {
    @lib.Return => true
    @lib.Br(label) => tail_call_label_exits_function(label, tail_labels)
    next if tail_call_pure_selector(next) && index + 2 == instrs.length() - 1 =>
      match instrs[index + 2] {
        @lib.BrIf(label) =>
          is_tail && tail_call_label_exits_function(label, tail_labels)
        @lib.BrTable(targets, default) => {
          if !is_tail || !tail_call_label_exits_function(default, tail_labels) {
            return false
          }
          targets.all(fn(label) {
            tail_call_label_exits_function(label, tail_labels)
          })
        }
        _ => false
      }
    _ => false
  }
}

///|
fn tail_call_rewrite_with_label(
  expr : @lib.Expr,
  is_tail : Bool,
  label_is_tail : Bool,
  tail_labels : Array[Bool],
  in_local_handler : Bool,
  changed : Array[Bool],
) -> @lib.Expr {
  tail_labels.push(label_is_tail)
  let rewritten = tail_call_rewrite_expr(
    expr, is_tail, tail_labels, in_local_handler, changed,
  )
  ignore(tail_labels.pop())
  rewritten
}

///|
fn tail_call_rewrite_expr(
  expr : @lib.Expr,
  is_tail : Bool,
  tail_labels : Array[Bool],
  in_local_handler : Bool,
  changed : Array[Bool],
) -> @lib.Expr {
  let @lib.Expr(instrs) = expr
  let rewritten : Array[@lib.Instruction] = []
  for index, instr in instrs {
    let child_is_tail = tail_call_transfer_after(
      instrs, index, is_tail, tail_labels,
    )
    let next = match instr {
      @lib.Block(block_type, body) =>
        @lib.Instruction::block(
          block_type,
          tail_call_rewrite_with_label(
            body, child_is_tail, child_is_tail, tail_labels, in_local_handler, changed,
          ),
        )
      @lib.Loop(block_type, body) =>
        @lib.Instruction::loop_(
          block_type,
          tail_call_rewrite_with_label(
            body, child_is_tail, false, tail_labels, in_local_handler, changed,
          ),
        )
      @lib.If(block_type, then_body, else_body) => {
        let @lib.Expr(next_then) = tail_call_rewrite_with_label(
          @lib.Expr::new(then_body),
          child_is_tail,
          child_is_tail,
          tail_labels,
          in_local_handler,
          changed,
        )
        let next_else = match else_body {
          Some(body) => {
            let @lib.Expr(next_body) = tail_call_rewrite_with_label(
              @lib.Expr::new(body),
              child_is_tail,
              child_is_tail,
              tail_labels,
              in_local_handler,
              changed,
            )
            Some(next_body)
          }
          None => None
        }
        @lib.Instruction::if_(block_type, next_then, next_else)
      }
      @lib.Try(block_type, body, catches, delegate) => {
        tail_labels.push(child_is_tail)
        // A tail call cannot skip a handler covering the protected body.
        // A catch body is covered only by handlers outside this try.
        let protected = in_local_handler ||
          !catches.is_empty() ||
          delegate is Some(_)
        let body = tail_call_rewrite_expr(
          body, child_is_tail, tail_labels, protected, changed,
        )
        let catches = catches.map(fn(c) {
          match c {
            @lib.LegacyCatch(tag, handler) =>
              @lib.LegacyCatch::tagged(
                tag,
                tail_call_rewrite_expr(
                  handler, child_is_tail, tail_labels, in_local_handler, changed,
                ),
              )
            @lib.LegacyCatchAll(handler) =>
              @lib.LegacyCatch::all(
                tail_call_rewrite_expr(
                  handler, child_is_tail, tail_labels, in_local_handler, changed,
                ),
              )
          }
        })
        ignore(tail_labels.pop())
        @lib.Instruction::try_(block_type, body, catches, delegate~)
      }
      @lib.TryTable(block_type, catches, body) => {
        tail_labels.push(child_is_tail)
        let rewritten = @lib.Instruction::try_table(
          block_type,
          catches,
          tail_call_rewrite_expr(
            body,
            child_is_tail,
            tail_labels,
            in_local_handler || !catches.is_empty(),
            changed,
          ),
        )
        ignore(tail_labels.pop())
        rewritten
      }
      @lib.Call(target) =>
        if !in_local_handler &&
          tail_call_transfer_after(instrs, index, is_tail, tail_labels) {
          changed[0] = true
          @lib.Instruction::return_call(target)
        } else {
          instr
        }
      @lib.CallIndirect(type_idx, table_idx) =>
        if !in_local_handler &&
          tail_call_transfer_after(instrs, index, is_tail, tail_labels) {
          changed[0] = true
          @lib.Instruction::return_call_indirect(type_idx, table_idx)
        } else {
          instr
        }
      @lib.CallRef(type_idx) =>
        if !in_local_handler &&
          tail_call_transfer_after(instrs, index, is_tail, tail_labels) {
          changed[0] = true
          @lib.Instruction::return_call_ref(type_idx)
        } else {
          instr
        }
      _ => instr
    }
    rewritten.push(next)
  }
  @lib.Expr::new(rewritten)
}

///|
fn tail_call_run_module_pass(mod_ : @lib.Module) -> @lib.Module {
  if !msf_module_allows_tail_call_thunks(mod_) {
    return mod_
  }
  let @lib.CodeSec(funcs) = match mod_.code_sec {
    Some(section) => section
    None => return mod_
  }
  let changed = [false]
  let next_funcs : Array[@lib.Func] = []
  for func in funcs {
    let @lib.Func(locals, body) = func
    next_funcs.push(
      @lib.Func::new(
        locals,
        tail_call_rewrite_expr(body, true, [], false, changed),
      ),
    )
  }
  if !changed[0] {
    return mod_
  }
  let candidate = mod_.with_code_sec(@lib.CodeSec::new(next_funcs))
  if @validate.validate_module(candidate).is_ok() {
    candidate
  } else {
    mod_
  }
}