Commit 8c9f468fdd

alichraghi <alichraghi@pm.me>
2022-07-13 15:00:51
std.ArrayList.ensureTotalCapacity: optimize and fix integer overflow
Fixes #12099
1 parent 78cd7b5
Changed files (1)
lib/std/array_list.zig
@@ -3,6 +3,7 @@ const debug = std.debug;
 const assert = debug.assert;
 const testing = std.testing;
 const mem = std.mem;
+const math = std.math;
 const Allocator = mem.Allocator;
 
 /// A contiguous, growable list of items in memory.
@@ -330,17 +331,17 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
         /// Invalidates pointers if additional memory is needed.
         pub fn ensureTotalCapacity(self: *Self, new_capacity: usize) Allocator.Error!void {
             if (@sizeOf(T) > 0) {
-                var better_capacity = self.capacity;
-                if (better_capacity >= new_capacity) return;
+                if (self.capacity >= new_capacity) return;
 
+                var better_capacity = self.capacity;
                 while (true) {
-                    better_capacity += better_capacity / 2 + 8;
+                    better_capacity +|= better_capacity / 2 + 8;
                     if (better_capacity >= new_capacity) break;
                 }
 
                 return self.ensureTotalCapacityPrecise(better_capacity);
             } else {
-                self.capacity = std.math.maxInt(usize);
+                self.capacity = math.maxInt(usize);
             }
         }
 
@@ -357,7 +358,7 @@ pub fn ArrayListAligned(comptime T: type, comptime alignment: ?u29) type {
                 self.items.ptr = new_memory.ptr;
                 self.capacity = new_memory.len;
             } else {
-                self.capacity = std.math.maxInt(usize);
+                self.capacity = math.maxInt(usize);
             }
         }
 
@@ -725,11 +726,11 @@ pub fn ArrayListAlignedUnmanaged(comptime T: type, comptime alignment: ?u29) typ
         /// Modify the array so that it can hold at least `new_capacity` items.
         /// Invalidates pointers if additional memory is needed.
         pub fn ensureTotalCapacity(self: *Self, allocator: Allocator, new_capacity: usize) Allocator.Error!void {
-            var better_capacity = self.capacity;
-            if (better_capacity >= new_capacity) return;
+            if (self.capacity >= new_capacity) return;
 
+            var better_capacity = self.capacity;
             while (true) {
-                better_capacity += better_capacity / 2 + 8;
+                better_capacity +|= better_capacity / 2 + 8;
                 if (better_capacity >= new_capacity) break;
             }