std: add new array list functions

1 files changed, 60 insertions, 0 deletions

diff --git a/lib/std/array_list.zig b/lib/std/array_list.zig
index a68c1fa9d6..d667bc4d17 100644
--- a/lib/std/array_list.zig
+++ b/lib/std/array_list.zig
@@ -257,6 +257,24 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
             return &self.items[self.items.len - 1];
         }
 
+        /// Resize the array, adding `n` new elements, which have `undefined` values.
+        /// The return value is an array pointing to the newly allocated elements.
+        pub fn addManyAsArray(self: *Self, comptime n: usize) !*[n]T {
+            const prev_len = self.items.len;
+            try self.resize(self.items.len + n);
+            return self.items[prev_len..][0..n];
+        }
+
+        /// Resize the array, adding `n` new elements, which have `undefined` values.
+        /// The return value is an array pointing to the newly allocated elements.
+        /// Asserts that there is already space for the new item without allocating more.
+        pub fn addManyAsArrayAssumeCapacity(self: *Self, comptime n: usize) *[n]T {
+            assert(self.items.len + n <= self.capacity);
+            const prev_len = self.items.len;
+            self.items.len += n;
+            return self.items[prev_len..][0..n];
+        }
+
         /// Remove and return the last element from the list.
         /// Asserts the list has at least one item.
         pub fn pop(self: *Self) T {
@@ -488,6 +506,24 @@ pub fn ArrayListAlignedUnmanaged(comptime T: type, comptime alignment: ?u29) typ
             return &self.items[self.items.len - 1];
         }
 
+        /// Resize the array, adding `n` new elements, which have `undefined` values.
+        /// The return value is an array pointing to the newly allocated elements.
+        pub fn addManyAsArray(self: *Self, allocator: *Allocator, comptime n: usize) !*[n]T {
+            const prev_len = self.items.len;
+            try self.resize(allocator, self.items.len + n);
+            return self.items[prev_len..][0..n];
+        }
+
+        /// Resize the array, adding `n` new elements, which have `undefined` values.
+        /// The return value is an array pointing to the newly allocated elements.
+        /// Asserts that there is already space for the new item without allocating more.
+        pub fn addManyAsArrayAssumeCapacity(self: *Self, comptime n: usize) *[n]T {
+            assert(self.items.len + n <= self.capacity);
+            const prev_len = self.items.len;
+            self.items.len += n;
+            return self.items[prev_len..][0..n];
+        }
+
         /// Remove and return the last element from the list.
         /// Asserts the list has at least one item.
         /// This operation does not invalidate any element pointers.
@@ -727,3 +763,27 @@ test "std.ArrayList.writer" {
     try writer.writeAll("efg");
     testing.expectEqualSlices(u8, list.items, "abcdefg");
 }
+
+test "addManyAsArray" {
+    const a = std.testing.allocator;
+    {
+        var list = ArrayList(u8).init(a);
+        defer list.deinit();
+
+        (try list.addManyAsArray(4)).* = "aoeu".*;
+        try list.ensureCapacity(8);
+        list.addManyAsArrayAssumeCapacity(4).* = "asdf".*;
+
+        testing.expectEqualSlices(u8, list.items, "aoeuasdf");
+    }
+    {
+        var list = ArrayListUnmanaged(u8){};
+        defer list.deinit(a);
+
+        (try list.addManyAsArray(a, 4)).* = "aoeu".*;
+        try list.ensureCapacity(a, 8);
+        list.addManyAsArrayAssumeCapacity(4).* = "asdf".*;
+
+        testing.expectEqualSlices(u8, list.items, "aoeuasdf");
+    }
+}