Iter 8b: lambdas with capture (Term::Lam + closure conversion)
Adds anonymous functions to AILang. A lambda value is constructed by
malloc'ing an env struct that holds its captured locals, plus a
closure pair `{ thunk_ptr, env_ptr }` (Iter 8a's ABI). The lambda's
body lifts to a top-level thunk `@ail_<m>_<def>_lam<id>(ptr %env,
params...)` that unpacks captures back into named locals before
running.
AST: new Term::Lam { params, paramTypes, retType, effects, body }.
Serde tag is "lam"; existing modules (no Lam) serialize identically,
so all current hashes stay stable.
Pretty-printer: renders `(\\ (x: T ...) -> R . body)`.
Typecheck (ailang-check): a lambda's type is the declared Type::Fn.
Body's effect set must be a subset of declared effects (no row
polymorphism). Constructing a lambda is pure — only calling it picks
up the declared effects, via the existing App branch.
Codegen (ailang-codegen):
- collect_captures: free-var walk; excludes builtins, current-module
top-level fns, and qualified names.
- lower_lambda: per-fn lam counter, save/restore emitter state, emit
thunk into a deferred queue (flushed after the parent fn), pack env
+ closure pair on heap, register the resulting closure-pair SSA in
the sidetable so subsequent indirect calls find it.
- Capture sigs propagate: a fn-typed capture keeps its FnSig in the
thunk's sidetable so `f(args)` inside the body still indirect-calls.
- Env layout: 8 bytes per capture (typed load/store reads only the
needed bytes; padding wasted but uniform).
Tests: 49 green (was 48). New examples/closure.ail.json + e2e
`closure_captures_let_n` exercises `let n = 3 in apply(\\x. x + n, 39)`
end-to-end and asserts the binary prints 42.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -22,7 +22,7 @@
|
||||
|
||||
use ailang_core::ast::*;
|
||||
use ailang_core::Workspace;
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum CodegenError {
|
||||
@@ -228,6 +228,15 @@ struct Emitter<'a> {
|
||||
/// function entry. Used by `Term::App` when the callee is not a
|
||||
/// statically-known top-level name.
|
||||
ssa_fn_sigs: BTreeMap<String, FnSig>,
|
||||
/// Iter 8b: name of the currently-emitted def (for lambda thunk
|
||||
/// naming `<def>_lam<n>`).
|
||||
current_def: String,
|
||||
/// Iter 8b: per-def counter for lambda thunks. Reset in emit_fn.
|
||||
lam_counter: u32,
|
||||
/// Iter 8b: thunk fn IR text for lambdas encountered during
|
||||
/// lowering. Flushed at the end of emit_fn (LLVM IR allows fns in
|
||||
/// any order).
|
||||
deferred_thunks: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -300,6 +309,9 @@ impl<'a> Emitter<'a> {
|
||||
ctor_index,
|
||||
current_block: String::new(),
|
||||
ssa_fn_sigs: BTreeMap::new(),
|
||||
current_def: String::new(),
|
||||
lam_counter: 0,
|
||||
deferred_thunks: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -388,6 +400,10 @@ impl<'a> Emitter<'a> {
|
||||
// Per-fn body: the sidetable starts empty. Top-level fn references
|
||||
// get registered on demand by `lower_term(Term::Var)`.
|
||||
self.ssa_fn_sigs.clear();
|
||||
// Iter 8b: lambda thunks live in the same module body but get
|
||||
// collected during lowering and appended after the parent fn.
|
||||
self.current_def = f.name.clone();
|
||||
self.lam_counter = 0;
|
||||
|
||||
let mut sig = format!(
|
||||
"define {ret} @ail_{module}_{name}(",
|
||||
@@ -429,6 +445,13 @@ impl<'a> Emitter<'a> {
|
||||
self.body
|
||||
.push_str(&format!(" ret {val_ty} {val}\n}}\n\n"));
|
||||
|
||||
// Iter 8b: flush lambda thunks collected while lowering this fn's
|
||||
// body. They go after the closing `}` of the parent fn, before
|
||||
// the adapter, so the parent fn is contiguous.
|
||||
for t in self.deferred_thunks.drain(..) {
|
||||
self.body.push_str(&t);
|
||||
}
|
||||
|
||||
// Iter 8a: emit closure-pair scaffold (adapter + static closure)
|
||||
// for this fn. The adapter takes an extra `ptr %_env` (ignored,
|
||||
// null sentinel for top-level fns) and forwards to the real fn.
|
||||
@@ -608,6 +631,9 @@ impl<'a> Emitter<'a> {
|
||||
Term::Do { op, args } => self.lower_effect_op(op, args),
|
||||
Term::Ctor { type_name, ctor, args } => self.lower_ctor(type_name, ctor, args),
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1009,6 +1035,312 @@ 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_<m>_<def>_lam<id>`, 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,
|
||||
) -> Result<(String, String)> {
|
||||
let llvm_param_tys: Vec<String> =
|
||||
lam_param_tys.iter().map(llvm_type).collect::<Result<_>>()?;
|
||||
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<String> = self
|
||||
.module_user_fns
|
||||
.get(self.module_name)
|
||||
.map(|m| m.keys().cloned().collect())
|
||||
.unwrap_or_default();
|
||||
let builtins_owned: Vec<String> = 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<String> = BTreeSet::new();
|
||||
for p in lam_params {
|
||||
bound.insert(p.clone());
|
||||
}
|
||||
let mut captures: Vec<String> = Vec::new();
|
||||
let mut captures_set: BTreeSet<String> = BTreeSet::new();
|
||||
Self::collect_captures(
|
||||
lam_body,
|
||||
&mut bound,
|
||||
&mut captures,
|
||||
&mut captures_set,
|
||||
&builtins,
|
||||
&top_level,
|
||||
&self.import_map,
|
||||
);
|
||||
|
||||
// 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.
|
||||
let mut cap_meta: Vec<(String, String, String, Option<FnSig>)> = Vec::new();
|
||||
// (name, outer_ssa, llvm_type, optional_sig_for_fn_typed_capture)
|
||||
for c in &captures {
|
||||
let (_, outer_ssa, lty) = 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, 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_sigs = std::mem::take(&mut self.ssa_fn_sigs);
|
||||
// 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, 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()));
|
||||
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()));
|
||||
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 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"));
|
||||
|
||||
// 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.ssa_fn_sigs = saved_sigs;
|
||||
|
||||
// 3. Emit allocation + capture filling + closure-pair packing
|
||||
// in the OUTER body. Captures use 8 bytes each; closure-pair
|
||||
// is 16 bytes.
|
||||
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 @malloc(i64 {env_size})\n"));
|
||||
for (i, (_cname, outer_ssa, cty, _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();
|
||||
self.body
|
||||
.push_str(&format!(" {clos} = call ptr @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"
|
||||
));
|
||||
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);
|
||||
|
||||
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<String>,
|
||||
captures: &mut Vec<String>,
|
||||
captures_set: &mut BTreeSet<String>,
|
||||
builtins: &BTreeSet<&str>,
|
||||
top_level: &BTreeSet<String>,
|
||||
import_map: &BTreeMap<String, String>,
|
||||
) {
|
||||
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, import_map);
|
||||
for a in args {
|
||||
Self::collect_captures(a, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
}
|
||||
}
|
||||
Term::Let { name, value, body } => {
|
||||
Self::collect_captures(value, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
let inserted = bound.insert(name.clone());
|
||||
Self::collect_captures(body, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
if inserted {
|
||||
bound.remove(name);
|
||||
}
|
||||
}
|
||||
Term::If { cond, then, else_ } => {
|
||||
Self::collect_captures(cond, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
Self::collect_captures(then, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
Self::collect_captures(else_, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
}
|
||||
Term::Do { args, .. } => {
|
||||
for a in args {
|
||||
Self::collect_captures(a, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
}
|
||||
}
|
||||
Term::Ctor { args, .. } => {
|
||||
for a in args {
|
||||
Self::collect_captures(a, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
}
|
||||
}
|
||||
Term::Match { scrutinee, arms } => {
|
||||
Self::collect_captures(scrutinee, bound, captures, captures_set, builtins, top_level, import_map);
|
||||
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, import_map);
|
||||
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, import_map);
|
||||
for n in newly_bound {
|
||||
bound.remove(&n);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
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
|
||||
/// current-module top-level fns, and qualified `prefix.def`.
|
||||
|
||||
Reference in New Issue
Block a user