std: ArenaAllocator tries to resize before allocating

Closes ziglang#5116
LemonBoy · Sep 25, 2020 · bd9003e · bd9003e
1 parent bd89bd6
commit bd9003e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git 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
@@ -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;
+                },
+            };
         }
     }