Brummel
bea5c92591
Fourth slice of the codegen split. Pulls the closure-emission
cluster into a dedicated module. No behaviour change.
Methods moved:
- lower_lambda (~290 lines, pub(crate)): capture analysis +
thunk fn lifting + heap-env construction + closure-pair
packing + per-pair drop-fn emission under --alloc=rc.
- collect_captures (~108 lines, private): recursive free-var
walk used only as lower_lambda's prelude.
- pattern_bound_names (~17 lines, private): pattern-binder
enumeration used by collect_captures's match arm.
Imports: super::escape (for analyze_fn_body), super::synth
(for fn_sig_from_type, llvm_type), the usual Emitter/Result
re-exports.
lib.rs drops from 3237 → 2825 lines. lambda.rs is 447 lines.
Codegen split summary (5295 → 2825 lib.rs, ~47% reduction):
lib.rs 2825 — Emitter struct, lifecycle (new,
start_block, emit_module, emit_*),
lower_term mass dispatcher, app/eq/effect
lowering, lookup helpers, synth helpers.
match_lower.rs 935 — Term::Ctor / ReuseAs / Match.
escape.rs 722 — escape analysis (predates the split).
drop.rs 696 — per-type drop fns + per-let-close drop.
lambda.rs 447 — Term::Lam + closure machinery.
subst.rs 319 — monomorphisation substitution.
synth.rs 216 — IR shaping helpers + builtin types.
cargo test --workspace clean across all four phases.
448 lines
19 KiB
Rust
448 lines
19 KiB
Rust
//! 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_<m>_<def>_lam<id>`,
|
|
//! 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_<m>_<def>_lam<id>`, 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<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,
|
|
);
|
|
|
|
// 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<FnSig>)> = 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_<thunk>_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_<thunk>_pair(p)` — load the env from offset 8,
|
|
// call `drop_<thunk>_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<String>,
|
|
captures: &mut Vec<String>,
|
|
captures_set: &mut BTreeSet<String>,
|
|
builtins: &BTreeSet<&str>,
|
|
top_level: &BTreeSet<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);
|
|
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<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);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|