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:
2026-05-07 13:03:20 +02:00
parent 99f68e89fa
commit ecde8fa7af
7 changed files with 514 additions and 1 deletions
+13
View File
@@ -843,6 +843,19 @@ fn walk_term(
}
}
}
Term::Lam { params, body, .. } => {
// Lambda params shadow outer scope inside the body.
let mut newly = Vec::new();
for p in params {
if scope.insert(p.clone()) {
newly.push(p.clone());
}
}
walk_term(body, out, builtins, scope);
for p in newly {
scope.remove(&p);
}
}
}
}
+10
View File
@@ -72,6 +72,16 @@ fn higher_order_apply_inc() {
assert_eq!(stdout.trim(), "42");
}
/// Iter 8b: lambdas with capture. `let n = 3 in apply(\x. x + n, 39)`
/// must print 42 — the lambda captures `n` from the enclosing `let`,
/// the env struct is malloc'd, the closure pair is passed to `apply`,
/// and apply's indirect call unpacks both halves.
#[test]
fn closure_captures_let_n() {
let stdout = build_and_run("closure.ail.json");
assert_eq!(stdout.trim(), "42");
}
/// Guards `ail diff`: a modified body changes the hash of `sum`, while
/// `main` stays unchanged. Expects exit code 1, `changed` contains exactly
/// `sum`, `unchanged` contains `main`, `added`/`removed` empty.
+42
View File
@@ -797,6 +797,48 @@ fn synth(
Ok(result_ty.expect("checked arms is non-empty"))
}
Term::Lam { params, param_tys, ret_ty, effects: lam_effects, body } => {
// Iter 8b: a lambda's type is the declared `Type::Fn`. Push
// params as locals, check the body's type matches `ret_ty`,
// and ensure body effects are a subset of the declared set.
// (We trust the JSON schema for params.len() == param_tys.len();
// serde would have rejected a malformed AST.)
let mut pushed: Vec<(String, Option<Type>)> = Vec::new();
for (n, t) in params.iter().zip(param_tys.iter()) {
let prev = locals.insert(n.clone(), t.clone());
pushed.push((n.clone(), prev));
}
let mut body_effects: BTreeSet<String> = BTreeSet::new();
let body_ty = synth(body, env, locals, &mut body_effects, in_def);
for (n, prev) in pushed.into_iter().rev() {
match prev {
Some(p) => {
locals.insert(n, p);
}
None => {
locals.shift_remove(&n);
}
}
}
let body_ty = body_ty?;
expect_eq(ret_ty, &body_ty)?;
// Constructing a lambda is pure — the body's effects are
// sealed into the lambda's type, not propagated to the
// outer effect set. Calling the lambda (Term::App) will
// pick those effects up via Type::Fn.effects.
// Body effects must be a subset of the declared lam_effects.
let declared: BTreeSet<String> = lam_effects.iter().cloned().collect();
for e in &body_effects {
if !declared.contains(e) {
return Err(CheckError::UndeclaredEffect(e.clone()));
}
}
Ok(Type::Fn {
params: param_tys.clone(),
ret: Box::new((**ret_ty).clone()),
effects: lam_effects.clone(),
})
}
}
}
+333 -1
View File
@@ -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`.
+14
View File
@@ -128,6 +128,20 @@ pub enum Term {
scrutinee: Box<Term>,
arms: Vec<Arm>,
},
/// Anonymous function (Iter 8b). Captures any free variables of
/// `body` from the enclosing scope. Param/return types are
/// declared inline so the typechecker stays HM-monomorphic on
/// the inferred shape.
Lam {
params: Vec<String>,
#[serde(rename = "paramTypes")]
param_tys: Vec<Type>,
#[serde(rename = "retType")]
ret_ty: Box<Type>,
#[serde(default)]
effects: Vec<String>,
body: Box<Term>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
+25
View File
@@ -198,6 +198,28 @@ fn term_block(t: &Term, indent: usize) -> String {
s.push(')');
s
}
Term::Lam { params, param_tys, ret_ty, effects, body } => {
// (\ [params] :: typed-sig . body)
let typed_params: Vec<String> = params
.iter()
.zip(param_tys.iter())
.map(|(n, t)| format!("{n}: {}", type_to_string(t)))
.collect();
let eff = if effects.is_empty() {
String::new()
} else {
format!(" !{}", effects.join(","))
};
let mut s = format!(
"{pad}(\\ ({}) -> {}{}\n",
typed_params.join(" "),
type_to_string(ret_ty),
eff,
);
s.push_str(&term_block(body, indent + 2));
s.push(')');
s
}
}
}
@@ -273,6 +295,9 @@ fn term_inline(t: &Term) -> String {
term_inline(else_)
)
}
Term::Lam { params, .. } => {
format!("(\\ {} ...)", params.join(" "))
}
}
}
+77
View File
@@ -0,0 +1,77 @@
{
"schema": "ailang/v0",
"name": "closure",
"imports": [],
"defs": [
{
"kind": "fn",
"name": "apply",
"type": {
"k": "fn",
"params": [
{
"k": "fn",
"params": [{ "k": "con", "name": "Int" }],
"ret": { "k": "con", "name": "Int" },
"effects": []
},
{ "k": "con", "name": "Int" }
],
"ret": { "k": "con", "name": "Int" },
"effects": []
},
"params": ["f", "x"],
"doc": "apply a fn-of-Int to an Int",
"body": {
"t": "app",
"fn": { "t": "var", "name": "f" },
"args": [{ "t": "var", "name": "x" }]
}
},
{
"kind": "fn",
"name": "main",
"type": {
"k": "fn",
"params": [],
"ret": { "k": "con", "name": "Unit" },
"effects": ["IO"]
},
"params": [],
"doc": "captures `n` from the enclosing let, returns 42",
"body": {
"t": "let",
"name": "n",
"value": { "t": "lit", "lit": { "kind": "int", "value": 3 } },
"body": {
"t": "do",
"op": "io/print_int",
"args": [
{
"t": "app",
"fn": { "t": "var", "name": "apply" },
"args": [
{
"t": "lam",
"params": ["x"],
"paramTypes": [{ "k": "con", "name": "Int" }],
"retType": { "k": "con", "name": "Int" },
"effects": [],
"body": {
"t": "app",
"fn": { "t": "var", "name": "+" },
"args": [
{ "t": "var", "name": "x" },
{ "t": "var", "name": "n" }
]
}
},
{ "t": "lit", "lit": { "kind": "int", "value": 39 } }
]
}
]
}
}
}
]
}