authorgravatar for aleksey.kladov@gmail.comAlex Kladov <aleksey.kladov@gmail.com> 2024-05-13 15:17:11+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-05-20 12:04:20-04:00
log9f4f43cf7fdd2b2a50e475a1e7d740c9bcc2227e
tree0f34d876e3552abbc898bd813cc895f0026d9e4f
parent8aae0d87b5808e3674f2d37c8de7b10e305de9da

std: align PriorityQueue and ArrayList API-wise

ArrayList uses `items` slice to store len initialized items, while PriorityQueue stores `capacity` potentially uninitialized items. This is a surprising difference in the API that leads to bugs! https://github.com/tigerbeetle/tigerbeetle/pull/1948

1 files changed, 52 insertions(+), 38 deletions(-)

lib/std/priority_queue.zig+52-38
...@@ -19,7 +19,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -19,7 +19,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
19 const Self = @This();19 const Self = @This();
2020
21 items: []T,21 items: []T,
22 len: usize,22 cap: usize,
23 allocator: Allocator,23 allocator: Allocator,
24 context: Context,24 context: Context,
2525
...@@ -27,7 +27,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -27,7 +27,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
27 pub fn init(allocator: Allocator, context: Context) Self {27 pub fn init(allocator: Allocator, context: Context) Self {
28 return Self{28 return Self{
29 .items = &[_]T{},29 .items = &[_]T{},
30 .len = 0,30 .cap = 0,
31 .allocator = allocator,31 .allocator = allocator,
32 .context = context,32 .context = context,
33 };33 };
...@@ -35,7 +35,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -35,7 +35,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
3535
36 /// Free memory used by the queue.36 /// Free memory used by the queue.
37 pub fn deinit(self: Self) void {37 pub fn deinit(self: Self) void {
38 self.allocator.free(self.items);38 self.allocator.free(self.allocatedSlice());
39 }39 }
4040
41 /// Insert a new element, maintaining priority.41 /// Insert a new element, maintaining priority.
...@@ -45,9 +45,9 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -45,9 +45,9 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
45 }45 }
4646
47 fn addUnchecked(self: *Self, elem: T) void {47 fn addUnchecked(self: *Self, elem: T) void {
48 self.items[self.len] = elem;48 self.items.len += 1;
49 siftUp(self, self.len);49 self.items[self.items.len - 1] = elem;
50 self.len += 1;50 siftUp(self, self.items.len - 1);
51 }51 }
5252
53 fn siftUp(self: *Self, start_index: usize) void {53 fn siftUp(self: *Self, start_index: usize) void {
...@@ -74,13 +74,13 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -74,13 +74,13 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
74 /// Look at the highest priority element in the queue. Returns74 /// Look at the highest priority element in the queue. Returns
75 /// `null` if empty.75 /// `null` if empty.
76 pub fn peek(self: *Self) ?T {76 pub fn peek(self: *Self) ?T {
77 return if (self.len > 0) self.items[0] else null;77 return if (self.items.len > 0) self.items[0] else null;
78 }78 }
7979
80 /// Pop the highest priority element from the queue. Returns80 /// Pop the highest priority element from the queue. Returns
81 /// `null` if empty.81 /// `null` if empty.
82 pub fn removeOrNull(self: *Self) ?T {82 pub fn removeOrNull(self: *Self) ?T {
83 return if (self.len > 0) self.remove() else null;83 return if (self.items.len > 0) self.remove() else null;
84 }84 }
8585
86 /// Remove and return the highest priority element from the86 /// Remove and return the highest priority element from the
...@@ -93,13 +93,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -93,13 +93,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
93 /// same order as iterator, which is not necessarily priority93 /// same order as iterator, which is not necessarily priority
94 /// order.94 /// order.
95 pub fn removeIndex(self: *Self, index: usize) T {95 pub fn removeIndex(self: *Self, index: usize) T {
96 assert(self.len > index);96 assert(self.items.len > index);
97 const last = self.items[self.len - 1];97 const last = self.items[self.items.len - 1];
98 const item = self.items[index];98 const item = self.items[index];
99 self.items[index] = last;99 self.items[index] = last;
100 self.len -= 1;100 self.items.len -= 1;
101101
102 if (index == 0) {102 if (index == self.items.len) {
103 // Last element removed, nothing more to do.
104 } else if (index == 0) {
103 siftDown(self, index);105 siftDown(self, index);
104 } else {106 } else {
105 const parent_index = ((index - 1) >> 1);107 const parent_index = ((index - 1) >> 1);
...@@ -117,13 +119,20 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -117,13 +119,20 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
117 /// Return the number of elements remaining in the priority119 /// Return the number of elements remaining in the priority
118 /// queue.120 /// queue.
119 pub fn count(self: Self) usize {121 pub fn count(self: Self) usize {
120 return self.len;122 return self.items.len;
121 }123 }
122124
123 /// Return the number of elements that can be added to the125 /// Return the number of elements that can be added to the
124 /// queue before more memory is allocated.126 /// queue before more memory is allocated.
125 pub fn capacity(self: Self) usize {127 pub fn capacity(self: Self) usize {
126 return self.items.len;128 return self.cap;
129 }
130
131 /// Returns a slice of all the items plus the extra capacity, whose memory
132 /// contents are `undefined`.
133 fn allocatedSlice(self: Self) []T {
134 // `items.len` is the length, not the capacity.
135 return self.items.ptr[0..self.cap];
127 }136 }
128137
129 fn siftDown(self: *Self, target_index: usize) void {138 fn siftDown(self: *Self, target_index: usize) void {
...@@ -131,10 +140,10 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -131,10 +140,10 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
131 var index = target_index;140 var index = target_index;
132 while (true) {141 while (true) {
133 var lesser_child_i = (std.math.mul(usize, index, 2) catch break) | 1;142 var lesser_child_i = (std.math.mul(usize, index, 2) catch break) | 1;
134 if (!(lesser_child_i < self.len)) break;143 if (!(lesser_child_i < self.items.len)) break;
135144
136 const next_child_i = lesser_child_i + 1;145 const next_child_i = lesser_child_i + 1;
137 if (next_child_i < self.len and compareFn(self.context, self.items[next_child_i], self.items[lesser_child_i]) == .lt) {146 if (next_child_i < self.items.len and compareFn(self.context, self.items[next_child_i], self.items[lesser_child_i]) == .lt) {
138 lesser_child_i = next_child_i;147 lesser_child_i = next_child_i;
139 }148 }
140149
...@@ -152,12 +161,12 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -152,12 +161,12 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
152 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {161 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
153 var self = Self{162 var self = Self{
154 .items = items,163 .items = items,
155 .len = items.len,164 .cap = items.len,
156 .allocator = allocator,165 .allocator = allocator,
157 .context = context,166 .context = context,
158 };167 };
159168
160 var i = self.len >> 1;169 var i = self.items.len >> 1;
161 while (i > 0) {170 while (i > 0) {
162 i -= 1;171 i -= 1;
163 self.siftDown(i);172 self.siftDown(i);
...@@ -167,39 +176,45 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -167,39 +176,45 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
167176
168 /// Ensure that the queue can fit at least `new_capacity` items.177 /// Ensure that the queue can fit at least `new_capacity` items.
169 pub fn ensureTotalCapacity(self: *Self, new_capacity: usize) !void {178 pub fn ensureTotalCapacity(self: *Self, new_capacity: usize) !void {
170 var better_capacity = self.capacity();179 var better_capacity = self.cap;
171 if (better_capacity >= new_capacity) return;180 if (better_capacity >= new_capacity) return;
172 while (true) {181 while (true) {
173 better_capacity += better_capacity / 2 + 8;182 better_capacity += better_capacity / 2 + 8;
174 if (better_capacity >= new_capacity) break;183 if (better_capacity >= new_capacity) break;
175 }184 }
176 self.items = try self.allocator.realloc(self.items, better_capacity);185 const old_memory = self.allocatedSlice();
186 const new_memory = try self.allocator.realloc(old_memory, better_capacity);
187 self.items.ptr = new_memory.ptr;
188 self.cap = new_memory.len;
177 }189 }
178190
179 /// Ensure that the queue can fit at least `additional_count` **more** item.191 /// Ensure that the queue can fit at least `additional_count` **more** item.
180 pub fn ensureUnusedCapacity(self: *Self, additional_count: usize) !void {192 pub fn ensureUnusedCapacity(self: *Self, additional_count: usize) !void {
181 return self.ensureTotalCapacity(self.len + additional_count);193 return self.ensureTotalCapacity(self.items.len + additional_count);
182 }194 }
183195
184 /// Reduce allocated capacity to `new_len`.196 /// Reduce allocated capacity to `new_capacity`.
185 pub fn shrinkAndFree(self: *Self, new_len: usize) void {197 pub fn shrinkAndFree(self: *Self, new_capacity: usize) void {
186 assert(new_len <= self.items.len);198 assert(new_capacity <= self.cap);
187199
188 // Cannot shrink to smaller than the current queue size without invalidating the heap property200 // Cannot shrink to smaller than the current queue size without invalidating the heap property
189 assert(new_len >= self.len);201 assert(new_capacity >= self.items.len);
190202
191 self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {203 const old_memory = self.allocatedSlice();
204 const new_memory = self.allocator.realloc(old_memory, new_capacity) catch |e| switch (e) {
192 error.OutOfMemory => { // no problem, capacity is still correct then.205 error.OutOfMemory => { // no problem, capacity is still correct then.
193 self.items.len = new_len;
194 return;206 return;
195 },207 },
196 };208 };
209
210 self.items.ptr = new_memory.ptr;
211 self.cap = new_memory.len;
197 }212 }
198213
199 pub fn update(self: *Self, elem: T, new_elem: T) !void {214 pub fn update(self: *Self, elem: T, new_elem: T) !void {
200 const update_index = blk: {215 const update_index = blk: {
201 var idx: usize = 0;216 var idx: usize = 0;
202 while (idx < self.len) : (idx += 1) {217 while (idx < self.items.len) : (idx += 1) {
203 const item = self.items[idx];218 const item = self.items[idx];
204 if (compareFn(self.context, item, elem) == .eq) break :blk idx;219 if (compareFn(self.context, item, elem) == .eq) break :blk idx;
205 }220 }
...@@ -219,7 +234,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -219,7 +234,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
219 count: usize,234 count: usize,
220235
221 pub fn next(it: *Iterator) ?T {236 pub fn next(it: *Iterator) ?T {
222 if (it.count >= it.queue.len) return null;237 if (it.count >= it.queue.items.len) return null;
223 const out = it.count;238 const out = it.count;
224 it.count += 1;239 it.count += 1;
225 return it.queue.items[out];240 return it.queue.items[out];
...@@ -244,16 +259,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -244,16 +259,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
244 const print = std.debug.print;259 const print = std.debug.print;
245 print("{{ ", .{});260 print("{{ ", .{});
246 print("items: ", .{});261 print("items: ", .{});
247 for (self.items, 0..) |e, i| {262 for (self.items) |e| {
248 if (i >= self.len) break;
249 print("{}, ", .{e});263 print("{}, ", .{e});
250 }264 }
251 print("array: ", .{});265 print("array: ", .{});
252 for (self.items) |e| {266 for (self.items) |e| {
253 print("{}, ", .{e});267 print("{}, ", .{e});
254 }268 }
255 print("len: {} ", .{self.len});269 print("len: {} ", .{self.items.len});
256 print("capacity: {}", .{self.capacity()});270 print("capacity: {}", .{self.cap});
257 print(" }}\n", .{});271 print(" }}\n", .{});
258 }272 }
259 };273 };
...@@ -369,7 +383,7 @@ test "fromOwnedSlice trivial case 0" {...@@ -369,7 +383,7 @@ test "fromOwnedSlice trivial case 0" {
369 const queue_items = try testing.allocator.dupe(u32, &items);383 const queue_items = try testing.allocator.dupe(u32, &items);
370 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});384 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
371 defer queue.deinit();385 defer queue.deinit();
372 try expectEqual(@as(usize, 0), queue.len);386 try expectEqual(@as(usize, 0), queue.count());
373 try expect(queue.removeOrNull() == null);387 try expect(queue.removeOrNull() == null);
374}388}
375389
...@@ -379,7 +393,7 @@ test "fromOwnedSlice trivial case 1" {...@@ -379,7 +393,7 @@ test "fromOwnedSlice trivial case 1" {
379 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});393 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
380 defer queue.deinit();394 defer queue.deinit();
381395
382 try expectEqual(@as(usize, 1), queue.len);396 try expectEqual(@as(usize, 1), queue.count());
383 try expectEqual(items[0], queue.remove());397 try expectEqual(items[0], queue.remove());
384 try expect(queue.removeOrNull() == null);398 try expect(queue.removeOrNull() == null);
385}399}
...@@ -500,11 +514,11 @@ test "shrinkAndFree" {...@@ -500,11 +514,11 @@ test "shrinkAndFree" {
500 try queue.add(2);514 try queue.add(2);
501 try queue.add(3);515 try queue.add(3);
502 try expect(queue.capacity() >= 4);516 try expect(queue.capacity() >= 4);
503 try expectEqual(@as(usize, 3), queue.len);517 try expectEqual(@as(usize, 3), queue.count());
504518
505 queue.shrinkAndFree(3);519 queue.shrinkAndFree(3);
506 try expectEqual(@as(usize, 3), queue.capacity());520 try expectEqual(@as(usize, 3), queue.capacity());
507 try expectEqual(@as(usize, 3), queue.len);521 try expectEqual(@as(usize, 3), queue.count());
508522
509 try expectEqual(@as(u32, 1), queue.remove());523 try expectEqual(@as(u32, 1), queue.remove());
510 try expectEqual(@as(u32, 2), queue.remove());524 try expectEqual(@as(u32, 2), queue.remove());
...@@ -589,7 +603,7 @@ test "siftUp in remove" {...@@ -589,7 +603,7 @@ test "siftUp in remove" {
589603
590 try queue.addSlice(&.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 });604 try queue.addSlice(&.{ 0, 1, 100, 2, 3, 101, 102, 4, 5, 6, 7, 103, 104, 105, 106, 8 });
591605
592 _ = queue.removeIndex(std.mem.indexOfScalar(u32, queue.items[0..queue.len], 102).?);606 _ = queue.removeIndex(std.mem.indexOfScalar(u32, queue.items[0..queue.count()], 102).?);
593607
594 const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };608 const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };
595 for (sorted_items) |e| {609 for (sorted_items) |e| {