Merge pull request #6413 from LemonBoy/fix-5116

Make ArenaAllocator try to resize first

2 files changed, 22 insertions, 6 deletions

diff --git a/lib/std/heap.zig b/lib/std/heap.zig
index 16de215cc2..cf32cff645 100644
--- a/lib/std/heap.zig
+++ b/lib/std/heap.zig
@@ -919,6 +919,13 @@ pub fn testAllocator(base_allocator: *mem.Allocator) !void {
     const zero_bit_ptr = try allocator.create(u0);
     zero_bit_ptr.* = 0;
     allocator.destroy(zero_bit_ptr);
+
+    const oversize = try allocator.allocAdvanced(u32, null, 5, .at_least);
+    testing.expect(oversize.len >= 5);
+    for (oversize) |*item| {
+        item.* = 0xDEADBEEF;
+    }
+    allocator.free(oversize);
 }
 
 pub fn testAllocatorAligned(base_allocator: *mem.Allocator, comptime alignment: u29) !void {
diff --git a/lib/std/heap/arena_allocator.zig b/lib/std/heap/arena_allocator.zig
index 0737cb2ef8..b7ee1d54c1 100644
--- a/lib/std/heap/arena_allocator.zig
+++ b/lib/std/heap/arena_allocator.zig
@@ -75,13 +75,22 @@ pub const ArenaAllocator = struct {
             const adjusted_addr = mem.alignForward(addr, ptr_align);
             const adjusted_index = self.state.end_index + (adjusted_addr - addr);
             const new_end_index = adjusted_index + n;
-            if (new_end_index > cur_buf.len) {
-                cur_node = try self.createNode(cur_buf.len, n + ptr_align);
-                continue;
+
+            if (new_end_index <= cur_buf.len) {
+                const result = cur_buf[adjusted_index..new_end_index];
+                self.state.end_index = new_end_index;
+                return result;
             }
-            const result = cur_buf[adjusted_index..new_end_index];
-            self.state.end_index = new_end_index;
-            return result;
+
+            const bigger_buf_size = @sizeOf(BufNode) + new_end_index;
+            // Try to grow the buffer in-place
+            cur_node.data = self.child_allocator.resize(cur_node.data, bigger_buf_size) catch |err| switch (err) {
+                error.OutOfMemory => {
+                    // Allocate a new node if that's not possible
+                    cur_node = try self.createNode(cur_buf.len, n + ptr_align);
+                    continue;
+                },
+            };
         }
     }