Lab 14: Closures and Upvalues (Milestone 10)

Background

Two implementations, in order. First the naive one: heap-allocate a box for every captured local at declaration. Get the semantics right. Then replace it with Lua's open/closed upvalues and benchmark the difference.

The diff between those two commits is the deliverable. Keep it.

Why This Lab Matters

• This is the conceptual peak. If you can explain how make_counter works after this lab, you can read any dynamic-language runtime's closure implementation.
• Doing it twice separates two problems. Version 1 answers "what must happen"; version 2 answers "how do you make it cheap". Attempting both at once is why people lose a week here.
• The differential test is about to earn its keep again. The tree walker captures by holding an Rc to a shared cell; the VM captures with stack indices. Two structurally different implementations of one semantic.

Prerequisites

• Lab 13 complete.
• Closures and Upvalues read — both, before starting.
• Warm-up Experiment 3 re-run, with your predictions in front of you.

Predict First

1. Two closures returned from one make_counter() call. Same count or two? What must be true of the implementation for your answer?
2. for i = 1, 3 do fs[i] = function() return i end end — 1 2 3 or 3 3 3? Which lab decided that, and did you get it right there?
3. Three-level nesting where the innermost function reads the outermost's local. How many upvalues does the middle function have?
4. local function f() return f() end — is f inside the body a local or an upvalue?
5. In the naive version, how many allocations does function() end (capturing nothing) cost?
6. In the open/closed version, how many allocations does a function that captures one variable but is never returned cost?

Step 1: Capture Analysis (Compiler)

This is the same in both versions. Write resolve_upvalue, add_upvalue, and mark_captured per the closures chapter.

#![allow(unused)]
fn main() {
fn resolve(&mut self, name: &str) -> Access {
    if let Some(s) = self.resolve_local(name)   { return Access::Local(s); }
    if let Some(u) = self.resolve_upvalue(name) { return Access::Upval(u); }
    Access::Global
}
}

The order is the semantic: local, then upvalue, then global. Check globals first and every closure silently breaks. Add a test that a variable shadowing a global is not a global access.

Then Expr::Name gains its third case, and Expr::Function builds a proto with an upvals list:

$ ember disassemble -e 'local x = 1 local f = function() return x end'
--- proto [0]: <anonymous> ---
upvalues: [0] x  (parent local, slot 0)
0000     1  GET_UPVAL    0          ; x
0001     |  RETURN       1

That upvalues: block in the disassembly is what tells you capture analysis worked, before any of it runs. Add it to the disassembler now, not after you have a bug.

Step 2: Version 1 — Boxes at Declaration

Goal. The counter example works. Nothing is fast.

#![allow(unused)]
fn main() {
// Every CAPTURED local becomes a heap cell at declaration. Uncaptured locals
// are untouched — the compiler already knows which is which, via `mark_captured`.
Op::NewBox(slot) => {            // emitted right after a captured local's initializer
    let v = self.stack[base + slot as usize];
    let b = self.heap.alloc_upvalue(Upvalue::Closed(v))?;
    self.stack[base + slot as usize] = Value::Upvalue(b);   // the SLOT holds the box
}
Op::GetLocal(s) if self.is_boxed(s) => { /* deref */ }
}

Two hours, and it is correct. Run the warm-up:

$ diff <(ember run tests/golden/closures/counter.ember) <(lua tests/golden/closures/counter.lua)

Commit here, with the message lab-14: closures, version 1 (boxes at declaration). You need this commit to exist for Step 4's diff.

Note: Version 1 needs a temporary Value::Upvalue variant or a per-slot boxed flag, and it makes GET_LOCAL conditional. Both are things version 2 deletes. That is fine — this is scaffolding, and building scaffolding you will remove is a normal engineering activity that teaching materials rarely admit to.

Step 3: Version 2 — Open and Closed

Now replace it. NewBox disappears; GET_LOCAL becomes unconditional again; three things arrive:

#![allow(unused)]
fn main() {
// 1. The open list, and the ONE == that implements sharing.
fn find_or_create_open_upvalue(&mut self, slot: usize) -> Result<GcRef<Upvalue>> { /* … */ }

// 2. Closing, called from exactly two places.
fn close_upvalues(&mut self, from: usize) { /* … */ }

// 3. CLOSURE, capturing per the proto's descriptors.
Op::Closure(p) => { /* … */ }
}

And the compiler's end_scope gains the CLOSE_UPVALS-instead-of-POP branch from the upvalues chapter.

Get these three orderings right, because none of them errors when wrong:

Step 4: Measure, and Keep the Diff

git diff <version-1-commit> HEAD -- src/ > docs/learning/10-upvalues.diff
cargo bench --bench closures | tee docs/learning/closures-v2.txt

Three benchmarks, and the third is the one that shows why version 2 exists:

That third row is the whole argument. Under v1, if rare then return function() return x end end boxes x on every call, including the ~100% of calls that take the other branch.

The Trace

$ ember trace --upvalues -e '
local function make_counter()
  local count = 0
  local function inc() count = count + 1 return count end
  local function get() return count end
  return inc, get
end
local inc, get = make_counter()
inc() inc() inc()
return get()'

dep  ip    op                        open upvalues        note
  2  0000  LOAD_INT   0                []                 count → slot 0 of frame 2
  2  0001  CLOSURE    0                [Open(12)]         inc captures slot 12 (= base+0)
  2  0003  CLOSURE    1                [Open(12)]         get: SAME slot → SAME box ★
  2  0006  GET_LOCAL  1
  2  0007  GET_LOCAL  2
  2  0008  RETURN     2                []                 ← CLOSE_UPVALS ran first:
                                                            Open(12) → Closed(0)
  1  ....  CALL (inc)                  []
  3  0002  GET_UPVAL  0                []                 reads Closed(0) → 0
  3  0005  SET_UPVAL  0                []                 writes Closed(1)
  ...
  1  ....  CALL (get)
  4  0000  GET_UPVAL  0                []                 reads Closed(3)  ← SHARED
3

The line marked ★ is the mechanism. The second CLOSURE asks for slot 12 and the open list already has an entry, so it gets the same GcRef. That single lookup is why inc and get share count. Delete the s == slot branch and this trace shows two different upvalue handles and the program prints 0.

And the RETURN line is the other half. The upvalue transitions Open(12) → Closed(0) before the stack truncates, so the value is rescued rather than read from a dead slot.

Now the loop case, which is the one people get wrong:

$ ember trace --upvalues -e '
local fs = {}
for i = 1, 3 do
  local doubled = i * 2
  fs[i] = function() return doubled end
end
return fs[1]() .. "," .. fs[2]() .. "," .. fs[3]()'

  1  ....  CLOSURE   0     [Open(9)]     iteration 1: captures slot 9
  1  ....  CLOSE_UPVALS 1  []            ← BLOCK END closes it: Closed(2)
  1  ....  CLOSURE   0     [Open(9)]     iteration 2: slot 9 REUSED, but the old
                                          upvalue is already closed → a NEW box
  1  ....  CLOSE_UPVALS 1  []            Closed(4)
  1  ....  CLOSURE   0     [Open(9)]     iteration 3
  1  ....  CLOSE_UPVALS 1  []            Closed(6)
2,4,6

If CLOSE_UPVALS only ran at function return, all three closures would still point at one open upvalue on slot 9, which would close once with the last value, and the answer would be 6,6,6. That is the bug the block-exit closing prevents, and this trace is how you see it.

Expected Output

$ ember run -e 'local function mk() local c=0
                  return function() c=c+1 return c end, function() return c end end
                local inc, get = mk()
                inc() inc() inc() return get()'
3
$ ember run -e 'local fs={} for i=1,3 do fs[i]=function() return i end end
                return fs[1]()..","..fs[2]()..","..fs[3]()'
1,2,3
$ ember run -e 'local x=1
                local f = function() return function() return x end end
                return f()()'
1
$ ember disassemble -e 'local x=1 local f=function() return function() return x end end' \
  | grep -A2 'proto \[0\]'
--- proto [0]: <anonymous> ---
upvalues: [0] x  (parent local, slot 0)      ← the MIDDLE function carries it

That last one is the transitive-capture check: the middle function has an upvalue for a variable it never mentions.

Debugging Steps

get() returns 0 instead of 3

Sharing is broken. Either find_or_create_open_upvalue is not checking s == slot, or the compiler created two upvalue descriptors for one name (add_upvalue is not deduplicating).

The loop case prints 6,6,6

CLOSE_UPVALS is only emitted at function return, not at block exit. end_scope needs the captured branch.

Closures capture garbage after a return

Closing happens after stack.truncate. Swap them.

The middle function of a three-level nest has no upvalue

resolve_upvalue's recursive case (case 2) is missing, so capture is not transitive.

local function f() return f() end — f compiles to GET_GLOBAL

The local is declared after the function expression is compiled. This was Lab 10's rule and it bites again here, because now the wrong answer is an upvalue question rather than a slot question.

--gc-stress fails only in closure tests

A newly allocated Upvalue between alloc_upvalue and being stored in the closure is reachable only from the open list. Is the open list in enumerate_roots? (It is root set 4.)

The open list is not sorted and closing misses entries

Add debug_assert!(self.open_upvalues.is_sorted_by_key(...)) after every mutation.

Experiment

CLAIM. Capturing a variable costs nothing until a closure is actually created, in version 2, and costs an allocation per call in version 1.

METHOD. Write the declare_no_closure benchmark:

local function f(flag)
  local x = 42
  if flag then return function() return x end end
  return 0
end
for i = 1, 1000000 do f(false) end     -- the closure is NEVER created

Run it under both versions. Count allocations with --stats.

PREDICTION. How many allocations under v1? Under v2? What is the wall-clock ratio?

RESULT. Record it in docs/learning/10-upvalues.md alongside the diff. This is the number that justifies eighty lines of open-list machinery, and having produced it yourself is the difference between knowing the design and believing it.

Test

#![allow(unused)]
fn main() {
#[test]
fn two_closures_over_one_variable_share_it() {
    // THE test. Warm-up Experiment 3, in our runtime.
    assert_eq!(run("local function mk() local c=0
                      return function() c=c+1 return c end, function() return c end end
                    local inc, get = mk()
                    inc() inc() inc() return get()"), "3");
}

#[test]
fn separate_calls_get_separate_variables() {
    assert_eq!(run("local function mk() local c=0 return function() c=c+1 return c end end
                    local a, b = mk(), mk()
                    a() a() return b()"), "1");
}

#[test]
fn loop_variables_are_fresh_per_iteration() {
    // Decided in Lab 6's scope handling; OBSERVABLE only now.
    assert_eq!(run("local fs={} for i=1,3 do fs[i]=function() return i end end
                    return fs[1]()..fs[2]()..fs[3]()"), "123");
    // And a block-scoped local inside the loop, which needs CLOSE_UPVALS at BLOCK exit:
    assert_eq!(run("local fs={} for i=1,3 do local d=i*2 fs[i]=function() return d end end
                    return fs[1]()..fs[2]()..fs[3]()"), "246");
}

#[test]
fn capture_is_transitive_through_a_middle_function() {
    assert_eq!(run("local x=1
                    local f=function() return function() return x end end
                    return f()()"), "1");
    let p = compile_proto("local x=1 local f=function() return function() return x end end");
    assert_eq!(p.protos[0].upvals.len(), 1, "the middle function must carry the capture");
}

#[test]
fn upvalues_are_deduplicated_per_proto() {
    // Two mentions of `x` must be ONE upvalue, or the read and the write use
    // different boxes and the write is invisible.
    let p = compile_proto("local x=1 local f=function() x = x + 1 return x end");
    assert_eq!(p.protos[0].upvals.len(), 1);
    assert_eq!(run("local x=1 local f=function() x=x+1 return x end f() return x"), "2");
}

#[test]
fn upvalues_close_before_the_stack_truncates() {
    // If close ran after truncate, this returns garbage or nil.
    assert_eq!(run("local function mk() local v='captured' return function() return v end end
                    return mk()()"), "captured");
}

#[test]
fn closures_survive_gc_stress() {
    with_gc_stress(|| {
        assert_eq!(run("local function mk() local c=0 return function() c=c+1 return c end end
                        local f=mk() f() f() return f()"), "3");
    });
}
}

Challenge Extensions

1. Intrusive open list. Replace the Vec with a linked list threaded through the Upvalue objects, as Lua does. No auxiliary allocation, no reallocation. Measure — at realistic list lengths it will probably lose to the Vec, and finding that out is the point.
2. Escape analysis. Detect closures that provably do not outlive their frame (created, called, and dropped within one expression) and keep their captures on the stack entirely. Report how often it fires on your corpus.
3. _ENV. Now that upvalues exist, implement Lua's real global design: x desugars to _ENV.x, _ENV is an upvalue, and load(chunk, env) becomes trivial. Then decide whether to keep it — this is the ADR-012 upgrade path, and superseding an ADR is good practice.
4. Upvalue debug names in tracebacks. in upvalue 'inner' is what Lua prints. Make yours match.
5. A closure-cycle demo. Write the smallest program whose closures form a cycle, show that collectgarbage() frees it, and put it in tests/golden/gc/. You will need it in Lab 15.

Deliverables

• Capture analysis: local / upvalue / global, in that order, with dedup and transitivity.
• The disassembler prints each proto's upvalues: block.
• Version 1 committed separately before version 2 exists.
• Version 2: open list with sharing, close_upvalues, CLOSE_UPVALS at block exit and return, and on break/return out of captured scopes.
• All three orderings correct, each with a test.
• docs/learning/10-upvalues.diff — the v1→v2 diff.
• The declare_no_closure benchmark under both versions, recorded.
• --trace-upvalues shows the open list and every close event.
• The open list is in enumerate_roots; closure tests pass under --gc-stress.
• Differential tests green.
• docs/learning/09-closures.md and 10-upvalues.md written.

Validation / Self-check

1. Which single line makes two closures share a variable? What happens if you delete it?
2. Why must closing happen before stack.truncate, and why is there no error if it does not?
3. Give the program that distinguishes "close at block exit" from "close at return only".
4. In a three-level nest, what does the middle function's proto contain, and why?
5. Why must add_upvalue deduplicate? Give the program that exposes it.
6. What does version 1 allocate that version 2 does not? Give your measured numbers.
7. Why is the open list a GC root set, and what breaks under --gc-stress without it?
8. Why is local function f different from local f = function, in terms of upvalues?

Next: Lab 15 — Mark and Sweep.