57. Closures and upvalues

The counter, again

You met this program in chapters 6 and 41.

fn makeCounter() {
  let n = 0;
  fn count() {
    n = n + 1;
    return n;
  }
  return count;
}
let a = makeCounter();
print a(); // 1
print a(); // 2

The tree-walker got this almost for free: every call made an Environment on the heap, and count kept a reference to it.

In the VM, n is a stack slot at a fixed offset from the frame pointer (chapter 54). When makeCounter returns, OP_RETURN sets the stack top back to the frame’s base, and the next call writes over the slot.

So count can’t just remember “slot 2”. By the time a() runs, slot 2 belongs to someone else.

Upvalues: found at compile time

Each closure gets a list of small objects, one per variable it uses from an enclosing function: its upvalues. The body says “my upvalue 0”, and the upvalue knows where the value is.

The compiler already tracks every function’s locals (chapter 54), so it can find these variables as it goes. When a name isn’t a local of the current function, namedVariable tries resolveUpvalue before falling back to a global:

function resolveUpvalue(c, name) {
  if (!c.enclosing) return -1;
  const local =
    resolveLocal(c.enclosing, name);
  if (local !== -1) {
    c.enclosing.locals[local]
      .isCaptured = true;
    return addUpvalue(c, local, true);
  }
  const up =
    resolveUpvalue(c.enclosing, name);
  if (up !== -1) {
    return addUpvalue(c, up, false);
  }
  return -1;
}

Two cases. If the name is a local of the function directly around this one, the upvalue captures that local’s slot (isLocal true). Otherwise the function recurses outward. If some outer function has it, the enclosing function gets an upvalue of its own, and this function’s upvalue copies that one (isLocal false). addUpvalue appends { index, isLocal } to the function’s list, unless an equal entry is already there, and returns its position.

When the function is done, the compiler writes that list into the bytecode as byte pairs after OP_CLOSURE. The companion’s disassembler (chapter 46) shows them:

== makeCounter ==
0000    2 OP_CONSTANT         0 '0'
0002    6 OP_CLOSURE          1 <fn count>
0004    |                     local 1
== count ==
0000    4 OP_GET_UPVALUE      0
0002    | OP_CONSTANT         0 '1'
0004    | OP_ADD
0005    | OP_SET_UPVALUE      0

local 1 means: capture slot 1 of the running frame, which is n (slot 0 holds the function itself). Inside count, n became OP_GET_UPVALUE 0 and OP_SET_UPVALUE 0. With three levels, say inner inside middle inside outer, where only inner uses outer’s x, the middle function passes it along: middle gets local 1, and inner gets upvalue 0, meaning “take my enclosing closure’s upvalue 0”.

Open: pointing into the stack

At run time, OP_CLOSURE builds the closure and fills in its upvalues, one pair at a time:

case OP.CLOSURE: {
  const cl = newClosure(readConstant());
  push(cl);
  const ups = cl.upvalues;
  for (let i = 0; i < ups.length; i++) {
    const isLocal = readByte();
    const index = readByte();
    ups[i] = isLocal
      ? captureUpvalue(frame.slots + index)
      : frame.closure.upvalues[index];
  }
  break;
}

frame.slots + index turns a frame-relative slot into an absolute stack index. The upvalue object (ObjUpvalue) has just three fields: slot, closed and nextOpen. While makeCounter runs, n really is in slot 2, so the upvalue records slot: 2, and reading through it reads the stack. That’s an open upvalue.

There is one subtlety. Two closures made in the same call must share n, not each get a copy. So the VM keeps every open upvalue on a list, openUpvalues, and captureUpvalue looks there first:

captureUpvalue(slot: number) {
  let prev = null;
  let u = this.openUpvalues;
  while (u && u.slot > slot) {
    prev = u;
    u = u.nextOpen;
  }
  if (u && u.slot === slot) return u;
  const made = heap.newUpvalue(slot);
  made.nextOpen = u;
  if (prev) prev.nextOpen = made;
  else this.openUpvalues = made;
  return made;
}

The list is sorted by slot, highest first. The search stops once it passes the slot it wants, and links a new upvalue in at that spot.

Closed: moving to the heap

Now step through what happens when makeCounter returns:

fn makeCounter() {
  let n = 0;
  fn count() {
    n = n + 1;
    return n;
  }
  return count;
}
let a = makeCounter();
print a();

The move happens at the last possible moment. OP_RETURN calls closeUpvalues(frame.slots) before it drops the frame:

closeUpvalues(last: number) {
  while (this.openUpvalues &&
         this.openUpvalues.slot >= last) {
    const u = this.openUpvalues;
    u.closed = this.stack[u.slot];
    u.slot = -1;
    this.openUpvalues = u.nextOpen;
  }
}

Every open upvalue at or above the frame’s base points into the dying frame. The loop copies each value into the upvalue’s own closed field, marks it closed, and unlinks it. Because the list is sorted, the first upvalue below last ends the loop: everything after it is lower still. From then on, readUpvalue sees slot === -1 and returns closed instead of stack[slot].

A function’s return isn’t the only time a local dies. A local declared in a block dies at the block’s }. That’s why resolveUpvalue set isCaptured: when endScope discards a local, it emits OP_POP for an ordinary one and OP_CLOSE_UPVALUE for a captured one, which closes it (closeUpvalues(sp - 1)) and then pops it. So a closure made in a loop body that captures a body local gets a fresh upvalue on every pass. The extra exercise at the end puts that to work.

Your turn

function capture(vm, slot) {
  // reuse an open upvalue for slot,
  // or make and remember a new one
}

function closeFrom(vm, last) {
  // close upvalues with slot >= last
}

function read(vm, u) {
  return u.slot === -1
    ? u.closed
    : vm.stack[u.slot];
}

let f1;
let f2;
let f3;
for (let i = 1; i <= 3; i = i + 1) {
  fn show() { print i; }
  if (i == 1) f1 = show;
  if (i == 2) f2 = show;
  if (i == 3) f3 = show;
}
f1();
f2();
f3();

What you learned

• In a stack VM, a local’s slot is reused once its call returns, so a closure needs an upvalue: a small object standing between it and the variable.
• The compiler resolves upvalues at compile time and emits (isLocal, index) pairs after OP_CLOSURE: capture a local of the enclosing frame, or reuse the enclosing closure’s upvalue.
• An open upvalue points at a live stack slot. The VM keeps open upvalues on one sorted list and reuses them, so closures share variables.
• A closed upvalue has copied its value into itself. Closing happens at OP_RETURN and at OP_CLOSE_UPVALUE, just before the slot dies.

Next: the VM has been allocating closures, upvalues, strings and instances, and freeing none of them. Time for a garbage collector.