tiny.memtrace.ledger

1const std = @import("std");23const Allocator = std.mem.Allocator;4pub const Sequence = u64;56pub const Record = struct {7    address: usize = 0,8    len: usize = 0,9    allocation_sequence: Sequence = 0,10    allocation_return_address: usize = 0,11    resize_return_address: usize = 0,12    remap_return_address: usize = 0,13    free_sequence: Sequence = 0,14    free_return_address: usize = 0,15    active: bool = false,16};1718pub const AllocationLedger = struct {19    backing: Allocator,20    records: []Record,21    count: usize = 0,22    sequence: Sequence = 0,2324    pub fn init(backing: Allocator, records: []Record) AllocationLedger {25        std.debug.assert(records.len != 0);26        return .{27            .backing = backing,28            .records = records,29        };30    }3132    pub fn allocator(self: *AllocationLedger) Allocator {33        return .{ .ptr = self, .vtable = &vtable };34    }3536    pub fn exhausted(self: *const AllocationLedger) bool {37        return self.count == self.records.len;38    }3940    pub fn allocationCount(self: *const AllocationLedger) usize {41        return self.count;42    }4344    pub fn allocation(45        self: *const AllocationLedger,46        id: usize,47    ) Record {48        std.debug.assert(id < self.count);49        return self.records[id];50    }5152    pub fn allocationContaining(53        self: *const AllocationLedger,54        pointer: anytype,55        len: usize,56    ) ?usize {57        const address = @intFromPtr(pointer);58        var index = self.count;59        while (index != 0) {60            index -= 1;61            if (rangeContains(self.records[index], address, len)) return index;62        }63        return null;64    }6566    pub fn liveAllocationContaining(67        self: *const AllocationLedger,68        pointer: anytype,69        len: usize,70    ) ?usize {71        const address = @intFromPtr(pointer);72        var index = self.count;73        while (index != 0) {74            index -= 1;75            const record = self.records[index];76            if (record.active and rangeContains(record, address, len)) {77                return index;78            }79        }80        return null;81    }8283    pub fn isActive(84        self: *const AllocationLedger,85        id: usize,86    ) bool {87        return self.allocation(id).active;88    }8990    pub fn freeSequence(91        self: *const AllocationLedger,92        id: usize,93    ) Sequence {94        const record = self.allocation(id);95        std.debug.assert(!record.active);96        return record.free_sequence;97    }9899    fn nextSequence(self: *AllocationLedger) Sequence {100        std.debug.assert(self.sequence != std.math.maxInt(Sequence));101        self.sequence += 1;102        return self.sequence;103    }104105    fn append(106        self: *AllocationLedger,107        pointer: [*]u8,108        len: usize,109        return_address: usize,110    ) void {111        std.debug.assert(self.count < self.records.len);112        self.records[self.count] = .{113            .address = @intFromPtr(pointer),114            .len = len,115            .allocation_sequence = self.nextSequence(),116            .allocation_return_address = return_address,117            .active = true,118        };119        self.count += 1;120    }121122    fn findActiveExact(123        self: *const AllocationLedger,124        pointer: [*]u8,125    ) usize {126        const address = @intFromPtr(pointer);127        var index = self.count;128        while (index != 0) {129            index -= 1;130            const record = self.records[index];131            if (record.active and record.address == address) return index;132        }133        unreachable;134    }135136    fn rawAlloc(137        context: *anyopaque,138        len: usize,139        alignment: std.mem.Alignment,140        return_address: usize,141    ) ?[*]u8 {142        const self: *AllocationLedger = @ptrCast(@alignCast(context));143        if (self.exhausted()) return null;144        const result = self.backing.rawAlloc(145            len,146            alignment,147            return_address,148        ) orelse return null;149        self.append(result, len, return_address);150        return result;151    }152153    fn rawResize(154        context: *anyopaque,155        memory: []u8,156        alignment: std.mem.Alignment,157        new_len: usize,158        return_address: usize,159    ) bool {160        const self: *AllocationLedger = @ptrCast(@alignCast(context));161        const resized = self.backing.rawResize(162            memory,163            alignment,164            new_len,165            return_address,166        );167        if (resized) {168            const id = self.findActiveExact(memory.ptr);169            _ = self.nextSequence();170            self.records[id].len = new_len;171            self.records[id].resize_return_address = return_address;172        }173        return resized;174    }175176    fn rawRemap(177        context: *anyopaque,178        memory: []u8,179        alignment: std.mem.Alignment,180        new_len: usize,181        return_address: usize,182    ) ?[*]u8 {183        const self: *AllocationLedger = @ptrCast(@alignCast(context));184        const id = self.findActiveExact(memory.ptr);185        const result = self.backing.rawRemap(186            memory,187            alignment,188            new_len,189            return_address,190        ) orelse return null;191        _ = self.nextSequence();192        self.records[id].address = @intFromPtr(result);193        self.records[id].len = new_len;194        self.records[id].remap_return_address = return_address;195        return result;196    }197198    fn rawFree(199        context: *anyopaque,200        memory: []u8,201        alignment: std.mem.Alignment,202        return_address: usize,203    ) void {204        const self: *AllocationLedger = @ptrCast(@alignCast(context));205        const id = self.findActiveExact(memory.ptr);206        self.backing.rawFree(memory, alignment, return_address);207        const sequence = self.nextSequence();208        self.records[id].active = false;209        self.records[id].free_sequence = sequence;210        self.records[id].free_return_address = return_address;211    }212213    const vtable: Allocator.VTable = .{214        .alloc = rawAlloc,215        .resize = rawResize,216        .remap = rawRemap,217        .free = rawFree,218    };219};220221fn rangeContains(record: Record, address: usize, len: usize) bool {222    if (address < record.address) return false;223    const end = std.math.add(usize, address, len) catch return false;224    const record_end = std.math.add(225        usize,226        record.address,227        record.len,228    ) catch return false;229    return end <= record_end;230}231232test "allocation ledger attributes ranges and teardown order" {233    const testing = std.testing;234    var records: [4]Record = undefined;235    var ledger = AllocationLedger.init(testing.allocator, &records);236    const allocator = ledger.allocator();237    const first = try allocator.alloc(u8, 64);238    const second = try allocator.alloc(u8, 128);239    const first_id = ledger.liveAllocationContaining(240        first.ptr + 7,241        8,242    ) orelse return error.TestUnexpectedResult;243    const second_id = ledger.liveAllocationContaining(244        second.ptr + 31,245        16,246    ) orelse return error.TestUnexpectedResult;247    try testing.expect(first_id != second_id);248    try testing.expect(249        ledger.allocation(first_id).allocation_return_address != 0,250    );251    allocator.free(second);252    allocator.free(first);253    try testing.expect(!ledger.isActive(first_id));254    try testing.expect(!ledger.isActive(second_id));255    try testing.expect(256        ledger.freeSequence(second_id) < ledger.freeSequence(first_id),257    );258}259260test "allocation ledger applies its caller-provided record bound" {261    const testing = std.testing;262    var records: [1]Record = undefined;263    var ledger = AllocationLedger.init(testing.allocator, &records);264    const allocator = ledger.allocator();265    const first = try allocator.alloc(u8, 1);266    defer allocator.free(first);267    try testing.expect(ledger.exhausted());268    try testing.expectError(error.OutOfMemory, allocator.alloc(u8, 1));269}270271test "allocation ledger tracks successful resize and remap operations" {272    const testing = std.testing;273    var backing_storage: [256]u8 = undefined;274    var backing = std.heap.FixedBufferAllocator.init(&backing_storage);275    var records: [1]Record = undefined;276    var ledger = AllocationLedger.init(backing.allocator(), &records);277    const allocator = ledger.allocator();278    var memory = try allocator.alloc(u8, 64);279    const id = ledger.liveAllocationContaining(280        memory.ptr,281        memory.len,282    ) orelse return error.TestUnexpectedResult;283284    try testing.expect(!allocator.resize(memory, backing_storage.len + 1));285    try testing.expectEqual(@as(usize, 64), ledger.allocation(id).len);286    try testing.expectEqual(@as(usize, 0), ledger.allocation(id).resize_return_address);287    try testing.expect(allocator.resize(memory, 96));288    memory = memory.ptr[0..96];289    try testing.expectEqual(@as(usize, 96), ledger.allocation(id).len);290    try testing.expect(ledger.allocation(id).resize_return_address != 0);291292    try testing.expect(allocator.remap(memory, backing_storage.len + 1) == null);293    try testing.expectEqual(@intFromPtr(memory.ptr), ledger.allocation(id).address);294    try testing.expectEqual(@as(usize, 96), ledger.allocation(id).len);295    try testing.expectEqual(@as(usize, 0), ledger.allocation(id).remap_return_address);296    memory = allocator.remap(memory, 128) orelse return error.TestUnexpectedResult;297    try testing.expectEqual(@intFromPtr(memory.ptr), ledger.allocation(id).address);298    try testing.expectEqual(@as(usize, 128), ledger.allocation(id).len);299    try testing.expect(ledger.allocation(id).remap_return_address != 0);300    try testing.expect(ledger.isActive(id));301302    allocator.free(memory);303    try testing.expect(!ledger.isActive(id));304}

1pub const ledger = ledger_mod;