1 files changed, 135 insertions, 0 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);
+}