Iter 17a: per-fn arena via alloca for non-escaping allocations
Adds a conservative escape analysis that flags Term::Ctor and
Term::Lam allocations as non-escaping when they are bound by a
let, never flow to a fn arg / ctor field / tail-return / lambda
capture, and remain inside the allocating fn's frame. Codegen
lowers flagged sites to LLVM `alloca` instead of `@GC_malloc`;
escaping sites continue to use Boehm.
- escape.rs (new): name-based taint propagation; let-only
candidates; tail-position / app-arg / ctor-field / lam-capture
all taint.
- codegen: three sites switched (lower_ctor, lam env, closure
pair). Save/restore non_escape across lambda thunk emit.
- examples/escape_local_demo.{ailx,ail.json}: focused fixture
exercising both branches (peek and recursive count).
- e2e + escape unit tests: 133 → 141 (+8).
Stop-gate: 17a closes the queue. Memory-management discussion
(GC scope) is the next user-driven step. Per-fn arena observations
appended to JOURNAL — including the headline finding that 0/270
shipped allocations are flagged non-escaping under this rule
(real AILang threads ctors directly between fns; build-locally /
consume-locally is not idiomatic).
All existing fixture stdouts and content hashes bit-identical;
IR snapshots unchanged.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,684 @@
|
||||
//! 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::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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -38,6 +38,9 @@ use ailang_core::ast::*;
|
||||
use ailang_core::Workspace;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
mod escape;
|
||||
use escape::NonEscapeSet;
|
||||
|
||||
/// Failure modes of [`emit_ir`] / [`lower_workspace`].
|
||||
///
|
||||
/// Most variants signal a compiler invariant violation rather than a
|
||||
@@ -455,6 +458,13 @@ struct Emitter<'a> {
|
||||
/// lowering. Flushed at the end of emit_fn (LLVM IR allows fns in
|
||||
/// any order).
|
||||
deferred_thunks: Vec<String>,
|
||||
/// Iter 17a: per-fn escape-analysis result. Set of pointer-as-usize
|
||||
/// addresses of `Term::Ctor` and `Term::Lam` nodes that the
|
||||
/// analysis proved do not escape the fn frame they are allocated
|
||||
/// in. Such allocations lower to `alloca` instead of `@GC_malloc`.
|
||||
/// Populated by `analyze_fn_body` at the start of `emit_fn` and at
|
||||
/// the start of every lambda thunk emission inside `lower_lambda`.
|
||||
non_escape: NonEscapeSet,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -557,6 +567,7 @@ impl<'a> Emitter<'a> {
|
||||
current_def: String::new(),
|
||||
lam_counter: 0,
|
||||
deferred_thunks: Vec::new(),
|
||||
non_escape: NonEscapeSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -738,6 +749,13 @@ impl<'a> Emitter<'a> {
|
||||
// collected during lowering and appended after the parent fn.
|
||||
self.current_def = f.name.clone();
|
||||
self.lam_counter = 0;
|
||||
// Iter 17a: run escape analysis over the fn body. The result
|
||||
// is queried at every `Term::Ctor` / `Term::Lam` lowering site
|
||||
// to decide between `alloca` (non-escaping) and `@GC_malloc`
|
||||
// (escaping). The analysis is purely additive — a stale or
|
||||
// empty result only loses optimisation opportunities, never
|
||||
// correctness.
|
||||
self.non_escape = escape::analyze_fn_body(&f.body);
|
||||
|
||||
let mut sig = format!(
|
||||
"define {ret} @ail_{module}_{name}(",
|
||||
@@ -1040,10 +1058,20 @@ impl<'a> Emitter<'a> {
|
||||
self.emit_indirect_call(&callee_ssa, &sig, args, *tail)
|
||||
}
|
||||
Term::Do { op, args, tail } => self.lower_effect_op(op, args, *tail),
|
||||
Term::Ctor { type_name, ctor, args } => self.lower_ctor(type_name, ctor, args),
|
||||
Term::Ctor { type_name, ctor, args } => {
|
||||
// Iter 17a: pass the term pointer so `lower_ctor` can
|
||||
// consult the escape-analysis result for this exact
|
||||
// allocation site.
|
||||
let term_ptr = (t as *const Term) as usize;
|
||||
self.lower_ctor(type_name, ctor, args, term_ptr)
|
||||
}
|
||||
Term::Match { scrutinee, arms } => self.lower_match(scrutinee, arms),
|
||||
Term::Lam { params, param_tys, ret_ty, effects: _, body } => {
|
||||
self.lower_lambda(params, param_tys, ret_ty, body)
|
||||
// Iter 17a: same as `Ctor` — pass the term pointer for
|
||||
// escape-analysis lookup. A non-escaping closure pair
|
||||
// (and its env) lower to `alloca`.
|
||||
let term_ptr = (t as *const Term) as usize;
|
||||
self.lower_lambda(params, param_tys, ret_ty, body, term_ptr)
|
||||
}
|
||||
Term::Seq { lhs, rhs } => {
|
||||
// Iter 10: lower lhs for its effects, discard the SSA;
|
||||
@@ -1185,11 +1213,19 @@ impl<'a> Emitter<'a> {
|
||||
/// Heap box layout: 8 bytes tag (i64) followed by 8 bytes per field.
|
||||
/// i1 and i8 fields also occupy a full 8-byte slot — the typed
|
||||
/// load/store instructions write/read only the required size.
|
||||
///
|
||||
/// Iter 17a: `term_ptr` is the pointer-as-usize of the lowered
|
||||
/// `Term::Ctor` node. If escape analysis flagged this site as
|
||||
/// non-escaping (i.e., the value cannot live past the current fn
|
||||
/// frame), allocation lowers to LLVM `alloca` instead of
|
||||
/// `@GC_malloc`. The rest of the lowering (tag store, field
|
||||
/// stores, ptr return) is identical.
|
||||
fn lower_ctor(
|
||||
&mut self,
|
||||
type_name: &str,
|
||||
ctor_name: &str,
|
||||
args: &[Term],
|
||||
term_ptr: usize,
|
||||
) -> Result<(String, String)> {
|
||||
let cref = self.lookup_ctor_by_type(type_name, ctor_name)?;
|
||||
if args.len() != cref.ail_fields.len() {
|
||||
@@ -1251,9 +1287,18 @@ impl<'a> Emitter<'a> {
|
||||
}
|
||||
let size_bytes = 8 + (compiled.len() * 8) as i64;
|
||||
let p = self.fresh_ssa();
|
||||
self.body.push_str(&format!(
|
||||
" {p} = call ptr @GC_malloc(i64 {size_bytes})\n"
|
||||
));
|
||||
// Iter 17a: pick allocator based on escape analysis. `alloca`
|
||||
// for non-escaping (stack-allocated, freed on fn return);
|
||||
// `@GC_malloc` for everything else.
|
||||
if self.non_escape.contains(&term_ptr) {
|
||||
self.body.push_str(&format!(
|
||||
" {p} = alloca i8, i64 {size_bytes}, align 8\n"
|
||||
));
|
||||
} else {
|
||||
self.body.push_str(&format!(
|
||||
" {p} = call ptr @GC_malloc(i64 {size_bytes})\n"
|
||||
));
|
||||
}
|
||||
// Write tag.
|
||||
self.body.push_str(&format!(
|
||||
" store i64 {tag}, ptr {p}, align 8\n",
|
||||
@@ -1867,6 +1912,7 @@ impl<'a> Emitter<'a> {
|
||||
lam_param_tys: &[Type],
|
||||
lam_ret_ty: &Type,
|
||||
lam_body: &Term,
|
||||
term_ptr: usize,
|
||||
) -> Result<(String, String)> {
|
||||
let llvm_param_tys: Vec<String> =
|
||||
lam_param_tys.iter().map(llvm_type).collect::<Result<_>>()?;
|
||||
@@ -1940,6 +1986,12 @@ impl<'a> Emitter<'a> {
|
||||
let saved_terminated = self.block_terminated;
|
||||
self.block_terminated = false;
|
||||
let saved_sigs = std::mem::take(&mut self.ssa_fn_sigs);
|
||||
// Iter 17a: a lambda thunk is its own fn frame for escape
|
||||
// analysis. Save the outer fn's non-escape set and recompute
|
||||
// for the lambda body. Restored at the end alongside the rest
|
||||
// of the saved emitter state.
|
||||
let saved_non_escape = std::mem::take(&mut self.non_escape);
|
||||
self.non_escape = escape::analyze_fn_body(lam_body);
|
||||
// Lambdas inside lambdas are fine: they get their own counter
|
||||
// namespace within the enclosing thunk. They share the
|
||||
// `deferred_thunks` queue (top-level body owns it).
|
||||
@@ -2020,15 +2072,30 @@ impl<'a> Emitter<'a> {
|
||||
self.current_block = saved_block;
|
||||
self.block_terminated = saved_terminated;
|
||||
self.ssa_fn_sigs = saved_sigs;
|
||||
self.non_escape = saved_non_escape;
|
||||
|
||||
// 3. Emit allocation + capture filling + closure-pair packing
|
||||
// in the OUTER body. Captures use 8 bytes each; closure-pair
|
||||
// is 16 bytes.
|
||||
// Iter 17a: env and closure-pair share the Lam term's escape
|
||||
// status — they have parallel lifetimes. If the closure pair
|
||||
// does not escape (only used as a callee in the outer body),
|
||||
// both can be `alloca`'d. The escape-analysis lookup runs
|
||||
// against the outer fn's `non_escape` set (already restored
|
||||
// above).
|
||||
let lam_local = self.non_escape.contains(&term_ptr);
|
||||
let env_size = (cap_meta.len() * 8) as i64;
|
||||
let env_ssa = if env_size > 0 {
|
||||
let env = self.fresh_ssa();
|
||||
self.body
|
||||
.push_str(&format!(" {env} = call ptr @GC_malloc(i64 {env_size})\n"));
|
||||
if lam_local {
|
||||
self.body.push_str(&format!(
|
||||
" {env} = alloca i8, i64 {env_size}, align 8\n"
|
||||
));
|
||||
} else {
|
||||
self.body.push_str(&format!(
|
||||
" {env} = call ptr @GC_malloc(i64 {env_size})\n"
|
||||
));
|
||||
}
|
||||
for (i, (_cname, outer_ssa, cty, _c_ail, _sig)) in cap_meta.iter().enumerate() {
|
||||
let offset = (i * 8) as i64;
|
||||
let slot = self.fresh_ssa();
|
||||
@@ -2045,8 +2112,13 @@ impl<'a> Emitter<'a> {
|
||||
};
|
||||
|
||||
let clos = self.fresh_ssa();
|
||||
self.body
|
||||
.push_str(&format!(" {clos} = call ptr @GC_malloc(i64 16)\n"));
|
||||
if lam_local {
|
||||
self.body
|
||||
.push_str(&format!(" {clos} = alloca i8, i64 16, align 8\n"));
|
||||
} else {
|
||||
self.body
|
||||
.push_str(&format!(" {clos} = call ptr @GC_malloc(i64 16)\n"));
|
||||
}
|
||||
let cs_t = self.fresh_ssa();
|
||||
self.body.push_str(&format!(
|
||||
" {cs_t} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 0\n"
|
||||
|
||||
Reference in New Issue
Block a user