diff --git a/crates/ailang-codegen/src/lambda.rs b/crates/ailang-codegen/src/lambda.rs new file mode 100644 index 0000000..02b48d7 --- /dev/null +++ b/crates/ailang-codegen/src/lambda.rs @@ -0,0 +1,447 @@ +//! Lowering of `Term::Lam` and the surrounding closure machinery. +//! +//! `lower_lambda` walks the body to find free vars (captures), lifts +//! the body into a top-level thunk fn `@ail___lam`, +//! allocates a heap env for the captures, and returns a closure +//! pair `{ thunk_ptr, env_ptr }`. Under `--alloc=rc` it also defers +//! per-pair drop-fn IR via `build_env_drop_fn` / +//! `build_pair_drop_fn` (in `drop.rs`) and registers the pair-drop +//! symbol in `closure_drops` so the surrounding `Term::Let` close +//! can call it. +//! +//! The two helpers `collect_captures` and `pattern_bound_names` are +//! free-variable analyses needed by the env-construction step. They +//! stay private to this module — `collect_captures` is recursive +//! and only ever entered through `lower_lambda`'s prelude; +//! `pattern_bound_names` only feeds the match-arm branch of +//! `collect_captures`. +//! +//! Cross-references back into the parent module: +//! - `start_block`, `lower_term`, `fresh_ssa`, +//! `build_env_drop_fn`, `build_pair_drop_fn`: lifecycle and +//! drop-emission primitives, all `pub(crate)`. +//! - `synth::llvm_type`, `synth::fn_sig_from_type`: free-function +//! helpers for IR shaping. + +use ailang_core::ast::*; +use std::collections::BTreeSet; + +use super::escape; +use super::synth::{fn_sig_from_type, llvm_type}; +use super::{AllocStrategy, CodegenError, Emitter, FnSig, Result}; + +impl<'a> Emitter<'a> { + /// Iter 8b: lower a `Term::Lam`. Walks the body to find captures + /// (free vars relative to the enclosing scope, minus builtins and + /// top-level fns), allocates a heap env for the captures, lifts the + /// body to a top-level thunk fn `@ail___lam`, and + /// returns a closure-pair pointer that pairs the thunk with the env. + pub(crate) fn lower_lambda( + &mut self, + lam_params: &[String], + lam_param_tys: &[Type], + lam_ret_ty: &Type, + lam_body: &Term, + term_ptr: usize, + ) -> Result<(String, String)> { + let llvm_param_tys: Vec = + lam_param_tys.iter().map(llvm_type).collect::>()?; + let llvm_ret = llvm_type(lam_ret_ty)?; + + // 1. Capture analysis. The "bound" set seeds with everything + // that's a local in the enclosing scope OR a lambda param — + // but params are bound only INSIDE the body, not before. So + // pass them through as part of the recursion. + let top_level: BTreeSet = self + .module_user_fns + .get(self.module_name) + .map(|m| m.keys().cloned().collect()) + .unwrap_or_default(); + let builtins_owned: Vec = ailang_check::builtins::value_names() + .into_iter() + .map(|s| s.to_string()) + .collect(); + let builtins: BTreeSet<&str> = + builtins_owned.iter().map(|s| s.as_str()).collect(); + + let mut bound: BTreeSet = BTreeSet::new(); + for p in lam_params { + bound.insert(p.clone()); + } + let mut captures: Vec = Vec::new(); + let mut captures_set: BTreeSet = BTreeSet::new(); + Self::collect_captures( + lam_body, + &mut bound, + &mut captures, + &mut captures_set, + &builtins, + &top_level, + ); + + // Outer scope provides every capture as a local — assert and + // pull SSA + LLVM type from `self.locals`. Unknown captures + // would mean the typechecker let through an unbound var, so a + // hard internal error is right. + // Tuple: (name, outer_ssa, llvm_type, ail_type, optional_fn_sig). + let mut cap_meta: Vec<(String, String, String, Type, Option)> = Vec::new(); + for c in &captures { + let (_, outer_ssa, lty, ail_ty) = self + .locals + .iter() + .rev() + .find(|(n, _, _, _)| n == c) + .ok_or_else(|| { + CodegenError::Internal(format!( + "lambda capture `{c}` not in scope (typechecker bug?)" + )) + })? + .clone(); + let sig = self.ssa_fn_sigs.get(&outer_ssa).cloned(); + cap_meta.push((c.clone(), outer_ssa, lty, ail_ty, sig)); + } + + // 2. Pick a thunk name and switch the emitter into "thunk + // emission" mode by saving and resetting per-fn state. + let lam_id = self.lam_counter; + self.lam_counter += 1; + let parent = self.current_def.clone(); + let thunk_name = format!("{parent}_lam{lam_id}"); + let thunk_symbol = format!("@ail_{m}_{thunk_name}", m = self.module_name); + + let saved_body = std::mem::take(&mut self.body); + let saved_locals = std::mem::take(&mut self.locals); + let saved_counter = self.counter; + let saved_block = std::mem::take(&mut self.current_block); + 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); + // Iter 18d.3: a lambda thunk is its own fn frame for move + // tracking — the outer fn's binders are not in scope inside + // the thunk body (only its captures + params). Save and reset. + let saved_moved = std::mem::take(&mut self.moved_slots); + // Iter 18d.4 fix: a lambda thunk is its own fn frame for + // param-mode lookup. The outer fn's params are not in scope + // inside the thunk; the thunk's own params are pushed below + // and (currently) carry no mode annotation, so they default + // to `Implicit` — `lower_match`'s Iter A gate will skip arm- + // close pattern-binder dec for matches on lambda params, + // mirroring the fn-level Implicit-param treatment. + let saved_param_modes = std::mem::take(&mut self.current_param_modes); + for pname in lam_params.iter() { + self.current_param_modes + .insert(pname.clone(), ParamMode::Implicit); + } + // 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). + self.counter = 0; + + // Thunk header: `(ptr %env, params...)`. + let mut thunk_sig = format!("define {ret} {thunk_symbol}(ptr %env", ret = llvm_ret); + for (pname, pty) in lam_params.iter().zip(llvm_param_tys.iter()) { + thunk_sig.push_str(&format!(", {pty} %arg_{pname}")); + } + thunk_sig.push_str(") {\n"); + self.body.push_str(&thunk_sig); + self.start_block("entry"); + + // Unpack captures back into named locals at the start of the + // body. Layout: 8 bytes per slot, regardless of LLVM type + // (typed load reads only the needed bytes; padding is wasted + // but uniform). + for (i, (cname, _outer, cty, c_ail, sig_opt)) in cap_meta.iter().enumerate() { + let offset = (i * 8) as i64; + let slot = self.fresh_ssa(); + let val = self.fresh_ssa(); + self.body.push_str(&format!( + " {slot} = getelementptr inbounds i8, ptr %env, i64 {offset}\n" + )); + self.body + .push_str(&format!(" {val} = load {cty}, ptr {slot}\n")); + self.locals.push(( + cname.clone(), + val.clone(), + cty.clone(), + c_ail.clone(), + )); + if let Some(sig) = sig_opt { + self.ssa_fn_sigs.insert(val, sig.clone()); + } + } + + // Push lambda params as locals (after captures so shadowing + // honors lexical order — params win). + for ((pname, pty), pty_ail) in lam_params + .iter() + .zip(llvm_param_tys.iter()) + .zip(lam_param_tys.iter()) + { + let pssa = format!("%arg_{pname}"); + self.locals.push(( + pname.clone(), + pssa.clone(), + pty.clone(), + pty_ail.clone(), + )); + if let Some(fs) = fn_sig_from_type(pty_ail) { + self.ssa_fn_sigs.insert(pssa, fs); + } + } + + let (body_v, body_ty) = self.lower_term(lam_body)?; + if !self.block_terminated { + if body_ty != llvm_ret { + return Err(CodegenError::Internal(format!( + "lambda `{thunk_name}`: body type {body_ty} != return type {llvm_ret}" + ))); + } + self.body + .push_str(&format!(" ret {body_ty} {body_v}\n}}\n\n")); + } else { + self.body.push_str("}\n\n"); + } + + // Park the thunk text in the deferred queue and restore outer + // emitter state. + let thunk_text = std::mem::take(&mut self.body); + self.deferred_thunks.push(thunk_text); + self.body = saved_body; + self.locals = saved_locals; + self.counter = saved_counter; + self.current_block = saved_block; + self.block_terminated = saved_terminated; + self.ssa_fn_sigs = saved_sigs; + self.non_escape = saved_non_escape; + self.moved_slots = saved_moved; + self.current_param_modes = saved_param_modes; + + // 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(); + 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 @{}(i64 {env_size})\n", + self.alloc.fn_name() + )); + } + 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(); + self.body.push_str(&format!( + " {slot} = getelementptr inbounds i8, ptr {env}, i64 {offset}\n" + )); + self.body.push_str(&format!( + " store {cty} {outer_ssa}, ptr {slot}\n" + )); + } + env + } else { + "null".into() + }; + + let clos = self.fresh_ssa(); + 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 @{}(i64 16)\n", + self.alloc.fn_name() + )); + } + let cs_t = self.fresh_ssa(); + self.body.push_str(&format!( + " {cs_t} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 0\n" + )); + self.body + .push_str(&format!(" store ptr {thunk_symbol}, ptr {cs_t}\n")); + let cs_e = self.fresh_ssa(); + self.body.push_str(&format!( + " {cs_e} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 1\n" + )); + self.body + .push_str(&format!(" store ptr {env_ssa}, ptr {cs_e}\n")); + + // Register sig so subsequent indirect calls on `clos` work. + let fs = FnSig { + params: llvm_param_tys, + ret: llvm_ret, + }; + self.ssa_fn_sigs.insert(clos.clone(), fs); + + // Iter 18c.4: emit per-pair drop fns for heap-allocated + // closures under `--alloc=rc`. `lam_local` closures live on + // the stack and need no drop fn — LLVM `alloca` reclaims the + // pair on fn return. For heap closures we emit two symbols: + // + // - `drop__env(env)` — dec each pointer-typed + // capture, then dec the env block itself. ADT captures + // cascade through the per-type drop fn; closure-typed + // captures fall back to plain `ailang_rc_dec` because we + // don't keep a per-pair drop pointer alongside the env + // cell (yet). Iter 18d/18e revisit this. + // + // - `drop__pair(p)` — load the env from offset 8, + // call `drop__env(env)`, then dec the pair box. + // + // The pair drop is the symbol `Term::Let` lowering calls + // when the binder is a closure; we record it in + // `closure_drops` keyed by the closure-pair SSA. + if matches!(self.alloc, AllocStrategy::Rc) && !lam_local { + let pair_drop = format!("drop_{m}_{thunk_name}_pair", m = self.module_name); + let env_drop = format!("drop_{m}_{thunk_name}_env", m = self.module_name); + let env_text = self.build_env_drop_fn(&env_drop, &cap_meta); + self.deferred_thunks.push(env_text); + let pair_text = self.build_pair_drop_fn(&pair_drop, &env_drop, !cap_meta.is_empty()); + self.deferred_thunks.push(pair_text); + self.closure_drops.insert(clos.clone(), pair_drop); + } + + Ok((clos, "ptr".into())) + } + + /// Iter 8b: walk a Term collecting free-variable names that should + /// be captured by an enclosing lambda. Builtins and top-level fns + /// are excluded — they're globally accessible. Qualified names + /// (`prefix.def`) are excluded too. + fn collect_captures( + t: &Term, + bound: &mut BTreeSet, + captures: &mut Vec, + captures_set: &mut BTreeSet, + builtins: &BTreeSet<&str>, + top_level: &BTreeSet, + ) { + match t { + Term::Lit { .. } => {} + Term::Var { name } => { + if name.contains('.') { + return; // qualified ref is module-level + } + if bound.contains(name) { + return; + } + if builtins.contains(name.as_str()) { + return; + } + if top_level.contains(name) { + return; + } + if captures_set.insert(name.clone()) { + captures.push(name.clone()); + } + } + Term::App { callee, args, .. } => { + Self::collect_captures(callee, bound, captures, captures_set, builtins, top_level); + for a in args { + Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); + } + } + Term::Let { name, value, body } => { + Self::collect_captures(value, bound, captures, captures_set, builtins, top_level); + let inserted = bound.insert(name.clone()); + Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); + if inserted { + bound.remove(name); + } + } + Term::If { cond, then, else_ } => { + Self::collect_captures(cond, bound, captures, captures_set, builtins, top_level); + Self::collect_captures(then, bound, captures, captures_set, builtins, top_level); + Self::collect_captures(else_, bound, captures, captures_set, builtins, top_level); + } + Term::Do { args, .. } => { + for a in args { + Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); + } + } + Term::Ctor { args, .. } => { + for a in args { + Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); + } + } + Term::Match { scrutinee, arms } => { + Self::collect_captures(scrutinee, bound, captures, captures_set, builtins, top_level); + for arm in arms { + let mut pat_bindings = Vec::new(); + Self::pattern_bound_names(&arm.pat, &mut pat_bindings); + let mut newly_bound = Vec::new(); + for n in &pat_bindings { + if bound.insert(n.clone()) { + newly_bound.push(n.clone()); + } + } + Self::collect_captures(&arm.body, bound, captures, captures_set, builtins, top_level); + for n in newly_bound { + bound.remove(&n); + } + } + } + Term::Lam { params, body, .. } => { + // Inner lambda's params don't escape; its free vars + // (relative to the outer) DO contribute to outer captures. + let mut newly_bound = Vec::new(); + for p in params { + if bound.insert(p.clone()) { + newly_bound.push(p.clone()); + } + } + Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); + for n in newly_bound { + bound.remove(&n); + } + } + Term::Seq { lhs, rhs } => { + Self::collect_captures(lhs, bound, captures, captures_set, builtins, top_level); + Self::collect_captures(rhs, bound, captures, captures_set, builtins, top_level); + } + Term::LetRec { .. } => { + // Iter 16b.1: eliminated by desugar before codegen. + unreachable!("Term::LetRec eliminated by desugar") + } + Term::Clone { value } => { + // Iter 18c.1: clone is identity for capture analysis — + // free vars of `(clone X)` are exactly the free vars of `X`. + Self::collect_captures(value, bound, captures, captures_set, builtins, top_level); + } + Term::ReuseAs { source, body } => { + // Iter 18d.1: free vars are the union of source and body. + Self::collect_captures(source, bound, captures, captures_set, builtins, top_level); + Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); + } + } + } + + fn pattern_bound_names(p: &Pattern, out: &mut Vec) { + match p { + Pattern::Wild | Pattern::Lit { .. } => {} + Pattern::Var { name } => out.push(name.clone()), + Pattern::Ctor { fields, .. } => { + for f in fields { + Self::pattern_bound_names(f, out); + } + } + } + } +} diff --git a/crates/ailang-codegen/src/lib.rs b/crates/ailang-codegen/src/lib.rs index 540ccbe..9368b7f 100644 --- a/crates/ailang-codegen/src/lib.rs +++ b/crates/ailang-codegen/src/lib.rs @@ -42,6 +42,7 @@ use ailang_check::uniqueness::{infer_module, UniquenessTable}; mod drop; mod escape; +mod lambda; mod match_lower; mod subst; mod synth; @@ -2018,419 +2019,6 @@ impl<'a> Emitter<'a> { Ok((dst, sig.ret.clone())) } - /// Iter 8b: lower a `Term::Lam`. Walks the body to find captures - /// (free vars relative to the enclosing scope, minus builtins and - /// top-level fns), allocates a heap env for the captures, lifts the - /// body to a top-level thunk fn `@ail___lam`, and - /// returns a closure-pair pointer that pairs the thunk with the env. - fn lower_lambda( - &mut self, - lam_params: &[String], - lam_param_tys: &[Type], - lam_ret_ty: &Type, - lam_body: &Term, - term_ptr: usize, - ) -> Result<(String, String)> { - let llvm_param_tys: Vec = - lam_param_tys.iter().map(llvm_type).collect::>()?; - let llvm_ret = llvm_type(lam_ret_ty)?; - - // 1. Capture analysis. The "bound" set seeds with everything - // that's a local in the enclosing scope OR a lambda param — - // but params are bound only INSIDE the body, not before. So - // pass them through as part of the recursion. - let top_level: BTreeSet = self - .module_user_fns - .get(self.module_name) - .map(|m| m.keys().cloned().collect()) - .unwrap_or_default(); - let builtins_owned: Vec = ailang_check::builtins::value_names() - .into_iter() - .map(|s| s.to_string()) - .collect(); - let builtins: BTreeSet<&str> = - builtins_owned.iter().map(|s| s.as_str()).collect(); - - let mut bound: BTreeSet = BTreeSet::new(); - for p in lam_params { - bound.insert(p.clone()); - } - let mut captures: Vec = Vec::new(); - let mut captures_set: BTreeSet = BTreeSet::new(); - Self::collect_captures( - lam_body, - &mut bound, - &mut captures, - &mut captures_set, - &builtins, - &top_level, - ); - - // Outer scope provides every capture as a local — assert and - // pull SSA + LLVM type from `self.locals`. Unknown captures - // would mean the typechecker let through an unbound var, so a - // hard internal error is right. - // Tuple: (name, outer_ssa, llvm_type, ail_type, optional_fn_sig). - let mut cap_meta: Vec<(String, String, String, Type, Option)> = Vec::new(); - for c in &captures { - let (_, outer_ssa, lty, ail_ty) = self - .locals - .iter() - .rev() - .find(|(n, _, _, _)| n == c) - .ok_or_else(|| { - CodegenError::Internal(format!( - "lambda capture `{c}` not in scope (typechecker bug?)" - )) - })? - .clone(); - let sig = self.ssa_fn_sigs.get(&outer_ssa).cloned(); - cap_meta.push((c.clone(), outer_ssa, lty, ail_ty, sig)); - } - - // 2. Pick a thunk name and switch the emitter into "thunk - // emission" mode by saving and resetting per-fn state. - let lam_id = self.lam_counter; - self.lam_counter += 1; - let parent = self.current_def.clone(); - let thunk_name = format!("{parent}_lam{lam_id}"); - let thunk_symbol = format!("@ail_{m}_{thunk_name}", m = self.module_name); - - let saved_body = std::mem::take(&mut self.body); - let saved_locals = std::mem::take(&mut self.locals); - let saved_counter = self.counter; - let saved_block = std::mem::take(&mut self.current_block); - 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); - // Iter 18d.3: a lambda thunk is its own fn frame for move - // tracking — the outer fn's binders are not in scope inside - // the thunk body (only its captures + params). Save and reset. - let saved_moved = std::mem::take(&mut self.moved_slots); - // Iter 18d.4 fix: a lambda thunk is its own fn frame for - // param-mode lookup. The outer fn's params are not in scope - // inside the thunk; the thunk's own params are pushed below - // and (currently) carry no mode annotation, so they default - // to `Implicit` — `lower_match`'s Iter A gate will skip arm- - // close pattern-binder dec for matches on lambda params, - // mirroring the fn-level Implicit-param treatment. - let saved_param_modes = std::mem::take(&mut self.current_param_modes); - for pname in lam_params.iter() { - self.current_param_modes - .insert(pname.clone(), ParamMode::Implicit); - } - // 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). - self.counter = 0; - - // Thunk header: `(ptr %env, params...)`. - let mut thunk_sig = format!("define {ret} {thunk_symbol}(ptr %env", ret = llvm_ret); - for (pname, pty) in lam_params.iter().zip(llvm_param_tys.iter()) { - thunk_sig.push_str(&format!(", {pty} %arg_{pname}")); - } - thunk_sig.push_str(") {\n"); - self.body.push_str(&thunk_sig); - self.start_block("entry"); - - // Unpack captures back into named locals at the start of the - // body. Layout: 8 bytes per slot, regardless of LLVM type - // (typed load reads only the needed bytes; padding is wasted - // but uniform). - for (i, (cname, _outer, cty, c_ail, sig_opt)) in cap_meta.iter().enumerate() { - let offset = (i * 8) as i64; - let slot = self.fresh_ssa(); - let val = self.fresh_ssa(); - self.body.push_str(&format!( - " {slot} = getelementptr inbounds i8, ptr %env, i64 {offset}\n" - )); - self.body - .push_str(&format!(" {val} = load {cty}, ptr {slot}\n")); - self.locals.push(( - cname.clone(), - val.clone(), - cty.clone(), - c_ail.clone(), - )); - if let Some(sig) = sig_opt { - self.ssa_fn_sigs.insert(val, sig.clone()); - } - } - - // Push lambda params as locals (after captures so shadowing - // honors lexical order — params win). - for ((pname, pty), pty_ail) in lam_params - .iter() - .zip(llvm_param_tys.iter()) - .zip(lam_param_tys.iter()) - { - let pssa = format!("%arg_{pname}"); - self.locals.push(( - pname.clone(), - pssa.clone(), - pty.clone(), - pty_ail.clone(), - )); - if let Some(fs) = fn_sig_from_type(pty_ail) { - self.ssa_fn_sigs.insert(pssa, fs); - } - } - - let (body_v, body_ty) = self.lower_term(lam_body)?; - if !self.block_terminated { - if body_ty != llvm_ret { - return Err(CodegenError::Internal(format!( - "lambda `{thunk_name}`: body type {body_ty} != return type {llvm_ret}" - ))); - } - self.body - .push_str(&format!(" ret {body_ty} {body_v}\n}}\n\n")); - } else { - self.body.push_str("}\n\n"); - } - - // Park the thunk text in the deferred queue and restore outer - // emitter state. - let thunk_text = std::mem::take(&mut self.body); - self.deferred_thunks.push(thunk_text); - self.body = saved_body; - self.locals = saved_locals; - self.counter = saved_counter; - self.current_block = saved_block; - self.block_terminated = saved_terminated; - self.ssa_fn_sigs = saved_sigs; - self.non_escape = saved_non_escape; - self.moved_slots = saved_moved; - self.current_param_modes = saved_param_modes; - - // 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(); - 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 @{}(i64 {env_size})\n", - self.alloc.fn_name() - )); - } - 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(); - self.body.push_str(&format!( - " {slot} = getelementptr inbounds i8, ptr {env}, i64 {offset}\n" - )); - self.body.push_str(&format!( - " store {cty} {outer_ssa}, ptr {slot}\n" - )); - } - env - } else { - "null".into() - }; - - let clos = self.fresh_ssa(); - 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 @{}(i64 16)\n", - self.alloc.fn_name() - )); - } - let cs_t = self.fresh_ssa(); - self.body.push_str(&format!( - " {cs_t} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 0\n" - )); - self.body - .push_str(&format!(" store ptr {thunk_symbol}, ptr {cs_t}\n")); - let cs_e = self.fresh_ssa(); - self.body.push_str(&format!( - " {cs_e} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 1\n" - )); - self.body - .push_str(&format!(" store ptr {env_ssa}, ptr {cs_e}\n")); - - // Register sig so subsequent indirect calls on `clos` work. - let fs = FnSig { - params: llvm_param_tys, - ret: llvm_ret, - }; - self.ssa_fn_sigs.insert(clos.clone(), fs); - - // Iter 18c.4: emit per-pair drop fns for heap-allocated - // closures under `--alloc=rc`. `lam_local` closures live on - // the stack and need no drop fn — LLVM `alloca` reclaims the - // pair on fn return. For heap closures we emit two symbols: - // - // - `drop__env(env)` — dec each pointer-typed - // capture, then dec the env block itself. ADT captures - // cascade through the per-type drop fn; closure-typed - // captures fall back to plain `ailang_rc_dec` because we - // don't keep a per-pair drop pointer alongside the env - // cell (yet). Iter 18d/18e revisit this. - // - // - `drop__pair(p)` — load the env from offset 8, - // call `drop__env(env)`, then dec the pair box. - // - // The pair drop is the symbol `Term::Let` lowering calls - // when the binder is a closure; we record it in - // `closure_drops` keyed by the closure-pair SSA. - if matches!(self.alloc, AllocStrategy::Rc) && !lam_local { - let pair_drop = format!("drop_{m}_{thunk_name}_pair", m = self.module_name); - let env_drop = format!("drop_{m}_{thunk_name}_env", m = self.module_name); - let env_text = self.build_env_drop_fn(&env_drop, &cap_meta); - self.deferred_thunks.push(env_text); - let pair_text = self.build_pair_drop_fn(&pair_drop, &env_drop, !cap_meta.is_empty()); - self.deferred_thunks.push(pair_text); - self.closure_drops.insert(clos.clone(), pair_drop); - } - - Ok((clos, "ptr".into())) - } - - /// Iter 8b: walk a Term collecting free-variable names that should - /// be captured by an enclosing lambda. Builtins and top-level fns - /// are excluded — they're globally accessible. Qualified names - /// (`prefix.def`) are excluded too. - fn collect_captures( - t: &Term, - bound: &mut BTreeSet, - captures: &mut Vec, - captures_set: &mut BTreeSet, - builtins: &BTreeSet<&str>, - top_level: &BTreeSet, - ) { - match t { - Term::Lit { .. } => {} - Term::Var { name } => { - if name.contains('.') { - return; // qualified ref is module-level - } - if bound.contains(name) { - return; - } - if builtins.contains(name.as_str()) { - return; - } - if top_level.contains(name) { - return; - } - if captures_set.insert(name.clone()) { - captures.push(name.clone()); - } - } - Term::App { callee, args, .. } => { - Self::collect_captures(callee, bound, captures, captures_set, builtins, top_level); - for a in args { - Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); - } - } - Term::Let { name, value, body } => { - Self::collect_captures(value, bound, captures, captures_set, builtins, top_level); - let inserted = bound.insert(name.clone()); - Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); - if inserted { - bound.remove(name); - } - } - Term::If { cond, then, else_ } => { - Self::collect_captures(cond, bound, captures, captures_set, builtins, top_level); - Self::collect_captures(then, bound, captures, captures_set, builtins, top_level); - Self::collect_captures(else_, bound, captures, captures_set, builtins, top_level); - } - Term::Do { args, .. } => { - for a in args { - Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); - } - } - Term::Ctor { args, .. } => { - for a in args { - Self::collect_captures(a, bound, captures, captures_set, builtins, top_level); - } - } - Term::Match { scrutinee, arms } => { - Self::collect_captures(scrutinee, bound, captures, captures_set, builtins, top_level); - for arm in arms { - let mut pat_bindings = Vec::new(); - Self::pattern_bound_names(&arm.pat, &mut pat_bindings); - let mut newly_bound = Vec::new(); - for n in &pat_bindings { - if bound.insert(n.clone()) { - newly_bound.push(n.clone()); - } - } - Self::collect_captures(&arm.body, bound, captures, captures_set, builtins, top_level); - for n in newly_bound { - bound.remove(&n); - } - } - } - Term::Lam { params, body, .. } => { - // Inner lambda's params don't escape; its free vars - // (relative to the outer) DO contribute to outer captures. - let mut newly_bound = Vec::new(); - for p in params { - if bound.insert(p.clone()) { - newly_bound.push(p.clone()); - } - } - Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); - for n in newly_bound { - bound.remove(&n); - } - } - Term::Seq { lhs, rhs } => { - Self::collect_captures(lhs, bound, captures, captures_set, builtins, top_level); - Self::collect_captures(rhs, bound, captures, captures_set, builtins, top_level); - } - Term::LetRec { .. } => { - // Iter 16b.1: eliminated by desugar before codegen. - unreachable!("Term::LetRec eliminated by desugar") - } - Term::Clone { value } => { - // Iter 18c.1: clone is identity for capture analysis — - // free vars of `(clone X)` are exactly the free vars of `X`. - Self::collect_captures(value, bound, captures, captures_set, builtins, top_level); - } - Term::ReuseAs { source, body } => { - // Iter 18d.1: free vars are the union of source and body. - Self::collect_captures(source, bound, captures, captures_set, builtins, top_level); - Self::collect_captures(body, bound, captures, captures_set, builtins, top_level); - } - } - } - - fn pattern_bound_names(p: &Pattern, out: &mut Vec) { - match p { - Pattern::Wild | Pattern::Lit { .. } => {} - Pattern::Var { name } => out.push(name.clone()), - Pattern::Ctor { fields, .. } => { - for f in fields { - Self::pattern_bound_names(f, out); - } - } - } - } /// Iter 7: is `name` a callee that can be resolved at compile time /// (no fn-pointer needed)? True for builtin operators, the