Rewrite FixedBufferAllocator and rename to BumpAllocator.

Fix overflow for certain allocation sizes.
Reduce struct size by 1 usize.
Add a few more tests.
Simplify logic.

Author: RetroDev256

lib/compiler/aro/aro/Driver/Filesystem.zig

@@ -23,10 +23,11 @@ fn findProgramByNameFake(entries: []const Filesystem.Entry, name: []const u8, pa
     }
     const path_env = path orelse return null;
     var fib = std.heap.FixedBufferAllocator.init(buf);
+    const fib_initial_state = fib.savestate();
 
     var it = mem.tokenizeScalar(u8, path_env, std.fs.path.delimiter);
     while (it.next()) |path_dir| {
-        defer fib.reset();
+        defer fib.restore(fib_initial_state);
         const full_path = std.fs.path.join(fib.allocator(), &.{ path_dir, name }) catch continue;
         if (canExecuteFake(entries, full_path)) return full_path;
     }
@@ -84,10 +85,11 @@ fn findProgramByNamePosix(name: []const u8, path: ?[]const u8, buf: []u8) ?[]con
     }
     const path_env = path orelse return null;
     var fib = std.heap.FixedBufferAllocator.init(buf);
+    const fib_initial_state = fib.savestate();
 
     var it = mem.tokenizeScalar(u8, path_env, std.fs.path.delimiter);
     while (it.next()) |path_dir| {
-        defer fib.reset();
+        defer fib.restore(fib_initial_state);
         const full_path = std.fs.path.join(fib.allocator(), &.{ path_dir, name }) catch continue;
         if (canExecutePosix(full_path)) return full_path;
     }

lib/compiler/aro/aro/Toolchain.zig

@@ -216,14 +216,15 @@ pub fn addFilePathLibArgs(tc: *const Toolchain, argv: *std.ArrayList([]const u8)
 fn getProgramPath(tc: *const Toolchain, name: []const u8, buf: []u8) []const u8 {
     var path_buf: [std.fs.max_path_bytes]u8 = undefined;
     var fib = std.heap.FixedBufferAllocator.init(&path_buf);
+    const fib_initial_state = fib.savestate();
 
     var tool_specific_buf: [64]u8 = undefined;
     var possible_name_buf: [2][]const u8 = undefined;
     const possible_names = possibleProgramNames(tc.driver.raw_target_triple, name, &tool_specific_buf, &possible_name_buf);
 
     for (possible_names) |tool_name| {
         for (tc.program_paths.items) |program_path| {
-            defer fib.reset();
+            defer fib.restore(fib_initial_state);
 
             const candidate = std.fs.path.join(fib.allocator(), &.{ program_path, tool_name }) catch continue;
 

lib/compiler/test_runner.zig

@@ -37,6 +37,7 @@ pub fn main() void {
         return mainSimple() catch @panic("test failure\n");
     }
 
+    const fba_initial_state = fba.savestate();
     const args = std.process.argsAlloc(fba.allocator()) catch
         @panic("unable to parse command line args");
 
@@ -61,7 +62,7 @@ pub fn main() void {
         fuzz_abi.fuzzer_init(.fromSlice(cache_dir));
     }
 
-    fba.reset();
+    fba.restore(fba_initial_state);
 
     if (listen) {
         return mainServer() catch @panic("internal test runner failure");

lib/std/heap.zig

@@ -10,8 +10,10 @@ const windows = std.os.windows;
 const Alignment = std.mem.Alignment;
 
 pub const ArenaAllocator = @import("heap/arena_allocator.zig").ArenaAllocator;
+pub const BumpAllocator = @import("heap/BumpAllocator.zig");
 pub const SmpAllocator = @import("heap/SmpAllocator.zig");
-pub const FixedBufferAllocator = @import("heap/FixedBufferAllocator.zig");
+/// Deprecated; to be removed after 0.16.0 is tagged.
+pub const FixedBufferAllocator = BumpAllocator;
 pub const PageAllocator = @import("heap/PageAllocator.zig");
 pub const SbrkAllocator = @import("heap/sbrk_allocator.zig").SbrkAllocator;
 pub const ThreadSafeAllocator = @import("heap/ThreadSafeAllocator.zig");
@@ -374,38 +376,36 @@ pub const wasm_allocator: Allocator = .{
 };
 
 /// Returns a `StackFallbackAllocator` allocating using either a
-/// `FixedBufferAllocator` on an array of size `size` and falling back to
+/// `BumpAllocator` 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,
+        .bump_allocator = undefined,
     };
 }
 
 /// An allocator that attempts to allocate using a
-/// `FixedBufferAllocator` using an array of size `size`. If the
+/// `BumpAllocator` 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();
-
         buffer: [size]u8,
         fallback_allocator: Allocator,
-        fixed_buffer_allocator: FixedBufferAllocator,
+        bump_allocator: BumpAllocator,
         get_called: if (std.debug.runtime_safety) bool else void =
             if (std.debug.runtime_safety) false else {},
 
         /// This function both fetches a `Allocator` interface to this
         /// allocator *and* resets the internal buffer allocator.
-        pub fn get(self: *Self) Allocator {
+        pub fn get(self: *@This()) 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..]);
+            self.bump_allocator = .init(self.buffer[0..]);
             return .{
                 .ptr = self,
                 .vtable = &.{
@@ -429,8 +429,8 @@ pub fn StackFallbackAllocator(comptime size: usize) type {
             alignment: Alignment,
             ra: usize,
         ) ?[*]u8 {
-            const self: *Self = @ptrCast(@alignCast(ctx));
-            return FixedBufferAllocator.alloc(&self.fixed_buffer_allocator, len, alignment, ra) orelse
+            const self: *@This() = @ptrCast(@alignCast(ctx));
+            return BumpAllocator.alloc(&self.bump_allocator, len, alignment, ra) orelse
                 return self.fallback_allocator.rawAlloc(len, alignment, ra);
         }
 
@@ -441,9 +441,9 @@ pub fn StackFallbackAllocator(comptime size: usize) type {
             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);
+            const self: *@This() = @ptrCast(@alignCast(ctx));
+            if (mem.sliceOwnsPtr(u8, &self.buffer, @ptrCast(buf.ptr))) {
+                return BumpAllocator.resize(&self.bump_allocator, buf, alignment, new_len, ra);
             } else {
                 return self.fallback_allocator.rawResize(buf, alignment, new_len, ra);
             }
@@ -456,9 +456,9 @@ pub fn StackFallbackAllocator(comptime size: usize) type {
             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);
+            const self: *@This() = @ptrCast(@alignCast(context));
+            if (mem.sliceOwnsPtr(u8, &self.buffer, @ptrCast(memory.ptr))) {
+                return BumpAllocator.remap(&self.bump_allocator, memory, alignment, new_len, return_address);
             } else {
                 return self.fallback_allocator.rawRemap(memory, alignment, new_len, return_address);
             }
@@ -470,9 +470,9 @@ pub fn StackFallbackAllocator(comptime size: usize) type {
             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);
+            const self: *@This() = @ptrCast(@alignCast(ctx));
+            if (mem.sliceOwnsPtr(u8, &self.buffer, @ptrCast(buf.ptr))) {
+                return BumpAllocator.free(&self.bump_allocator, buf, alignment, ra);
             } else {
                 return self.fallback_allocator.rawFree(buf, alignment, ra);
             }
@@ -666,7 +666,7 @@ pub fn testAllocatorAlignedShrink(base_allocator: mem.Allocator) !void {
     const allocator = validationAllocator.allocator();
 
     var debug_buffer: [1000]u8 = undefined;
-    var fib = FixedBufferAllocator.init(&debug_buffer);
+    var fib: BumpAllocator = .init(&debug_buffer);
     const debug_allocator = fib.allocator();
 
     const alloc_size = pageSize() * 2 + 50;
@@ -990,7 +990,7 @@ test {
     _ = @import("heap/memory_pool.zig");
     _ = ArenaAllocator;
     _ = GeneralPurposeAllocator;
-    _ = FixedBufferAllocator;
+    _ = BumpAllocator;
     _ = ThreadSafeAllocator;
     _ = SbrkAllocator;
     if (builtin.target.cpu.arch.isWasm()) {

lib/std/heap/BumpAllocator.zig

@@ -0,0 +1,212 @@
+const std = @import("../std.zig");
+const Alignment = std.mem.Alignment;
+const Allocator = std.mem.Allocator;
+
+base: usize,
+limit: usize,
+
+pub fn init(buffer: []u8) @This() {
+    const base: usize = @intFromPtr(buffer.ptr);
+    const limit: usize = base + buffer.len;
+    return .{ .base = base, .limit = limit };
+}
+
+pub fn allocator(self: *@This()) Allocator {
+    return .{
+        .ptr = self,
+        .vtable = &.{
+            .alloc = alloc,
+            .resize = resize,
+            .remap = remap,
+            .free = free,
+        },
+    };
+}
+
+/// Save the current state of the allocator
+pub fn savestate(self: *@This()) usize {
+    return self.base;
+}
+
+/// Restore a previously saved allocator state
+pub fn restore(self: *@This(), state: usize) void {
+    self.base = state;
+}
+
+pub fn alloc(
+    ctx: *anyopaque,
+    length: usize,
+    alignment: Alignment,
+    _: usize,
+) ?[*]u8 {
+    const self: *@This() = @alignCast(@ptrCast(ctx));
+
+    // Only allocate if we have enough space
+    const aligned = alignment.forward(self.base);
+    const end_addr = @addWithOverflow(aligned, length);
+    if ((end_addr[1] == 1) | (end_addr[0] > self.limit)) return null;
+
+    self.base = end_addr[0];
+    return @ptrFromInt(aligned);
+}
+
+pub fn resize(
+    ctx: *anyopaque,
+    memory: []u8,
+    _: Alignment,
+    new_length: usize,
+    _: usize,
+) bool {
+    const self: *@This() = @alignCast(@ptrCast(ctx));
+
+    const alloc_base = @intFromPtr(memory.ptr);
+    const next_alloc = alloc_base + memory.len;
+
+    // Prior allocations can be shrunk, but not grown
+    const shrinking = memory.len >= new_length;
+    if (next_alloc != self.base) return shrinking;
+
+    // Grow allocations only if we have enough space
+    const end_addr = @addWithOverflow(alloc_base, new_length);
+    const overflow = (end_addr[1] == 1) | (end_addr[0] > self.limit);
+    if (!shrinking and overflow) return false;
+
+    self.base = end_addr[0];
+    return true;
+}
+
+pub fn remap(
+    ctx: *anyopaque,
+    memory: []u8,
+    _: Alignment,
+    new_length: usize,
+    _: usize,
+) ?[*]u8 {
+    if (resize(ctx, memory, undefined, new_length, undefined)) {
+        return memory.ptr;
+    } else {
+        return null;
+    }
+}
+
+pub fn free(
+    ctx: *anyopaque,
+    memory: []u8,
+    _: Alignment,
+    _: usize,
+) void {
+    const self: *@This() = @alignCast(@ptrCast(ctx));
+
+    // Only free the immediate last allocation
+    const alloc_base = @intFromPtr(memory.ptr);
+    const next_alloc = alloc_base + memory.len;
+    if (next_alloc != self.base) return;
+
+    self.base = self.base - memory.len;
+}
+
+test "BumpAllocator" {
+    var buffer: [1 << 20]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.allocator();
+
+    try std.heap.testAllocator(gpa);
+    try std.heap.testAllocatorAligned(gpa);
+    try std.heap.testAllocatorAlignedShrink(gpa);
+    try std.heap.testAllocatorLargeAlignment(gpa);
+}
+
+test "savestate and restore" {
+    var buffer: [256]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.allocator();
+
+    const state_before = bump_allocator.savestate();
+    _ = try gpa.alloc(u8, buffer.len);
+
+    bump_allocator.restore(state_before);
+    _ = try gpa.alloc(u8, buffer.len);
+}
+
+test "reuse memory on realloc" {
+    var buffer: [10]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.allocator();
+
+    const slice_0 = try gpa.alloc(u8, 5);
+    const slice_1 = try gpa.realloc(slice_0, 10);
+    try std.testing.expect(slice_1.ptr == slice_0.ptr);
+}
+
+test "don't grow one allocation into another" {
+    var buffer: [10]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.allocator();
+
+    const slice_0 = try gpa.alloc(u8, 3);
+    const slice_1 = try gpa.alloc(u8, 3);
+    const slice_2 = try gpa.realloc(slice_0, 4);
+    try std.testing.expect(slice_2.ptr == slice_1.ptr + 3);
+}
+
+test "avoid integer overflow for obscene allocations" {
+    var buffer: [10]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.allocator();
+
+    _ = try gpa.alloc(u8, 5);
+    const problem = gpa.alloc(u8, std.math.maxInt(usize));
+    try std.testing.expectError(error.OutOfMemory, problem);
+}
+
+/// Deprecated; to be removed after 0.16.0 is tagged.
+/// Provides a lock free thread safe `Allocator` interface to the underlying `FixedBufferAllocator`
+/// Using this at the same time as the interface returned by `allocator` is not thread safe.
+pub fn threadSafeAllocator(self: *@This()) Allocator {
+    return .{
+        .ptr = self,
+        .vtable = &.{
+            .alloc = threadSafeAlloc,
+            .resize = Allocator.noResize,
+            .remap = Allocator.noRemap,
+            .free = Allocator.noFree,
+        },
+    };
+}
+
+// Remove after 0.16.0 is tagged.
+fn threadSafeAlloc(
+    ctx: *anyopaque,
+    length: usize,
+    alignment: Alignment,
+    _: usize,
+) ?[*]u8 {
+    const self: *@This() = @alignCast(@ptrCast(ctx));
+
+    var old_base = @atomicLoad(usize, &self.base, .seq_cst);
+    while (true) {
+        // Only allocate if we have enough space
+        const aligned = alignment.forward(old_base);
+        const end_addr = @addWithOverflow(aligned, length);
+        if ((end_addr[1] == 1) | (end_addr[0] > self.limit)) return null;
+
+        if (@cmpxchgWeak(usize, &self.base, old_base, @intCast(end_addr[0]), .seq_cst, .seq_cst)) |prev| {
+            old_base = prev;
+            continue;
+        }
+
+        return @ptrFromInt(aligned);
+    }
+}
+
+// Remove after 0.16.0 is tagged.
+test "thread safe version" {
+    var buffer: [1 << 20]u8 = undefined;
+    var bump_allocator: @This() = .init(&buffer);
+    const gpa = bump_allocator.threadSafeAllocator();
+
+    try std.heap.testAllocator(gpa);
+    try std.heap.testAllocatorAligned(gpa);
+    try std.heap.testAllocatorAlignedShrink(gpa);
+    try std.heap.testAllocatorLargeAlignment(gpa);
+}