Chapter 19: Closure Conversion

New files

src/closure.rs

src/flat.rs

src/flat_machine.rs

src/flat_lint.rs

Changed files

src/lib.rs

src/main.rs

Motivation

A function can outlive the call that created it. That is easy to write in Fun, but it leaves our back end with a question. Where do the variables from that old call live now?

Our interpreters answered by storing an environment with each function. We’ll make the same idea explicit in the program we compile. Each lambda becomes a function in a program-wide table. Its closure carries the values that function needs from the surrounding scope.

We’ll also work out how to call these functions. When we know the target, we can call it directly. When a function arrives as an argument or comes back from another call, we need to inspect its closure. Partial application adds one more interesting case. We’ll model all of this in a Flat interpreter before we start laying out the objects in native memory.

The Language, Now

Here’s a small function that remembers x and can be passed to map.

make x = \y -> x+y

main = print (map (make 20) [1,2,3])

Every call of the returned function needs the supplied value 20. Here ANF uncurrying merges the adjacent lambdas into one global function, so its partial application retains that value in a PAP. The overloaded addition also needs a Num dictionary. Local lambdas separated by a binding or branch instead require explicit captured environments. Partial application, oversaturated application, recursive local functions, and class dictionaries used as values must all keep working. --engine=flat is now the default. Core and ANF remain selectable.

A closure environment is compiler-created storage. The source language gains no mutable field or reference operation. The only mutation introduced here initializes capture slots for a recursive group before user code can observe the closures.

Theory

Free-variable analysis identifies the values a lambda needs beyond its own parameters. Globals need no local capture slot because the program can address them directly. For each remaining free variable, conversion chooses a stable slot, stores the value when constructing the closure, and replaces lexical access with an environment load in the lifted function.

Closure conversion and lambda lifting are related but different choices. Closure conversion packages captured values into an environment. Lambda lifting can instead add them as explicit function parameters when call sites permit it (Appel 1992). This chapter keeps one environment mechanism for local functions, even non-escaping ones. Optional lambda lifting is deferred. A known local call can still bypass general apply while passing its environment explicitly.

A source function can be called a little at a time. Our lifted machine function expects a fixed number of arguments. How do we reconcile those two views?

If the function’s syntactic arity is n and we have fewer arguments, we save them in a partial application, or PAP. If we have exactly n, we enter the function. If we have more, we call it with n and apply the leftover arguments to the function it returns. That’s the eval/apply protocol in miniature (Marlow and Peyton Jones 2004). A single calling mechanism now handles all three cases.

Mutual recursion gives us a chicken-and-egg problem. If even captures odd and odd captures even, which closure can we build first? We’ll allocate both with placeholder capture slots, then fill the slots once both addresses exist. No source code can call them until this initialization is finished. Requiring recursive bindings to be functions makes this possible without trying to evaluate an unfinished recursive value.

Implementation

The Flat Grammar

Flat keeps ANF’s explicit control flow, but functions now live in a separate function table. There is no lambda form inside an executable body. Instead, we allocate a closure that names a lifted function and stores its captures.

$e = let x = r in e ∣ case a of {\overline{q \overline{x} \to e}} ∣ join j (\overline{x}) = e_{1} in e_{2} ∣ jump j (\overline{a}) ∣ return a ∣ tail f [ρ] (\overline{a}) ∣ tailapply a (\overline{a}) ∣ fail_{s} strict binding branch on an atom local continuation jump to a join return a value known tail call unknown tail call pattern failure at span s$

We keep the ANF notation for atoms $a$ , join names $j$ , alternative tags $q$ , and sequences marked with an overbar. The name $f$ identifies a lifted function. Its optional environment operand $ρ$ is either an atom or absent. A known call names its code directly. An unknown call must inspect a closure or PAP to find out what to do.

$a ρ r = x ∣ g ∣ ℓ = a ∣ none = a ∣ call f [ρ] (\overline{a}) ∣ apply a (\overline{a}) ∣ con K (\overline{a}) ∣ closure f [\overline{a}] ∣ env x [i] ∣ field a [i] ∣ p (\overline{a}) ∣ patch a [i] = a^{'} local, global, or literal optional environment atom known call unknown call allocate a constructor allocate a closure read a capture read a constructor field primitive operation fill a closure capture$

Here $K$ is a constructor, $p$ a primitive, and $i$ a slot index. Primitive source spans are omitted. patch is an internal operation used while tying recursive closures together, not mutable state available to a Fun program. The LetRec from ANF has disappeared into allocations followed by these patches.

Lifted Functions and Calls

Flat IR removes lambda expressions from executable bodies. Each lifted function records its parameters, optional environment parameter, body, and ordered capture list:

#![allow(unused)]
fn main() {
pub use crate::anf::{Atom, Var};
use crate::{
    ast::Op,
    core::{self, AltCon},
    data::ConId,
    source::Span,
};
use std::rc::Rc;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
/// An index into this program's lifted function table.
pub struct FuncId(pub u32);
pub type Body = Rc<Flat>;
#[derive(Clone, Debug)]
/// Function globals are installed before ordered strict value initializers run.
pub struct GlobalInit {
    pub name: Var,
    pub function: Option<FuncId>,
    pub body: Body,
}
#[derive(Clone, Debug)]
pub struct Program {
    pub funcs: Vec<Func>,
    pub globals: Vec<GlobalInit>,
    pub metadata: core::Program,
}
#[derive(Clone, Debug)]
/// PatchClosure is a compiler-only operation used after all letrec allocations.
pub enum Rhs {
    Atom(Atom),
    CallKnown(FuncId, Option<Atom>, Vec<Atom>),
    CallUnknown(Atom, Vec<Atom>),
    AllocCon(ConId, Vec<Atom>),
    AllocClosure(FuncId, Vec<Atom>),
    EnvField(Var, u32),
    Field(Atom, u32),
    Prim(Op, Vec<Atom>, Span),
    PatchClosure(Atom, u32, Atom),
}
#[derive(Clone, Debug)]
/// Control transfers are explicit and joins remain local to each function.
pub enum Flat {
    Let(Var, Rhs, Body),
    Case(Atom, Vec<(AltCon, Vec<Var>, Body)>),
    Join(Var, Vec<Var>, Body, Body),
    Jump(Var, Vec<Atom>),
    Ret(Atom),
    TailCallKnown(FuncId, Option<Atom>, Vec<Atom>),
    TailCallUnknown(Atom, Vec<Atom>),
    Fail(Span),
}
impl std::fmt::Display for FuncId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "f{}", self.0)
    }
}
fn list<T: std::fmt::Display>(xs: &[T]) -> String {
    xs.iter()
        .map(ToString::to_string)
        .collect::<Vec<_>>()
        .join(", ")
}
impl std::fmt::Display for Rhs {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Atom(a) => write!(f, "{a}"),
            Self::CallKnown(id, env, args) => write!(
                f,
                "call {id}[{}]({})",
                env.as_ref().map(ToString::to_string).unwrap_or_default(),
                list(args)
            ),
            Self::CallUnknown(a, args) => {
                write!(f, "apply {a}({})", list(args))
            }
            Self::AllocCon(c, args) => {
                write!(f, "{}({})", c.name, list(args))
            }
            Self::AllocClosure(id, args) => {
                write!(f, "closure {id}[{}]", list(args))
            }
            Self::EnvField(v, i) => write!(f, "env {v}[{i}]"),
            Self::Field(a, i) => write!(f, "field {a}[{i}]"),
            Self::Prim(p, args, _) => {
                write!(f, "{}({})", p.text(), list(args))
            }
            Self::PatchClosure(a, i, v) => {
                write!(f, "patch {a}[{i}] = {v}")
            }
        }
    }
}
impl std::fmt::Display for Flat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Let(v, r, b) => write!(f, "let {v} = {r};\n{b}"),
            Self::Join(j, ps, b, r) => {
                write!(f, "join {j}({}) = {{ {b} }};\n{r}", list(ps))
            }
            Self::Jump(j, args) => write!(f, "jump {j}({})", list(args)),
            Self::Ret(a) => write!(f, "return {a}"),
            Self::TailCallKnown(id, env, args) => write!(
                f,
                "tail {id}[{}]({})",
                env.as_ref().map(ToString::to_string).unwrap_or_default(),
                list(args)
            ),
            Self::TailCallUnknown(a, args) => {
                write!(f, "tail-apply {a}({})", list(args))
            }
            Self::Fail(s) => write!(f, "fail {s:?}"),
            Self::Case(a, alts) => {
                writeln!(f, "case {a} {{")?;
                for (c, ps, b) in alts {
                    let head = match c {
                        AltCon::Con(c) => c.name.clone(),
                        AltCon::Lit(l) => l.to_string(),
                        AltCon::Default => "_".into(),
                    };
                    writeln!(f, "{head}({}) -> {{ {b} }}", list(ps))?;
                }
                write!(f, "}}")
            }
        }
    }
}
#[derive(Clone, Debug)]
/// Only lifted local lambdas receive an environment parameter.
pub struct Func {
    pub name: FuncId,
    pub source: Var,
    pub env: Option<Var>,
    pub params: Vec<Var>,
    pub body: Body,
    pub captures: Vec<Var>,
}
}

closure::convert collects lambda definitions, computes free variables, and analyzes uses. Capture order follows unique-name order, so dumps and slot indexes are deterministic. Top-level function entries have no environment parameter. Lifted local entries do, even when their environment is empty. Prefix EnvField bindings load captures at the start of each local function.

A call is known only when its argument count matches the syntactic arity and its head is a global function or a non-escaping local lambda. Merely having a function type is insufficient. The actual decision is small enough to inspect:

#![allow(unused)]
fn main() {
fn known(&self, a: &Atom, n: usize) -> Option<(FuncId, Option<Atom>)> {
    let (Atom::Var(v) | Atom::Global(v)) = a else {
        return None;
    };
    let info = self.infos.get(v)?;
    (n == info.lambda.params.len() && (info.global || !info.escapes))
        .then(|| (info.id, (!info.global).then(|| a.clone())))
}
}

Escape analysis regards saturated use as a callee differently from returning, storing, or passing the function as an argument. A known local call still carries its closure as an environment operand. Other calls use CallUnknown. Tail position selects the corresponding tail-call form. Partial application therefore continues through apply even when the source spelling names a familiar function.

The right-hand-side language makes allocation and access visible:

#![allow(unused)]
fn main() {
pub use crate::anf::{Atom, Var};
use crate::{
    ast::Op,
    core::{self, AltCon},
    data::ConId,
    source::Span,
};
use std::rc::Rc;
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
/// An index into this program's lifted function table.
pub struct FuncId(pub u32);
pub type Body = Rc<Flat>;
#[derive(Clone, Debug)]
/// Only lifted local lambdas receive an environment parameter.
pub struct Func {
    pub name: FuncId,
    pub source: Var,
    pub env: Option<Var>,
    pub params: Vec<Var>,
    pub body: Body,
    pub captures: Vec<Var>,
}
#[derive(Clone, Debug)]
/// Function globals are installed before ordered strict value initializers run.
pub struct GlobalInit {
    pub name: Var,
    pub function: Option<FuncId>,
    pub body: Body,
}
#[derive(Clone, Debug)]
pub struct Program {
    pub funcs: Vec<Func>,
    pub globals: Vec<GlobalInit>,
    pub metadata: core::Program,
}
#[derive(Clone, Debug)]
/// Control transfers are explicit and joins remain local to each function.
pub enum Flat {
    Let(Var, Rhs, Body),
    Case(Atom, Vec<(AltCon, Vec<Var>, Body)>),
    Join(Var, Vec<Var>, Body, Body),
    Jump(Var, Vec<Atom>),
    Ret(Atom),
    TailCallKnown(FuncId, Option<Atom>, Vec<Atom>),
    TailCallUnknown(Atom, Vec<Atom>),
    Fail(Span),
}
impl std::fmt::Display for FuncId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "f{}", self.0)
    }
}
fn list<T: std::fmt::Display>(xs: &[T]) -> String {
    xs.iter()
        .map(ToString::to_string)
        .collect::<Vec<_>>()
        .join(", ")
}
impl std::fmt::Display for Rhs {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Atom(a) => write!(f, "{a}"),
            Self::CallKnown(id, env, args) => write!(
                f,
                "call {id}[{}]({})",
                env.as_ref().map(ToString::to_string).unwrap_or_default(),
                list(args)
            ),
            Self::CallUnknown(a, args) => {
                write!(f, "apply {a}({})", list(args))
            }
            Self::AllocCon(c, args) => {
                write!(f, "{}({})", c.name, list(args))
            }
            Self::AllocClosure(id, args) => {
                write!(f, "closure {id}[{}]", list(args))
            }
            Self::EnvField(v, i) => write!(f, "env {v}[{i}]"),
            Self::Field(a, i) => write!(f, "field {a}[{i}]"),
            Self::Prim(p, args, _) => {
                write!(f, "{}({})", p.text(), list(args))
            }
            Self::PatchClosure(a, i, v) => {
                write!(f, "patch {a}[{i}] = {v}")
            }
        }
    }
}
impl std::fmt::Display for Flat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Let(v, r, b) => write!(f, "let {v} = {r};\n{b}"),
            Self::Join(j, ps, b, r) => {
                write!(f, "join {j}({}) = {{ {b} }};\n{r}", list(ps))
            }
            Self::Jump(j, args) => write!(f, "jump {j}({})", list(args)),
            Self::Ret(a) => write!(f, "return {a}"),
            Self::TailCallKnown(id, env, args) => write!(
                f,
                "tail {id}[{}]({})",
                env.as_ref().map(ToString::to_string).unwrap_or_default(),
                list(args)
            ),
            Self::TailCallUnknown(a, args) => {
                write!(f, "tail-apply {a}({})", list(args))
            }
            Self::Fail(s) => write!(f, "fail {s:?}"),
            Self::Case(a, alts) => {
                writeln!(f, "case {a} {{")?;
                for (c, ps, b) in alts {
                    let head = match c {
                        AltCon::Con(c) => c.name.clone(),
                        AltCon::Lit(l) => l.to_string(),
                        AltCon::Default => "_".into(),
                    };
                    writeln!(f, "{head}({}) -> {{ {b} }}", list(ps))?;
                }
                write!(f, "}}")
            }
        }
    }
}
#[derive(Clone, Debug)]
/// PatchClosure is a compiler-only operation used after all letrec allocations.
pub enum Rhs {
    Atom(Atom),
    CallKnown(FuncId, Option<Atom>, Vec<Atom>),
    CallUnknown(Atom, Vec<Atom>),
    AllocCon(ConId, Vec<Atom>),
    AllocClosure(FuncId, Vec<Atom>),
    EnvField(Var, u32),
    Field(Atom, u32),
    Prim(Op, Vec<Atom>, Span),
    PatchClosure(Atom, u32, Atom),
}
}

For a LetRec, conversion emits every AllocClosure before any PatchClosure. The initial capture slots contain placeholder integer literals in the IR. After patching, each contains its actual captured value. Patch operations return unit into unused temporary binders. All patches precede the converted body, so no source expression can call a partially initialized closure.

Flat lint checks each lifted function independently. It rejects an implicit free variable, unknown function identity, incorrect known-call arity, missing or unexpected environment operand, out-of-bounds environment load, or invalid join transfer. Capture binder identities can recur in different functions, since each environment load rebinds the name in a separate function scope.

The Flat machine owns an indexed arena of closure records. A closure value is an arena index. Environment slots contain ordinary interpreter values, including other closure indexes. That permits recursive cycles without reference-count cycles. PAPs store a function value and supplied arguments. The arena lives for one execution and models object relationships. It does not yet model reclamation.

Execution is an iterative loop with separate code, known-call, apply, and return states. Known calls validate the entry’s arity and environment. Apply flattens PAP arguments, handles short and exact calls, and keeps excess arguments in an explicit continuation. Joins remain lexical control transfers within a function. Primitive operands cross a data adapter, but evaluation never delegates a Flat body to the Core or ANF interpreter. Agreement is meaningful because the three engines implement control flow independently.

Running It

$ cd crates/ch19_closure
$ cargo run -q -- run tests/programs/closure/capture-return.fun --lint
[21,22,23]
$ cargo run -q -- dump --pass=types tests/programs/closure/capture-return.fun
main : IO ()

Dump this file with --pass=flat to find the partial application of make. The global entry has no environment. Apply saves its dictionary and integer arguments in a PAP. In contrast, escape-recursive.fun returns a local recursive closure capturing an outer value, and known-local.fun uses direct local calls. pap-chain.fun supplies arguments in stages. over-application.fun supplies arguments to a function whose result is itself a function. The small Flat snapshots in tests/anf.rs omit unrelated Prelude functions by selecting reachable code.

Try supplying arguments in stages, returning a local function, and making two local functions call each other while both capture an outer value. These all look natural in source, but exercise different parts of the closure machinery. Core, ANF, and Flat should agree on their results. The Flat machine also checks that a long tail-recursive sum keeps a bounded number of pending continuations.

Exercises

Warm-up. Trace how uncurrying and partial application preserve the example’s integer argument and Num dictionary.
Intermediate. Change a local function from saturated-only use to returning it as a value. Explain the resulting call classification.
Intermediate. Trace a PAP receiving another argument, followed by an oversaturated call whose first result is a closure.
Advanced. Construct mutually recursive local functions that also capture an outer integer. Draw their environment graph after patching.
Advanced. Prototype lambda lifting for non-escaping functions and compare allocation counts while retaining differential agreement.

Write You a Haskell