From a4d8e9608e0cb2704bad00b2f87af078c6d4ed59 Mon Sep 17 00:00:00 2001 From: Mason Remaley Date: Mon, 13 Apr 2026 02:55:07 -0700 Subject: [PATCH] Reworks the stack fallback allocator The previous approach had a few downsides: 1. You couldn't set the alignment of the internal buffer. Many callers in the standard library trying to use this for a small vec style optimization worked around this by setting the alignment for the struct itself, this ends up very verbose and also assumes a specific layout for the struct which isn't guaranteed. 2. It was generic over the size of the buffer. This type is used a lot in std with various sizes. 3. It has an awkward API where you had to call get which mutated the type unlike all other allocators, and then had a runtime check to make sure you didn't get this wrong. The new approach resolves all of these issues by just taking the buf as an argument. This is particularly amenable to smallvec style optimizations: you can just declare the buf as an array of the item you want to allocate to get the exact minimum size. --- lib/std/heap.zig | 187 +++++++++++++++++++++-------------------------- 1 file changed, 82 insertions(+), 105 deletions(-) diff --git a/lib/std/heap.zig b/lib/std/heap.zig index 1b979fd0241bf7123df885624d2f6dc78c438e29..1e234ac4a2e8a1cc7050413bc22d46c1c4d08374 100644 --- a/lib/std/heap.zig +++ b/lib/std/heap.zig @@ -368,112 +368,88 @@ pub const brk_allocator: Allocator = .{ .vtable = &BrkAllocator.vtable, }; -/// Returns a `StackFallbackAllocator` allocating using either a -/// `FixedBufferAllocator` on an array of size `size` and falling back to -/// `fallback_allocator` if that fails. -pub fn stackFallback(comptime size: usize, fallback_allocator: Allocator) StackFallbackAllocator(size) { - return StackFallbackAllocator(size){ - .buffer = undefined, - .fallback_allocator = fallback_allocator, - .fixed_buffer_allocator = undefined, - }; -} +/// An allocator that attempts to allocate from the given buffer, falling back to +/// `fallback_allocator` if this fails. +pub const StackFallbackAllocator = struct { + const Self = @This(); -/// An allocator that attempts to allocate using a -/// `FixedBufferAllocator` using an array of size `size`. If the -/// allocation fails, it will fall back to using -/// `fallback_allocator`. Easily created with `stackFallback`. -pub fn StackFallbackAllocator(comptime size: usize) type { - return struct { - const Self = @This(); + fallback_allocator: Allocator, + fixed_buffer_allocator: FixedBufferAllocator, - buffer: [size]u8, - fallback_allocator: Allocator, - fixed_buffer_allocator: FixedBufferAllocator, - get_called: if (std.debug.runtime_safety) bool else void = - if (std.debug.runtime_safety) false else {}, + pub fn init(buf: []u8, fallback_allocator: Allocator) Self { + return .{ + .fallback_allocator = fallback_allocator, + .fixed_buffer_allocator = .init(buf), + }; + } - /// This function both fetches a `Allocator` interface to this - /// allocator *and* resets the internal buffer allocator. - pub fn get(self: *Self) Allocator { - if (std.debug.runtime_safety) { - assert(!self.get_called); // `get` called multiple times; instead use `const allocator = stackFallback(N).get();` - self.get_called = true; - } - self.fixed_buffer_allocator = FixedBufferAllocator.init(self.buffer[0..]); - return .{ - .ptr = self, - .vtable = &.{ - .alloc = alloc, - .resize = resize, - .remap = remap, - .free = free, - }, - }; - } - - /// Unlike most std allocators `StackFallbackAllocator` modifies - /// its internal state before returning an implementation of - /// the`Allocator` interface and therefore also doesn't use - /// the usual `.allocator()` method. - pub const allocator = @compileError("use 'const allocator = stackFallback(N).get();' instead"); + pub fn allocator(self: *Self) Allocator { + return .{ + .ptr = self, + .vtable = &.{ + .alloc = alloc, + .resize = resize, + .remap = remap, + .free = free, + }, + }; + } - fn alloc( - ctx: *anyopaque, - len: usize, - alignment: Alignment, - ra: usize, - ) ?[*]u8 { - const self: *Self = @ptrCast(@alignCast(ctx)); - return FixedBufferAllocator.alloc(&self.fixed_buffer_allocator, len, alignment, ra) orelse - return self.fallback_allocator.rawAlloc(len, alignment, ra); - } + fn alloc( + ctx: *anyopaque, + len: usize, + alignment: Alignment, + ra: usize, + ) ?[*]u8 { + const self: *Self = @ptrCast(@alignCast(ctx)); + return FixedBufferAllocator.alloc(&self.fixed_buffer_allocator, len, alignment, ra) orelse + return self.fallback_allocator.rawAlloc(len, alignment, ra); + } - fn resize( - ctx: *anyopaque, - buf: []u8, - alignment: Alignment, - new_len: usize, - ra: usize, - ) bool { - const self: *Self = @ptrCast(@alignCast(ctx)); - if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) { - return FixedBufferAllocator.resize(&self.fixed_buffer_allocator, buf, alignment, new_len, ra); - } else { - return self.fallback_allocator.rawResize(buf, alignment, new_len, ra); - } + fn resize( + ctx: *anyopaque, + buf: []u8, + alignment: Alignment, + new_len: usize, + ra: usize, + ) bool { + const self: *Self = @ptrCast(@alignCast(ctx)); + if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) { + return FixedBufferAllocator.resize(&self.fixed_buffer_allocator, buf, alignment, new_len, ra); + } else { + return self.fallback_allocator.rawResize(buf, alignment, new_len, ra); } + } - fn remap( - context: *anyopaque, - memory: []u8, - alignment: Alignment, - new_len: usize, - return_address: usize, - ) ?[*]u8 { - const self: *Self = @ptrCast(@alignCast(context)); - if (self.fixed_buffer_allocator.ownsPtr(memory.ptr)) { - return FixedBufferAllocator.remap(&self.fixed_buffer_allocator, memory, alignment, new_len, return_address); - } else { - return self.fallback_allocator.rawRemap(memory, alignment, new_len, return_address); - } + fn remap( + context: *anyopaque, + memory: []u8, + alignment: Alignment, + new_len: usize, + return_address: usize, + ) ?[*]u8 { + const self: *Self = @ptrCast(@alignCast(context)); + if (self.fixed_buffer_allocator.ownsPtr(memory.ptr)) { + return FixedBufferAllocator.remap(&self.fixed_buffer_allocator, memory, alignment, new_len, return_address); + } else { + return self.fallback_allocator.rawRemap(memory, alignment, new_len, return_address); } + } - fn free( - ctx: *anyopaque, - buf: []u8, - alignment: Alignment, - ra: usize, - ) void { - const self: *Self = @ptrCast(@alignCast(ctx)); - if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) { - return FixedBufferAllocator.free(&self.fixed_buffer_allocator, buf, alignment, ra); - } else { - return self.fallback_allocator.rawFree(buf, alignment, ra); - } + fn free( + ctx: *anyopaque, + buf: []u8, + alignment: Alignment, + ra: usize, + ) void { + const self: *Self = @ptrCast(@alignCast(ctx)); + if (self.fixed_buffer_allocator.ownsPtr(buf.ptr)) { + return FixedBufferAllocator.free(&self.fixed_buffer_allocator, buf, alignment, ra); + } else { + return self.fallback_allocator.rawFree(buf, alignment, ra); } - }; -} + } +}; test c_allocator { if (builtin.link_libc) { @@ -525,22 +501,23 @@ test ArenaAllocator { try testAllocatorAlignedShrink(allocator); } -test "StackFallbackAllocator" { +test StackFallbackAllocator { + var buf: [4096]u8 = undefined; { - var stack_allocator = stackFallback(4096, std.testing.allocator); - try testAllocator(stack_allocator.get()); + var stack_allocator: StackFallbackAllocator = .init(&buf, std.testing.allocator); + try testAllocator(stack_allocator.allocator()); } { - var stack_allocator = stackFallback(4096, std.testing.allocator); - try testAllocatorAligned(stack_allocator.get()); + var stack_allocator: StackFallbackAllocator = .init(&buf, std.testing.allocator); + try testAllocatorAligned(stack_allocator.allocator()); } { - var stack_allocator = stackFallback(4096, std.testing.allocator); - try testAllocatorLargeAlignment(stack_allocator.get()); + var stack_allocator: StackFallbackAllocator = .init(&buf, std.testing.allocator); + try testAllocatorLargeAlignment(stack_allocator.allocator()); } { - var stack_allocator = stackFallback(4096, std.testing.allocator); - try testAllocatorAlignedShrink(stack_allocator.get()); + var stack_allocator: StackFallbackAllocator = .init(&buf, std.testing.allocator); + try testAllocatorAlignedShrink(stack_allocator.allocator()); } } -- 2.54.0