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
1919 const Self = @This();
2020
2121 items: []T,
22 len: usize,
22 cap: usize,
2323 allocator: Allocator,
2424 context: Context,
2525
......@@ -27,7 +27,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
2727 pub fn init(allocator: Allocator, context: Context) Self {
2828 return Self{
2929 .items = &[_]T{},
30 .len = 0,
30 .cap = 0,
3131 .allocator = allocator,
3232 .context = context,
3333 };
......@@ -35,7 +35,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
3535
3636 /// Free memory used by the queue.
3737 pub fn deinit(self: Self) void {
38 self.allocator.free(self.items);
38 self.allocator.free(self.allocatedSlice());
3939 }
4040
4141 /// Insert a new element, maintaining priority.
......@@ -45,9 +45,9 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
4545 }
4646
4747 fn addUnchecked(self: *Self, elem: T) void {
48 self.items[self.len] = elem;
49 siftUp(self, self.len);
50 self.len += 1;
48 self.items.len += 1;
49 self.items[self.items.len - 1] = elem;
50 siftUp(self, self.items.len - 1);
5151 }
5252
5353 fn siftUp(self: *Self, start_index: usize) void {
......@@ -74,13 +74,13 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
7474 /// Look at the highest priority element in the queue. Returns
7575 /// `null` if empty.
7676 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;
7878 }
7979
8080 /// Pop the highest priority element from the queue. Returns
8181 /// `null` if empty.
8282 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;
8484 }
8585
8686 /// Remove and return the highest priority element from the
......@@ -93,13 +93,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
9393 /// same order as iterator, which is not necessarily priority
9494 /// order.
9595 pub fn removeIndex(self: *Self, index: usize) T {
96 assert(self.len > index);
97 const last = self.items[self.len - 1];
96 assert(self.items.len > index);
97 const last = self.items[self.items.len - 1];
9898 const item = self.items[index];
9999 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) {
103105 siftDown(self, index);
104106 } else {
105107 const parent_index = ((index - 1) >> 1);
......@@ -117,13 +119,20 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
117119 /// Return the number of elements remaining in the priority
118120 /// queue.
119121 pub fn count(self: Self) usize {
120 return self.len;
122 return self.items.len;
121123 }
122124
123125 /// Return the number of elements that can be added to the
124126 /// queue before more memory is allocated.
125127 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];
127136 }
128137
129138 fn siftDown(self: *Self, target_index: usize) void {
......@@ -131,10 +140,10 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
131140 var index = target_index;
132141 while (true) {
133142 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
136145 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) {
138147 lesser_child_i = next_child_i;
139148 }
140149
......@@ -152,12 +161,12 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
152161 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
153162 var self = Self{
154163 .items = items,
155 .len = items.len,
164 .cap = items.len,
156165 .allocator = allocator,
157166 .context = context,
158167 };
159168
160 var i = self.len >> 1;
169 var i = self.items.len >> 1;
161170 while (i > 0) {
162171 i -= 1;
163172 self.siftDown(i);
......@@ -167,39 +176,45 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
167176
168177 /// Ensure that the queue can fit at least `new_capacity` items.
169178 pub fn ensureTotalCapacity(self: *Self, new_capacity: usize) !void {
170 var better_capacity = self.capacity();
179 var better_capacity = self.cap;
171180 if (better_capacity >= new_capacity) return;
172181 while (true) {
173182 better_capacity += better_capacity / 2 + 8;
174183 if (better_capacity >= new_capacity) break;
175184 }
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;
177189 }
178190
179191 /// Ensure that the queue can fit at least `additional_count` **more** item.
180192 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);
182194 }
183195
184 /// Reduce allocated capacity to `new_len`.
185 pub fn shrinkAndFree(self: *Self, new_len: usize) void {
186 assert(new_len <= self.items.len);
196 /// Reduce allocated capacity to `new_capacity`.
197 pub fn shrinkAndFree(self: *Self, new_capacity: usize) void {
198 assert(new_capacity <= self.cap);
187199
188200 // 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) {
192205 error.OutOfMemory => { // no problem, capacity is still correct then.
193 self.items.len = new_len;
194206 return;
195207 },
196208 };
209
210 self.items.ptr = new_memory.ptr;
211 self.cap = new_memory.len;
197212 }
198213
199214 pub fn update(self: *Self, elem: T, new_elem: T) !void {
200215 const update_index = blk: {
201216 var idx: usize = 0;
202 while (idx < self.len) : (idx += 1) {
217 while (idx < self.items.len) : (idx += 1) {
203218 const item = self.items[idx];
204219 if (compareFn(self.context, item, elem) == .eq) break :blk idx;
205220 }
......@@ -219,7 +234,7 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
219234 count: usize,
220235
221236 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;
223238 const out = it.count;
224239 it.count += 1;
225240 return it.queue.items[out];
......@@ -244,16 +259,15 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
244259 const print = std.debug.print;
245260 print("{{ ", .{});
246261 print("items: ", .{});
247 for (self.items, 0..) |e, i| {
248 if (i >= self.len) break;
262 for (self.items) |e| {
249263 print("{}, ", .{e});
250264 }
251265 print("array: ", .{});
252266 for (self.items) |e| {
253267 print("{}, ", .{e});
254268 }
255 print("len: {} ", .{self.len});
256 print("capacity: {}", .{self.capacity()});
269 print("len: {} ", .{self.items.len});
270 print("capacity: {}", .{self.cap});
257271 print(" }}\n", .{});
258272 }
259273 };
......@@ -369,7 +383,7 @@ test "fromOwnedSlice trivial case 0" {
369383 const queue_items = try testing.allocator.dupe(u32, &items);
370384 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
371385 defer queue.deinit();
372 try expectEqual(@as(usize, 0), queue.len);
386 try expectEqual(@as(usize, 0), queue.count());
373387 try expect(queue.removeOrNull() == null);
374388}
375389
......@@ -379,7 +393,7 @@ test "fromOwnedSlice trivial case 1" {
379393 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
380394 defer queue.deinit();
381395
382 try expectEqual(@as(usize, 1), queue.len);
396 try expectEqual(@as(usize, 1), queue.count());
383397 try expectEqual(items[0], queue.remove());
384398 try expect(queue.removeOrNull() == null);
385399}
......@@ -500,11 +514,11 @@ test "shrinkAndFree" {
500514 try queue.add(2);
501515 try queue.add(3);
502516 try expect(queue.capacity() >= 4);
503 try expectEqual(@as(usize, 3), queue.len);
517 try expectEqual(@as(usize, 3), queue.count());
504518
505519 queue.shrinkAndFree(3);
506520 try expectEqual(@as(usize, 3), queue.capacity());
507 try expectEqual(@as(usize, 3), queue.len);
521 try expectEqual(@as(usize, 3), queue.count());
508522
509523 try expectEqual(@as(u32, 1), queue.remove());
510524 try expectEqual(@as(u32, 2), queue.remove());
......@@ -589,7 +603,7 @@ test "siftUp in remove" {
589603
590604 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
594608 const sorted_items = [_]u32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 100, 101, 103, 104, 105, 106 };
595609 for (sorted_items) |e| {