159. A garbage collector in Rust

Who decides to collect?

Chapter 58’s TypeScript heap collects from inside allocate. It finds the roots through heap.vm, the compiler’s functions, and an extraRoots stack that heap.push(v) fills to protect a half-built value. Chapter 145 showed why drip-rs can’t: the heap is part of the VM and can’t borrow it back. So the caller decides. Adapted from src/vm.rs:

fn alloc<T: HeapObject>(&mut self, obj: T)
    -> Gc<T> {
    if self.heap.should_collect() {
        self.collect_garbage();
    }
    self.heap.alloc(obj)
}

Heap::alloc itself never collects. should_collect() is true once bytes_allocated passes next_gc, or always in stress mode. The price is one rule, in the comment on Vm::alloc: everything the new object refers to must already be reachable from a root. That is what chapter 158’s closure argument checked.

Roots are trait objects

Anything that can report the handles it holds implements one trait, from src/gc.rs:

pub trait Trace {
    fn trace(&self, marks: &mut Marks);
}

impl Trace for Value {
    fn trace(&self, m: &mut Marks) {
        match *self {
            Value::Closure(c) => m.mark(c),
            // ...an arm per object kind,
            // none for nil, bool, number
        }
    }
}

There are impls for Gc<T>, StrRef, Table, Option<T>, [T] and Vec<T>, so self.stack.trace(marks) marks a whole stack. The collector takes its roots as a slice of trait objects ^[2]:

pub fn collect(&mut self,
               roots: &[&dyn Trace])

dyn is needed because one call mixes types. The VM passes &[&roots], its VmRoots of borrowed fields (stack, frames, globals, open upvalues, the "init" string). The compiler passes &[self.roots, &self.functions]: whatever the VM gave it, plus its own Vec<FunctionState> of half-built functions (chapter 152). --disassemble, with no VM at all, passes &().

struct Marks {
    bits: Vec<bool>,
}

trait Trace {
    fn trace(&self, m: &mut Marks);
}

// A handle: an arena index.
struct Handle(usize);

impl Trace for Handle {
    fn trace(&self, m: &mut Marks) {
        m.bits[self.0] = true;
    }
}

impl<T: Trace> Trace for Vec<T> {
    fn trace(&self, m: &mut Marks) {
        for t in self {
            t.trace(m);
        }
    }
}

// No roots at all.
impl Trace for () {
    fn trace(&self, _: &mut Marks) {}
}

// Borrowed fields, as in VmRoots.
struct VmRoots<'a> {
    stack: &'a Vec<Handle>,
    globals: &'a Vec<Handle>,
}

impl Trace for VmRoots<'_> {
    fn trace(&self, m: &mut Marks) {
        self.stack.trace(m);
        self.globals.trace(m);
    }
}

fn mark(
    n: usize,
    roots: &[&dyn Trace],
) -> Vec<bool> {
    let mut m = Marks {
        bits: vec![false; n],
    };
    for r in roots {
        r.trace(&mut m);
    }
    m.bits
}

fn main() {
    let s = vec![Handle(0), Handle(2)];
    let g = vec![Handle(5)];
    let vm = VmRoots {
        stack: &s,
        globals: &g,
    };
    // The compiler's half-built work:
    let compiling = vec![Handle(3)];
    println!("{:?}",
        mark(6, &[&vm, &compiling]));
    println!("{:?}", mark(6, &[&()]));
}

Mark bits outside the objects

TypeScript objects have an isMarked field in their header. Put a mark bit in the arena and try to mark what an object refers to:

fn blacken(&mut self, i: usize) {
    for &r in &self.nodes[i].refs {
        self.mark(r); // sets a mark bit
    }
}

rustc 1.97 (trimmed):

error[E0502]: cannot borrow `*self` as
  mutable because it is also borrowed
  as immutable
13 | for &r in &self.nodes[i].refs {
   |           ------------------- immutable
14 |     self.mark(r);
   |     ^^^^^^^^^^^^ mutable borrow

Tracing reads objects; marking writes bits. With both in one struct, the read borrow and the write borrow overlap. So drip-rs keeps the bits in a separate Marks value, one Vec<bool> per arena plus the gray stack, and lifts it out of the heap while it works. Adapted from Heap::collect:

let mut marks =
    std::mem::take(&mut self.marks);
marks.reset(self);
for root in roots {           // 1. mark
    root.trace(&mut marks);
}
while let Some(o) = marks.gray.pop() {
    self.blacken(o, &mut marks); // 2.
}
self.interned.retain_keys(    // 3.
    |s| marks.strings[s.handle.index()]);
self.sweep(&marks);           // 4.
self.marks = marks;

mem::take swaps an empty Marks in and hands back the real one ^[3]. Now blacken(&self, o, &mut marks) borrows the heap shared and the bitmap mutably, two separate values. reset resizes every bit vector to its arena and clears it, so each collection starts from a fresh bitmap, and sweep never has to unmark survivors, as the TypeScript sweep does. The vectors’ memory is reused from one collection to the next.

The gray stack holds a Gray, an enum with one variant per object kind that can refer to others. Strings and natives have none: marking one turns it black at once, where TypeScript pushes every object gray.

Sweep and the free list

Sweep runs the same loop over all eight arenas (stamped out by a small macro_rules!), freeing with Arena::remove, from src/object.rs:

pub fn remove(&mut self, index: usize,
              reuse: bool) -> Option<T> {
    let object = self.slots[index].take();
    if reuse && object.is_some() {
        self.free.push(index as u32);
    }
    object
}

The slot becomes None and the object is dropped, with its Box<str> or its tables. The TypeScript sweep unlinks the object from its list instead. Pacing is the same in both VMs: subtract what was freed from bytes_allocated, and set next_gc to twice what is left, but at least 1 MiB.

A collection before every allocation

A missing root is the classic collector bug, and it hides: it bites only if a collection runs at the wrong moment. Stress mode makes every moment that moment: should_collect() is always true, so the heap collects before every allocation. It also sets poison_freed, so remove never puts a slot on the free list. A stale handle then finds None and panics with “use after free”, instead of quietly reading whatever object moved in. That does the job of the TypeScript heap’s poisoning, which deletes a dead object’s fields.

tests/stress_gc.rs runs the 95 corpus programs in closure/, class/, inheritance/, string/, function/, super/, method/ and field/ this way. To see it catch something, I deleted one line from a scratch copy of drip-rs, self.open_upvalues.trace(marks) in VmRoots. The normal corpus suite still passed, 246 of 246. The stress suite panicked on its sixth program, class/local_reference_self.drip (reformatted a little):

{
  class Foo {
    returnSelf() { return Foo; }
  }
  print Foo().returnSelf();
}

test class/local_reference_self.drip
  ... FAILED
---- class/local_reference_self.drip
  panicked: use after free: an
  object of type ObjUpvalue was
  collected while still in use

The method captures the local Foo. At OP_CLOSURE, the new upvalue was reachable only from the open list, now unrooted, so the closure’s own allocation freed it. The backtrace showed the panic later, inside blacken, when a collection traced the method closure’s dangling handle.

Your turn

struct Obj {
    refs: Vec<usize>,
}

struct Heap {
    objs: Vec<Option<Obj>>,
    free: Vec<usize>,
    // Never reuse a freed slot.
    poison: bool,
    freed: usize,
}

impl Heap {
    fn new(poison: bool) -> Heap {
        Heap {
            objs: Vec::new(),
            free: Vec::new(),
            poison,
            freed: 0,
        }
    }

    fn alloc(
        &mut self,
        refs: Vec<usize>,
    ) -> usize {
        let obj = Some(Obj { refs });
        let top = self.free.pop();
        if let Some(i) = top {
            self.objs[i] = obj;
            return i;
        }
        self.objs.push(obj);
        self.objs.len() - 1
    }

    fn get(&self, h: usize) -> &Obj {
        let o = self.objs[h].as_ref();
        o.expect("use after free")
    }

    // Free every object the roots
    // can't reach.
    fn collect(
        &mut self,
        roots: &[usize],
    ) {
        todo!()
    }
}

What you learned

• In drip-rs, allocation never collects. The VM asks should_collect() and calls collect with its roots, so a new object’s references must already be reachable.
• Roots are &dyn Trace values: the VM’s borrowed fields, the compiler’s functions, or ().
• Mark bits live in a separate bitmap, taken out of the heap with mem::take and cleared at each collection, so tracing (reads) and marking (writes) never borrow the same thing.
• Sweep frees unmarked slots and pushes them on per-arena free lists; the weak string table drops unmarked strings first.
• Stress mode collects before every allocation and never reuses a slot, so a missing root panics at once. 95 programs run this way.