Hand every C library's allocator hook the same C heap

C's malloc, calloc, realloc and free over a Zig allocator were written three
times: syntax.zig's for tree-sitter, pdf.zig's for MuPDF, and gui.zig's for
FreeType and HarfBuzz. They are now one, src/c_heap.zig's Heap. Each hook
names the allocator variable its heap reads; gui.zig exports its heap under
the ui_malloc names font.c and HarfBuzz call. It is a module of its own
because pdf.zig is the MuPDF module, and a file cannot belong to that module
and the core's at once.

A block starts with a 16-byte header holding its size and the allocator that
made it, so it goes back where it came from even after the heap is repointed,
as the fonts' copy did. The fonts' copy carried the 16-byte Allocator itself
in a 32-byte header; here the header holds the allocator's index in a table of
every allocator a heap has used (entries are published once and never change,
and a heap remembers where its current one was last found). The difference is
measurable: with the 32-byte header, pardes-pdf-bench's filtered page render
at 96 dpi took 1.1% longer than before this change, slower in 15 of 16 pinned
rounds in two separate runs, and highlighting all of src/pardes.zig (305k
tree-sitter blocks) up to 0.6% longer, against an A/A pair within 0.7%. With
the 16-byte header both are back within the A/A pair's noise.

Two other things change. For tree-sitter and MuPDF, a block made before the
allocator was repointed now goes back to the one that made it rather than
through the current one. For the fonts, realloc(block, 0) now frees and
returns null, as glibc's does and as the other two copies already did;
neither FreeType nor HarfBuzz asks for it. The heap's test replaces the
tests of the old copies.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

4 files changed, 165 insertions, 182 deletions

diff --git a/src/c_heap.zig b/src/c_heap.zig
new file mode 100644
index 00000000..2a3e865f
--- /dev/null
+++ b/src/c_heap.zig
@@ -0,0 +1,135 @@
+//! C's malloc, calloc, realloc and free over a Zig allocator: what every C
+//! library that takes an allocator hook is handed (tree-sitter, MuPDF,
+//! FreeType, HarfBuzz, SDL), so what they allocate shows up in a Debug
+//! build's leak report and under a test's std.testing.allocator like the
+//! rest of pardes. A module of its own because pdf.zig, the MuPDF module,
+//! uses it as well as the core.
+const std = @import("std");
+
+/// C's free and realloc pass no size, so a block starts with its size and
+/// the allocator that made it: a block goes back where it came from even
+/// after the heap is repointed (at a test's std.testing.allocator, or away
+/// from the default heap once a library's process-wide objects exist). The
+/// allocator is named by its place in `allocators` rather than carried
+/// whole, which keeps the header at 16 bytes, max_align_t's alignment:
+/// carrying the 16-byte Allocator itself made a 32-byte header, and that
+/// measurably slowed tree-sitter, which makes a block per syntax node.
+const Block = struct { size: usize, allocator: usize };
+const header = 16;
+comptime {
+    std.debug.assert(@sizeOf(Block) <= header);
+}
+
+/// Every allocator any heap has made a block from. An entry is written once,
+/// before `allocators_len` publishes it, and never changes, so a block's
+/// index stays good for the life of the process. Few allocators ever back a
+/// heap: the gpa, the subsystem allocators, a test's.
+var allocators: [64]std.mem.Allocator = undefined;
+var allocators_len: std.atomic.Value(usize) = .init(0);
+var allocators_lock: std.atomic.Mutex = .unlocked;
+
+/// The four functions over whatever `current` holds when a block is made.
+/// A hook is installed once and `current` may be repointed later, but not
+/// while another thread may be allocating from it.
+pub fn Heap(comptime current: *const std.mem.Allocator) type {
+    return struct {
+        /// Where `current` was last found in `allocators`: checked, not
+        /// trusted, on every malloc, so a repointed `current` is looked up
+        /// (and added) again.
+        var last: std.atomic.Value(usize) = .init(0);
+
+        pub fn malloc(size: usize) callconv(.c) ?*anyopaque {
+            const total = std.math.add(usize, header, size) catch return null;
+            const bytes = current.alignedAlloc(u8, .@"16", total) catch return null;
+            var index = last.load(.monotonic);
+            if (index >= allocators_len.load(.acquire) or
+                allocators[index].ptr != current.ptr or allocators[index].vtable != current.vtable)
+            {
+                while (!allocators_lock.tryLock()) std.atomic.spinLoopHint();
+                defer allocators_lock.unlock();
+                const len = allocators_len.load(.monotonic);
+                index = for (allocators[0..len], 0..) |known, at| {
+                    if (known.ptr == current.ptr and known.vtable == current.vtable) break at;
+                } else len;
+                if (index == len) {
+                    if (len == allocators.len) @panic("c_heap: more distinct allocators than it can name");
+                    allocators[len] = current.*;
+                    allocators_len.store(len + 1, .release);
+                }
+                last.store(index, .monotonic);
+            }
+            @as(*Block, @ptrCast(bytes.ptr)).* = .{ .size = size, .allocator = index };
+            return bytes.ptr + header;
+        }
+
+        pub fn calloc(count: usize, size: usize) callconv(.c) ?*anyopaque {
+            const total = std.math.mul(usize, count, size) catch return null;
+            const block: [*]u8 = @ptrCast(malloc(total) orelse return null);
+            @memset(block[0..total], 0);
+            return block;
+        }
+
+        /// realloc(block, 0) frees the block and returns null, as glibc's
+        /// does: every library handed this heap was written against glibc.
+        pub fn realloc(block: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque {
+            const old = block orelse return malloc(size);
+            if (size == 0) {
+                free(old);
+                return null;
+            }
+            const base: [*]align(16) u8 = @alignCast(@as([*]u8, @ptrCast(old)) - header);
+            const made: *Block = @ptrCast(base);
+            const total = std.math.add(usize, header, size) catch return null;
+            const bytes = allocators[made.allocator].realloc(base[0 .. header + made.size], total) catch return null;
+            @as(*Block, @ptrCast(bytes.ptr)).size = size;
+            return bytes.ptr + header;
+        }
+
+        pub fn free(block: ?*anyopaque) callconv(.c) void {
+            const old = block orelse return;
+            const base: [*]align(16) u8 = @alignCast(@as([*]u8, @ptrCast(old)) - header);
+            const made: *Block = @ptrCast(base);
+            allocators[made.allocator].free(base[0 .. header + made.size]);
+        }
+    };
+}
+
+var test_allocator: std.mem.Allocator = undefined;
+
+test "a C heap block keeps its bytes, frees exactly, and goes back where it came from" {
+    const heap = Heap(&test_allocator);
+    test_allocator = std.testing.allocator;
+
+    var live: ?*anyopaque = heap.calloc(4, 1) orelse return error.OutOfMemory;
+    defer if (live) |block| heap.free(block);
+    const original: [*]u8 = @ptrCast(live.?);
+    try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, original[0..4]);
+    @memcpy(original[0..4], "data");
+    try std.testing.expectEqual(@as(usize, 0), @intFromPtr(original) % 16);
+
+    // Repointed: the live block still grows and is freed through the testing
+    // allocator that made it, which is what reports a leak or a bad free.
+    var other: std.heap.DebugAllocator(.{}) = .init;
+    defer if (other.deinit() != .ok) @panic("the other allocator leaked");
+    test_allocator = other.allocator();
+    const fresh = heap.malloc(24) orelse return error.OutOfMemory;
+    live = heap.realloc(live, 4096) orelse return error.OutOfMemory;
+    const grown: [*]u8 = @ptrCast(live.?);
+    try std.testing.expectEqualSlices(u8, "data", grown[0..4]);
+    heap.free(fresh);
+
+    // ...and back: the testing allocator is found again, not added twice.
+    test_allocator = std.testing.allocator;
+    const known = allocators_len.load(.monotonic);
+    const again = heap.malloc(8) orelse return error.OutOfMemory;
+    heap.free(again);
+    try std.testing.expectEqual(known, allocators_len.load(.monotonic));
+
+    const zero = heap.malloc(0) orelse return error.OutOfMemory;
+    heap.free(zero);
+    try std.testing.expect(heap.malloc(std.math.maxInt(usize)) == null);
+    try std.testing.expect(heap.calloc(std.math.maxInt(usize), 2) == null);
+    try std.testing.expect(heap.realloc(live, 0) == null);
+    live = null;
+    heap.free(null);
+}
diff --git a/src/gui/gui.zig b/src/gui/gui.zig
index dd74ae58..73f17a22 100644
--- a/src/gui/gui.zig
+++ b/src/gui/gui.zig
@@ -15,6 +15,7 @@ const crt = @import("crt.zig");
 const pet = @import("pet.zig");
 const fonts = @import("../fonts.zig"); // the Font builtin's half of the seam
 const selection_pipe = @import("../selection_pipe.zig");
+const c_heap = @import("c_heap");
 
 const ninep_io = @import("../9p_io.zig");
 
@@ -133,53 +134,20 @@ fn taglineRaster(font: *c.UIFont, body_px: f32, body_cell_w: u32, body_cell_h: u
         .height = @intCast(band_h),
     };
 }
-/// What FreeType and HarfBuzz allocate from, through ui_malloc and friends
-/// below. runNative points it at the gui's gpa once HarfBuzz's process-wide
-/// objects exist (ui_font_prime), so a Debug build's leak report is about
-/// fonts; a test that counts what fonts leak points it at
-/// std.testing.allocator the same way. Only the render thread calls into
+/// What FreeType and HarfBuzz allocate from, through ui_malloc and friends:
+/// the shared C heap, exported under the names font.c and HarfBuzz's
+/// hb_malloc_impl expect. runNative points it at the gui's gpa once
+/// HarfBuzz's process-wide objects exist (ui_font_prime), so a Debug build's
+/// leak report is about fonts; a test that counts what fonts leak points it
+/// at std.testing.allocator the same way. Only the render thread calls into
 /// fonts, and every allocator it is pointed at is thread-safe anyway.
 var font_allocator: std.mem.Allocator = std.heap.smp_allocator;
-
-/// C's free and realloc pass no size, so a block starts with its size and the
-/// allocator that made it: a block made before font_allocator was repointed
-/// goes back where it came from. 32 bytes keep the caller's pointer 16-byte
-/// (max_align_t) aligned.
-const FontBlock = struct { size: usize, allocator: std.mem.Allocator };
-const font_block_header = 32;
+const font_heap = c_heap.Heap(&font_allocator);
 comptime {
-    std.debug.assert(@sizeOf(FontBlock) <= font_block_header);
-}
-
-export fn ui_malloc(size: usize) ?*anyopaque {
-    const total = std.math.add(usize, font_block_header, size) catch return null;
-    const bytes = font_allocator.alignedAlloc(u8, .@"16", total) catch return null;
-    @as(*FontBlock, @ptrCast(bytes.ptr)).* = .{ .size = size, .allocator = font_allocator };
-    return bytes.ptr + font_block_header;
-}
-
-export fn ui_calloc(count: usize, size: usize) ?*anyopaque {
-    const total = std.math.mul(usize, count, size) catch return null;
-    const block: [*]u8 = @ptrCast(ui_malloc(total) orelse return null);
-    @memset(block[0..total], 0);
-    return block;
-}
-
-export fn ui_realloc(block: ?*anyopaque, size: usize) ?*anyopaque {
-    const old = block orelse return ui_malloc(size);
-    const base: [*]align(16) u8 = @alignCast(@as([*]u8, @ptrCast(old)) - font_block_header);
-    const header: *FontBlock = @ptrCast(base);
-    const total = std.math.add(usize, font_block_header, size) catch return null;
-    const bytes = header.allocator.realloc(base[0 .. font_block_header + header.size], total) catch return null;
-    @as(*FontBlock, @ptrCast(bytes.ptr)).size = size;
-    return bytes.ptr + font_block_header;
-}
-
-export fn ui_free(block: ?*anyopaque) void {
-    const old = block orelse return;
-    const base: [*]align(16) u8 = @alignCast(@as([*]u8, @ptrCast(old)) - font_block_header);
-    const header: *FontBlock = @ptrCast(base);
-    header.allocator.free(base[0 .. font_block_header + header.size]);
+    @export(&font_heap.malloc, .{ .name = "ui_malloc" });
+    @export(&font_heap.calloc, .{ .name = "ui_calloc" });
+    @export(&font_heap.realloc, .{ .name = "ui_realloc" });
+    @export(&font_heap.free, .{ .name = "ui_free" });
 }
 
 test "FreeType and HarfBuzz allocate from the gui allocator and give it all back" {
diff --git a/src/pdf.zig b/src/pdf.zig
index de5e6d86..8ddca497 100644
--- a/src/pdf.zig
+++ b/src/pdf.zig
@@ -2,66 +2,30 @@
 //! setjmp/longjmp; pdf_bridge.c catches every such throw before control returns
 //! to Zig. This module is present in native builds unless `-Dmupdf=false`.
 const std = @import("std");
+const c_heap = @import("c_heap");
 
 const c = @cImport({
     @cInclude("pdf_bridge.h");
 });
 
-const allocation_header_size = 16;
+/// What MuPDF allocates from, itself and the libraries it bundles (they all
+/// allocate through its context): `start`'s allocator, through the shared C
+/// heap. MuPDF's allocator functions carry a user pointer the heap does not
+/// need.
 var bridge_allocator: std.mem.Allocator = undefined;
+const bridge_heap = c_heap.Heap(&bridge_allocator);
 var bridge_started = false;
 
-comptime {
-    if (allocation_header_size < @sizeOf(usize))
-        @compileError("PDF allocation header is too small");
-}
-
-fn bridgeAlloc(user: ?*anyopaque, size_arg: usize) callconv(.c) ?*anyopaque {
-    const user_ptr = user orelse return null;
-    const allocator: *std.mem.Allocator = @ptrCast(@alignCast(user_ptr));
-    const size = @max(size_arg, 1);
-    const total = std.math.add(usize, allocation_header_size, size) catch
-        return null;
-    const bytes = allocator.alignedAlloc(u8, .@"16", total) catch return null;
-    const header: *align(16) usize = @ptrCast(bytes.ptr);
-    header.* = total;
-    return @ptrCast(bytes.ptr + allocation_header_size);
+fn bridgeAlloc(_: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque {
+    return bridge_heap.malloc(size);
 }
 
-fn bridgeFree(user: ?*anyopaque, pointer: ?*anyopaque) callconv(.c) void {
-    const data = pointer orelse return;
-    const user_ptr = user orelse @panic("PDF allocator callback has no allocator");
-    const allocator: *std.mem.Allocator = @ptrCast(@alignCast(user_ptr));
-    const data_bytes: [*]u8 = @ptrCast(data);
-    const base: [*]align(16) u8 =
-        @ptrCast(@alignCast(data_bytes - allocation_header_size));
-    const header: *align(16) usize = @ptrCast(base);
-    allocator.free(base[0..header.*]);
+fn bridgeRealloc(_: ?*anyopaque, block: ?*anyopaque, size: usize) callconv(.c) ?*anyopaque {
+    return bridge_heap.realloc(block, size);
 }
 
-fn bridgeRealloc(
-    user: ?*anyopaque,
-    pointer: ?*anyopaque,
-    size_arg: usize,
-) callconv(.c) ?*anyopaque {
-    if (pointer == null) return bridgeAlloc(user, size_arg);
-    if (size_arg == 0) {
-        bridgeFree(user, pointer);
-        return null;
-    }
-
-    const user_ptr = user orelse return null;
-    const allocator: *std.mem.Allocator = @ptrCast(@alignCast(user_ptr));
-    const total = std.math.add(usize, allocation_header_size, size_arg) catch
-        return null;
-    const data_bytes: [*]u8 = @ptrCast(pointer.?);
-    const base: [*]align(16) u8 =
-        @ptrCast(@alignCast(data_bytes - allocation_header_size));
-    const header: *align(16) usize = @ptrCast(base);
-    const resized = allocator.realloc(base[0..header.*], total) catch return null;
-    const resized_header: *align(16) usize = @ptrCast(resized.ptr);
-    resized_header.* = total;
-    return @ptrCast(resized.ptr + allocation_header_size);
+fn bridgeFree(_: ?*anyopaque, block: ?*anyopaque) callconv(.c) void {
+    bridge_heap.free(block);
 }
 
 /// Installs the allocator used by all subsequently opened PDF documents and
@@ -70,7 +34,7 @@ pub fn start(allocator: std.mem.Allocator) void {
     if (bridge_started) @panic("PDF allocator is already started");
     bridge_allocator = allocator;
     const descriptor = c.pardes_pdf_allocator{
-        .user = &bridge_allocator,
+        .user = null,
         .alloc = bridgeAlloc,
         .realloc = bridgeRealloc,
         .free = bridgeFree,
@@ -90,31 +54,6 @@ pub fn stop() void {
     bridge_allocator = undefined;
 }
 
-test "PDF C allocator callbacks preserve and free exact allocations" {
-    var allocator = std.testing.allocator;
-
-    var live: ?*anyopaque = bridgeAlloc(&allocator, 4) orelse
-        return error.OutOfMemory;
-    defer if (live) |pointer| bridgeFree(&allocator, pointer);
-    const original: [*]u8 = @ptrCast(live.?);
-    @memcpy(original[0..4], "data");
-
-    live = bridgeRealloc(&allocator, live, 32) orelse
-        return error.OutOfMemory;
-    const grown: [*]u8 = @ptrCast(live.?);
-    try std.testing.expectEqualSlices(u8, "data", grown[0..4]);
-
-    try std.testing.expect(bridgeAlloc(
-        &allocator,
-        std.math.maxInt(usize),
-    ) == null);
-
-    const zero = bridgeAlloc(&allocator, 0) orelse return error.OutOfMemory;
-    bridgeFree(&allocator, zero);
-    try std.testing.expect(bridgeRealloc(&allocator, live, 0) == null);
-    live = null;
-}
-
 test "custom allocator owns a complete PDF document lifecycle" {
     start(std.testing.allocator);
     defer stop();
diff --git a/src/syntax.zig b/src/syntax.zig
index e637bf67..5c8b9c0f 100644
--- a/src/syntax.zig
+++ b/src/syntax.zig
@@ -3,6 +3,7 @@ const config = @import("pardes_config");
 const tracy = @import("tracy.zig");
 const grammar_manifest = @import("grammar_manifest.zig");
 const look = @import("look.zig");
+const c_heap = @import("c_heap");
 
 pub const enabled = config.syntax_highlighting;
 const zig_grammar = config.syntax_zig_grammar;
@@ -139,8 +140,10 @@ fn initSpecs() [specCount()]Spec {
 
 var specs = initSpecs();
 
-const allocation_header_size = 16;
+/// What the tree-sitter runtime allocates from: `start`'s allocator, through
+/// the shared C heap.
 var syntax_allocator: std.mem.Allocator = undefined;
+const syntax_heap = c_heap.Heap(&syntax_allocator);
 var syntax_started = false;
 
 extern fn ts_set_allocator(
@@ -150,56 +153,13 @@ extern fn ts_set_allocator(
     new_free: ?*const fn (ptr: ?*anyopaque) callconv(.c) void,
 ) void;
 
-fn syntaxAlloc(size_arg: usize) callconv(.c) ?*anyopaque {
-    const size = @max(size_arg, 1);
-    const total = std.math.add(usize, allocation_header_size, size) catch return null;
-    const bytes = syntax_allocator.alignedAlloc(u8, .@"16", total) catch return null;
-    const header: *align(16) usize = @ptrCast(bytes.ptr);
-    header.* = total;
-    return @ptrCast(bytes.ptr + allocation_header_size);
-}
-
-fn syntaxCalloc(count: usize, size: usize) callconv(.c) ?*anyopaque {
-    const len = std.math.mul(usize, count, size) catch return null;
-    const pointer = syntaxAlloc(len) orelse return null;
-    const bytes: [*]u8 = @ptrCast(pointer);
-    @memset(bytes[0..len], 0);
-    return pointer;
-}
-
-fn syntaxRealloc(ptr: ?*anyopaque, new_size: usize) callconv(.c) ?*anyopaque {
-    const pointer = ptr orelse return syntaxAlloc(new_size);
-    if (new_size == 0) {
-        syntaxFree(pointer);
-        return null;
-    }
-    const user: [*]u8 = @ptrCast(pointer);
-    const header: *align(16) usize = @ptrCast(@alignCast(user - allocation_header_size));
-    const old_total = header.*;
-    const old_bytes: []align(16) u8 = @as([*]align(16) u8, @ptrCast(header))[0..old_total];
-    const new_total = std.math.add(usize, allocation_header_size, new_size) catch return null;
-    const new_bytes = syntax_allocator.realloc(old_bytes, new_total) catch return null;
-    const new_header: *align(16) usize = @ptrCast(new_bytes.ptr);
-    new_header.* = new_total;
-    return @ptrCast(new_bytes.ptr + allocation_header_size);
-}
-
-fn syntaxFree(ptr: ?*anyopaque) callconv(.c) void {
-    const pointer = ptr orelse return;
-    const user: [*]u8 = @ptrCast(pointer);
-    const header: *align(16) usize = @ptrCast(@alignCast(user - allocation_header_size));
-    const total = header.*;
-    const bytes: []align(16) u8 = @as([*]align(16) u8, @ptrCast(header))[0..total];
-    syntax_allocator.free(bytes);
-}
-
 pub fn start(gpa: std.mem.Allocator) void {
     if (comptime enabled) {
         std.debug.assert(!syntax_started);
         stop();
         syntax_allocator = gpa;
         syntax_started = true;
-        ts_set_allocator(syntaxAlloc, syntaxCalloc, syntaxRealloc, syntaxFree);
+        ts_set_allocator(syntax_heap.malloc, syntax_heap.calloc, syntax_heap.realloc, syntax_heap.free);
     }
 }
 
@@ -935,25 +895,6 @@ test "syntax codeAfterLocation takes whole-token locations and nothing else" {
     try std.testing.expect(codeAfterLocation("   leading space") == null);
 }
 
-test "syntax tree-sitter allocator callbacks preserve and free exact allocations" {
-    syntax_allocator = std.testing.allocator;
-    defer syntax_allocator = undefined;
-
-    var live: ?*anyopaque = syntaxCalloc(4, 1) orelse return error.OutOfMemory;
-    defer if (live) |pointer| syntaxFree(pointer);
-    const original: [*]u8 = @ptrCast(live.?);
-    try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, original[0..4]);
-    @memcpy(original[0..4], "data");
-
-    live = syntaxRealloc(live, 32) orelse return error.OutOfMemory;
-    const grown: [*]u8 = @ptrCast(live.?);
-    try std.testing.expectEqualSlices(u8, "data", grown[0..4]);
-
-    try std.testing.expect(syntaxCalloc(std.math.maxInt(usize), 2) == null);
-    try std.testing.expect(syntaxRealloc(live, 0) == null);
-    live = null;
-}
-
 test "syntax default full grammar set highlights Typst source" {
     if (!enabled or !full_grammars) return;
     start(std.testing.allocator);