Garbage collection with Zig's Allocator interface

Zig provides an interface for memory allocators in the form of std.mem.Allocator. Some allocators implement this interface directly (see for example std.heap.PageAllocator), but you can also use it to wrap an existing allocator. This is what std.heap.debug_allocator does, for example.

The way this generally works is you create your own custom allocator struct and then write a function that takes your allocator and outputs a std.mem.Allocator, which is a struct in this format:

  ptr: *anyopaque,
  vtable: *const VTable,

  pub const VTable = struct {
      alloc: *const fn (...) ?[*]u8,
      resize: *const fn (...) bool,
      remap: *const fn (...) ?[*]u8,
      free: *const fn (...) void,
  };

(Arguments omitted here to make it easier to read.) The underlying allocator will already have such a VTable, whose functions are conveniently also available in the interface under rawAlloc, rawResize, rawRemap, rawFree.

The example of std.head.ThreadSafeAllocator is instructive:

  child_allocator: Allocator,
  mutex: std.Thread.Mutex = .{},

  pub fn allocator(self: *ThreadSafeAllocator) Allocator {
      return .{
          .ptr = self,
          .vtable = &.{
              .alloc = alloc,
              .resize = resize,
              .remap = remap,
              .free = free,
          },
      };
  }

  fn alloc(ctx: *anyopaque, n: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
      const self: *ThreadSafeAllocator = @ptrCast(@alignCast(ctx));
      self.mutex.lock();
      defer self.mutex.unlock();

      return self.child_allocator.rawAlloc(n, alignment, ra);
  }

  // (...)

Other than the pointer casting this is pretty straightforward: acquire a lock, run rawAlloc() in the underlying allocator (which in turn will call alloc()), release the lock.

This interface proved useful to me while working my way through Crafting Interpreters, specifically the Garbage Collection chapter. In the book the VM is implemented in C, and memory management is implemented as a few wrappers to realloc and free (see memory.c, memory.h). Additionally there's a few macros for handling resizable arrays. Before reaching this chapter, my Zig implementation didn't have any wrappers, although I did stick to the allocator interface instead of hardcoding a particular allocator. I also didn't feel the need to implement resizable arrays since Zig's ArrayList can already handle that.

For Lox's garbage collector, we need two things: first, we need to hold a reference to the VM (including the compiler), so we can traverse the objects it allocates and perform the mark-and-sweep algorithm. We also need to keep track of how much memory it's allocating so we can decide when to perform garbage collection. Essentially it looks like this:

  fn alloc(ctx: *anyopaque, n: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
      const self: *Self = @ptrCast(@alignCast(ctx));

      if (self.shouldGC()) {
          self.collectGarbage();
      }

      return self.child_allocator.rawAlloc(n, alignment, ra);
  }

  fn shouldGC(self: Self) bool {
      return self.debug_stress or self.bytes_allocated > self.next_gc;
  }

  fn collectGarbage(self: *Self) void {
      self.markRoots(self.vm);
      self.markCompilerRoots(self.compiler);
      self.traceReferences();
      self.tableRemoveWhite(compiler.vm, &compiler.vm.strings);
      self.sweep(compiler.vm);
      self.next_gc = self.bytes_allocated * gc_heap_growth_factor;
  }

resize() and remap() can also allocate memory, so we'll want to run collectGarbage() there as well. One note about keeping track of the allocated memory: these functions (including alloc()) don't always allocate, so we need to check the return value before updating the counter. For example, alloc():

  fn alloc(ctx: *anyopaque, n: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
      const self: *Self = @ptrCast(@alignCast(ctx));

      if (self.shouldGC()) {
          self.collectGarbage();
      }

      const ptr = self.child_allocator.rawAlloc(n, alignment, ra);

      if (ptr != null) {
          self.bytes_allocated += len;
      }

      return ptr;
  }

Another bit of nuance (which is covered in the book) is that at times we need to prevent some memory from being garbage-collected, for example if we're in the middle of a partly initialized object; in these cases the allocated memory won't be marked so it'd get swept if GC runs. The book uses a few tricks where it temporarily puts values in the stack so the memory gets marked, which I suppose is fine but I've found it more convenient to temporarily disable the GC in such situations. This works quite well in conjunction with Zig's defer; whenever I disable GC I can put a defer gc.is_disabled = false and be sure it'll be enabled again when the function returns.