158. Closures and upvalues in Rust

A counter that outlives its frame

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

n is a local of makeCounter, in a stack slot. By the time c() runs, that frame has returned and its slots are gone, yet count still reads and writes n. Chapter 57’s answer was the upvalue: an object that points at the variable’s slot while the slot exists, and holds the variable itself afterwards. drip-rs prints 1 and 2, as the TypeScript VM does. What changes is how an upvalue is written down.

Two states, one enum

The TypeScript upvalue, from object.ts:

interface ObjUpvalue extends ObjHeader {
  type: 'upvalue';
  slot: number;   // -1 once closed
  closed: Value;
  nextOpen: ObjUpvalue | null;
}

And drip-rs’s, from src/object.rs:

pub enum ObjUpvalue {
    /// Still on the stack, in this slot.
    Open(usize),
    /// Its frame returned: it holds the
    /// variable itself.
    Closed(Value),
}

The TypeScript object has three fields for two states, and the rule “closed means something only when slot is -1” lives in your head. The enum makes the states exclusive: an open upvalue has no stale closed value, and a closed one has no slot (chapter 128). Reading one is a match, adapted from the GetUpvalue and SetUpvalue arms:

let c = self.closure_register();
let up = self.heap[c].upvalues[i];
// GetUpvalue:
let value = match self.heap[up] {
    ObjUpvalue::Open(slot) =>
        self.stack[slot],
    ObjUpvalue::Closed(v) => v,
};
// SetUpvalue, with value = peek(0):
match &mut self.heap[up] {
    ObjUpvalue::Open(slot) =>
        self.stack[*slot] = value,
    ObjUpvalue::Closed(v) => *v = value,
}

closure_register is the top frame’s closure handle. A TypeScript closure holds its upvalues directly; here it is two lookups, the closure and then the upvalue. In SetUpvalue, &mut self.heap[up] and self.stack[*slot] are borrowed at once: different fields, so that’s fine.

#[derive(Clone, Copy, Debug)]
enum Up {
    Open(usize),
    Closed(f64),
}

struct Closure {
    name: &'static str,
    ups: Vec<usize>, // handles
}

fn main() {
    let mut stack = vec![0.0, 5.0];
    // The upvalue arena.
    let mut ups = vec![Up::Open(1)];
    // Both capture slot 1: handle 0.
    let inc = Closure {
        name: "inc",
        ups: vec![0],
    };
    let get = Closure {
        name: "get",
        ups: vec![0],
    };
    println!("{:?}", ups[get.ups[0]]);

    // The frame returns: close it.
    if let Up::Open(s) = ups[0] {
        ups[0] = Up::Closed(stack[s]);
    }
    stack.truncate(1);

    // inc writes through its handle...
    let h = inc.ups[0];
    if let Up::Closed(v) = &mut ups[h] {
        *v += 1.0;
    }
    // ...and get sees the write.
    let h = get.ups[0];
    println!("{} sees {:?}", get.name,
        ups[h]);
    println!("{} wrote it", inc.name);
}

The open list is a Vec

The VM must find an existing upvalue for a slot, so that two closures over n share it, and it must close the right ones when a frame ends. TypeScript links the open upvalues through nextOpen, highest slot first. drip-rs keeps their handles in a Vec sorted by slot, lowest first, so the end of the Vec matches the top of the stack. Adapted from src/vm.rs:

fn capture_upvalue(&mut self, slot: usize)
    -> Gc<ObjUpvalue> {
    let mut i = self.open_upvalues.len();
    while i > 0 {
        let up = self.open_upvalues[i - 1];
        let s = self.open_slot(up);
        if s == slot { return up; }
        if s < slot { break; }
        i -= 1;
    }
    let new = self.alloc(Open(slot));
    self.open_upvalues.insert(i, new);
    new
}

fn close_upvalues(&mut self, last: usize) {
    while let Some(&up) =
        self.open_upvalues.last()
    {
        let slot = self.open_slot(up);
        if slot < last { break; }
        let v = self.stack[slot];
        self.heap[up] = Closed(v);
        self.open_upvalues.pop();
    }
}

(Open and Closed are short for the ObjUpvalue:: variants.) open_slot reads the slot of an Open upvalue; a Closed one in the open list would be a VM bug, so it panics. close_upvalues(last) runs at Return, with the frame’s base, and at CloseUpvalue, with the top slot of a block that ends (chapter 155). The assignment replaces the whole object in its arena slot. The handle stays the same, so every closure holding it now reads the Closed value. Step through the counter:

Capture first, then allocate

The TypeScript CLOSURE case allocates the closure first, with an array of nulls, pushes it so the collector can see it, and then fills each entry with captureUpvalue. Try the same order in a toy Rust VM:

self.closures.push(Closure::default());
let c = self.closures.last_mut().unwrap();
for &s in slots {
    c.upvalues.push(self.capture(s));
}

rustc 1.97 (trimmed):

error[E0499]: cannot borrow `*self` as
  mutable more than once at a time
16 | let c = self.closures.last_mut()...
   |         ------------- first mutable
   |            borrow occurs here
18 | c.upvalues.push(self.capture(s));
   | ----------      ^^^^ second mutable
   | |                    borrow occurs here
   | first borrow later used here

capture needs all of self, because it may allocate, while c still borrows the closure. So drip-rs reverses the order, adapted from Vm::closure:

let n = self.heap[func].upvalue_count;
let mut upvalues = Vec::with_capacity(n);
for _ in 0..n {
    let is_local = self.read_byte() == 1;
    let index = self.read_byte() as usize;
    upvalues.push(if is_local {
        let s = self.slots + index;
        self.capture_upvalue(s)
    } else {
        let c = self.closure_register();
        self.heap[c].upvalues[index]
    });
}
let closure = self.alloc(ObjClosure {
    function: func, upvalues });
self.push(Value::Closure(closure));

The type gets simpler too: Vec<Gc<ObjUpvalue>>, with no nulls, because the list is complete before the closure exists.

Is it safe? upvalues is a local Vec, invisible to the collector, and both capture_upvalue and alloc may collect. The comment in vm.rs gives the argument. Every handle in upvalues is either in open_upvalues, which is a root, or in the enclosing closure, which the running frame keeps alive. The function comes from the running chunk’s constants, which are traced too. Chapter 159 shows what happens when the first half of that argument breaks.

Your turn

#[derive(Clone, Copy, Debug, PartialEq)]
enum Up {
    Open(usize),
    Closed(f64),
}

struct Vm {
    stack: Vec<f64>,
    // The upvalue arena: a handle is
    // an index into it.
    ups: Vec<Up>,
    // Open handles, sorted by slot,
    // lowest first.
    open: Vec<usize>,
}

impl Vm {
    fn new(stack: Vec<f64>) -> Vm {
        Vm {
            stack,
            ups: Vec::new(),
            open: Vec::new(),
        }
    }

    // The slot of an open upvalue.
    fn slot(&self, h: usize) -> usize {
        match self.ups[h] {
            Up::Open(s) => s,
            _ => unreachable!(),
        }
    }

    // The open upvalue for slot: the
    // existing one, if there is one.
    fn capture(
        &mut self,
        slot: usize,
    ) -> usize {
        todo!()
    }

    // Close every open upvalue at or
    // above stack slot last.
    fn close(&mut self, last: usize) {
        todo!()
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
enum Up {
    Open(usize),
    Closed(f64),
}

struct Vm {
    stack: Vec<f64>,
    ups: Vec<Up>,
}

impl Vm {
    // GetUpvalue: read the variable.
    fn get(&self, h: usize) -> f64 {
        todo!()
    }

    // SetUpvalue: write the variable.
    fn set(
        &mut self,
        h: usize,
        v: f64,
    ) {
        todo!()
    }
}

What you learned

• drip-rs’s ObjUpvalue is an enum: Open(slot) while the variable is on the stack, Closed(value) once its frame returns.
• Closures hold upvalue handles. Closing overwrites the object in its arena slot, so every closure sharing the handle sees the change.
• The open upvalues are a Vec of handles sorted by slot, lowest first: capture searches from the end, and close pops from the end.
• drip-rs captures every upvalue before allocating the closure: the borrow checker forbids the other order, and every captured handle is already reachable from a root.