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