AILang/crates/ailang-core/src/desugar.rs

//! AST → AST rewriter that runs **after** [`crate::load_module`] and
//! **before** anything that consumes a [`Module`] (typecheck, codegen).
//!
//! ## Iter 16a: nested constructor patterns in `match`
//!
//! Through Iter 15e, a [`Pattern::Ctor`]'s sub-patterns were restricted to
//! [`Pattern::Var`] and [`Pattern::Wild`] — nested ctors like
//! `(pat-ctor Cons a (pat-ctor Cons b _))` were rejected by the
//! checker (`nested-ctor-pattern-not-allowed`) and silently mishandled
//! by the codegen. Iter 16a lifts the gate by **desugaring** nested
//! ctor sub-patterns into chains of single-level matches at the AST
//! level, before either the checker or codegen sees the module. The
//! rewrite is pure (no side effects, no I/O), runs in memory, and
//! produces a [`Module`] semantically equivalent to its input.
//!
//! Lit sub-patterns inside a Ctor remain rejected — that is a separate
//! iter. The [`crate::ast::Pattern`] enum, the JSON schema, and every
//! on-disk hash are unchanged: the desugar runs **after** `load_module`,
//! so canonical bytes computed from the source file are untouched. The
//! checker / codegen consume the desugared form.
//!
//! ## Algorithm
//!
//! For a [`Term::Match`] whose arms contain a non-flat ctor pattern:
//!
//! 1. Bottom-up: recursively desugar the scrutinee and every arm body
//!    first (children are normalised before parents).
//! 2. Let-bind the scrutinee to a fresh name `$mp_N` (so that we don't
//!    re-evaluate effectful scrutinees per arm).
//! 3. Build a chain of single-level matches via
//!    `build_chain(s_var, arms, default)` where `default` is a unit
//!    literal — a placeholder for the unreachable fallthrough. Valid
//!    programs never reach it because the checker requires
//!    exhaustiveness (the catch-all arm dominates the chain).
//! 4. Each arm is lowered via `desugar_one_arm`:
//!    - `Pattern::Wild` → arm body (catch-all; later arms are dropped).
//!    - `Pattern::Var { name }` → `Let { name = scrutinee_var; body }`.
//!    - `Pattern::Lit { lit }` → `Term::If { cond = (== s_var lit),
//!      then = arm.body, else_ = fall_k }` (Iter 16c). After this
//!      rewrite, no `Pattern::Lit` reaches typecheck or codegen.
//!    - `Pattern::Ctor { ctor, fields }`: lift each field to a fresh
//!      var, build a flat outer pattern, walk fields right-to-left
//!      wrapping inner sub-patterns via `wrap_sub`. Emit a
//!      `Term::Match` with two arms: the flat ctor and a wildcard
//!      fall-through into the rest of the chain.
//!
//! `wrap_sub(fv, sub, body, fall_k)` recurses on `sub`:
//! - `Var(n)` → `Let(n, Var(fv), body)`.
//! - `Wild` → `body`.
//! - `Lit` (Iter 16c) → `Term::If { cond = (== fv lit), then = body,
//!   else_ = fall_k }`. Same shape as the top-level case but on the
//!   field-bound fresh variable.
//! - `Ctor` → recursively `desugar_match(Var(fv), [(sub, body),
//!   (Wild, fall_k)])`. The recursion is what handles arbitrary
//!   nesting depth (`Cons a (Cons b (Cons c _))`).
//!
//! The fall-through term is cloned per inner match because each branch
//! that fails must continue to the same fallthrough. Worst-case size is
//! O(arms × depth); acceptable for typical patterns.
//!
//! ## Fresh-name safety
//!
//! `$` is a valid identifier character in form (A) (the lexer's `Ident`
//! token is "anything not paren/int/string"), so a user *could* write
//! `$mp_0` themselves. To avoid collisions, the fresh-name generator
//! reads a pre-collected set of every name appearing in any
//! [`crate::ast::Term::Var`] or [`crate::ast::Pattern::Var`] of the
//! module and bumps the counter until it lands on a name not in that
//! set.
//!
//! ## Iter 16b.1: local recursive `let` (no-capture)
//!
//! [`Term::LetRec`] is a surface-only AST node introduced in 16b.1. The
//! desugar pass eliminates it before typecheck/codegen by **lifting**
//! the LetRec to a synthetic top-level fn in the same module. The
//! lifted name has the form `<hint>$lr_N`, fresh against both the
//! existing module-top-level def names and the [`Desugarer::used`] set.
//! Inside the LetRec's `body` and `in_term`, every reference to the
//! original local name is rewritten via [`subst_var`] to the lifted
//! name. The lifted [`FnDef`] is appended to the module's `defs` so
//! the typechecker / codegen see it as if it had been written
//! top-level by hand.
//!
//! Capture is **not** supported in 16b.1: if the LetRec's body's free
//! variables (excluding `{name} ∪ params`) intersect the enclosing
//! lexical scope (params of the surrounding fn, let-bound names, or
//! pattern-bound names), the desugar pass panics with a diagnostic
//! that points the author at 16b.2 (closure conversion). Free vars
//! that resolve to module-top-level def names, qualified import names
//! (containing `.`), effect-op names (containing `/`), and builtin
//! operators are all ignored — those reach a lifted top-level fn
//! through the same path they reach any other top-level fn.
//!
//! ## What this module deliberately does not do
//!
//! - It does not alter [`crate::ast::Pattern`] or [`crate::ast::Term`].
//! - It does not introduce a new IR or hashable form.
//! - It does not implement Maranget-style decision trees; the rewrite
//!   is a literal-translation chain that the existing single-level
//!   match codegen already handles.
//! - It does not implement closure conversion for [`Term::LetRec`]
//!   bodies that capture from the enclosing scope (queued as 16b.2).

use crate::ast::*;
use std::collections::BTreeSet;

/// Rewrite `m` so that every [`Pattern::Ctor`] sub-pattern is a
/// [`Pattern::Var`] or [`Pattern::Wild`] (i.e. flat) and every
/// [`Term::LetRec`] is replaced by a reference to a synthetic
/// top-level fn.
///
/// Pure / total: returns a new [`Module`]; does not mutate `m`.
/// Idempotent: a module that is already flat and LetRec-free is
/// returned as-is modulo cloning.
pub fn desugar_module(m: &Module) -> Module {
    let mut used: BTreeSet<String> = BTreeSet::new();
    for def in &m.defs {
        match def {
            Def::Fn(f) => {
                used.insert(f.name.clone());
                for p in &f.params {
                    used.insert(p.clone());
                }
                collect_used_in_term(&f.body, &mut used);
            }
            Def::Const(c) => {
                used.insert(c.name.clone());
                collect_used_in_term(&c.value, &mut used);
            }
            Def::Type(td) => {
                used.insert(td.name.clone());
            }
        }
    }
    // Iter 16b.1: pre-collect every top-level def name. The LetRec
    // lifter consults this set to (a) know whether a free var of a
    // LetRec body resolves to a top-level def (no capture) and
    // (b) keep its fresh-name generator from colliding with an
    // existing def — including ones lifted earlier in the same pass.
    let mut module_top_names: BTreeSet<String> = BTreeSet::new();
    for def in &m.defs {
        module_top_names.insert(def.name().to_string());
    }
    let mut d = Desugarer {
        counter: 0,
        used,
        lifted: Vec::new(),
        module_top_names,
    };
    let mut out = m.clone();
    for def in &mut out.defs {
        match def {
            Def::Fn(f) => {
                let scope: BTreeSet<String> = f.params.iter().cloned().collect();
                f.body = d.desugar_term(&f.body, &scope);
            }
            Def::Const(c) => {
                let scope: BTreeSet<String> = BTreeSet::new();
                c.value = d.desugar_term(&c.value, &scope);
            }
            Def::Type(_) => {}
        }
    }
    // Iter 16b.1: append every lifted fn to the desugared module.
    // Order: original defs first, then lifts in the order they were
    // produced by the bottom-up walk. Typecheck/codegen are order-
    // insensitive at the def list level (they index by name), so this
    // is purely cosmetic.
    out.defs.extend(d.lifted.into_iter());
    out
}

/// Walks a term and inserts every [`Term::Var.name`] and every
/// [`Pattern::Var.name`] into `used`. Used to seed
/// [`Desugarer::fresh`] so a generated name `$mp_N` cannot shadow a
/// source-level identifier.
fn collect_used_in_term(t: &Term, used: &mut BTreeSet<String>) {
    match t {
        Term::Lit { .. } => {}
        Term::Var { name } => {
            used.insert(name.clone());
        }
        Term::App { callee, args, .. } => {
            collect_used_in_term(callee, used);
            for a in args {
                collect_used_in_term(a, used);
            }
        }
        Term::Let { name, value, body } => {
            used.insert(name.clone());
            collect_used_in_term(value, used);
            collect_used_in_term(body, used);
        }
        Term::If { cond, then, else_ } => {
            collect_used_in_term(cond, used);
            collect_used_in_term(then, used);
            collect_used_in_term(else_, used);
        }
        Term::Do { args, .. } => {
            for a in args {
                collect_used_in_term(a, used);
            }
        }
        Term::Ctor { args, .. } => {
            for a in args {
                collect_used_in_term(a, used);
            }
        }
        Term::Match { scrutinee, arms } => {
            collect_used_in_term(scrutinee, used);
            for arm in arms {
                collect_used_in_pattern(&arm.pat, used);
                collect_used_in_term(&arm.body, used);
            }
        }
        Term::Lam { params, body, .. } => {
            for p in params {
                used.insert(p.clone());
            }
            collect_used_in_term(body, used);
        }
        Term::Seq { lhs, rhs } => {
            collect_used_in_term(lhs, used);
            collect_used_in_term(rhs, used);
        }
        Term::LetRec { name, params, body, in_term, .. } => {
            used.insert(name.clone());
            for p in params {
                used.insert(p.clone());
            }
            collect_used_in_term(body, used);
            collect_used_in_term(in_term, used);
        }
    }
}

/// Walks a pattern and inserts every [`Pattern::Var.name`] into `used`.
fn collect_used_in_pattern(p: &Pattern, used: &mut BTreeSet<String>) {
    match p {
        Pattern::Wild => {}
        Pattern::Var { name } => {
            used.insert(name.clone());
        }
        Pattern::Lit { .. } => {}
        Pattern::Ctor { fields, .. } => {
            for sub in fields {
                collect_used_in_pattern(sub, used);
            }
        }
    }
}

/// State carried across a module-wide desugar pass.
///
/// `counter` is incremented for every fresh-name request; `used`
/// contains every source-level identifier (var-bind or var-reference)
/// in the module so that [`fresh`](Self::fresh) cannot shadow one.
///
/// Iter 16b.1 fields:
/// - `lifted` accumulates synthetic top-level fns produced by
///   [`Term::LetRec`] desugaring. Appended to `Module.defs` once the
///   per-def walk finishes.
/// - `module_top_names` mirrors every name reachable as a top-level
///   def at this point in the pass (originals + already-lifted). The
///   capture analyser uses it to classify free vars; the fresh-name
///   generator [`fresh_lifted`](Self::fresh_lifted) consults it so
///   later lifts cannot collide with earlier ones.
struct Desugarer {
    counter: u64,
    used: BTreeSet<String>,
    lifted: Vec<Def>,
    module_top_names: BTreeSet<String>,
}

impl Desugarer {
    /// Returns a name of the form `$mp_N` not present in `used`.
    /// Increments the counter until it lands on a free name; the
    /// generated name itself is added to `used` so a subsequent call
    /// returns a distinct one.
    fn fresh(&mut self) -> String {
        loop {
            let n = format!("$mp_{}", self.counter);
            self.counter += 1;
            if !self.used.contains(&n) {
                self.used.insert(n.clone());
                return n;
            }
        }
    }

    /// Iter 16b.1: returns a name of the form `<hint>$lr_N` not in
    /// `used` and not in `module_top_names`. Bumps both sets so a
    /// later lift cannot collide. The `hint` is the source-level
    /// LetRec name; it makes lifted bindings traceable through the
    /// pipeline (typecheck errors / IR mangling / panic messages).
    fn fresh_lifted(&mut self, hint: &str) -> String {
        let mut n = 0u64;
        loop {
            let candidate = format!("{hint}$lr_{n}");
            n += 1;
            if !self.used.contains(&candidate) && !self.module_top_names.contains(&candidate) {
                self.used.insert(candidate.clone());
                self.module_top_names.insert(candidate.clone());
                return candidate;
            }
        }
    }

    /// Recursively rewrites `t`: descends into every child first
    /// (bottom-up), then dispatches a [`Term::Match`] to
    /// [`desugar_match`](Self::desugar_match) and a [`Term::LetRec`]
    /// to the 16b.1 lifter.
    ///
    /// `scope` is the set of names lexically bound at this point —
    /// initialised at the def boundary (fn params for [`Def::Fn`],
    /// empty for [`Def::Const`]) and extended locally by every
    /// binder ([`Term::Let`], [`Term::Lam`], [`Term::Match`] arms,
    /// [`Term::LetRec`]). Used by the LetRec lifter to detect
    /// captures.
    fn desugar_term(&mut self, t: &Term, scope: &BTreeSet<String>) -> Term {
        match t {
            Term::Lit { .. } | Term::Var { .. } => t.clone(),
            Term::App { callee, args, tail } => Term::App {
                callee: Box::new(self.desugar_term(callee, scope)),
                args: args.iter().map(|a| self.desugar_term(a, scope)).collect(),
                tail: *tail,
            },
            Term::Let { name, value, body } => {
                let v = self.desugar_term(value, scope);
                let mut inner = scope.clone();
                inner.insert(name.clone());
                let b = self.desugar_term(body, &inner);
                Term::Let {
                    name: name.clone(),
                    value: Box::new(v),
                    body: Box::new(b),
                }
            }
            Term::If { cond, then, else_ } => Term::If {
                cond: Box::new(self.desugar_term(cond, scope)),
                then: Box::new(self.desugar_term(then, scope)),
                else_: Box::new(self.desugar_term(else_, scope)),
            },
            Term::Do { op, args, tail } => Term::Do {
                op: op.clone(),
                args: args.iter().map(|a| self.desugar_term(a, scope)).collect(),
                tail: *tail,
            },
            Term::Ctor { type_name, ctor, args } => Term::Ctor {
                type_name: type_name.clone(),
                ctor: ctor.clone(),
                args: args.iter().map(|a| self.desugar_term(a, scope)).collect(),
            },
            Term::Match { scrutinee, arms } => {
                // Recurse into children first (bottom-up).
                let scrutinee = self.desugar_term(scrutinee, scope);
                let arms: Vec<Arm> = arms
                    .iter()
                    .map(|a| {
                        let mut inner = scope.clone();
                        pattern_binds(&a.pat, &mut inner);
                        Arm {
                            pat: a.pat.clone(),
                            body: self.desugar_term(&a.body, &inner),
                        }
                    })
                    .collect();
                self.desugar_match(scrutinee, arms)
            }
            Term::Lam {
                params,
                param_tys,
                ret_ty,
                effects,
                body,
            } => {
                let mut inner = scope.clone();
                for p in params {
                    inner.insert(p.clone());
                }
                Term::Lam {
                    params: params.clone(),
                    param_tys: param_tys.clone(),
                    ret_ty: ret_ty.clone(),
                    effects: effects.clone(),
                    body: Box::new(self.desugar_term(body, &inner)),
                }
            }
            Term::Seq { lhs, rhs } => Term::Seq {
                lhs: Box::new(self.desugar_term(lhs, scope)),
                rhs: Box::new(self.desugar_term(rhs, scope)),
            },
            Term::LetRec { name, ty, params, body, in_term } => {
                // Iter 16b.1: lift to a synthetic top-level fn.
                //
                // Body's scope: outer ∪ {name} ∪ params (recursive
                // self-reference is legal; params shadow outer names).
                let mut body_scope = scope.clone();
                body_scope.insert(name.clone());
                for p in params {
                    body_scope.insert(p.clone());
                }
                let desugared_body = self.desugar_term(body, &body_scope);
                // in_term's scope: outer ∪ {name} (params are lambda-
                // local to the LetRec's body, not visible in `in`).
                let mut in_scope = scope.clone();
                in_scope.insert(name.clone());
                let desugared_in = self.desugar_term(in_term, &in_scope);

                // Capture detection: free vars of the desugared body
                // wrt {name} ∪ params, intersected with the outer
                // `scope`. Anything in that intersection is a capture
                // (16b.2 territory).
                let mut local_bound: BTreeSet<String> = BTreeSet::new();
                local_bound.insert(name.clone());
                for p in params {
                    local_bound.insert(p.clone());
                }
                let mut frees: BTreeSet<String> = BTreeSet::new();
                free_vars_in_term(&desugared_body, &local_bound, &mut frees);
                let captures: Vec<String> = frees.intersection(scope).cloned().collect();
                if !captures.is_empty() {
                    panic!(
                        "Iter 16b.1: local recursive `let` `{}` would capture {:?} \
                         from enclosing scope; capture is queued for 16b.2. Lift \
                         the helper to top-level for now.",
                        name, captures
                    );
                }

                // Lift: synthesise a fresh top-level def name and
                // substitute every reference to the local `name`
                // (in body and in_term) to the lifted name.
                let lifted_name = self.fresh_lifted(name);
                let body_lifted = subst_var(&desugared_body, name, &lifted_name);
                let in_lifted = subst_var(&desugared_in, name, &lifted_name);
                self.lifted.push(Def::Fn(FnDef {
                    name: lifted_name,
                    ty: ty.clone(),
                    params: params.clone(),
                    body: body_lifted,
                    doc: None,
                }));
                in_lifted
            }
        }
    }

    /// Lowers a `match` whose children have already been desugared.
    ///
    /// If every arm is already flat ([`is_flat`] returns true), the
    /// match is reconstructed unchanged (no let-binding or chain
    /// allocation overhead). Otherwise the scrutinee is let-bound to
    /// a fresh `$mp_N` and the arms are translated to a chain of
    /// single-level matches over that variable.
    fn desugar_match(&mut self, scrutinee: Term, arms: Vec<Arm>) -> Term {
        if arms.iter().all(|a| is_flat(&a.pat)) {
            return Term::Match {
                scrutinee: Box::new(scrutinee),
                arms,
            };
        }
        let s = self.fresh();
        let s_var = Term::Var { name: s.clone() };
        // Document: `default` is unreachable for valid programs because
        // the typechecker requires a catch-all arm (`pat-wild` or
        // exhaustive ctor coverage). It only exists to give
        // `build_chain` a well-typed terminator.
        let default = Term::Lit { lit: Literal::Unit };
        let chain = self.build_chain(&s_var, &arms, &default);
        Term::Let {
            name: s,
            value: Box::new(scrutinee),
            body: Box::new(chain),
        }
    }

    /// Recursively builds a chain of single-arm matches with a shared
    /// fall-through. Empty arms ⇒ `default`; otherwise the first arm
    /// is desugared with the rest of the chain as its fall-through.
    fn build_chain(&mut self, s_var: &Term, arms: &[Arm], default: &Term) -> Term {
        if arms.is_empty() {
            return default.clone();
        }
        let head = &arms[0];
        let rest = &arms[1..];
        let fall_k = self.build_chain(s_var, rest, default);
        self.desugar_one_arm(s_var, head, fall_k)
    }

    /// Lowers one arm into a term. Wild/Var arms drop the chain (the
    /// arm matches everything); Lit and Ctor arms emit a `Term::Match`
    /// with the desugared head pattern as the first arm and a
    /// wildcard fall-through to `fall_k`.
    fn desugar_one_arm(&mut self, s_var: &Term, arm: &Arm, fall_k: Term) -> Term {
        match &arm.pat {
            Pattern::Wild => arm.body.clone(),
            Pattern::Var { name } => Term::Let {
                name: name.clone(),
                value: Box::new(s_var.clone()),
                body: Box::new(arm.body.clone()),
            },
            Pattern::Lit { lit } => {
                // Iter 16c: a top-level lit arm desugars to `if (== s_var lit)
                // then arm.body else fall_k`. Eliminates `Pattern::Lit` from
                // the desugared output entirely; codegen never sees it.
                let cmp = build_eq(s_var.clone(), lit);
                Term::If {
                    cond: Box::new(cmp),
                    then: Box::new(arm.body.clone()),
                    else_: Box::new(fall_k),
                }
            }
            Pattern::Ctor { ctor, fields } => {
                // Lift each field to a fresh var; build the flat outer
                // pattern, then walk right-to-left wrapping each inner
                // sub-pattern via `wrap_sub` so the deepest field is
                // matched first inside-out.
                let fresh_vars: Vec<String> = fields.iter().map(|_| self.fresh()).collect();
                let flat_fields: Vec<Pattern> = fresh_vars
                    .iter()
                    .map(|n| Pattern::Var { name: n.clone() })
                    .collect();
                let mut inner = arm.body.clone();
                for (sub, fv) in fields.iter().zip(fresh_vars.iter()).rev() {
                    inner = self.wrap_sub(fv, sub, inner, &fall_k);
                }
                Term::Match {
                    scrutinee: Box::new(s_var.clone()),
                    arms: vec![
                        Arm {
                            pat: Pattern::Ctor {
                                ctor: ctor.clone(),
                                fields: flat_fields,
                            },
                            body: inner,
                        },
                        Arm {
                            pat: Pattern::Wild,
                            body: fall_k,
                        },
                    ],
                }
            }
        }
    }

    /// Wraps `body` so that it only runs when the value bound to `fv`
    /// matches `sub`. For `Var` / `Wild` the wrap is a let-bind / no-op;
    /// for nested `Ctor` / `Lit` the wrap is a recursive
    /// [`desugar_match`](Self::desugar_match) — that recursion is what
    /// flattens arbitrarily-nested patterns.
    fn wrap_sub(&mut self, fv: &str, sub: &Pattern, body: Term, fall_k: &Term) -> Term {
        match sub {
            Pattern::Wild => body,
            Pattern::Var { name } => Term::Let {
                name: name.clone(),
                value: Box::new(Term::Var {
                    name: fv.to_string(),
                }),
                body: Box::new(body),
            },
            Pattern::Lit { lit } => {
                // Iter 16c: a lit sub-pattern desugars to `if (== fv lit) body
                // else fall_k`. Same shape as the top-level case, but on the
                // field-bound fresh variable rather than the original
                // scrutinee.
                let cmp = build_eq(Term::Var { name: fv.to_string() }, lit);
                Term::If {
                    cond: Box::new(cmp),
                    then: Box::new(body),
                    else_: Box::new(fall_k.clone()),
                }
            }
            Pattern::Ctor { .. } => self.desugar_match(
                Term::Var {
                    name: fv.to_string(),
                },
                vec![
                    Arm {
                        pat: sub.clone(),
                        body,
                    },
                    Arm {
                        pat: Pattern::Wild,
                        body: fall_k.clone(),
                    },
                ],
            ),
        }
    }
}

/// True iff `p` is fully shallow: a [`Pattern::Var`], [`Pattern::Wild`],
/// or a [`Pattern::Ctor`] all of whose fields are `Var` or `Wild`. Used
/// by [`Desugarer::desugar_match`] to skip the let-binding and chain
/// construction for already-flat matches.
///
/// Iter 16c: [`Pattern::Lit`] is **not** flat — it always desugars to a
/// [`Term::If`] via [`build_eq`], so it must take the chain path even
/// when no other arm needs flattening. Otherwise the early-return in
/// [`Desugarer::desugar_match`] would leak a `Pattern::Lit` arm to
/// typecheck/codegen, which the codegen rejects with an internal error.
fn is_flat(p: &Pattern) -> bool {
    match p {
        Pattern::Wild | Pattern::Var { .. } => true,
        // Iter 16c: lit patterns are not flat; they desugar to If.
        Pattern::Lit { .. } => false,
        Pattern::Ctor { fields, .. } => fields
            .iter()
            .all(|f| matches!(f, Pattern::Var { .. } | Pattern::Wild)),
    }
}

/// Iter 16c: Build an equality-test term for a lit pattern. The shape is
/// `(app == scrutinee lit_term)` for Int/Bool/Str. Unit literals are
/// degenerate — every Unit value is equal — so an emitted `true` literal
/// stands in.
///
/// Note: as of Iter 16c, the `==` builtin is typed `(Int, Int) -> Bool`
/// only. Bool- and Str-lit patterns therefore reach typecheck as a
/// type error rather than a clean rewrite. They remain authorable at the
/// AST level (the desugar never refuses) but are not yet usable end-to-
/// end. Lifting that gate is bound to extending `==` to those types.
fn build_eq(scrutinee: Term, lit: &Literal) -> Term {
    match lit {
        Literal::Unit => Term::Lit {
            lit: Literal::Bool { value: true },
        },
        Literal::Int { .. } | Literal::Bool { .. } | Literal::Str { .. } => Term::App {
            callee: Box::new(Term::Var {
                name: "==".to_string(),
            }),
            args: vec![scrutinee, Term::Lit { lit: lit.clone() }],
            tail: false,
        },
    }
}

/// Iter 16b.1: collect free variable names of `t`, with `bound` being
/// the set of names lexically bound at this point. A
/// [`Term::Var`] `{ name }` is "free" iff `name` is not in `bound`.
/// Compound binders ([`Term::Let`], [`Term::Lam`], [`Term::Match`]
/// arms, [`Term::LetRec`]) extend `bound` for the sub-walk.
///
/// Used by the [`Term::LetRec`] desugar to decide whether a local
/// recursive let would capture any name from the enclosing scope
/// (16b.2 has not shipped yet, so capture is a panic at desugar time).
fn free_vars_in_term(t: &Term, bound: &BTreeSet<String>, out: &mut BTreeSet<String>) {
    match t {
        Term::Lit { .. } => {}
        Term::Var { name } => {
            if !bound.contains(name) {
                out.insert(name.clone());
            }
        }
        Term::App { callee, args, .. } => {
            free_vars_in_term(callee, bound, out);
            for a in args {
                free_vars_in_term(a, bound, out);
            }
        }
        Term::Let { name, value, body } => {
            free_vars_in_term(value, bound, out);
            let mut b = bound.clone();
            b.insert(name.clone());
            free_vars_in_term(body, &b, out);
        }
        Term::If { cond, then, else_ } => {
            free_vars_in_term(cond, bound, out);
            free_vars_in_term(then, bound, out);
            free_vars_in_term(else_, bound, out);
        }
        Term::Do { args, .. } => {
            for a in args {
                free_vars_in_term(a, bound, out);
            }
        }
        Term::Ctor { args, .. } => {
            for a in args {
                free_vars_in_term(a, bound, out);
            }
        }
        Term::Match { scrutinee, arms } => {
            free_vars_in_term(scrutinee, bound, out);
            for arm in arms {
                let mut b = bound.clone();
                pattern_binds(&arm.pat, &mut b);
                free_vars_in_term(&arm.body, &b, out);
            }
        }
        Term::Lam { params, body, .. } => {
            let mut b = bound.clone();
            for p in params {
                b.insert(p.clone());
            }
            free_vars_in_term(body, &b, out);
        }
        Term::Seq { lhs, rhs } => {
            free_vars_in_term(lhs, bound, out);
            free_vars_in_term(rhs, bound, out);
        }
        Term::LetRec { name, params, body, in_term, .. } => {
            // Inside body: name + params are bound.
            let mut b_body = bound.clone();
            b_body.insert(name.clone());
            for p in params {
                b_body.insert(p.clone());
            }
            free_vars_in_term(body, &b_body, out);
            // Inside in_term: only name is bound.
            let mut b_in = bound.clone();
            b_in.insert(name.clone());
            free_vars_in_term(in_term, &b_in, out);
        }
    }
}

/// Iter 16b.1: collect every name a pattern binds into `out`. Mirrors
/// `Pattern::pattern_bound_names` in `ailang-codegen`; duplicated here
/// because `ailang-core` cannot depend on the codegen crate.
fn pattern_binds(p: &Pattern, out: &mut BTreeSet<String>) {
    match p {
        Pattern::Wild | Pattern::Lit { .. } => {}
        Pattern::Var { name } => {
            out.insert(name.clone());
        }
        Pattern::Ctor { fields, .. } => {
            for sub in fields {
                pattern_binds(sub, out);
            }
        }
    }
}

/// Iter 16b.1: rewrite every free [`Term::Var`] `{ name == from }`
/// reference in `t` to [`Term::Var`] `{ name = to }`. Respects
/// shadowing: if a binder rebinds `from`, occurrences inside that
/// binder's scope stop being free and are left alone.
///
/// Used after lifting a [`Term::LetRec`] body to a synthetic top-level
/// fn — every recursive self-call site in the lifted body and every
/// reference in the in-term needs to point at the new global name.
fn subst_var(t: &Term, from: &str, to: &str) -> Term {
    match t {
        Term::Lit { .. } => t.clone(),
        Term::Var { name } => {
            if name == from {
                Term::Var { name: to.to_string() }
            } else {
                t.clone()
            }
        }
        Term::App { callee, args, tail } => Term::App {
            callee: Box::new(subst_var(callee, from, to)),
            args: args.iter().map(|a| subst_var(a, from, to)).collect(),
            tail: *tail,
        },
        Term::Let { name, value, body } => {
            let v = subst_var(value, from, to);
            // If this `let` rebinds `from`, leave its body alone.
            let b = if name == from {
                (**body).clone()
            } else {
                subst_var(body, from, to)
            };
            Term::Let {
                name: name.clone(),
                value: Box::new(v),
                body: Box::new(b),
            }
        }
        Term::If { cond, then, else_ } => Term::If {
            cond: Box::new(subst_var(cond, from, to)),
            then: Box::new(subst_var(then, from, to)),
            else_: Box::new(subst_var(else_, from, to)),
        },
        Term::Do { op, args, tail } => Term::Do {
            op: op.clone(),
            args: args.iter().map(|a| subst_var(a, from, to)).collect(),
            tail: *tail,
        },
        Term::Ctor { type_name, ctor, args } => Term::Ctor {
            type_name: type_name.clone(),
            ctor: ctor.clone(),
            args: args.iter().map(|a| subst_var(a, from, to)).collect(),
        },
        Term::Match { scrutinee, arms } => Term::Match {
            scrutinee: Box::new(subst_var(scrutinee, from, to)),
            arms: arms
                .iter()
                .map(|a| {
                    // Pattern-bound vars shadow `from` inside the arm.
                    let mut binds: BTreeSet<String> = BTreeSet::new();
                    pattern_binds(&a.pat, &mut binds);
                    let body = if binds.contains(from) {
                        a.body.clone()
                    } else {
                        subst_var(&a.body, from, to)
                    };
                    Arm { pat: a.pat.clone(), body }
                })
                .collect(),
        },
        Term::Lam { params, param_tys, ret_ty, effects, body } => {
            let body = if params.iter().any(|p| p == from) {
                (**body).clone()
            } else {
                subst_var(body, from, to)
            };
            Term::Lam {
                params: params.clone(),
                param_tys: param_tys.clone(),
                ret_ty: ret_ty.clone(),
                effects: effects.clone(),
                body: Box::new(body),
            }
        }
        Term::Seq { lhs, rhs } => Term::Seq {
            lhs: Box::new(subst_var(lhs, from, to)),
            rhs: Box::new(subst_var(rhs, from, to)),
        },
        Term::LetRec { name, ty, params, body, in_term } => {
            let body_shadowed = name == from || params.iter().any(|p| p == from);
            let body_rw = if body_shadowed {
                (**body).clone()
            } else {
                subst_var(body, from, to)
            };
            let in_shadowed = name == from;
            let in_rw = if in_shadowed {
                (**in_term).clone()
            } else {
                subst_var(in_term, from, to)
            };
            Term::LetRec {
                name: name.clone(),
                ty: ty.clone(),
                params: params.clone(),
                body: Box::new(body_rw),
                in_term: Box::new(in_rw),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: walk a term, return true iff any [`Pattern::Ctor`] in
    /// any reachable [`Term::Match`] has a non-flat field.
    fn any_nested_ctor(t: &Term) -> bool {
        match t {
            Term::Lit { .. } | Term::Var { .. } => false,
            Term::App { callee, args, .. } => {
                any_nested_ctor(callee) || args.iter().any(any_nested_ctor)
            }
            Term::Let { value, body, .. } => any_nested_ctor(value) || any_nested_ctor(body),
            Term::If { cond, then, else_ } => {
                any_nested_ctor(cond) || any_nested_ctor(then) || any_nested_ctor(else_)
            }
            Term::Do { args, .. } => args.iter().any(any_nested_ctor),
            Term::Ctor { args, .. } => args.iter().any(any_nested_ctor),
            Term::Match { scrutinee, arms } => {
                if any_nested_ctor(scrutinee) {
                    return true;
                }
                for a in arms {
                    if !is_flat(&a.pat) {
                        return true;
                    }
                    if any_nested_ctor(&a.body) {
                        return true;
                    }
                }
                false
            }
            Term::Lam { body, .. } => any_nested_ctor(body),
            Term::Seq { lhs, rhs } => any_nested_ctor(lhs) || any_nested_ctor(rhs),
            Term::LetRec { body, in_term, .. } => {
                any_nested_ctor(body) || any_nested_ctor(in_term)
            }
        }
    }

    /// Helper: walk a term, return true iff any [`Term::LetRec`] is
    /// reachable. After 16b.1 desugaring this should be `false`.
    fn any_let_rec(t: &Term) -> bool {
        match t {
            Term::Lit { .. } | Term::Var { .. } => false,
            Term::App { callee, args, .. } => {
                any_let_rec(callee) || args.iter().any(any_let_rec)
            }
            Term::Let { value, body, .. } => any_let_rec(value) || any_let_rec(body),
            Term::If { cond, then, else_ } => {
                any_let_rec(cond) || any_let_rec(then) || any_let_rec(else_)
            }
            Term::Do { args, .. } => args.iter().any(any_let_rec),
            Term::Ctor { args, .. } => args.iter().any(any_let_rec),
            Term::Match { scrutinee, arms } => {
                any_let_rec(scrutinee) || arms.iter().any(|a| any_let_rec(&a.body))
            }
            Term::Lam { body, .. } => any_let_rec(body),
            Term::Seq { lhs, rhs } => any_let_rec(lhs) || any_let_rec(rhs),
            Term::LetRec { .. } => true,
        }
    }

    /// Iter 16a: a `match` containing `(Cons a (Cons b _))` desugars to
    /// a tree with no nested ctor sub-patterns. Property: the rewriter's
    /// output never has a `Pattern::Ctor` whose field is itself a
    /// `Pattern::Ctor`.
    #[test]
    fn nested_cons_cons_flattens() {
        // Build: match xs of
        //   Cons a (Cons b _) -> ...
        //   _                 -> ...
        let body_match = Term::Match {
            scrutinee: Box::new(Term::Var { name: "xs".into() }),
            arms: vec![
                Arm {
                    pat: Pattern::Ctor {
                        ctor: "Cons".into(),
                        fields: vec![
                            Pattern::Var { name: "a".into() },
                            Pattern::Ctor {
                                ctor: "Cons".into(),
                                fields: vec![
                                    Pattern::Var { name: "b".into() },
                                    Pattern::Wild,
                                ],
                            },
                        ],
                    },
                    body: Term::App {
                        callee: Box::new(Term::Var { name: "+".into() }),
                        args: vec![
                            Term::Var { name: "a".into() },
                            Term::Var { name: "b".into() },
                        ],
                        tail: false,
                    },
                },
                Arm {
                    pat: Pattern::Wild,
                    body: Term::Lit { lit: Literal::Int { value: 0 } },
                },
            ],
        };
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "f".into(),
                ty: Type::Fn {
                    params: vec![Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec!["xs".into()],
                body: body_match,
                doc: None,
            })],
        };
        let out = desugar_module(&m);
        let body = match &out.defs[0] {
            Def::Fn(f) => &f.body,
            _ => unreachable!(),
        };
        assert!(
            !any_nested_ctor(body),
            "desugarer must remove every nested-ctor sub-pattern; got: {body:#?}"
        );
    }

    /// A flat-only match round-trips structurally: there should be no
    /// extra let-binding or chain wrapping.
    #[test]
    fn flat_match_is_unchanged_shape() {
        let original = Term::Match {
            scrutinee: Box::new(Term::Var { name: "xs".into() }),
            arms: vec![
                Arm {
                    pat: Pattern::Ctor {
                        ctor: "Nil".into(),
                        fields: vec![],
                    },
                    body: Term::Lit { lit: Literal::Int { value: 0 } },
                },
                Arm {
                    pat: Pattern::Ctor {
                        ctor: "Cons".into(),
                        fields: vec![
                            Pattern::Var { name: "h".into() },
                            Pattern::Var { name: "t".into() },
                        ],
                    },
                    body: Term::Var { name: "h".into() },
                },
            ],
        };
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "f".into(),
                ty: Type::Fn {
                    params: vec![Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec!["xs".into()],
                body: original.clone(),
                doc: None,
            })],
        };
        let out = desugar_module(&m);
        let body = match &out.defs[0] {
            Def::Fn(f) => f.body.clone(),
            _ => unreachable!(),
        };
        // Already-flat match must round-trip with the same Match shape
        // at the top level (no Let-wrap).
        assert!(matches!(body, Term::Match { .. }));
    }

    /// Iter 16b.1: a [`Term::LetRec`] whose body has no captures from
    /// the enclosing scope is lifted to a synthetic top-level fn. The
    /// resulting module's `defs` grows by one and the original LetRec
    /// is gone from the term tree.
    fn fact_letrec_term() -> Term {
        // (let-rec fact (params n) (type ...) (body ...) (in (app fact 5)))
        // Body just returns `n` so the test doesn't depend on `<=`/etc.
        Term::LetRec {
            name: "fact".into(),
            ty: Type::Fn {
                params: vec![Type::int()],
                ret: Box::new(Type::int()),
                effects: vec![],
            },
            params: vec!["n".into()],
            body: Box::new(Term::Var { name: "n".into() }),
            in_term: Box::new(Term::App {
                callee: Box::new(Term::Var { name: "fact".into() }),
                args: vec![Term::Lit { lit: Literal::Int { value: 5 } }],
                tail: false,
            }),
        }
    }

    #[test]
    fn let_rec_no_capture_lifts_to_top_level() {
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "main".into(),
                ty: Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec![],
                body: fact_letrec_term(),
                doc: None,
            })],
        };
        let out = desugar_module(&m);
        // One original + one lifted = two defs.
        assert_eq!(out.defs.len(), 2, "expected one lifted fn appended");
        // The new def is a fn whose name starts with the LetRec hint.
        let lifted = match &out.defs[1] {
            Def::Fn(f) => f,
            _ => panic!("expected a lifted FnDef"),
        };
        assert!(
            lifted.name.starts_with("fact$lr_"),
            "lifted name `{}` should start with `fact$lr_`",
            lifted.name
        );
        assert_eq!(lifted.params, vec!["n".to_string()]);
        // The original main's body must no longer contain a LetRec.
        let main_body = match &out.defs[0] {
            Def::Fn(f) => &f.body,
            _ => unreachable!(),
        };
        assert!(
            !any_let_rec(main_body),
            "desugarer must remove every Term::LetRec from the term tree; got: {main_body:#?}"
        );
        // The in-term of the LetRec was `(app fact 5)`; after the lift
        // it should be `(app <lifted_name> 5)`.
        match main_body {
            Term::App { callee, .. } => match callee.as_ref() {
                Term::Var { name } => assert_eq!(name, &lifted.name),
                other => panic!("expected lifted Var as callee, got {other:?}"),
            },
            other => panic!("expected App in main body, got {other:?}"),
        }
    }

    /// Iter 16b.1: a LetRec whose body references a var bound in the
    /// enclosing scope (here, `outer` is a parameter of the
    /// surrounding fn) panics at desugar time. 16b.2 will lift this.
    #[test]
    #[should_panic(expected = "would capture")]
    fn let_rec_with_capture_panics() {
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "main".into(),
                ty: Type::Fn {
                    params: vec![Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec!["outer".into()],
                // (let-rec helper (params x) ... (body (app + x outer)) ...)
                body: Term::LetRec {
                    name: "helper".into(),
                    ty: Type::Fn {
                        params: vec![Type::int()],
                        ret: Box::new(Type::int()),
                        effects: vec![],
                    },
                    params: vec!["x".into()],
                    body: Box::new(Term::App {
                        callee: Box::new(Term::Var { name: "+".into() }),
                        args: vec![
                            Term::Var { name: "x".into() },
                            Term::Var { name: "outer".into() },
                        ],
                        tail: false,
                    }),
                    in_term: Box::new(Term::App {
                        callee: Box::new(Term::Var { name: "helper".into() }),
                        args: vec![Term::Lit { lit: Literal::Int { value: 1 } }],
                        tail: false,
                    }),
                },
                doc: None,
            })],
        };
        let _ = desugar_module(&m);
    }

    /// Iter 16c: walks a term, returns true iff any [`Pattern::Lit`] is
    /// reachable in any [`Term::Match`] arm. After 16c desugaring, the
    /// invariant is `false` for every desugared output — `Pattern::Lit`
    /// is gone before typecheck/codegen runs.
    fn any_lit_pattern(t: &Term) -> bool {
        fn pat_has_lit(p: &Pattern) -> bool {
            match p {
                Pattern::Wild | Pattern::Var { .. } => false,
                Pattern::Lit { .. } => true,
                Pattern::Ctor { fields, .. } => fields.iter().any(pat_has_lit),
            }
        }
        match t {
            Term::Lit { .. } | Term::Var { .. } => false,
            Term::App { callee, args, .. } => {
                any_lit_pattern(callee) || args.iter().any(any_lit_pattern)
            }
            Term::Let { value, body, .. } => any_lit_pattern(value) || any_lit_pattern(body),
            Term::If { cond, then, else_ } => {
                any_lit_pattern(cond) || any_lit_pattern(then) || any_lit_pattern(else_)
            }
            Term::Do { args, .. } => args.iter().any(any_lit_pattern),
            Term::Ctor { args, .. } => args.iter().any(any_lit_pattern),
            Term::Match { scrutinee, arms } => {
                if any_lit_pattern(scrutinee) {
                    return true;
                }
                for a in arms {
                    if pat_has_lit(&a.pat) {
                        return true;
                    }
                    if any_lit_pattern(&a.body) {
                        return true;
                    }
                }
                false
            }
            Term::Lam { body, .. } => any_lit_pattern(body),
            Term::Seq { lhs, rhs } => any_lit_pattern(lhs) || any_lit_pattern(rhs),
            Term::LetRec { body, in_term, .. } => {
                any_lit_pattern(body) || any_lit_pattern(in_term)
            }
        }
    }

    /// Iter 16c: a top-level lit arm desugars to a `Term::If` whose
    /// condition is `(== s_var lit)`. The desugared body must contain
    /// no `Pattern::Lit` anywhere.
    #[test]
    fn top_level_lit_desugars_to_if() {
        // (match n
        //   (case (pat-lit 0) 100)
        //   (case _           999))
        let body_match = Term::Match {
            scrutinee: Box::new(Term::Var { name: "n".into() }),
            arms: vec![
                Arm {
                    pat: Pattern::Lit {
                        lit: Literal::Int { value: 0 },
                    },
                    body: Term::Lit {
                        lit: Literal::Int { value: 100 },
                    },
                },
                Arm {
                    pat: Pattern::Wild,
                    body: Term::Lit {
                        lit: Literal::Int { value: 999 },
                    },
                },
            ],
        };
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "f".into(),
                ty: Type::Fn {
                    params: vec![Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec!["n".into()],
                body: body_match,
                doc: None,
            })],
        };
        let out = desugar_module(&m);
        let body = match &out.defs[0] {
            Def::Fn(f) => &f.body,
            _ => unreachable!(),
        };
        assert!(
            !any_lit_pattern(body),
            "desugarer must remove every Pattern::Lit from the term tree; got: {body:#?}"
        );
        // The desugar wraps the match in `let $mp_N = scrutinee in <chain>`,
        // and the chain head must be a `Term::If` (the lit arm).
        let chain = match body {
            Term::Let { body, .. } => body.as_ref(),
            other => panic!("expected outer Let from chain machinery, got {other:?}"),
        };
        assert!(
            matches!(chain, Term::If { .. }),
            "expected Term::If at the chain head, got {chain:?}"
        );
    }

    /// Iter 16c: a `(pat-ctor Cons (pat-lit 0) _)` sub-pattern desugars
    /// to a tree with no `Pattern::Lit` anywhere — the lit is rewritten
    /// to a `Term::If` against the field-bound fresh variable.
    #[test]
    fn nested_lit_in_ctor_desugars_to_if() {
        // (match xs
        //   (case (pat-ctor Cons (pat-lit 0) _) 0)
        //   (case _ -1))
        let body_match = Term::Match {
            scrutinee: Box::new(Term::Var { name: "xs".into() }),
            arms: vec![
                Arm {
                    pat: Pattern::Ctor {
                        ctor: "Cons".into(),
                        fields: vec![
                            Pattern::Lit {
                                lit: Literal::Int { value: 0 },
                            },
                            Pattern::Wild,
                        ],
                    },
                    body: Term::Lit {
                        lit: Literal::Int { value: 0 },
                    },
                },
                Arm {
                    pat: Pattern::Wild,
                    body: Term::Lit {
                        lit: Literal::Int { value: -1 },
                    },
                },
            ],
        };
        let m = Module {
            schema: crate::SCHEMA.to_string(),
            name: "t".into(),
            imports: vec![],
            defs: vec![Def::Fn(FnDef {
                name: "f".into(),
                ty: Type::Fn {
                    params: vec![Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                params: vec!["xs".into()],
                body: body_match,
                doc: None,
            })],
        };
        let out = desugar_module(&m);
        let body = match &out.defs[0] {
            Def::Fn(f) => &f.body,
            _ => unreachable!(),
        };
        assert!(
            !any_lit_pattern(body),
            "desugarer must remove every Pattern::Lit from the term tree; got: {body:#?}"
        );
        assert!(
            !any_nested_ctor(body),
            "desugarer must also produce no nested ctor sub-patterns; got: {body:#?}"
        );
    }

    /// Iter 16c regression guard: `is_flat` must classify
    /// `Pattern::Lit` as **not** flat, so the chain machinery in
    /// [`Desugarer::desugar_match`] is always invoked for lit-arms
    /// (rather than the early-return that would leak the lit pattern
    /// through to typecheck/codegen).
    #[test]
    fn flat_arm_with_lit_is_no_longer_flat() {
        let p = Pattern::Lit {
            lit: Literal::Int { value: 0 },
        };
        assert!(
            !is_flat(&p),
            "Pattern::Lit must be classified as non-flat after 16c"
        );
    }
}