From bd9003ed5b16a6e187999fb1190d89eb80bd587b Mon Sep 17 00:00:00 2001
From: LemonBoy <thatlemon@gmail.com>
Date: Fri, 25 Sep 2020 00:21:57 +0200
Subject: [PATCH] std: ArenaAllocator tries to resize before allocating

Closes #5116
---
 lib/std/heap.zig                 |  7 +++++++
 lib/std/heap/arena_allocator.zig | 21 +++++++++++++++------
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/lib/std/heap.zig b/lib/std/heap.zig
index 16de215cc2e1..cf32cff645a8 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 0737cb2ef8ab..b7ee1d54c12b 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;
+                },
+            };
         }
     }