AILang/crates/ailang-codegen/src/escape.rs

//! Per-fn escape analysis for ADT and closure allocations (Iter 17a).
//!
//! Identifies `Term::Ctor` and `Term::Lam` allocations whose value
//! does not escape the function in which they are allocated. Such
//! allocations can be lowered to LLVM `alloca` instead of `@GC_malloc`,
//! producing stack-allocated boxes that are auto-freed at fn return —
//! semantically equivalent to the all-GC version but bypassing the
//! collector entirely for fully-local data.
//!
//! The analysis is *purely an optimisation*. A pessimistic answer
//! (claiming an allocation escapes when it does not) only loses
//! efficiency, never correctness.
//!
//! ## Rule (conservative, name-based taint propagation)
//!
//! For each allocation site that is the value of a `Term::Let`:
//!
//!   `let X = ALLOC in B`
//!
//! we ask whether the body `B` makes the value of `X` (or any value
//! derived from it via projection) reachable past the fn return.
//!
//! Taint propagation rule:
//! - The bound name `X` is **tainted** in `B`.
//! - A `Term::Match` whose scrutinee is a `Var` referring to a tainted
//!   name propagates taint to every pattern-bound name in every arm.
//!   (Pattern bindings hold field projections of the scrutinee, which
//!   live inside the same allocation.)
//! - A `Term::Let { name = Y, value = Var(t), ... }` where `t` is
//!   tainted makes `Y` tainted too in the let's body.
//!
//! Escape rule (any tainted name appearing in any of these positions
//! marks the allocation as escaping):
//! - **Tail position** of `B` (the value of `B` is the value of the
//!   enclosing region; the let-body's value flows up).
//! - **`Term::App` arg position** (arg flows into a callee, which may
//!   keep the value).
//! - **`Term::Do` arg position** (effect ops are runtime calls).
//! - **`Term::Ctor` field position** (storing into another box; the
//!   parent box may itself escape).
//! - **`Term::Lam` capture** (free var of an enclosed lambda — the
//!   lambda's env, possibly heap-allocated, holds the value past the
//!   current fn frame).
//! - **`Term::App` callee position** when the callee is **not** the
//!   simple Var-tainted call (we only special-case calling a
//!   tainted-Var-bound closure as non-escaping; any other indirect
//!   callee derived from the tainted value is escape).
//!
//! Allowed (non-escaping) positions:
//! - **Scrutinee of `Term::Match`** (scrutinee is read-then-projected;
//!   the box itself doesn't escape unless an arm's body escapes a
//!   pattern binding, handled by the taint propagation above).
//! - **Callee of `Term::App` when the callee is a bare `Var` to the
//!   tainted name** (the closure pair is loaded and called locally).
//!
//! ## Lambda-specific note
//!
//! Closures often escape (passed as args to higher-order fns). The
//! same rule applies. A `let f = \x. body in B` allocation is
//! non-escaping iff `f` appears in `B` exclusively in the callee
//! position of `Term::App` (so the closure pair is loaded and
//! invoked locally).
//!
//! The Lam's *body* is itself a fresh fn frame from the analyzer's
//! perspective — the Lam body's captures already cross a fn
//! boundary, so the analyzer does not recurse into it. (The Lam's
//! body would be analyzed in its own top-level pass over the lifted
//! thunk, but in this codebase lambdas are not lifted to top-level
//! before codegen — they are emitted as deferred thunks during the
//! same fn-emission pass. We handle this by running the analysis on
//! every fn body and on every Lam body recursively.)
//!
//! ## What this analysis is NOT
//!
//! - Not a region inference system. Per-fn only.
//! - Not field-sensitive (we taint pattern bindings wholesale).
//! - Not flow-sensitive within an arm (we just check if any tainted
//!   name appears in any escape position).
//!
//! Precision can be improved later; correctness is the priority for
//! this iter.

use ailang_core::ast::{Pattern, Term};
use std::collections::BTreeSet;

/// Set of allocation sites (identified by raw pointer address cast to
/// `usize`) that the escape analysis has proven do not escape their
/// allocating fn.
///
/// Pointer identity is stable for the duration of codegen because the
/// `Module` AST is borrowed immutably throughout.
pub type NonEscapeSet = BTreeSet<usize>;

/// Run the analysis over a fn body. Returns the set of allocation
/// sites (pointers to `Term::Ctor` / `Term::Lam` nodes) that may
/// be safely lowered with `alloca` instead of `@GC_malloc`.
pub fn analyze_fn_body(body: &Term) -> NonEscapeSet {
    let mut out = NonEscapeSet::new();
    walk(body, &mut out);
    out
}

/// Walk the term, looking for `Let { name, value: Ctor | Lam, body }`
/// shapes. For each such shape, run the escape check on `body` with
/// the bound name as the initial tainted seed. If the check returns
/// false (no escape), record the value's pointer in `out`.
///
/// Then recurse into all sub-terms so that nested let-allocations and
/// lambdas are analyzed too.
fn walk(t: &Term, out: &mut NonEscapeSet) {
    match t {
        Term::Let { name, value, body } => {
            // Candidate shape: let X = Ctor|Lam in B.
            let alloc_kind = match value.as_ref() {
                Term::Ctor { .. } => Some("ctor"),
                Term::Lam { .. } => Some("lam"),
                _ => None,
            };
            if alloc_kind.is_some() {
                let mut tainted = BTreeSet::new();
                tainted.insert(name.clone());
                // The let's body is in tail position iff the let is —
                // but we don't know the outer context here. Treat the
                // body as if it were tail-of-the-fn: safer (more
                // conservative). If body's tail value contains the
                // tainted name, it escapes the let, which transitively
                // escapes the fn.
                if !escapes(body, &tainted, true) {
                    let ptr = (value.as_ref() as *const Term) as usize;
                    out.insert(ptr);
                }
            }
            // Recurse into value and body for any nested allocations.
            walk(value, out);
            walk(body, out);
        }
        Term::App { callee, args, .. } => {
            walk(callee, out);
            for a in args {
                walk(a, out);
            }
        }
        Term::Do { args, .. } => {
            for a in args {
                walk(a, out);
            }
        }
        Term::Ctor { args, .. } => {
            for a in args {
                walk(a, out);
            }
        }
        Term::Match { scrutinee, arms } => {
            walk(scrutinee, out);
            for arm in arms {
                walk(&arm.body, out);
            }
        }
        Term::If { cond, then, else_ } => {
            walk(cond, out);
            walk(then, out);
            walk(else_, out);
        }
        Term::Seq { lhs, rhs } => {
            walk(lhs, out);
            walk(rhs, out);
        }
        Term::Lam { body, .. } => {
            // Recurse into the lambda body — it is itself a fn frame
            // for our analysis purposes (each lifted thunk runs the
            // same analysis when emit_fn is called for it; for
            // un-lifted closures we still want to flag inner local
            // allocations).
            walk(body, out);
        }
        Term::LetRec { body, in_term, .. } => {
            // Iter 16b.x: LetRec is normally eliminated before codegen.
            // Still walk for completeness.
            walk(body, out);
            walk(in_term, out);
        }
        Term::Clone { value } => {
            // Iter 18c.1: clone is identity at IR level — only walk
            // sub-terms looking for Let-allocation candidates.
            walk(value, out);
        }
        Term::ReuseAs { source, body } => {
            // Iter 18d.1: identity at IR level (codegen lowers `body`,
            // ignores `source`). Walk both children for completeness.
            walk(source, out);
            walk(body, out);
        }
        Term::Loop { binders, body } => {
            for b in binders {
                walk(&b.init, out);
            }
            walk(body, out);
        }
        Term::Recur { args } => {
            for a in args {
                walk(a, out);
            }
        }
        Term::Lit { .. } | Term::Var { .. } => {}
    }
}

/// Returns `true` if any tainted name escapes from the term `t`.
///
/// `in_tail` indicates whether `t`'s value flows up to the fn return
/// (or to the let-body whose value flows up, transitively). At the
/// initial call from `walk`, we set `in_tail = true` for the let's
/// body because the body's value is the let's value, which we cannot
/// see beyond — better to assume escape.
fn escapes(t: &Term, tainted: &BTreeSet<String>, in_tail: bool) -> bool {
    match t {
        Term::Lit { .. } => false,
        Term::Var { name } => {
            // A bare Var reference whose value flows to tail = escape.
            in_tail && tainted.contains(name)
        }
        Term::App { callee, args, .. } => {
            // Callee position. Special case: a direct call to a
            // tainted Var (calling the bound closure) is not escape;
            // the closure pair is loaded and invoked locally.
            match callee.as_ref() {
                Term::Var { name: cname } if tainted.contains(cname) => {
                    // Calling locally is fine; don't flag the callee as escape.
                }
                _ => {
                    if escapes(callee, tainted, false) {
                        return true;
                    }
                }
            }
            // Args: any tainted Var as arg = escape. Sub-expressions
            // recurse with in_tail=false.
            for a in args {
                if let Term::Var { name } = a {
                    if tainted.contains(name) {
                        return true;
                    }
                }
                if escapes(a, tainted, false) {
                    return true;
                }
            }
            false
        }
        Term::Do { args, .. } => {
            for a in args {
                if let Term::Var { name } = a {
                    if tainted.contains(name) {
                        return true;
                    }
                }
                if escapes(a, tainted, false) {
                    return true;
                }
            }
            false
        }
        Term::Ctor { args, .. } => {
            for a in args {
                if let Term::Var { name } = a {
                    if tainted.contains(name) {
                        return true;
                    }
                }
                if escapes(a, tainted, false) {
                    return true;
                }
            }
            false
        }
        Term::Let { name, value, body } => {
            if escapes(value, tainted, false) {
                return true;
            }
            // If the let's value is a tainted Var, the new binding
            // inherits taint. (Same for `Var`-only sub-expr; for any
            // other shape we'd need a stronger analysis to detect
            // taint flowing through computation. Conservative cut:
            // only direct Var aliasing propagates.)
            let mut local = tainted.clone();
            if let Term::Var { name: vn } = value.as_ref() {
                if tainted.contains(vn) {
                    local.insert(name.clone());
                }
            }
            // Body: tail position propagates from the enclosing let.
            escapes(body, &local, in_tail)
        }
        Term::Match { scrutinee, arms } => {
            if escapes(scrutinee, tainted, false) {
                return true;
            }
            // If the scrutinee is a tainted Var, propagate taint to
            // pattern bindings (each arm independently).
            let scrutinee_tainted = match scrutinee.as_ref() {
                Term::Var { name } => tainted.contains(name),
                _ => false,
            };
            for arm in arms {
                let mut arm_tainted = tainted.clone();
                if scrutinee_tainted {
                    let mut binds = Vec::new();
                    pattern_bound_names(&arm.pat, &mut binds);
                    for b in binds {
                        arm_tainted.insert(b);
                    }
                }
                if escapes(&arm.body, &arm_tainted, in_tail) {
                    return true;
                }
            }
            false
        }
        Term::If { cond, then, else_ } => {
            if escapes(cond, tainted, false) {
                return true;
            }
            if escapes(then, tainted, in_tail) {
                return true;
            }
            if escapes(else_, tainted, in_tail) {
                return true;
            }
            false
        }
        Term::Seq { lhs, rhs } => {
            if escapes(lhs, tainted, false) {
                return true;
            }
            escapes(rhs, tainted, in_tail)
        }
        Term::Lam { params, body, .. } => {
            // A lambda captures free vars. If any free var is tainted,
            // the closure escapes the value via its env. (Even if the
            // closure pair itself doesn't escape, capturing creates a
            // store of the tainted value into a heap env, which the
            // collector keeps. For correctness we just need: tainted
            // value's lifetime <= current fn frame.)
            //
            // The lambda's own body is analyzed independently when its
            // turn comes (via `walk`); here we only check captures.
            let mut frees = BTreeSet::new();
            let mut bound_in_lam: BTreeSet<String> = params.iter().cloned().collect();
            collect_free_vars(body, &mut bound_in_lam, &mut frees);
            for f in &frees {
                if tainted.contains(f) {
                    return true;
                }
            }
            false
        }
        Term::LetRec { .. } => {
            // LetRec is normally eliminated by the lift-pass before
            // codegen. If one slips through, be conservative and
            // declare escape.
            true
        }
        Term::Clone { value } => {
            // Iter 18c.1: clone is identity. The wrapper's escape
            // verdict is exactly that of its inner term, threaded
            // through with the same `in_tail` context.
            escapes(value, tainted, in_tail)
        }
        Term::ReuseAs { source, body } => {
            // Iter 18d.1: identity at IR level. The whole expression's
            // value is `body`, so its escape verdict is the body's
            // (threaded through `in_tail`). The `source` is evaluated
            // but its value is discarded; it cannot escape via the
            // tail position, but a tainted name reachable inside it
            // could still escape — walk it in non-tail mode.
            if escapes(source, tainted, false) {
                return true;
            }
            escapes(body, tainted, in_tail)
        }
        Term::Loop { binders, body } => {
            for b in binders {
                if escapes(&b.init, tainted, false) {
                    return true;
                }
            }
            escapes(body, tainted, in_tail)
        }
        Term::Recur { args } => {
            for a in args {
                if escapes(a, tainted, false) {
                    return true;
                }
            }
            false
        }
    }
}

/// Collect pattern-bound names (variables only; wildcards and
/// literals bind nothing).
fn pattern_bound_names(p: &Pattern, out: &mut Vec<String>) {
    match p {
        Pattern::Wild | Pattern::Lit { .. } => {}
        Pattern::Var { name } => out.push(name.clone()),
        Pattern::Ctor { fields, .. } => {
            for f in fields {
                pattern_bound_names(f, out);
            }
        }
    }
}

/// Collect free variables of a term relative to a `bound` set. Used
/// only to check whether a Lam captures any tainted name. Mirrors
/// `Emitter::collect_captures` in lib.rs but without the builtin /
/// top-level filter (those don't matter for taint).
fn collect_free_vars(t: &Term, bound: &mut 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, .. } => {
            collect_free_vars(callee, bound, out);
            for a in args {
                collect_free_vars(a, bound, out);
            }
        }
        Term::Do { args, .. } => {
            for a in args {
                collect_free_vars(a, bound, out);
            }
        }
        Term::Ctor { args, .. } => {
            for a in args {
                collect_free_vars(a, bound, out);
            }
        }
        Term::Let { name, value, body } => {
            collect_free_vars(value, bound, out);
            let inserted = bound.insert(name.clone());
            collect_free_vars(body, bound, out);
            if inserted {
                bound.remove(name);
            }
        }
        Term::If { cond, then, else_ } => {
            collect_free_vars(cond, bound, out);
            collect_free_vars(then, bound, out);
            collect_free_vars(else_, bound, out);
        }
        Term::Match { scrutinee, arms } => {
            collect_free_vars(scrutinee, bound, out);
            for arm in arms {
                let mut binds = Vec::new();
                pattern_bound_names(&arm.pat, &mut binds);
                let mut newly = Vec::new();
                for n in &binds {
                    if bound.insert(n.clone()) {
                        newly.push(n.clone());
                    }
                }
                collect_free_vars(&arm.body, bound, out);
                for n in newly {
                    bound.remove(&n);
                }
            }
        }
        Term::Lam { params, body, .. } => {
            let mut newly = Vec::new();
            for p in params {
                if bound.insert(p.clone()) {
                    newly.push(p.clone());
                }
            }
            collect_free_vars(body, bound, out);
            for n in newly {
                bound.remove(&n);
            }
        }
        Term::Seq { lhs, rhs } => {
            collect_free_vars(lhs, bound, out);
            collect_free_vars(rhs, bound, out);
        }
        Term::LetRec { body, in_term, .. } => {
            collect_free_vars(body, bound, out);
            collect_free_vars(in_term, bound, out);
        }
        Term::Clone { value } => {
            // Iter 18c.1: clone is identity for free-var collection.
            collect_free_vars(value, bound, out);
        }
        Term::ReuseAs { source, body } => {
            // Iter 18d.1: free vars are the union of source and body.
            collect_free_vars(source, bound, out);
            collect_free_vars(body, bound, out);
        }
        Term::Loop { binders, body } => {
            let mut newly: Vec<String> = Vec::new();
            for b in binders {
                collect_free_vars(&b.init, bound, out);
                if bound.insert(b.name.clone()) {
                    newly.push(b.name.clone());
                }
            }
            collect_free_vars(body, bound, out);
            for n in newly {
                bound.remove(&n);
            }
        }
        Term::Recur { args } => {
            for a in args {
                collect_free_vars(a, bound, out);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ailang_core::ast::{Arm, Literal, Type};

    fn lit_int(v: i64) -> Term {
        Term::Lit {
            lit: Literal::Int { value: v },
        }
    }
    fn var(s: &str) -> Term {
        Term::Var { name: s.into() }
    }
    fn ctor(ty: &str, c: &str, args: Vec<Term>) -> Term {
        Term::Ctor {
            type_name: ty.into(),
            ctor: c.into(),
            args,
        }
    }

    /// Match-only consumption: `let xs = Cons(1, Nil) in match xs of _ -> 0`.
    /// xs only used as match scrutinee; arm body returns a literal.
    /// Expected: alloca-eligible.
    #[test]
    fn local_ctor_match_only() {
        let alloc = ctor(
            "L",
            "Cons",
            vec![lit_int(1), ctor("L", "Nil", vec![])],
        );
        let alloc_ptr = (&alloc as *const Term) as usize;
        let body = Term::Let {
            name: "xs".into(),
            value: Box::new(alloc.clone()),
            body: Box::new(Term::Match {
                scrutinee: Box::new(var("xs")),
                arms: vec![Arm {
                    pat: Pattern::Wild,
                    body: lit_int(0),
                }],
            }),
        };
        let neset = analyze_fn_body(&body);
        // The allocation site we care about is the `value` field of
        // the Let. Since `body` was constructed inline (the Let's
        // value is a Box copy of `alloc`), use the boxed copy's addr.
        // Walk the AST to find it.
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            neset.contains(&let_value_ptr),
            "expected let-bound ctor to be non-escaping; neset={:?}",
            neset
        );
        // Avoid unused warning in case ptr identity changed.
        let _ = alloc_ptr;
    }

    /// Returned value: `let xs = Cons(1, Nil) in xs`. xs in tail
    /// position. Expected: NOT alloca-eligible.
    #[test]
    fn returned_ctor_escapes() {
        let alloc = ctor(
            "L",
            "Cons",
            vec![lit_int(1), ctor("L", "Nil", vec![])],
        );
        let body = Term::Let {
            name: "xs".into(),
            value: Box::new(alloc),
            body: Box::new(var("xs")),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            !neset.contains(&let_value_ptr),
            "let-bound ctor that flows to tail must escape"
        );
    }

    /// Passed to a fn: `let xs = Cons(1, Nil) in length(xs)`. xs in
    /// arg position. Expected: NOT alloca-eligible.
    #[test]
    fn ctor_passed_as_arg_escapes() {
        let alloc = ctor("L", "Nil", vec![]);
        let body = Term::Let {
            name: "xs".into(),
            value: Box::new(alloc),
            body: Box::new(Term::App {
                callee: Box::new(var("length")),
                args: vec![var("xs")],
                tail: false,
            }),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            !neset.contains(&let_value_ptr),
            "let-bound ctor passed as arg must escape"
        );
    }

    /// Stored into another ctor: `let h = Mk(1) in Cons(h, Nil)`.
    /// h is a field of the outer Cons. Expected: NOT alloca-eligible.
    #[test]
    fn ctor_stored_in_ctor_field_escapes() {
        let h_alloc = ctor("X", "Mk", vec![lit_int(1)]);
        let body = Term::Let {
            name: "h".into(),
            value: Box::new(h_alloc),
            body: Box::new(ctor(
                "L",
                "Cons",
                vec![var("h"), ctor("L", "Nil", vec![])],
            )),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            !neset.contains(&let_value_ptr),
            "let-bound ctor stored as ctor field must escape"
        );
    }

    /// Match arm returns a pattern-bound name: `let xs = Cons(1, Nil)
    /// in match xs of Cons h _ -> h | Nil -> 0`. h is tainted via the
    /// scrutinee, returned in arm tail. Expected: NOT alloca-eligible.
    #[test]
    fn pattern_binding_returned_escapes() {
        let alloc = ctor(
            "L",
            "Cons",
            vec![lit_int(1), ctor("L", "Nil", vec![])],
        );
        let body = Term::Let {
            name: "xs".into(),
            value: Box::new(alloc),
            body: Box::new(Term::Match {
                scrutinee: Box::new(var("xs")),
                arms: vec![
                    Arm {
                        pat: Pattern::Ctor {
                            ctor: "Cons".into(),
                            fields: vec![
                                Pattern::Var { name: "h".into() },
                                Pattern::Wild,
                            ],
                        },
                        body: var("h"),
                    },
                    Arm {
                        pat: Pattern::Ctor {
                            ctor: "Nil".into(),
                            fields: vec![],
                        },
                        body: lit_int(0),
                    },
                ],
            }),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            !neset.contains(&let_value_ptr),
            "scrutinee whose pattern binding flows to arm tail must escape"
        );
    }

    /// Lambda call-only: `let f = \x. x in f(42)`. f only used as
    /// callee. Expected: alloca-eligible.
    #[test]
    fn local_lam_call_only() {
        let lam = Term::Lam {
            params: vec!["x".into()],
            param_tys: vec![Type::int()],
            ret_ty: Box::new(Type::int()),
            effects: vec![],
            body: Box::new(var("x")),
        };
        let body = Term::Let {
            name: "f".into(),
            value: Box::new(lam),
            body: Box::new(Term::App {
                callee: Box::new(var("f")),
                args: vec![lit_int(42)],
                tail: false,
            }),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            neset.contains(&let_value_ptr),
            "let-bound lambda used only as callee must be non-escaping; got {:?}",
            neset
        );
    }

    /// Lambda passed to HOF: `let f = \x. x in map(f, xs)`. f passed
    /// as arg. Expected: NOT alloca-eligible.
    #[test]
    fn lam_passed_as_arg_escapes() {
        let lam = Term::Lam {
            params: vec!["x".into()],
            param_tys: vec![Type::int()],
            ret_ty: Box::new(Type::int()),
            effects: vec![],
            body: Box::new(var("x")),
        };
        let body = Term::Let {
            name: "f".into(),
            value: Box::new(lam),
            body: Box::new(Term::App {
                callee: Box::new(var("map")),
                args: vec![var("f"), var("xs")],
                tail: false,
            }),
        };
        let neset = analyze_fn_body(&body);
        let let_value_ptr = match &body {
            Term::Let { value, .. } => (value.as_ref() as *const Term) as usize,
            _ => unreachable!(),
        };
        assert!(
            !neset.contains(&let_value_ptr),
            "let-bound lambda passed as arg must escape"
        );
    }
}