From 16caea38d1fe0fb1ea7f18e6344f21f8c11d8dfe Mon Sep 17 00:00:00 2001
From: Andrew Kelley <andrew@ziglang.org>
Date: Wed, 30 Nov 2022 15:42:59 -0700
Subject: [PATCH] std.ArrayList: fix shrinkAndFree

Fixes a regression introduced in
e35f297aeb993ec956ae80379ddf7f86069e109b.

Now there is test coverage for ArrayList.shrinkAndFree in the case when
resizing fails.
---
 lib/std/array_list.zig           | 62 ++++++++++++++++----------------
 lib/std/heap/arena_allocator.zig |  5 +--
 2 files changed, 33 insertions(+), 34 deletions(-)

diff --git a/lib/std/array_list.zig b/lib/std/array_list.zig
index f1455c86bb..64e973ae53 100644
--- a/lib/std/array_list.zig
+++ b/lib/std/array_list.zig
@@ -314,32 +314,9 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
         /// Reduce allocated capacity to `new_len`.
         /// May invalidate element pointers.
         pub fn shrinkAndFree(self: *Self, new_len: usize) void {
-            assert(new_len <= self.items.len);
-
-            if (@sizeOf(T) == 0) {
-                self.items.len = new_len;
-                return;
-            }
-
-            const old_memory = self.allocatedSlice();
-            if (self.allocator.resize(old_memory, new_len)) {
-                self.capacity = new_len;
-                self.items.len = new_len;
-                return;
-            }
-
-            const new_memory = self.allocator.alignedAlloc(T, alignment, new_len) catch |e| switch (e) {
-                error.OutOfMemory => {
-                    // No problem, capacity is still correct then.
-                    self.items.len = new_len;
-                    return;
-                },
-            };
-
-            mem.copy(T, new_memory, self.items);
-            self.allocator.free(old_memory);
-            self.items = new_memory;
-            self.capacity = new_memory.len;
+            var unmanaged = self.moveToUnmanaged();
+            unmanaged.shrinkAndFree(self.allocator, new_len);
+            self.* = unmanaged.toManaged(self.allocator);
         }
 
         /// Reduce length to `new_len`.
@@ -782,7 +759,7 @@ pub fn ArrayListAlignedUnmanaged(comptime T: type, comptime alignment: ?u29) typ
                 },
             };
 
-            mem.copy(T, new_memory, self.items);
+            mem.copy(T, new_memory, self.items[0..new_len]);
             allocator.free(old_memory);
             self.items = new_memory;
             self.capacity = new_memory.len;
@@ -1488,14 +1465,18 @@ test "std.ArrayListUnmanaged(u8) implements writer" {
     }
 }
 
-test "std.ArrayList/ArrayListUnmanaged.shrink still sets length on error.OutOfMemory" {
-    // use an arena allocator to make sure realloc returns error.OutOfMemory
-    var arena = std.heap.ArenaAllocator.init(testing.allocator);
-    defer arena.deinit();
-    const a = arena.allocator();
+test "shrink still sets length when resizing is disabled" {
+    // Use the testing allocator but with resize disabled.
+    var a = testing.allocator;
+    a.vtable = &.{
+        .alloc = a.vtable.alloc,
+        .resize = Allocator.noResize,
+        .free = a.vtable.free,
+    };
 
     {
         var list = ArrayList(i32).init(a);
+        defer list.deinit();
 
         try list.append(1);
         try list.append(2);
@@ -1506,6 +1487,7 @@ test "std.ArrayList/ArrayListUnmanaged.shrink still sets length on error.OutOfMe
     }
     {
         var list = ArrayListUnmanaged(i32){};
+        defer list.deinit(a);
 
         try list.append(a, 1);
         try list.append(a, 2);
@@ -1516,6 +1498,22 @@ test "std.ArrayList/ArrayListUnmanaged.shrink still sets length on error.OutOfMe
     }
 }
 
+test "shrinkAndFree with a copy" {
+    // Use the testing allocator but with resize disabled.
+    var a = testing.allocator;
+    a.vtable = &.{
+        .alloc = a.vtable.alloc,
+        .resize = Allocator.noResize,
+        .free = a.vtable.free,
+    };
+    var list = ArrayList(i32).init(a);
+    defer list.deinit();
+
+    try list.appendNTimes(3, 16);
+    list.shrinkAndFree(4);
+    try testing.expect(mem.eql(i32, list.items, &.{ 3, 3, 3, 3 }));
+}
+
 test "std.ArrayList/ArrayListUnmanaged.addManyAsArray" {
     const a = std.testing.allocator;
     {
diff --git a/lib/std/heap/arena_allocator.zig b/lib/std/heap/arena_allocator.zig
index d6cb2bf5f1..ec9783185f 100644
--- a/lib/std/heap/arena_allocator.zig
+++ b/lib/std/heap/arena_allocator.zig
@@ -107,8 +107,9 @@ pub const ArenaAllocator = struct {
         const cur_node = self.state.buffer_list.first orelse return false;
         const cur_buf = cur_node.data[@sizeOf(BufNode)..];
         if (@ptrToInt(cur_buf.ptr) + self.state.end_index != @ptrToInt(buf.ptr) + buf.len) {
-            if (new_len > buf.len) return false;
-            return true;
+            // It's not the most recent allocation, so it cannot be expanded,
+            // but it's fine if they want to make it smaller.
+            return new_len <= buf.len;
         }
 
         if (buf.len >= new_len) {