//! 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);
}