Skip to main content

aver/codegen/recursion/
mod.rs

1//! Shared proof-mode recursion analysis.
2//!
3//! Classifies each recursive pure fn into a [`RecursionPlan`] that tells
4//! the proof backends (Lean, Dafny) how to emit a fuel-guarded helper
5//! plus a wrapper with an appropriate fuel metric. The same classifier
6//! feeds both backends so supported shapes stay consistent.
7//!
8//! Emission is backend-specific (syntax, termination-proof mechanism,
9//! default-value for fuel exhaustion), but the recognition pass and the
10//! AST transform that rewrites recursive calls into helper calls are
11//! shared.
12
13pub mod detect;
14
15use std::collections::HashSet;
16
17use crate::ast::{Expr, FnBody, MatchArm, Spanned, Stmt, StrPart, TailCallData};
18use crate::codegen::common::expr_to_dotted_name;
19
20pub use detect::analyze_plans_in_scope;
21
22/// Classification for a single recursive fn (or a whole mutual-recursion
23/// SCC, in which case every fn in the SCC gets its own plan from the
24/// same family).
25///
26/// `Eq` is deliberately omitted — `IntAscending` holds an AST expression
27/// which only implements `PartialEq` (float literals inside it are
28/// partially ordered). `PartialEq` still works for the uses this enum
29/// sees (pattern matching, `matches!`, equality via `.eq`).
30#[derive(Clone, Debug, PartialEq)]
31pub enum RecursionPlan {
32    /// Single-fn recursion where an `Int` parameter decreases by 1.
33    /// The wrapper supplies `n.natAbs + 1` fuel so the helper terminates.
34    IntCountdown { param_index: usize },
35    /// Same shape as `IntCountdown` but body is `match p { 0 -> BASE; _ -> rec(p-1, ...) }`
36    /// and the fn is closed-world: every external callsite either passes a
37    /// non-negative literal or sits in a guard branch that proves `p ≥ 0`.
38    /// Proof backends emit a native def with an injected `p ≥ 0` precondition
39    /// plus a wrapper handling the `p < 0` case from the source's `0` arm, skipping
40    /// fuel entirely so Lean can `decide` / unfold for symbolic proofs. Carries
41    /// both arm bodies so the Lean emitter can switch the elaborated shape to
42    /// `if h_zero : p = 0 then BASE else REC` — needed because Lean's `match`
43    /// elaborator does not expose the case-split hypothesis to `omega` for the
44    /// recursive callsite's `p - 1 ≥ 0` discharge.
45    IntCountdownGuarded {
46        param_index: usize,
47        /// Literal int from the body's literal-match arm — the value the
48        /// `if h_zero : p = L then base else rec(p-1, ...)` aux splits on.
49        base_arm_literal: i64,
50        base_arm_body: Spanned<Expr>,
51        wildcard_arm_body: Spanned<Expr>,
52        /// Path-constraint chain extracted from the single external
53        /// caller's surrounding `match (Bool) { true/false -> CALL ... }`
54        /// (and `if/then/else`) guards. Already normalised to positive
55        /// form by flipping the comparison BinOp on `false`-arm guards
56        /// (`Lt ↔ Gte`, `Gt ↔ Lte`, `Eq ↔ Neq`), so every clause is a
57        /// plain Aver Bool expression. Substituted into callee's
58        /// variable space (caller's arg-binding renamed to callee's
59        /// param name). Conjunction of all clauses is the aux's
60        /// precondition.
61        ///
62        /// Empty means "no single external caller in `ctx`": the Lean
63        /// emitter falls back to `(h_dom : p ≥ 0)` so a free-standing
64        /// fibTR-shape (no caller in the proof artifact) keeps the
65        /// legacy precondition. Same `Spanned<Expr>`-as-predicate
66        /// representation as opaque types' smart-constructor predicate
67        /// (`refinement_info_for`) and verify-law `when` clauses — the
68        /// three sources differ only in where the predicate comes from.
69        precondition: Vec<Spanned<Expr>>,
70    },
71    /// Single-fn recursion where an `Int` param increases by 1 up to a
72    /// bound. The bound is kept as an Aver AST expression so each
73    /// backend renders it in its own idiom; the wrapper supplies
74    /// `(bound - n).natAbs + 1` fuel.
75    IntAscending {
76        param_index: usize,
77        bound: Spanned<Expr>,
78    },
79    /// Single-fn recursion where an `Int` parameter shrinks by a
80    /// literal-divisor floor division at every self-call —
81    /// `Result.withDefault(Int.div(p, k), d)` with literal `k >= 2`,
82    /// either inlined at the call site or through a unary wrapper fn
83    /// (`half(p)`) whose body is exactly that expression. Validated,
84    /// never guessed: the classifier additionally proves that the
85    /// guard chain enclosing every self-call site implies `p >= 1`
86    /// (so `p / k < p` and `p.toNat` strictly decreases). Backends
87    /// emit a native well-founded def — Lean
88    /// `termination_by p.toNat` (kernel re-checks the measure),
89    /// Dafny `decreases if p >= 0 then p else 0` with no synthesized
90    /// `requires`.
91    IntFloorDivCountdown {
92        param_index: usize,
93        /// The literal divisor (>= 2).
94        divisor: i64,
95        /// `Some(name)` when the shrink goes through a unary wrapper
96        /// fn; `None` for the inlined form.
97        helper_fn: Option<String>,
98    },
99    /// Affine second-order recurrence like `fib(n) = fib(n-1) + fib(n-2)`
100    /// with `0 / 1` bases and an `n < 0` guard. Emitted through a
101    /// private Nat helper (pair-state), not a fuel helper.
102    LinearRecurrence2,
103    /// Single-fn structural recursion on a `List<_>` parameter; proof
104    /// backends emit as structural recursion directly (no fuel).
105    ListStructural { param_index: usize },
106    /// Single-fn structural recursion on a recursive user ADT; proof
107    /// backends emit through a sizeOf-guarded fuel helper.
108    SizeOfStructural,
109    /// Single-fn recursion where the first `String` is preserved and
110    /// the second `Int` position parameter strictly advances (`pos +
111    /// k`, k ≥ 1). Wrapper fuel is derived from `s.length - pos`.
112    StringPosAdvance,
113    /// Mutual recursion SCC where the first `Int` parameter decreases
114    /// by 1 across every inter-fn call.
115    MutualIntCountdown,
116    /// Mutual recursion SCC where the first `String` is preserved and
117    /// the second `Int` either advances or stays the same across
118    /// rank-decreasing edges.
119    MutualStringPosAdvance { rank: usize },
120    /// Generic mutual recursion SCC using `sizeOf` on structural
121    /// parameters plus rank for same-measure edges.
122    MutualSizeOfRanked { rank: usize },
123}
124
125/// Classifier-side diagnostic: a recursive fn whose shape falls
126/// outside every supported pattern. `proof_lower::populate_fn_
127/// contracts` translates these into `ProofIR.unclassified_fns` —
128/// consumers should read `ctx.proof_ir.unclassified_fns` (typed
129/// `Vec<UnclassifiedFn>`) instead of reaching for this type.
130#[derive(Clone, Debug, Eq, PartialEq)]
131pub struct ProofModeIssue {
132    pub line: usize,
133    pub message: String,
134}
135
136/// Canonical suffix for a fuel-guarded helper fn. Deliberately contains
137/// only lowercase ASCII + underscores so both Lean and Dafny accept it
138/// as an identifier without renaming.
139pub fn fuel_helper_name(name: &str) -> String {
140    format!("{}__fuel", name)
141}
142
143/// Suffix for the native-emit auxiliary fn that carries the explicit
144/// precondition parameter. Mirrors [`fuel_helper_name`]'s ASCII-only
145/// shape so Lean accepts it verbatim.
146pub fn native_aux_name(name: &str) -> String {
147    format!("{}__aux", name)
148}
149
150/// Sentinel identifier injected as an extra synthetic argument at every
151/// recursive callsite inside an `IntCountdownGuarded` body. Lean's expr
152/// emitter recognises this name and renders it as `(by omega)`; Dafny's
153/// codegen never sees it because Dafny discharges preconditions via
154/// auto-inference at the existing fn-def emit path.
155pub const OMEGA_PROOF_SENTINEL: &str = "__aver_omega_proof__";
156
157/// Flip a comparison `BinOp` to its logical negation so a caller's
158/// `match (a OP b) { false -> ... }` arm normalises into a positive
159/// predicate. Returns `None` for non-comparison operators (`Add`,
160/// `Sub`, etc.) — those can't legally land as a `match … { Bool ->
161/// ... }` subject anyway. Used by the issue-84 caller-guard extractor
162/// so every collected predicate is the same `Spanned<Expr>` shape
163/// opaque types and verify `when` already use, avoiding a parallel
164/// `(expr, negated)` representation.
165pub fn flip_comparison_binop(expr: &Spanned<Expr>) -> Option<Spanned<Expr>> {
166    let Expr::BinOp(op, left, right) = &expr.node else {
167        return None;
168    };
169    use crate::ast::BinOp::*;
170    let flipped = match op {
171        Lt => Gte,
172        Gt => Lte,
173        Lte => Gt,
174        Gte => Lt,
175        Eq => Neq,
176        Neq => Eq,
177        _ => return None,
178    };
179    Some(Spanned::new(
180        Expr::BinOp(flipped, left.clone(), right.clone()),
181        expr.line,
182    ))
183}
184
185// Re-export `substitute_ident_in_expr` from `codegen::common` so the
186// recursion module's public API doesn't lose the symbol existing
187// callers (issue 84 caller-guard walker, etc.) reach for. The
188// definition lives in common because three predicate sources (opaque
189// constructor, caller guard, verify `when`-redundancy check) all use
190// the same substitution.
191pub use crate::codegen::common::substitute_ident_in_expr;
192
193/// True iff `expr` (recursively) mentions `name` as an `Ident`/
194/// `Resolved` reference. Used by the caller-guard extractor to filter
195/// out enclosing predicates that don't constrain the variable passed
196/// at the countdown-param position — those predicates name caller
197/// locals that aren't in scope inside the callee.
198pub fn expr_references_ident(expr: &Spanned<Expr>, name: &str) -> bool {
199    match &expr.node {
200        Expr::Ident(n) | Expr::Resolved { name: n, .. } => n == name,
201        Expr::Literal(_) => false,
202        Expr::Attr(obj, _) => expr_references_ident(obj, name),
203        Expr::FnCall(callee, args) => {
204            expr_references_ident(callee, name)
205                || args.iter().any(|a| expr_references_ident(a, name))
206        }
207        Expr::BinOp(_, l, r) => expr_references_ident(l, name) || expr_references_ident(r, name),
208        Expr::Neg(inner) => expr_references_ident(inner, name),
209        Expr::Match { subject, arms } => {
210            expr_references_ident(subject, name)
211                || arms.iter().any(|a| expr_references_ident(&a.body, name))
212        }
213        Expr::Constructor(_, arg) => arg
214            .as_deref()
215            .is_some_and(|a| expr_references_ident(a, name)),
216        Expr::ErrorProp(inner) => expr_references_ident(inner, name),
217        Expr::InterpolatedStr(parts) => parts.iter().any(|p| match p {
218            StrPart::Parsed(inner) => expr_references_ident(inner, name),
219            _ => false,
220        }),
221        Expr::List(items) | Expr::Tuple(items) | Expr::IndependentProduct(items, _) => {
222            items.iter().any(|i| expr_references_ident(i, name))
223        }
224        Expr::MapLiteral(entries) => entries
225            .iter()
226            .any(|(k, v)| expr_references_ident(k, name) || expr_references_ident(v, name)),
227        Expr::RecordCreate { fields, .. } => {
228            fields.iter().any(|(_, v)| expr_references_ident(v, name))
229        }
230        Expr::RecordUpdate { base, updates, .. } => {
231            expr_references_ident(base, name)
232                || updates.iter().any(|(_, v)| expr_references_ident(v, name))
233        }
234        Expr::TailCall(boxed) => boxed.args.iter().any(|a| expr_references_ident(a, name)),
235    }
236}
237
238/// AST transform: walk `expr` and replace every recursive call to a fn
239/// in `targets` with `fn__fuel(fuel_var, …args)`. Inter-fn mutual calls
240/// in the same SCC are rewritten identically (the fuel parameter is
241/// threaded through the whole group).
242pub fn rewrite_recursive_calls_expr(
243    expr: &Spanned<Expr>,
244    targets: &HashSet<String>,
245    fuel_var: &str,
246) -> Spanned<Expr> {
247    let line = expr.line;
248    let new_node = match &expr.node {
249        Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => return expr.clone(),
250        Expr::Attr(obj, field) => Expr::Attr(
251            Box::new(rewrite_recursive_calls_expr(obj, targets, fuel_var)),
252            field.clone(),
253        ),
254        Expr::FnCall(callee, args) => {
255            let rewritten_args: Vec<Spanned<Expr>> = args
256                .iter()
257                .map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
258                .collect();
259            if let Some(name) = expr_to_dotted_name(&callee.node)
260                && targets.contains(&name)
261            {
262                let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
263                call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
264                call_args.extend(rewritten_args);
265                Expr::FnCall(
266                    Box::new(Spanned::new(Expr::Ident(fuel_helper_name(&name)), line)),
267                    call_args,
268                )
269            } else {
270                Expr::FnCall(
271                    Box::new(rewrite_recursive_calls_expr(callee, targets, fuel_var)),
272                    rewritten_args,
273                )
274            }
275        }
276        Expr::BinOp(op, left, right) => Expr::BinOp(
277            *op,
278            Box::new(rewrite_recursive_calls_expr(left, targets, fuel_var)),
279            Box::new(rewrite_recursive_calls_expr(right, targets, fuel_var)),
280        ),
281        Expr::Neg(inner) => Expr::Neg(Box::new(rewrite_recursive_calls_expr(
282            inner, targets, fuel_var,
283        ))),
284        Expr::Match { subject, arms } => Expr::Match {
285            subject: Box::new(rewrite_recursive_calls_expr(subject, targets, fuel_var)),
286            arms: arms
287                .iter()
288                .map(|arm| MatchArm {
289                    pattern: arm.pattern.clone(),
290                    body: Box::new(rewrite_recursive_calls_expr(&arm.body, targets, fuel_var)),
291                    binding_slots: std::sync::OnceLock::new(),
292                })
293                .collect(),
294        },
295        Expr::Constructor(name, arg) => Expr::Constructor(
296            name.clone(),
297            arg.as_ref()
298                .map(|inner| Box::new(rewrite_recursive_calls_expr(inner, targets, fuel_var))),
299        ),
300        Expr::ErrorProp(inner) => Expr::ErrorProp(Box::new(rewrite_recursive_calls_expr(
301            inner, targets, fuel_var,
302        ))),
303        Expr::InterpolatedStr(parts) => Expr::InterpolatedStr(
304            parts
305                .iter()
306                .map(|part| match part {
307                    StrPart::Literal(_) => part.clone(),
308                    StrPart::Parsed(inner) => StrPart::Parsed(Box::new(
309                        rewrite_recursive_calls_expr(inner, targets, fuel_var),
310                    )),
311                })
312                .collect(),
313        ),
314        Expr::List(items) => Expr::List(
315            items
316                .iter()
317                .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
318                .collect(),
319        ),
320        Expr::Tuple(items) => Expr::Tuple(
321            items
322                .iter()
323                .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
324                .collect(),
325        ),
326        Expr::IndependentProduct(items, flag) => Expr::IndependentProduct(
327            items
328                .iter()
329                .map(|item| rewrite_recursive_calls_expr(item, targets, fuel_var))
330                .collect(),
331            *flag,
332        ),
333        Expr::MapLiteral(entries) => Expr::MapLiteral(
334            entries
335                .iter()
336                .map(|(k, v)| {
337                    (
338                        rewrite_recursive_calls_expr(k, targets, fuel_var),
339                        rewrite_recursive_calls_expr(v, targets, fuel_var),
340                    )
341                })
342                .collect(),
343        ),
344        Expr::RecordCreate { type_name, fields } => Expr::RecordCreate {
345            type_name: type_name.clone(),
346            fields: fields
347                .iter()
348                .map(|(name, value)| {
349                    (
350                        name.clone(),
351                        rewrite_recursive_calls_expr(value, targets, fuel_var),
352                    )
353                })
354                .collect(),
355        },
356        Expr::RecordUpdate {
357            type_name,
358            base,
359            updates,
360        } => Expr::RecordUpdate {
361            type_name: type_name.clone(),
362            base: Box::new(rewrite_recursive_calls_expr(base, targets, fuel_var)),
363            updates: updates
364                .iter()
365                .map(|(name, value)| {
366                    (
367                        name.clone(),
368                        rewrite_recursive_calls_expr(value, targets, fuel_var),
369                    )
370                })
371                .collect(),
372        },
373        Expr::TailCall(boxed) => {
374            let TailCallData { target, args, .. } = boxed.as_ref();
375            let rewritten_args: Vec<Spanned<Expr>> = args
376                .iter()
377                .map(|arg| rewrite_recursive_calls_expr(arg, targets, fuel_var))
378                .collect();
379            if targets.contains(target) {
380                let mut call_args = Vec::with_capacity(rewritten_args.len() + 1);
381                call_args.push(Spanned::new(Expr::Ident(fuel_var.to_string()), line));
382                call_args.extend(rewritten_args);
383                Expr::FnCall(
384                    Box::new(Spanned::new(Expr::Ident(fuel_helper_name(target)), line)),
385                    call_args,
386                )
387            } else {
388                Expr::TailCall(Box::new(TailCallData::new(target.clone(), rewritten_args)))
389            }
390        }
391    };
392    Spanned::new(new_node, line)
393}
394
395/// Walk `expr` and rewrite every recursive call to `fn_name` into a
396/// call to `aux_name` carrying an extra `OMEGA_PROOF_SENTINEL` ident
397/// at the end of the argument list. Used by the proof-mode
398/// `IntCountdownGuarded` lowering — the synthetic argument lets Lean
399/// discharge the precondition at every recursive site without
400/// touching the original Aver source.
401pub fn rewrite_native_guarded_calls_expr(
402    expr: &Spanned<Expr>,
403    fn_name: &str,
404    aux_name: &str,
405) -> Spanned<Expr> {
406    let line = expr.line;
407    let new_node = match &expr.node {
408        Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. } => return expr.clone(),
409        Expr::Attr(obj, field) => Expr::Attr(
410            Box::new(rewrite_native_guarded_calls_expr(obj, fn_name, aux_name)),
411            field.clone(),
412        ),
413        Expr::FnCall(callee, args) => {
414            let rewritten_args: Vec<Spanned<Expr>> = args
415                .iter()
416                .map(|arg| rewrite_native_guarded_calls_expr(arg, fn_name, aux_name))
417                .collect();
418            if let Some(name) = expr_to_dotted_name(&callee.node)
419                && name == fn_name
420            {
421                let mut call_args = rewritten_args;
422                call_args.push(Spanned::new(
423                    Expr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
424                    line,
425                ));
426                Expr::FnCall(
427                    Box::new(Spanned::new(Expr::Ident(aux_name.to_string()), line)),
428                    call_args,
429                )
430            } else {
431                Expr::FnCall(
432                    Box::new(rewrite_native_guarded_calls_expr(callee, fn_name, aux_name)),
433                    rewritten_args,
434                )
435            }
436        }
437        Expr::BinOp(op, left, right) => Expr::BinOp(
438            *op,
439            Box::new(rewrite_native_guarded_calls_expr(left, fn_name, aux_name)),
440            Box::new(rewrite_native_guarded_calls_expr(right, fn_name, aux_name)),
441        ),
442        Expr::Neg(inner) => Expr::Neg(Box::new(rewrite_native_guarded_calls_expr(
443            inner, fn_name, aux_name,
444        ))),
445        Expr::Match { subject, arms } => Expr::Match {
446            subject: Box::new(rewrite_native_guarded_calls_expr(
447                subject, fn_name, aux_name,
448            )),
449            arms: arms
450                .iter()
451                .map(|arm| MatchArm {
452                    pattern: arm.pattern.clone(),
453                    body: Box::new(rewrite_native_guarded_calls_expr(
454                        &arm.body, fn_name, aux_name,
455                    )),
456                    binding_slots: std::sync::OnceLock::new(),
457                })
458                .collect(),
459        },
460        Expr::Constructor(name, arg) => Expr::Constructor(
461            name.clone(),
462            arg.as_ref()
463                .map(|inner| Box::new(rewrite_native_guarded_calls_expr(inner, fn_name, aux_name))),
464        ),
465        Expr::ErrorProp(inner) => Expr::ErrorProp(Box::new(rewrite_native_guarded_calls_expr(
466            inner, fn_name, aux_name,
467        ))),
468        Expr::InterpolatedStr(parts) => Expr::InterpolatedStr(
469            parts
470                .iter()
471                .map(|part| match part {
472                    StrPart::Literal(_) => part.clone(),
473                    StrPart::Parsed(inner) => StrPart::Parsed(Box::new(
474                        rewrite_native_guarded_calls_expr(inner, fn_name, aux_name),
475                    )),
476                })
477                .collect(),
478        ),
479        Expr::List(items) => Expr::List(
480            items
481                .iter()
482                .map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
483                .collect(),
484        ),
485        Expr::Tuple(items) => Expr::Tuple(
486            items
487                .iter()
488                .map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
489                .collect(),
490        ),
491        Expr::IndependentProduct(items, flag) => Expr::IndependentProduct(
492            items
493                .iter()
494                .map(|item| rewrite_native_guarded_calls_expr(item, fn_name, aux_name))
495                .collect(),
496            *flag,
497        ),
498        Expr::MapLiteral(entries) => Expr::MapLiteral(
499            entries
500                .iter()
501                .map(|(k, v)| {
502                    (
503                        rewrite_native_guarded_calls_expr(k, fn_name, aux_name),
504                        rewrite_native_guarded_calls_expr(v, fn_name, aux_name),
505                    )
506                })
507                .collect(),
508        ),
509        Expr::RecordCreate { type_name, fields } => Expr::RecordCreate {
510            type_name: type_name.clone(),
511            fields: fields
512                .iter()
513                .map(|(name, value)| {
514                    (
515                        name.clone(),
516                        rewrite_native_guarded_calls_expr(value, fn_name, aux_name),
517                    )
518                })
519                .collect(),
520        },
521        Expr::RecordUpdate {
522            type_name,
523            base,
524            updates,
525        } => Expr::RecordUpdate {
526            type_name: type_name.clone(),
527            base: Box::new(rewrite_native_guarded_calls_expr(base, fn_name, aux_name)),
528            updates: updates
529                .iter()
530                .map(|(name, value)| {
531                    (
532                        name.clone(),
533                        rewrite_native_guarded_calls_expr(value, fn_name, aux_name),
534                    )
535                })
536                .collect(),
537        },
538        Expr::TailCall(boxed) => {
539            let TailCallData { target, args, .. } = boxed.as_ref();
540            let rewritten_args: Vec<Spanned<Expr>> = args
541                .iter()
542                .map(|arg| rewrite_native_guarded_calls_expr(arg, fn_name, aux_name))
543                .collect();
544            if target == fn_name {
545                let mut call_args = rewritten_args;
546                call_args.push(Spanned::new(
547                    Expr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
548                    line,
549                ));
550                Expr::FnCall(
551                    Box::new(Spanned::new(Expr::Ident(aux_name.to_string()), line)),
552                    call_args,
553                )
554            } else {
555                Expr::TailCall(Box::new(TailCallData::new(target.clone(), rewritten_args)))
556            }
557        }
558    };
559    Spanned::new(new_node, line)
560}
561
562/// `ResolvedExpr` mirror of [`rewrite_native_guarded_calls_expr`].
563/// The proof-mode `IntCountdownGuarded` shape stores the rewritten
564/// arms on `ProofIR` (now resolved), so the rewriter that lifts
565/// `fn(args)` to `aux_name(args, OMEGA_PROOF_SENTINEL)` works in
566/// resolved space too.
567///
568/// Target detection: callers pass `target_fn_id` — the opaque
569/// [`crate::ir::FnId`] of the recursive fn the lowerer pinned. Every
570/// `ResolvedCallee::Fn(callee_id)` / `TailCall { target, .. }` site
571/// is compared by id, so two same-bare-name fns across modules can't
572/// cross-rewrite by accident — same anti-collision rule as the rest
573/// of #147 phase E.
574///
575/// Synthesised aux call: emitted as
576/// `ResolvedCallee::Unresolved { callee: Ident(aux_name) }` rather
577/// than `Builtin(...)` because `aux_name` is a codegen-only helper
578/// (no entry in the program's symbol table); `Unresolved` is exactly
579/// the variant the resolver uses for names it can't classify, and
580/// the Lean expr emitter already renders the inner ident verbatim.
581pub fn rewrite_native_guarded_calls_resolved_expr(
582    expr: &Spanned<crate::ir::hir::ResolvedExpr>,
583    target_fn_id: crate::ir::FnId,
584    aux_name: &str,
585) -> Spanned<crate::ir::hir::ResolvedExpr> {
586    use crate::ir::hir::{ResolvedCallee, ResolvedExpr, ResolvedMatchArm, ResolvedStrPart};
587    let line = expr.line;
588    let rec = |e: &Spanned<ResolvedExpr>| {
589        rewrite_native_guarded_calls_resolved_expr(e, target_fn_id, aux_name)
590    };
591    let synth_aux_callee = |line: usize| -> ResolvedCallee {
592        ResolvedCallee::Unresolved {
593            callee: Box::new(Spanned::new(
594                ResolvedExpr::Ident(aux_name.to_string()),
595                line,
596            )),
597        }
598    };
599    let new_node = match &expr.node {
600        ResolvedExpr::Literal(_) | ResolvedExpr::Ident(_) | ResolvedExpr::Resolved { .. } => {
601            return expr.clone();
602        }
603        ResolvedExpr::Attr(obj, field) => ResolvedExpr::Attr(Box::new(rec(obj)), field.clone()),
604        ResolvedExpr::Call(callee, args) => {
605            let rewritten_args: Vec<Spanned<ResolvedExpr>> = args.iter().map(&rec).collect();
606            let target_matches =
607                matches!(callee, ResolvedCallee::Fn(callee_id) if *callee_id == target_fn_id);
608            if target_matches {
609                let mut call_args = rewritten_args;
610                call_args.push(Spanned::new(
611                    ResolvedExpr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
612                    line,
613                ));
614                ResolvedExpr::Call(synth_aux_callee(line), call_args)
615            } else {
616                ResolvedExpr::Call(callee.clone(), rewritten_args)
617            }
618        }
619        ResolvedExpr::BinOp(op, left, right) => {
620            ResolvedExpr::BinOp(*op, Box::new(rec(left)), Box::new(rec(right)))
621        }
622        ResolvedExpr::Neg(inner) => ResolvedExpr::Neg(Box::new(rec(inner))),
623        ResolvedExpr::Match { subject, arms } => ResolvedExpr::Match {
624            subject: Box::new(rec(subject)),
625            arms: arms
626                .iter()
627                .map(|arm| ResolvedMatchArm {
628                    pattern: arm.pattern.clone(),
629                    body: Box::new(rec(&arm.body)),
630                    binding_slots: std::sync::OnceLock::new(),
631                })
632                .collect(),
633        },
634        ResolvedExpr::Ctor(ctor, args) => {
635            ResolvedExpr::Ctor(ctor.clone(), args.iter().map(&rec).collect())
636        }
637        ResolvedExpr::ErrorProp(inner) => ResolvedExpr::ErrorProp(Box::new(rec(inner))),
638        ResolvedExpr::InterpolatedStr(parts) => ResolvedExpr::InterpolatedStr(
639            parts
640                .iter()
641                .map(|p| match p {
642                    ResolvedStrPart::Literal(_) => p.clone(),
643                    ResolvedStrPart::Parsed(inner) => ResolvedStrPart::Parsed(Box::new(rec(inner))),
644                })
645                .collect(),
646        ),
647        ResolvedExpr::List(items) => ResolvedExpr::List(items.iter().map(&rec).collect()),
648        ResolvedExpr::Tuple(items) => ResolvedExpr::Tuple(items.iter().map(&rec).collect()),
649        ResolvedExpr::IndependentProduct(items, flag) => {
650            ResolvedExpr::IndependentProduct(items.iter().map(&rec).collect(), *flag)
651        }
652        ResolvedExpr::MapLiteral(entries) => {
653            ResolvedExpr::MapLiteral(entries.iter().map(|(k, v)| (rec(k), rec(v))).collect())
654        }
655        ResolvedExpr::RecordCreate {
656            type_id,
657            type_name,
658            fields,
659        } => ResolvedExpr::RecordCreate {
660            type_id: *type_id,
661            type_name: type_name.clone(),
662            fields: fields.iter().map(|(n, v)| (n.clone(), rec(v))).collect(),
663        },
664        ResolvedExpr::RecordUpdate {
665            type_id,
666            type_name,
667            base,
668            updates,
669        } => ResolvedExpr::RecordUpdate {
670            type_id: *type_id,
671            type_name: type_name.clone(),
672            base: Box::new(rec(base)),
673            updates: updates.iter().map(|(n, v)| (n.clone(), rec(v))).collect(),
674        },
675        ResolvedExpr::TailCall { target, args } => {
676            let rewritten_args: Vec<Spanned<ResolvedExpr>> = args.iter().map(&rec).collect();
677            if *target == target_fn_id {
678                let mut call_args = rewritten_args;
679                call_args.push(Spanned::new(
680                    ResolvedExpr::Ident(OMEGA_PROOF_SENTINEL.to_string()),
681                    line,
682                ));
683                ResolvedExpr::Call(synth_aux_callee(line), call_args)
684            } else {
685                ResolvedExpr::TailCall {
686                    target: *target,
687                    args: rewritten_args,
688                }
689            }
690        }
691    };
692    Spanned::new(new_node, line)
693}
694
695/// Body-level wrapper around [`rewrite_native_guarded_calls_expr`].
696pub fn rewrite_native_guarded_calls_body(body: &FnBody, fn_name: &str, aux_name: &str) -> FnBody {
697    FnBody::Block(
698        body.stmts()
699            .iter()
700            .map(|stmt| match stmt {
701                Stmt::Binding(name, ty, expr) => Stmt::Binding(
702                    name.clone(),
703                    ty.clone(),
704                    rewrite_native_guarded_calls_expr(expr, fn_name, aux_name),
705                ),
706                Stmt::Expr(expr) => {
707                    Stmt::Expr(rewrite_native_guarded_calls_expr(expr, fn_name, aux_name))
708                }
709            })
710            .collect(),
711    )
712}
713
714/// Body-level wrapper around [`rewrite_recursive_calls_expr`].
715pub fn rewrite_recursive_calls_body(
716    body: &FnBody,
717    targets: &HashSet<String>,
718    fuel_var: &str,
719) -> FnBody {
720    FnBody::Block(
721        body.stmts()
722            .iter()
723            .map(|stmt| match stmt {
724                Stmt::Binding(name, ty, expr) => Stmt::Binding(
725                    name.clone(),
726                    ty.clone(),
727                    rewrite_recursive_calls_expr(expr, targets, fuel_var),
728                ),
729                Stmt::Expr(expr) => {
730                    Stmt::Expr(rewrite_recursive_calls_expr(expr, targets, fuel_var))
731                }
732            })
733            .collect(),
734    )
735}