Merge pull request #3769 from MCRusher/initcapacity-for-buffer-arraylist

Add initCapacity for buffer & arraylist

lib/std/array_list.zig

@@ -40,6 +40,14 @@ pub fn AlignedArrayList(comptime T: type, comptime alignment: ?u29) type {
                 .allocator = allocator,
             };
         }
+        
+        /// Initialize with capacity to hold at least num elements.
+        /// Deinitialize with `deinit` or use `toOwnedSlice`.
+        pub fn initCapacity(allocator: *Allocator, num: usize) !Self {
+            var self = Self.init(allocator);
+            try self.ensureCapacity(num);
+            return self;
+        }
 
         /// Release all allocated memory.
         pub fn deinit(self: Self) void {
@@ -271,6 +279,15 @@ test "std.ArrayList.init" {
     testing.expect(list.capacity() == 0);
 }
 
+test "std.ArrayList.initCapacity" {
+    var bytes: [1024]u8 = undefined;
+    const allocator = &std.heap.FixedBufferAllocator.init(bytes[0..]).allocator;
+    var list = try ArrayList(i8).initCapacity(allocator, 200);
+    defer list.deinit();
+    testing.expect(list.count() == 0);
+    testing.expect(list.capacity() >= 200);
+}
+
 test "std.ArrayList.basic" {
     var bytes: [1024]u8 = undefined;
     const allocator = &std.heap.FixedBufferAllocator.init(bytes[0..]).allocator;

lib/std/buffer.zig

@@ -16,13 +16,22 @@ pub const Buffer = struct {
         mem.copy(u8, self.list.items, m);
         return self;
     }
-
+    
+    /// Initialize memory to size bytes of undefined values.
     /// Must deinitialize with deinit.
     pub fn initSize(allocator: *Allocator, size: usize) !Buffer {
         var self = initNull(allocator);
         try self.resize(size);
         return self;
     }
+    
+    /// Initialize with capacity to hold at least num bytes.
+    /// Must deinitialize with deinit.
+    pub fn initCapacity(allocator: *Allocator, num: usize) !Buffer {
+        var self = Buffer{ .list = try ArrayList(u8).initCapacity(allocator, num + 1) };
+        self.list.appendAssumeCapacity(0);
+        return self;
+    }
 
     /// Must deinitialize with deinit.
     /// None of the other operations are valid until you do one of these:
@@ -98,6 +107,13 @@ pub const Buffer = struct {
     pub fn len(self: Buffer) usize {
         return self.list.len - 1;
     }
+    
+    pub fn capacity(self: Buffer) usize {
+        return if (self.list.items.len > 0)
+            self.list.items.len - 1
+        else
+            0;
+    }
 
     pub fn append(self: *Buffer, m: []const u8) !void {
         const old_len = self.len();
@@ -151,3 +167,21 @@ test "simple Buffer" {
     try buf2.resize(4);
     testing.expect(buf.startsWith(buf2.toSlice()));
 }
+
+test "Buffer.initSize" {
+    var buf = try Buffer.initSize(debug.global_allocator, 3);
+    testing.expect(buf.len() == 3);
+    try buf.append("hello");
+    testing.expect(mem.eql(u8, buf.toSliceConst()[3..], "hello"));
+}
+
+test "Buffer.initCapacity" {
+    var buf = try Buffer.initCapacity(debug.global_allocator, 10);
+    testing.expect(buf.len() == 0);
+    testing.expect(buf.capacity() >= 10);
+    const old_cap = buf.capacity();
+    try buf.append("hello");
+    testing.expect(buf.len() == 5);
+    testing.expect(buf.capacity() == old_cap);
+    testing.expect(mem.eql(u8, buf.toSliceConst(), "hello"));
+}