From f49d42729a22846ec54b6610db415ac8cdaa31db Mon Sep 17 00:00:00 2001
From: Ryan Liptak <squeek502@hotmail.com>
Date: Mon, 1 Nov 2021 00:38:11 -0700
Subject: std.ArrayList: add ensureTotalCapacityPrecise and update doc comments

initCapacity did and still does use the ensureTotalCapacityPrecise logic because the initial capacity of an ArrayList is not important in terms of how it grows, so allocating a more exact slice up-front allows for saving memory when the array list never exceeds that initial allocation size. There are use cases where this precise capacity is useful outside of the `init` function, though, like in instances where the user does not call the `init` function themselves but otherwise knows that an ArrayList is empty so calling `ensureTotalCapacityPrecise` can give the same memory savings that `initCapacity` would have.

Closes #9775
---
 lib/std/array_list.zig | 47 ++++++++++++++++++++++++++++++-----------------
 1 file changed, 30 insertions(+), 17 deletions(-)

(limited to 'lib/std/array_list.zig')

diff --git a/lib/std/array_list.zig b/lib/std/array_list.zig
index 77ba0646cf..8dbd9f8dae 100644
--- a/lib/std/array_list.zig
+++ b/lib/std/array_list.zig
@@ -56,19 +56,11 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
         }
 
         /// Initialize with capacity to hold at least `num` elements.
+        /// The resulting capacity is likely to be equal to `num`.
         /// Deinitialize with `deinit` or use `toOwnedSlice`.
         pub fn initCapacity(allocator: *Allocator, num: usize) !Self {
             var self = Self.init(allocator);
-
-            if (@sizeOf(T) > 0) {
-                const new_memory = try self.allocator.allocAdvanced(T, alignment, num, .at_least);
-                self.items.ptr = new_memory.ptr;
-                self.capacity = new_memory.len;
-            } else {
-                // If `T` is a zero-sized type, then we do not need to allocate memory.
-                self.capacity = std.math.maxInt(usize);
-            }
-
+            try self.ensureTotalCapacityPrecise(num);
             return self;
         }
 
@@ -330,8 +322,22 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
                     if (better_capacity >= new_capacity) break;
                 }
 
+                return self.ensureTotalCapacityPrecise(better_capacity);
+            } else {
+                self.capacity = std.math.maxInt(usize);
+            }
+        }
+
+        /// Modify the array so that it can hold at least `new_capacity` items.
+        /// Like `ensureTotalCapacity`, but the resulting capacity is much more likely
+        /// (but not guaranteed) to be equal to `new_capacity`.
+        /// Invalidates pointers if additional memory is needed.
+        pub fn ensureTotalCapacityPrecise(self: *Self, new_capacity: usize) !void {
+            if (@sizeOf(T) > 0) {
+                if (self.capacity >= new_capacity) return;
+
                 // TODO This can be optimized to avoid needlessly copying undefined memory.
-                const new_memory = try self.allocator.reallocAtLeast(self.allocatedSlice(), better_capacity);
+                const new_memory = try self.allocator.reallocAtLeast(self.allocatedSlice(), new_capacity);
                 self.items.ptr = new_memory.ptr;
                 self.capacity = new_memory.len;
             } else {
@@ -464,14 +470,11 @@ pub fn ArrayListAlignedUnmanaged(comptime T: type, comptime alignment: ?u29) typ
         pub const Slice = if (alignment) |a| ([]align(a) T) else []T;
 
         /// Initialize with capacity to hold at least num elements.
+        /// The resulting capacity is likely to be equal to `num`.
         /// Deinitialize with `deinit` or use `toOwnedSlice`.
         pub fn initCapacity(allocator: *Allocator, num: usize) !Self {
             var self = Self{};
-
-            const new_memory = try allocator.allocAdvanced(T, alignment, num, .at_least);
-            self.items.ptr = new_memory.ptr;
-            self.capacity = new_memory.len;
-
+            try self.ensureTotalCapacityPrecise(allocator, num);
             return self;
         }
 
@@ -685,7 +688,17 @@ pub fn ArrayListAlignedUnmanaged(comptime T: type, comptime alignment: ?u29) typ
                 if (better_capacity >= new_capacity) break;
             }
 
-            const new_memory = try allocator.reallocAtLeast(self.allocatedSlice(), better_capacity);
+            return self.ensureTotalCapacityPrecise(allocator, better_capacity);
+        }
+
+        /// Modify the array so that it can hold at least `new_capacity` items.
+        /// Like `ensureTotalCapacity`, but the resulting capacity is much more likely
+        /// (but not guaranteed) to be equal to `new_capacity`.
+        /// Invalidates pointers if additional memory is needed.
+        pub fn ensureTotalCapacityPrecise(self: *Self, allocator: *Allocator, new_capacity: usize) !void {
+            if (self.capacity >= new_capacity) return;
+
+            const new_memory = try allocator.reallocAtLeast(self.allocatedSlice(), new_capacity);
             self.items.ptr = new_memory.ptr;
             self.capacity = new_memory.len;
         }
-- 
cgit v1.2.3