Implementing Closures In a Compiled Lisp

Jan 2021

If writing a compiler is a final boss, implementing closures are one of its most powerful attacks. Closures are complicated. Suddenly functions transform from mere labels you can jump to into heap allocated objects with associated data and pointers of their own.

Seemingly as a result of this I’ve noticed that many intro texts and blog posts about writing a compiler tend to avoid closures. Even Ghuloum’s venerable paper An Incremental Approach to Compiler Construction doesn’t seem to give them as much time as they deserve.

Here, I hope to help right this by walking through how I added closures to my compiled lisp in as much detail as is reasonable. The process of implementing closures is pretty nicely split into two tasks. The first task is to determine what variables each function in your program closes over. The second task is then to emit code to close over those values and construct closure objects. We'll call these two steps static analysis and compilation.

Static Analysis

In the lisp that this is based on there are two ways that variables can be created: as function arguments, or inside of a let expression. Here’s what that looks like:

;; As function arguments
(fn (variable) variable)

;; Inside a let expression
(let variable 10)

The goal of our static analysis pass is for each function to determine what variables that are used inside of it that it does not define. We say that those variables are free variables and that the function closes over them. For example, in the function below the variable free is a free variable in the function bound to add-free.

(let free 10)
(let add-free (fn (n)
                  (add free n)))

We must take care too that our static analysis pass can handle the case where a variable is free for part of a function’s body and bound for the rest. For example, the following program where free is in such a situation:

(let free 10)
(let add-half-free (fn (n)
                       (let free (div free 2))
                       (add free n)))

Once we understand those constraints, we can write a function to perform that analysis with some care. Here’s what I used in my lisp:

pub(crate) fn annotate_free_variables(f: &mut LustFn) {
    let mut bound: HashSet<&String> = f.params.iter().collect();
    let mut free = HashSet::<&String>::new();

    for e in &f.body {
        let (newbound, newfree) = analyze_variables(e);
        // Extend with the free variables found that do not have
        // bindings in our scope yet.
        free.extend(newfree.difference(&bound));
        // Once that is done add the new bound variables to our
        // scope. This needs to happen second so variables that are
        // free and bound in analyze_variables are added properly.
        bound.extend(newbound);
    }

    f.free_variables = free.into_iter().cloned().collect();
}

annotate_free_variables takes a function as an argument and for each statement in its body calls analyze_variables to find the bound and unbound variables in it. Then, it merges them with its own bound and unbound variables and stores the result in the function. If you’re interested in how analyze_variables works you can take a look at it here.

Compilation

In the Lisp that this is based on all functions are anonymous and leading up to the compilation of closures they are compiled as follows:

1. Each function is collected into an object containing a unique name, its arguments, and its body.
2. The functions are lifted to the top of the program where they are compiled individually.
3. The source code is transformed so that each anonymous function compiled this way is replaced with a symbol containing its unique name.

The result is that a program like this:

(let I (fn (a) a))

Is transformed into a program like this:

(let I __anon_fn_0)

In order to make these functions usable, whenever we encounter a symbol in the form __anon_fn_0 we need to replace what would normally be a variable access with code to construct a closure containing a pointer to that function.

A closure is a heap allocated object that contains a pointer to the function it refers to and after that copies of all the values that function captures at the time of its construction. This memory representation is similar to the way an array is represented in C except instead of holding data it holds a pointer followed by the values of variables that are captured.

[fn_ptr, freevar_1, ..., freevar_N]

I’ll start by describing how closures are compiled and then describe how they can be called.

Compiling closures

Whenever we encounter a symbol that references a function we need to emit the code to create a closure on the heap. This is actually reasonably straightforward once we’ve done free variable annotation. In psudocode it looks something like this:

# Lookup the free variables for this function
free_vars = lookup_free_vars(fn_name)

# Malloc enough space to store the closure
closure_ptr = emit_malloc_words(len(free_vars) + 1)

# Get the function pointer and store it as position 0 in the closure.
fn_ptr = emit_get_fn_ptr(fn_name)
emit_store(fn_ptr, closure_ptr, 0)

for i, var in enumerate(free_vars):
    # Get and store the value at its position in the closure.
    val = emit_var_access(var)
    emit_store(val, closure_ptr, i)

# Return the closure pointer for the caller to use.
return closure_ptr

Admittedly this is an oversimplification of how this works, but I assume that if you’re at the point where you’re compiling closures you have variable bindings and heap allocations in your language and the emit_ functions should already exist in some form or another.

Calling closures

Once we’re compiling closures we need some way to call them. To do this we’ll modify all of the functions that we compile to take an additional argument that is their closure. As part of the preamble of the function it’ll move over the values in its closure and store them in variables corresponding to their names.

In a lower level sense this is to say: in the body of functions lookup free variables inside of the closure pointer passed to the function instead of the stack or wherever function arguments and local variables are normally stored.

Once that is done, the code for emitting a function call is as simple as this:

fn_ptr = emit_load(closure_ptr, 0)

# Push the functions arguments to the stack
emit_push_fn_args()
# Push the closure pointer to the stack
emit_push_closure_ptr()

emit_call_indirect(fn_ptr)

Conclusion

This stuff can be daunting. It took me a ton of experimentation, tutorials and an entire compilers course to work out how it all works. Even then I'm just starting to dip my toes into it all. If you're interested in writing your own compiler I highly reccomend taking a look at An Incremental Approach to Compiler Construction which helped me add some structure to my efforts.

The Lisp I'm working on that this is based on is built on Cranelift's JIT. Thus far I'm enjoying it quite a lot and it has been a nice excuse to really try and learn Rust. The performance seems to be quite good as well thus far which is an added bonus.

While pretty high level hopefully this has helped to demystify some of the methodology for adding closures to your compiled language. If you’d like to browse through some real code for doing this, you can take a look at these three commits in my Lisp compiler or browse the entire repository here.