From 47597a6d7cb7bbea313c15ba4c9f5dcecb972e64 Mon Sep 17 00:00:00 2001 From: "andrew.kraevskiii" Date: Thu, 12 Mar 2026 20:50:27 +0100 Subject: [PATCH] std.ArrayList: add toOwnedSliceAssert, shrinkToLen. (#30769) Adds a set of new APIs to ArrayList with goal of making it easier to avoid toOwnedSlice foot gun: ```zig fn wrongUsageOfCurrentApi(gpa: std.mem.Allocator) ![2][]const u8 { var a: std.ArrayList(u8) = .empty; defer a.deinit(gpa); var b: std.ArrayList(u8) = .empty; defer b.deinit(gpa); // ... return .{ try a.toOwnedSlice(gpa), try b.toOwnedSlice(gpa), // oom here causes leak :( }; } fn correctUsageOfCurrentApi(gpa: std.mem.Allocator) ![2][]const u8 { var a: std.ArrayList(u8) = .empty; defer a.deinit(gpa); var b: std.ArrayList(u8) = .empty; defer b.deinit(gpa); // ... const a_slice = try a.toOwnedSlice(gpa); errdefer gpa.free(a_slice); const b_slice = try b.toOwnedSlice(gpa); // may be omited but good if error appear bellow later in development. errdefer gpa.free(b_slice); return .{ a_slice, b_slice, }; } fn proposedApi(gpa: std.mem.Allocator) ![2][]const u8 { var a: std.ArrayList(u8) = .empty; defer a.deinit(gpa); var b: std.ArrayList(u8) = .empty; defer b.deinit(gpa); // ... try a.shrinkToLen(gpa); try b.shrinkToLen(gpa); return .{ a.toOwnedSliceAssert(), b.toOwnedSliceAssert(), }; } ``` Reviewed-on: https://codeberg.org/ziglang/zig/pulls/30769 Reviewed-by: Andrew Kelley Co-authored-by: andrew.kraevskiii Co-committed-by: andrew.kraevskiii --- lib/std/array_list.zig | 147 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 132 insertions(+), 15 deletions(-) diff --git a/lib/std/array_list.zig b/lib/std/array_list.zig index d763233beda0ec7805ed71af07932bfe88fecf8a..c51675030b60c7b3f729de74d1df2272d15518cb 100644 --- a/lib/std/array_list.zig +++ b/lib/std/array_list.zig @@ -674,6 +674,26 @@ pub fn Aligned(comptime T: type, comptime alignment: ?mem.Alignment) type { return result[0 .. result.len - 1 :sentinel]; } + /// The caller owns the returned memory. Empties this ArrayList. + /// Its capacity is cleared, making deinit() safe but unnecessary to call. + /// + /// Asserts what the capacity is equal to the length. + pub fn toOwnedSliceAssert(self: *Self) Slice { + assert(self.items.len == self.capacity); + const items = self.items; + self.* = .empty; + return items; + } + + /// The caller owns the returned memory. ArrayList becomes empty. + /// Asserts what the capacity is equal to the length + 1. + pub fn toOwnedSliceSentinelAssert(self: *Self, comptime sentinel: T) SentinelSlice(sentinel) { + std.debug.assert(self.items.len + 1 == self.capacity); + self.appendAssumeCapacity(sentinel); + const result = self.toOwnedSliceAssert(); + return result[0 .. result.len - 1 :sentinel]; + } + /// Creates a copy of this ArrayList. pub fn clone(self: Self, gpa: Allocator) Allocator.Error!Self { var cloned = try Self.initCapacity(gpa, self.capacity); @@ -1106,27 +1126,35 @@ pub fn Aligned(comptime T: type, comptime alignment: ?mem.Alignment) type { /// May invalidate element pointers. /// Asserts that the new length is less than or equal to the previous length. pub fn shrinkAndFree(self: *Self, gpa: Allocator, 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 (gpa.remap(old_memory, new_len)) |new_items| { - self.capacity = new_items.len; - self.items = new_items; - return; - } - - const new_memory = gpa.alignedAlloc(T, alignment, new_len) catch |e| switch (e) { + self.shrinkAndFreePrecise(gpa, new_len) catch |e| switch (e) { error.OutOfMemory => { // No problem, capacity is still correct then. self.items.len = new_len; return; }, }; + } + + /// Reduce allocated capacity to `new_len`. + /// May invalidate element pointers. + /// Asserts that the new length is less than or equal to the previous length. + /// If succeds capacity is guaranteed to be equal to the length. + pub fn shrinkAndFreePrecise(self: *Self, gpa: Allocator, new_len: usize) Allocator.Error!void { + assert(new_len <= self.items.len); + + if (@sizeOf(T) == 0) { + self.items.len = new_len; + return; + } + + const old_memory = self.allocatedSlice(); + if (gpa.remap(old_memory, new_len)) |new_items| { + self.capacity = new_items.len; + self.items = new_items; + return; + } + + const new_memory = try gpa.alignedAlloc(T, alignment, new_len); @memcpy(new_memory, self.items[0..new_len]); gpa.free(old_memory); @@ -1134,6 +1162,32 @@ pub fn Aligned(comptime T: type, comptime alignment: ?mem.Alignment) type { self.capacity = new_memory.len; } + /// Shrinks capacity to match length. + /// May invalidate element pointers. + /// If succeds it is safe to call toOwnedSliceAssert(). + pub fn shrinkToLen(self: *Self, gpa: Allocator) Allocator.Error!void { + try self.shrinkAndFreePrecise(gpa, self.items.len); + } + + /// Shrinks or expands capacity to match length + 1. + /// May invalidate element pointers. + /// If succeds it is safe to call toOwnedSliceSentinelAssert(). + pub fn shrinkToLenSentinel(self: *Self, gpa: Allocator) Allocator.Error!void { + std.debug.assert(self.items.len <= self.capacity); + const required_len = self.items.len + 1; + switch (std.math.order(required_len, self.capacity)) { + .eq => return, + .gt => { + try self.ensureTotalCapacityPrecise(gpa, required_len); + }, + .lt => { + self.items.len += 1; + defer self.items.len -= 1; + try self.shrinkToLen(gpa); + }, + } + } + /// Reduce length to `new_len`. /// Invalidates pointers to elements `items[new_len..]`. /// Keeps capacity the same. @@ -2089,6 +2143,30 @@ test "shrinkAndFree with a copy" { try testing.expect(mem.eql(i32, list.items, &.{ 3, 3, 3, 3 })); } +test "shrinkAndFreePrecise without resize succeeds" { + var failing_allocator = testing.FailingAllocator.init(testing.allocator, .{ .resize_fail_index = 0 }); + const a = failing_allocator.allocator(); + + var list: Aligned(i32, null) = .empty; + defer list.deinit(a); + + try list.appendNTimes(a, 3, 16); + try list.shrinkAndFreePrecise(a, 4); + try testing.expectEqualSlices(i32, &.{ 3, 3, 3, 3 }, list.items); + try testing.expectEqual(list.items.len, list.capacity); +} + +test "shrinkAndFreePrecise without resize and no copy failes" { + var failing_allocator = testing.FailingAllocator.init(testing.allocator, .{ .resize_fail_index = 0, .fail_index = 1 }); + const a = failing_allocator.allocator(); + + var list: Aligned(i32, null) = .empty; + defer list.deinit(a); + + try list.appendNTimes(a, 3, 16); + try std.testing.expectError(error.OutOfMemory, list.shrinkAndFreePrecise(a, 4)); +} + test "addManyAsArray" { const a = std.testing.allocator; { @@ -2219,6 +2297,45 @@ test "toOwnedSliceSentinel" { } } +test "toOwnedSliceAssert" { + var failing_allocator: testing.FailingAllocator = .init(testing.allocator, .{ + .fail_index = 2, + }); + const a = failing_allocator.allocator(); + + var list: Aligned(u8, null) = try .initCapacity(a, 6); // first alloc + list.appendSliceAssumeCapacity(&.{ 1, 2, 3 }); + + try list.shrinkToLen(a); // first resize + try std.testing.expectEqual(list.items.len, list.capacity); + try list.shrinkToLen(a); // no alloc or resize + + const slice = list.toOwnedSliceAssert(); + defer a.free(slice); + + try std.testing.expectEqual(Aligned(u8, null).empty, list); + try std.testing.expectEqualSlices(u8, &.{ 1, 2, 3 }, slice); +} + +test "toOwnedSliceSentinelAssert" { + const a = testing.allocator; + + var list: Aligned(u8, null) = try .initCapacity(a, 6); + list.appendSliceAssumeCapacity(&.{ 1, 2, 3 }); + + // shrinkToLenSentinel shrinks array + try list.shrinkToLenSentinel(a); + + // shrinkToLenSentinel expands array + try list.shrinkToLen(a); + try list.shrinkToLenSentinel(a); + + const slice = list.toOwnedSliceSentinelAssert(10); + defer a.free(slice); + + try std.testing.expectEqualSentinel(u8, 10, &.{ 1, 2, 3 }, slice); +} + test "accepts unaligned slices" { const a = testing.allocator; { -- 2.54.0