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()); } }