tiny.stabilizer.ShuffleAllocator

1pub const ShuffleAllocator = struct {2    backing: Allocator,3    config: HeapConfig,4    rng: Marsaglia,5    pools: std.ArrayListUnmanaged(Pool) = .empty,6    known: std.AutoHashMapUnmanaged(usize, AllocationState) = .empty,78    pub fn init(backing: Allocator, heap_config: HeapConfig, seed: u64) ShuffleAllocator {9        return .{10            .backing = backing,11            .config = normalizeHeapConfig(heap_config),12            .rng = Marsaglia.init(seed),13        };14    }1516    pub fn deinit(self: *ShuffleAllocator) void {17        for (self.pools.items) |*pool| pool.deinit(self);18        self.pools.deinit(self.backing);19        self.known.deinit(self.backing);20        self.* = undefined;21    }2223    pub fn allocator(self: *ShuffleAllocator) Allocator {24        return .{ .ptr = self, .vtable = &vtable };25    }2627    pub fn malloc(self: *ShuffleAllocator, len: usize) ![*]u8 {28        return try self.alloc(len, Alignment.fromByteUnits(1), true, @returnAddress());29    }3031    pub fn calloc(self: *ShuffleAllocator, count: usize, len: usize) ![*]u8 {32        const total = std.math.mul(usize, count, len) catch return error.OutOfMemory;33        const memory = try self.malloc(total);34        @memset(memory[0..total], 0);35        return memory;36    }3738    pub fn realloc(self: *ShuffleAllocator, ptr: ?[*]u8, new_len: usize) !?[*]u8 {39        const memory = ptr orelse return try self.malloc(new_len);40        const state = self.known.get(@intFromPtr(memory)) orelse return error.InvalidAllocation;41        if (state == .inactive) return error.InvalidAllocation;42        const header = checkedHeader(memory) orelse return error.InvalidAllocation;43        if (!header.allocated) return error.InvalidAllocation;44        const alignment = Alignment.fromByteUnits(header.alignment);45        const old_len = header.requested_len;46        const old_memory = memory[0..old_len];4748        if (new_len == 0) {49            self.free(old_memory, alignment, @returnAddress());50            return null;51        }5253        if (self.resize(old_memory, alignment, new_len, @returnAddress())) return memory;5455        const new_memory = try self.alloc(new_len, alignment, true, @returnAddress());56        @memcpy(new_memory[0..@min(old_len, new_len)], old_memory[0..@min(old_len, new_len)]);57        self.free(old_memory, alignment, @returnAddress());58        return new_memory;59    }6061    pub fn freePointer(self: *ShuffleAllocator, ptr: ?[*]u8) void {62        _ = self.freePointerIfOwned(ptr);63    }6465    pub fn freePointerIfOwned(self: *ShuffleAllocator, ptr: ?[*]u8) bool {66        const memory = ptr orelse return true;67        const state = self.known.get(@intFromPtr(memory)) orelse return false;68        if (state == .inactive) reportDoubleFreeError();69        const header = checkedHeader(memory) orelse return false;70        self.free(memory[0..header.requested_len], Alignment.fromByteUnits(header.alignment), @returnAddress());71        return true;72    }7374    pub fn ownsPointer(self: *const ShuffleAllocator, ptr: ?[*]u8) bool {75        const memory = ptr orelse return false;76        return self.known.get(@intFromPtr(memory)) != null;77    }7879    pub fn requestedLength(ptr: [*]u8) ?usize {80        const header = checkedHeader(ptr) orelse return null;81        if (!header.allocated) return null;82        return header.requested_len;83    }8485    fn normalizedSlots(self: *const ShuffleAllocator) usize {86        return if (self.config.shuffle_slots == 0) 1 else self.config.shuffle_slots;87    }8889    fn alloc(90        self: *ShuffleAllocator,91        len: usize,92        alignment: Alignment,93        pointer_validation: bool,94        ret_addr: usize,95    ) ![*]u8 {96        if (!self.config.enabled or len > self.config.max_shuffled_size) {97            const memory = try self.createUserMemory(98                len,99                alignment,100                0,101                pointer_validation,102                .active,103            );104            _ = ret_addr;105            return memory.ptr;106        }107108        const class_size = try sizeClass(len, self.config.min_class_size);109        const pool = try self.getPool(class_size, alignment, pointer_validation);110        const index = self.rng.bounded(pool.slots.len);111        const replacement = Slot{ .memory = try self.createUserMemory(112            class_size,113            alignment,114            class_size,115            pointer_validation,116            .inactive,117        ) };118        const selected = pool.slots[index];119        const header = headerFromUser(selected.memory.ptr);120        if (header.allocated) reportDoubleFreeError();121        if (pointer_validation) self.activatePointer(selected.memory.ptr);122        pool.slots[index] = replacement;123124        header.requested_len = len;125        header.allocated = true;126        return selected.memory.ptr;127    }128129    fn resize(self: *ShuffleAllocator, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {130        _ = self;131        _ = ret_addr;132        const header = checkedHeader(memory.ptr) orelse return false;133        if (!header.allocated) return false;134        if (header.alignment != alignment.toByteUnits()) return false;135        const current_class = header.class_size;136        if (current_class == 0) {137            if (new_len <= header.base_len - headerOffset(memory.ptr, header)) {138                header.requested_len = new_len;139                return true;140            }141            return false;142        }143        if (new_len <= current_class) {144            header.requested_len = new_len;145            return true;146        }147        return false;148    }149150    fn remap(self: *ShuffleAllocator, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {151        _ = self;152        _ = memory;153        _ = alignment;154        _ = new_len;155        _ = ret_addr;156        return null;157    }158159    fn free(self: *ShuffleAllocator, memory: []u8, alignment: Alignment, ret_addr: usize) void {160        const header = checkedHeader(memory.ptr) orelse return;161        if (header.alignment != alignment.toByteUnits()) return;162        if (!header.allocated) reportDoubleFreeError();163        header.allocated = false;164        if (header.tracked) {165            const state = self.known.getPtr(@intFromPtr(memory.ptr)) orelse unreachable;166            std.debug.assert(state.* == .active);167            state.* = .inactive;168        }169        if (header.class_size == 0 or !self.config.enabled) {170            self.destroyUserMemory(memory.ptr, ret_addr);171            return;172        }173174        const pool = self.findPool(header.class_size, alignment, header.tracked) orelse {175            self.destroyUserMemory(memory.ptr, ret_addr);176            return;177        };178        const index = self.rng.bounded(pool.slots.len);179        const evicted = pool.slots[index];180        pool.slots[index] = .{ .memory = memory.ptr[0..header.class_size] };181        self.destroySlot(evicted);182    }183184    fn getPool(185        self: *ShuffleAllocator,186        class_size: usize,187        alignment: Alignment,188        pointer_validation: bool,189    ) !*Pool {190        if (self.findPool(class_size, alignment, pointer_validation)) |pool| return pool;191        var pool = try Pool.init(self, class_size, alignment, pointer_validation);192        errdefer pool.deinit(self);193        try self.pools.append(self.backing, pool);194        return &self.pools.items[self.pools.items.len - 1];195    }196197    fn findPool(198        self: *ShuffleAllocator,199        class_size: usize,200        alignment: Alignment,201        pointer_validation: bool,202    ) ?*Pool {203        for (self.pools.items) |*pool| {204            if (pool.class_size == class_size and205                pool.alignment == alignment and206                pool.pointer_validation == pointer_validation)207            {208                return pool;209            }210        }211        return null;212    }213214    fn createUserMemory(215        self: *ShuffleAllocator,216        len: usize,217        alignment: Alignment,218        class_size: usize,219        pointer_validation: bool,220        state: AllocationState,221    ) ![]u8 {222        const user_alignment = Alignment.max(alignment, .of(Header));223        const alignment_bytes = user_alignment.toByteUnits();224        const overhead = std.math.add(usize, @sizeOf(Header), alignment_bytes - 1) catch225            return error.OutOfMemory;226        const base_len = std.math.add(usize, len, overhead) catch return error.OutOfMemory;227        const base = try self.backing.alloc(u8, base_len);228        const user_addr = user_alignment.forward(@intFromPtr(base.ptr) + @sizeOf(Header));229        const header: *Header = @ptrFromInt(user_addr - @sizeOf(Header));230        header.* = .{231            .magic = allocation_magic,232            .base_addr = @intFromPtr(base.ptr),233            .base_len = base_len,234            .class_size = class_size,235            .requested_len = len,236            .alignment = alignment.toByteUnits(),237            .allocated = class_size == 0,238            .tracked = false,239        };240        errdefer self.backing.free(base);241        if (pointer_validation) {242            try self.known.put(self.backing, user_addr, state);243            header.tracked = true;244        }245        return @as([*]u8, @ptrFromInt(user_addr))[0..len];246    }247248    fn activatePointer(self: *ShuffleAllocator, ptr: [*]u8) void {249        const header = headerFromUser(ptr);250        std.debug.assert(header.tracked);251        const state = self.known.getPtr(@intFromPtr(ptr)) orelse unreachable;252        std.debug.assert(state.* == .inactive);253        state.* = .active;254    }255256    fn destroySlot(self: *ShuffleAllocator, slot: Slot) void {257        self.destroyUserMemory(slot.memory.ptr, @returnAddress());258    }259260    fn destroyUserMemory(self: *ShuffleAllocator, ptr: [*]u8, ret_addr: usize) void {261        _ = ret_addr;262        const header = checkedHeader(ptr) orelse return;263        const base: [*]u8 = @ptrFromInt(header.base_addr);264        if (header.tracked) std.debug.assert(self.known.remove(@intFromPtr(ptr)));265        self.backing.free(base[0..header.base_len]);266    }267268    const vtable: Allocator.VTable = .{269        .alloc = rawAlloc,270        .resize = rawResize,271        .remap = rawRemap,272        .free = rawFree,273    };274275    fn rawAlloc(ctx: *anyopaque, len: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 {276        const self: *ShuffleAllocator = @ptrCast(@alignCast(ctx));277        return self.alloc(len, alignment, self.config.pointer_validation, ret_addr) catch null;278    }279280    fn rawResize(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {281        const self: *ShuffleAllocator = @ptrCast(@alignCast(ctx));282        return self.resize(memory, alignment, new_len, ret_addr);283    }284285    fn rawRemap(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {286        const self: *ShuffleAllocator = @ptrCast(@alignCast(ctx));287        return self.remap(memory, alignment, new_len, ret_addr);288    }289290    fn rawFree(ctx: *anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void {291        const self: *ShuffleAllocator = @ptrCast(@alignCast(ctx));292        self.free(memory, alignment, ret_addr);293    }294};

1pub const ShuffleAllocator = heap_mod.ShuffleAllocator;