Merge pull request #13666 from ziglang/allocator-interface

std.mem.Allocator: allow shrink to fail

1 files changed, 44 insertions, 23 deletions

diff --git a/lib/std/segmented_list.zig b/lib/std/segmented_list.zig
index 5b227b8c50..7ec29fe8d6 100644
--- a/lib/std/segmented_list.zig
+++ b/lib/std/segmented_list.zig
@@ -1,6 +1,7 @@
 const std = @import("std.zig");
 const assert = std.debug.assert;
 const testing = std.testing;
+const mem = std.mem;
 const Allocator = std.mem.Allocator;
 
 // Imagine that `fn at(self: *Self, index: usize) &T` is a customer asking for a box
@@ -177,24 +178,32 @@ pub fn SegmentedList(comptime T: type, comptime prealloc_item_count: usize) type
             return self.growCapacity(allocator, new_capacity);
         }
 
-        /// Only grows capacity, or retains current capacity
+        /// Only grows capacity, or retains current capacity.
         pub fn growCapacity(self: *Self, allocator: Allocator, new_capacity: usize) Allocator.Error!void {
             const new_cap_shelf_count = shelfCount(new_capacity);
             const old_shelf_count = @intCast(ShelfIndex, self.dynamic_segments.len);
-            if (new_cap_shelf_count > old_shelf_count) {
-                self.dynamic_segments = try allocator.realloc(self.dynamic_segments, new_cap_shelf_count);
-                var i = old_shelf_count;
-                errdefer {
-                    self.freeShelves(allocator, i, old_shelf_count);
-                    self.dynamic_segments = allocator.shrink(self.dynamic_segments, old_shelf_count);
-                }
-                while (i < new_cap_shelf_count) : (i += 1) {
-                    self.dynamic_segments[i] = (try allocator.alloc(T, shelfSize(i))).ptr;
-                }
+            if (new_cap_shelf_count <= old_shelf_count) return;
+
+            const new_dynamic_segments = try allocator.alloc([*]T, new_cap_shelf_count);
+            errdefer allocator.free(new_dynamic_segments);
+
+            var i: ShelfIndex = 0;
+            while (i < old_shelf_count) : (i += 1) {
+                new_dynamic_segments[i] = self.dynamic_segments[i];
             }
+            errdefer while (i > old_shelf_count) : (i -= 1) {
+                allocator.free(new_dynamic_segments[i][0..shelfSize(i)]);
+            };
+            while (i < new_cap_shelf_count) : (i += 1) {
+                new_dynamic_segments[i] = (try allocator.alloc(T, shelfSize(i))).ptr;
+            }
+
+            allocator.free(self.dynamic_segments);
+            self.dynamic_segments = new_dynamic_segments;
         }
 
-        /// Only shrinks capacity or retains current capacity
+        /// Only shrinks capacity or retains current capacity.
+        /// It may fail to reduce the capacity in which case the capacity will remain unchanged.
         pub fn shrinkCapacity(self: *Self, allocator: Allocator, new_capacity: usize) void {
             if (new_capacity <= prealloc_item_count) {
                 const len = @intCast(ShelfIndex, self.dynamic_segments.len);
@@ -207,12 +216,24 @@ pub fn SegmentedList(comptime T: type, comptime prealloc_item_count: usize) type
             const new_cap_shelf_count = shelfCount(new_capacity);
             const old_shelf_count = @intCast(ShelfIndex, self.dynamic_segments.len);
             assert(new_cap_shelf_count <= old_shelf_count);
-            if (new_cap_shelf_count == old_shelf_count) {
-                return;
-            }
+            if (new_cap_shelf_count == old_shelf_count) return;
 
+            // freeShelves() must be called before resizing the dynamic
+            // segments, but we don't know if resizing the dynamic segments
+            // will work until we try it. So we must allocate a fresh memory
+            // buffer in order to reduce capacity.
+            const new_dynamic_segments = allocator.alloc([*]T, new_cap_shelf_count) catch return;
             self.freeShelves(allocator, old_shelf_count, new_cap_shelf_count);
-            self.dynamic_segments = allocator.shrink(self.dynamic_segments, new_cap_shelf_count);
+            if (allocator.resize(self.dynamic_segments, new_cap_shelf_count)) {
+                // We didn't need the new memory allocation after all.
+                self.dynamic_segments = self.dynamic_segments[0..new_cap_shelf_count];
+                allocator.free(new_dynamic_segments);
+            } else {
+                // Good thing we allocated that new memory slice.
+                mem.copy([*]T, new_dynamic_segments, self.dynamic_segments[0..new_cap_shelf_count]);
+                allocator.free(self.dynamic_segments);
+                self.dynamic_segments = new_dynamic_segments;
+            }
         }
 
         pub fn shrink(self: *Self, new_len: usize) void {
@@ -227,10 +248,10 @@ pub fn SegmentedList(comptime T: type, comptime prealloc_item_count: usize) type
 
             var i = start;
             if (end <= prealloc_item_count) {
-                std.mem.copy(T, dest[i - start ..], self.prealloc_segment[i..end]);
+                mem.copy(T, dest[i - start ..], self.prealloc_segment[i..end]);
                 return;
             } else if (i < prealloc_item_count) {
-                std.mem.copy(T, dest[i - start ..], self.prealloc_segment[i..]);
+                mem.copy(T, dest[i - start ..], self.prealloc_segment[i..]);
                 i = prealloc_item_count;
             }
 
@@ -239,7 +260,7 @@ pub fn SegmentedList(comptime T: type, comptime prealloc_item_count: usize) type
                 const copy_start = boxIndex(i, shelf_index);
                 const copy_end = std.math.min(shelfSize(shelf_index), copy_start + end - i);
 
-                std.mem.copy(
+                mem.copy(
                     T,
                     dest[i - start ..],
                     self.dynamic_segments[shelf_index][copy_start..copy_end],
@@ -480,13 +501,13 @@ fn testSegmentedList(comptime prealloc: usize) !void {
             control[@intCast(usize, i)] = i + 1;
         }
 
-        std.mem.set(i32, dest[0..], 0);
+        mem.set(i32, dest[0..], 0);
         list.writeToSlice(dest[0..], 0);
-        try testing.expect(std.mem.eql(i32, control[0..], dest[0..]));
+        try testing.expect(mem.eql(i32, control[0..], dest[0..]));
 
-        std.mem.set(i32, dest[0..], 0);
+        mem.set(i32, dest[0..], 0);
         list.writeToSlice(dest[50..], 50);
-        try testing.expect(std.mem.eql(i32, control[50..], dest[50..]));
+        try testing.expect(mem.eql(i32, control[50..], dest[50..]));
     }
 
     try list.setCapacity(testing.allocator, 0);