From ecde8fa7af156a769842a6dcdfd161bfc3af2c5c Mon Sep 17 00:00:00 2001 From: Brummel Date: Thu, 7 May 2026 13:03:20 +0200 Subject: [PATCH] Iter 8b: lambdas with capture (Term::Lam + closure conversion) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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___lam(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) --- crates/ail/src/main.rs | 13 ++ crates/ail/tests/e2e.rs | 10 + crates/ailang-check/src/lib.rs | 42 ++++ crates/ailang-codegen/src/lib.rs | 334 ++++++++++++++++++++++++++++++- crates/ailang-core/src/ast.rs | 14 ++ crates/ailang-core/src/pretty.rs | 25 +++ examples/closure.ail.json | 77 +++++++ 7 files changed, 514 insertions(+), 1 deletion(-) create mode 100644 examples/closure.ail.json diff --git a/crates/ail/src/main.rs b/crates/ail/src/main.rs index 8b39d8e..556bb35 100644 --- a/crates/ail/src/main.rs +++ b/crates/ail/src/main.rs @@ -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); + } + } } } diff --git a/crates/ail/tests/e2e.rs b/crates/ail/tests/e2e.rs index 1f1c282..b48060d 100644 --- a/crates/ail/tests/e2e.rs +++ b/crates/ail/tests/e2e.rs @@ -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. diff --git a/crates/ailang-check/src/lib.rs b/crates/ailang-check/src/lib.rs index 32d7f2d..a4d027e 100644 --- a/crates/ailang-check/src/lib.rs +++ b/crates/ailang-check/src/lib.rs @@ -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)> = 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 = 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 = 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(), + }) + } } } diff --git a/crates/ailang-codegen/src/lib.rs b/crates/ailang-codegen/src/lib.rs index 471d85d..b6a87e9 100644 --- a/crates/ailang-codegen/src/lib.rs +++ b/crates/ailang-codegen/src/lib.rs @@ -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, + /// Iter 8b: name of the currently-emitted def (for lambda thunk + /// naming `_lam`). + 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, } #[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___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, + ) -> 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, + &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)> = 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, + captures: &mut Vec, + captures_set: &mut BTreeSet, + builtins: &BTreeSet<&str>, + top_level: &BTreeSet, + import_map: &BTreeMap, + ) { + 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) { + 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`. diff --git a/crates/ailang-core/src/ast.rs b/crates/ailang-core/src/ast.rs index 5ba1086..fd84b85 100644 --- a/crates/ailang-core/src/ast.rs +++ b/crates/ailang-core/src/ast.rs @@ -128,6 +128,20 @@ pub enum Term { scrutinee: Box, arms: Vec, }, + /// 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, + #[serde(rename = "paramTypes")] + param_tys: Vec, + #[serde(rename = "retType")] + ret_ty: Box, + #[serde(default)] + effects: Vec, + body: Box, + }, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/crates/ailang-core/src/pretty.rs b/crates/ailang-core/src/pretty.rs index 3b0bcbe..27a9428 100644 --- a/crates/ailang-core/src/pretty.rs +++ b/crates/ailang-core/src/pretty.rs @@ -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 = 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(" ")) + } } } diff --git a/examples/closure.ail.json b/examples/closure.ail.json new file mode 100644 index 0000000..3eac26f --- /dev/null +++ b/examples/closure.ail.json @@ -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 } } + ] + } + ] + } + } + } + ] +}