authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2021-04-02 10:52:23-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2021-04-02 10:52:23-07:00
logd780848ae404884b2f4c5d863dab980cb153cd8b
treed60c08b92369b2963d091698447d461211b56f24
parent833f258297fddd699be48b125bcfd8511dc2b470
parentce22c70586137bd20de101f3b1ec3a44fbdb54e7
signature Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #7792 from zanderxyz/zanderxyz/priority-dequeue

std: Add Priority Dequeue

3 files changed, 1068 insertions(+), 26 deletions(-)

lib/std/priority_dequeue.zig created+972
...@@ -0,0 +1,972 @@
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2015-2021 Zig Contributors
3// This file is part of [zig](https://ziglang.org/), which is MIT licensed.
4// The MIT license requires this copyright notice to be included in all copies
5// and substantial portions of the software.
6const std = @import("std.zig");
7const Allocator = std.mem.Allocator;
8const assert = std.debug.assert;
9const warn = std.debug.warn;
10const Order = std.math.Order;
11const testing = std.testing;
12const expect = testing.expect;
13const expectEqual = testing.expectEqual;
14const expectError = testing.expectError;
15
16/// Priority Dequeue for storing generic data. Initialize with `init`.
17pub fn PriorityDequeue(comptime T: type) type {
18 return struct {
19 const Self = @This();
20
21 items: []T,
22 len: usize,
23 allocator: *Allocator,
24 compareFn: fn (a: T, b: T) Order,
25
26 /// Initialize and return a new priority dequeue. Provide `compareFn`
27 /// that returns `Order.lt` when its first argument should
28 /// get min-popped before its second argument, `Order.eq` if the
29 /// arguments are of equal priority, or `Order.gt` if the second
30 /// argument should be min-popped first. Popping the max element works
31 /// in reverse. For example, to make `popMin` return the smallest
32 /// number, provide
33 ///
34 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
35 pub fn init(allocator: *Allocator, compareFn: fn (T, T) Order) Self {
36 return Self{
37 .items = &[_]T{},
38 .len = 0,
39 .allocator = allocator,
40 .compareFn = compareFn,
41 };
42 }
43
44 /// Free memory used by the dequeue.
45 pub fn deinit(self: Self) void {
46 self.allocator.free(self.items);
47 }
48
49 /// Insert a new element, maintaining priority.
50 pub fn add(self: *Self, elem: T) !void {
51 try ensureCapacity(self, self.len + 1);
52 addUnchecked(self, elem);
53 }
54
55 /// Add each element in `items` to the dequeue.
56 pub fn addSlice(self: *Self, items: []const T) !void {
57 try self.ensureCapacity(self.len + items.len);
58 for (items) |e| {
59 self.addUnchecked(e);
60 }
61 }
62
63 fn addUnchecked(self: *Self, elem: T) void {
64 self.items[self.len] = elem;
65
66 if (self.len > 0) {
67 const start = self.getStartForSiftUp(elem, self.len);
68 self.siftUp(start);
69 }
70
71 self.len += 1;
72 }
73
74 fn isMinLayer(index: usize) bool {
75 // In the min-max heap structure:
76 // The first element is on a min layer;
77 // next two are on a max layer;
78 // next four are on a min layer, and so on.
79 const leading_zeros = @clz(usize, index + 1);
80 const highest_set_bit = @bitSizeOf(usize) - 1 - leading_zeros;
81 return (highest_set_bit & 1) == 0;
82 }
83
84 fn nextIsMinLayer(self: Self) bool {
85 return isMinLayer(self.len);
86 }
87
88 const StartIndexAndLayer = struct {
89 index: usize,
90 min_layer: bool,
91 };
92
93 fn getStartForSiftUp(self: Self, child: T, index: usize) StartIndexAndLayer {
94 var child_index = index;
95 var parent_index = parentIndex(child_index);
96 const parent = self.items[parent_index];
97
98 const min_layer = self.nextIsMinLayer();
99 const order = self.compareFn(child, parent);
100 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
101 // We must swap the item with it's parent if it is on the "wrong" layer
102 self.items[parent_index] = child;
103 self.items[child_index] = parent;
104 return .{
105 .index = parent_index,
106 .min_layer = !min_layer,
107 };
108 } else {
109 return .{
110 .index = child_index,
111 .min_layer = min_layer,
112 };
113 }
114 }
115
116 fn siftUp(self: *Self, start: StartIndexAndLayer) void {
117 if (start.min_layer) {
118 doSiftUp(self, start.index, .lt);
119 } else {
120 doSiftUp(self, start.index, .gt);
121 }
122 }
123
124 fn doSiftUp(self: *Self, start_index: usize, target_order: Order) void {
125 var child_index = start_index;
126 while (child_index > 2) {
127 var grandparent_index = grandparentIndex(child_index);
128 const child = self.items[child_index];
129 const grandparent = self.items[grandparent_index];
130
131 // If the grandparent is already better or equal, we have gone as far as we need to
132 if (self.compareFn(child, grandparent) != target_order) break;
133
134 // Otherwise swap the item with it's grandparent
135 self.items[grandparent_index] = child;
136 self.items[child_index] = grandparent;
137 child_index = grandparent_index;
138 }
139 }
140
141 /// Look at the smallest element in the dequeue. Returns
142 /// `null` if empty.
143 pub fn peekMin(self: *Self) ?T {
144 return if (self.len > 0) self.items[0] else null;
145 }
146
147 /// Look at the largest element in the dequeue. Returns
148 /// `null` if empty.
149 pub fn peekMax(self: *Self) ?T {
150 if (self.len == 0) return null;
151 if (self.len == 1) return self.items[0];
152 if (self.len == 2) return self.items[1];
153 return self.bestItemAtIndices(1, 2, .gt).item;
154 }
155
156 fn maxIndex(self: Self) ?usize {
157 if (self.len == 0) return null;
158 if (self.len == 1) return 0;
159 if (self.len == 2) return 1;
160 return self.bestItemAtIndices(1, 2, .gt).index;
161 }
162
163 /// Pop the smallest element from the dequeue. Returns
164 /// `null` if empty.
165 pub fn removeMinOrNull(self: *Self) ?T {
166 return if (self.len > 0) self.removeMin() else null;
167 }
168
169 /// Remove and return the smallest element from the
170 /// dequeue.
171 pub fn removeMin(self: *Self) T {
172 return self.removeIndex(0);
173 }
174
175 /// Pop the largest element from the dequeue. Returns
176 /// `null` if empty.
177 pub fn removeMaxOrNull(self: *Self) ?T {
178 return if (self.len > 0) self.removeMax() else null;
179 }
180
181 /// Remove and return the largest element from the
182 /// dequeue.
183 pub fn removeMax(self: *Self) T {
184 return self.removeIndex(self.maxIndex().?);
185 }
186
187 /// Remove and return element at index. Indices are in the
188 /// same order as iterator, which is not necessarily priority
189 /// order.
190 pub fn removeIndex(self: *Self, index: usize) T {
191 assert(self.len > index);
192 const item = self.items[index];
193 const last = self.items[self.len - 1];
194
195 self.items[index] = last;
196 self.len -= 1;
197 siftDown(self, index);
198
199 return item;
200 }
201
202 fn siftDown(self: *Self, index: usize) void {
203 if (isMinLayer(index)) {
204 self.doSiftDown(index, .lt);
205 } else {
206 self.doSiftDown(index, .gt);
207 }
208 }
209
210 fn doSiftDown(self: *Self, start_index: usize, target_order: Order) void {
211 var index = start_index;
212 const half = self.len >> 1;
213 while (true) {
214 const first_grandchild_index = firstGrandchildIndex(index);
215 const last_grandchild_index = first_grandchild_index + 3;
216
217 const elem = self.items[index];
218
219 if (last_grandchild_index < self.len) {
220 // All four grandchildren exist
221 const index2 = first_grandchild_index + 1;
222 const index3 = index2 + 1;
223
224 // Find the best grandchild
225 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
226 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
227 const best_grandchild = self.bestItem(best_left, best_right, target_order);
228
229 // If the item is better than or equal to its best grandchild, we are done
230 if (self.compareFn(best_grandchild.item, elem) != target_order) return;
231
232 // Otherwise, swap them
233 self.items[best_grandchild.index] = elem;
234 self.items[index] = best_grandchild.item;
235 index = best_grandchild.index;
236
237 // We might need to swap the element with it's parent
238 self.swapIfParentIsBetter(elem, index, target_order);
239 } else {
240 // The children or grandchildren are the last layer
241 const first_child_index = firstChildIndex(index);
242 if (first_child_index > self.len) return;
243
244 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
245
246 // If the item is better than or equal to its best descendant, we are done
247 if (self.compareFn(best_descendent.item, elem) != target_order) return;
248
249 // Otherwise swap them
250 self.items[best_descendent.index] = elem;
251 self.items[index] = best_descendent.item;
252 index = best_descendent.index;
253
254 // If we didn't swap a grandchild, we are done
255 if (index < first_grandchild_index) return;
256
257 // We might need to swap the element with it's parent
258 self.swapIfParentIsBetter(elem, index, target_order);
259 return;
260 }
261
262 // If we are now in the last layer, we are done
263 if (index >= half) return;
264 }
265 }
266
267 fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, target_order: Order) void {
268 const parent_index = parentIndex(child_index);
269 const parent = self.items[parent_index];
270
271 if (self.compareFn(parent, child) == target_order) {
272 self.items[parent_index] = child;
273 self.items[child_index] = parent;
274 }
275 }
276
277 const ItemAndIndex = struct {
278 item: T,
279 index: usize,
280 };
281
282 fn getItem(self: Self, index: usize) ItemAndIndex {
283 return .{
284 .item = self.items[index],
285 .index = index,
286 };
287 }
288
289 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
290 if (self.compareFn(item1.item, item2.item) == target_order) {
291 return item1;
292 } else {
293 return item2;
294 }
295 }
296
297 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
298 var item1 = self.getItem(index1);
299 var item2 = self.getItem(index2);
300 return self.bestItem(item1, item2, target_order);
301 }
302
303 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
304 const second_child_index = first_child_index + 1;
305 if (first_grandchild_index >= self.len) {
306 // No grandchildren, find the best child (second may not exist)
307 if (second_child_index >= self.len) {
308 return .{
309 .item = self.items[first_child_index],
310 .index = first_child_index,
311 };
312 } else {
313 return self.bestItemAtIndices(first_child_index, second_child_index, target_order);
314 }
315 }
316
317 const second_grandchild_index = first_grandchild_index + 1;
318 if (second_grandchild_index >= self.len) {
319 // One grandchild, so we know there is a second child. Compare first grandchild and second child
320 return self.bestItemAtIndices(first_grandchild_index, second_child_index, target_order);
321 }
322
323 const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, target_order).index;
324 const third_grandchild_index = second_grandchild_index + 1;
325 if (third_grandchild_index >= self.len) {
326 // Two grandchildren, and we know the best. Compare this to second child.
327 return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, target_order);
328 } else {
329 // Three grandchildren, compare the min of the first two with the third
330 return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, target_order);
331 }
332 }
333
334 /// Return the number of elements remaining in the dequeue
335 pub fn count(self: Self) usize {
336 return self.len;
337 }
338
339 /// Return the number of elements that can be added to the
340 /// dequeue before more memory is allocated.
341 pub fn capacity(self: Self) usize {
342 return self.items.len;
343 }
344
345 /// Dequeue takes ownership of the passed in slice. The slice must have been
346 /// allocated with `allocator`.
347 /// De-initialize with `deinit`.
348 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (T, T) Order, items: []T) Self {
349 var queue = Self{
350 .items = items,
351 .len = items.len,
352 .allocator = allocator,
353 .compareFn = compareFn,
354 };
355
356 if (queue.len <= 1) return queue;
357
358 const half = (queue.len >> 1) - 1;
359 var i: usize = 0;
360 while (i <= half) : (i += 1) {
361 const index = half - i;
362 queue.siftDown(index);
363 }
364 return queue;
365 }
366
367 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
368 var better_capacity = self.capacity();
369 if (better_capacity >= new_capacity) return;
370 while (true) {
371 better_capacity += better_capacity / 2 + 8;
372 if (better_capacity >= new_capacity) break;
373 }
374 self.items = try self.allocator.realloc(self.items, better_capacity);
375 }
376
377 /// Reduce allocated capacity to `new_len`.
378 pub fn shrinkAndFree(self: *Self, new_len: usize) void {
379 assert(new_len <= self.items.len);
380
381 // Cannot shrink to smaller than the current queue size without invalidating the heap property
382 assert(new_len >= self.len);
383
384 self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
385 error.OutOfMemory => { // no problem, capacity is still correct then.
386 self.items.len = new_len;
387 return;
388 },
389 };
390 self.len = new_len;
391 }
392
393 /// Reduce length to `new_len`.
394 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
395 assert(new_len <= self.items.len);
396
397 // Cannot shrink to smaller than the current queue size without invalidating the heap property
398 assert(new_len >= self.len);
399
400 self.len = new_len;
401 }
402
403 pub fn update(self: *Self, elem: T, new_elem: T) !void {
404 var old_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
405 _ = self.removeIndex(old_index);
406 self.addUnchecked(new_elem);
407 }
408
409 pub const Iterator = struct {
410 queue: *PriorityDequeue(T),
411 count: usize,
412
413 pub fn next(it: *Iterator) ?T {
414 if (it.count >= it.queue.len) return null;
415 const out = it.count;
416 it.count += 1;
417 return it.queue.items[out];
418 }
419
420 pub fn reset(it: *Iterator) void {
421 it.count = 0;
422 }
423 };
424
425 /// Return an iterator that walks the queue without consuming
426 /// it. Invalidated if the queue is modified.
427 pub fn iterator(self: *Self) Iterator {
428 return Iterator{
429 .queue = self,
430 .count = 0,
431 };
432 }
433
434 fn dump(self: *Self) void {
435 warn("{{ ", .{});
436 warn("items: ", .{});
437 for (self.items) |e, i| {
438 if (i >= self.len) break;
439 warn("{}, ", .{e});
440 }
441 warn("array: ", .{});
442 for (self.items) |e, i| {
443 warn("{}, ", .{e});
444 }
445 warn("len: {} ", .{self.len});
446 warn("capacity: {}", .{self.capacity()});
447 warn(" }}\n", .{});
448 }
449
450 fn parentIndex(index: usize) usize {
451 return (index - 1) >> 1;
452 }
453
454 fn grandparentIndex(index: usize) usize {
455 return parentIndex(parentIndex(index));
456 }
457
458 fn firstChildIndex(index: usize) usize {
459 return (index << 1) + 1;
460 }
461
462 fn firstGrandchildIndex(index: usize) usize {
463 return firstChildIndex(firstChildIndex(index));
464 }
465 };
466}
467
468fn lessThanComparison(a: u32, b: u32) Order {
469 return std.math.order(a, b);
470}
471
472const PDQ = PriorityDequeue(u32);
473
474test "std.PriorityDequeue: add and remove min" {
475 var queue = PDQ.init(testing.allocator, lessThanComparison);
476 defer queue.deinit();
477
478 try queue.add(54);
479 try queue.add(12);
480 try queue.add(7);
481 try queue.add(23);
482 try queue.add(25);
483 try queue.add(13);
484
485 expectEqual(@as(u32, 7), queue.removeMin());
486 expectEqual(@as(u32, 12), queue.removeMin());
487 expectEqual(@as(u32, 13), queue.removeMin());
488 expectEqual(@as(u32, 23), queue.removeMin());
489 expectEqual(@as(u32, 25), queue.removeMin());
490 expectEqual(@as(u32, 54), queue.removeMin());
491}
492
493test "std.PriorityDequeue: add and remove min structs" {
494 const S = struct {
495 size: u32,
496 };
497 var queue = PriorityDequeue(S).init(testing.allocator, struct {
498 fn order(a: S, b: S) Order {
499 return std.math.order(a.size, b.size);
500 }
501 }.order);
502 defer queue.deinit();
503
504 try queue.add(.{ .size = 54 });
505 try queue.add(.{ .size = 12 });
506 try queue.add(.{ .size = 7 });
507 try queue.add(.{ .size = 23 });
508 try queue.add(.{ .size = 25 });
509 try queue.add(.{ .size = 13 });
510
511 expectEqual(@as(u32, 7), queue.removeMin().size);
512 expectEqual(@as(u32, 12), queue.removeMin().size);
513 expectEqual(@as(u32, 13), queue.removeMin().size);
514 expectEqual(@as(u32, 23), queue.removeMin().size);
515 expectEqual(@as(u32, 25), queue.removeMin().size);
516 expectEqual(@as(u32, 54), queue.removeMin().size);
517}
518
519test "std.PriorityDequeue: add and remove max" {
520 var queue = PDQ.init(testing.allocator, lessThanComparison);
521 defer queue.deinit();
522
523 try queue.add(54);
524 try queue.add(12);
525 try queue.add(7);
526 try queue.add(23);
527 try queue.add(25);
528 try queue.add(13);
529
530 expectEqual(@as(u32, 54), queue.removeMax());
531 expectEqual(@as(u32, 25), queue.removeMax());
532 expectEqual(@as(u32, 23), queue.removeMax());
533 expectEqual(@as(u32, 13), queue.removeMax());
534 expectEqual(@as(u32, 12), queue.removeMax());
535 expectEqual(@as(u32, 7), queue.removeMax());
536}
537
538test "std.PriorityDequeue: add and remove same min" {
539 var queue = PDQ.init(testing.allocator, lessThanComparison);
540 defer queue.deinit();
541
542 try queue.add(1);
543 try queue.add(1);
544 try queue.add(2);
545 try queue.add(2);
546 try queue.add(1);
547 try queue.add(1);
548
549 expectEqual(@as(u32, 1), queue.removeMin());
550 expectEqual(@as(u32, 1), queue.removeMin());
551 expectEqual(@as(u32, 1), queue.removeMin());
552 expectEqual(@as(u32, 1), queue.removeMin());
553 expectEqual(@as(u32, 2), queue.removeMin());
554 expectEqual(@as(u32, 2), queue.removeMin());
555}
556
557test "std.PriorityDequeue: add and remove same max" {
558 var queue = PDQ.init(testing.allocator, lessThanComparison);
559 defer queue.deinit();
560
561 try queue.add(1);
562 try queue.add(1);
563 try queue.add(2);
564 try queue.add(2);
565 try queue.add(1);
566 try queue.add(1);
567
568 expectEqual(@as(u32, 2), queue.removeMax());
569 expectEqual(@as(u32, 2), queue.removeMax());
570 expectEqual(@as(u32, 1), queue.removeMax());
571 expectEqual(@as(u32, 1), queue.removeMax());
572 expectEqual(@as(u32, 1), queue.removeMax());
573 expectEqual(@as(u32, 1), queue.removeMax());
574}
575
576test "std.PriorityDequeue: removeOrNull empty" {
577 var queue = PDQ.init(testing.allocator, lessThanComparison);
578 defer queue.deinit();
579
580 expect(queue.removeMinOrNull() == null);
581 expect(queue.removeMaxOrNull() == null);
582}
583
584test "std.PriorityDequeue: edge case 3 elements" {
585 var queue = PDQ.init(testing.allocator, lessThanComparison);
586 defer queue.deinit();
587
588 try queue.add(9);
589 try queue.add(3);
590 try queue.add(2);
591
592 expectEqual(@as(u32, 2), queue.removeMin());
593 expectEqual(@as(u32, 3), queue.removeMin());
594 expectEqual(@as(u32, 9), queue.removeMin());
595}
596
597test "std.PriorityDequeue: edge case 3 elements max" {
598 var queue = PDQ.init(testing.allocator, lessThanComparison);
599 defer queue.deinit();
600
601 try queue.add(9);
602 try queue.add(3);
603 try queue.add(2);
604
605 expectEqual(@as(u32, 9), queue.removeMax());
606 expectEqual(@as(u32, 3), queue.removeMax());
607 expectEqual(@as(u32, 2), queue.removeMax());
608}
609
610test "std.PriorityDequeue: peekMin" {
611 var queue = PDQ.init(testing.allocator, lessThanComparison);
612 defer queue.deinit();
613
614 expect(queue.peekMin() == null);
615
616 try queue.add(9);
617 try queue.add(3);
618 try queue.add(2);
619
620 expect(queue.peekMin().? == 2);
621 expect(queue.peekMin().? == 2);
622}
623
624test "std.PriorityDequeue: peekMax" {
625 var queue = PDQ.init(testing.allocator, lessThanComparison);
626 defer queue.deinit();
627
628 expect(queue.peekMin() == null);
629
630 try queue.add(9);
631 try queue.add(3);
632 try queue.add(2);
633
634 expect(queue.peekMax().? == 9);
635 expect(queue.peekMax().? == 9);
636}
637
638test "std.PriorityDequeue: sift up with odd indices" {
639 var queue = PDQ.init(testing.allocator, lessThanComparison);
640 defer queue.deinit();
641 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
642 for (items) |e| {
643 try queue.add(e);
644 }
645
646 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
647 for (sorted_items) |e| {
648 expectEqual(e, queue.removeMin());
649 }
650}
651
652test "std.PriorityDequeue: sift up with odd indices" {
653 var queue = PDQ.init(testing.allocator, lessThanComparison);
654 defer queue.deinit();
655 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
656 for (items) |e| {
657 try queue.add(e);
658 }
659
660 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
661 for (sorted_items) |e| {
662 expectEqual(e, queue.removeMax());
663 }
664}
665
666test "std.PriorityDequeue: addSlice min" {
667 var queue = PDQ.init(testing.allocator, lessThanComparison);
668 defer queue.deinit();
669 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
670 try queue.addSlice(items[0..]);
671
672 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
673 for (sorted_items) |e| {
674 expectEqual(e, queue.removeMin());
675 }
676}
677
678test "std.PriorityDequeue: addSlice max" {
679 var queue = PDQ.init(testing.allocator, lessThanComparison);
680 defer queue.deinit();
681 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
682 try queue.addSlice(items[0..]);
683
684 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
685 for (sorted_items) |e| {
686 expectEqual(e, queue.removeMax());
687 }
688}
689
690test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
691 const items = [0]u32{};
692 const queue_items = try testing.allocator.dupe(u32, &items);
693 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);
694 defer queue.deinit();
695 expectEqual(@as(usize, 0), queue.len);
696 expect(queue.removeMinOrNull() == null);
697}
698
699test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
700 const items = [1]u32{1};
701 const queue_items = try testing.allocator.dupe(u32, &items);
702 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);
703 defer queue.deinit();
704
705 expectEqual(@as(usize, 1), queue.len);
706 expectEqual(items[0], queue.removeMin());
707 expect(queue.removeMinOrNull() == null);
708}
709
710test "std.PriorityDequeue: fromOwnedSlice" {
711 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
712 const queue_items = try testing.allocator.dupe(u32, items[0..]);
713 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);
714 defer queue.deinit();
715
716 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
717 for (sorted_items) |e| {
718 expectEqual(e, queue.removeMin());
719 }
720}
721
722test "std.PriorityDequeue: update min queue" {
723 var queue = PDQ.init(testing.allocator, lessThanComparison);
724 defer queue.deinit();
725
726 try queue.add(55);
727 try queue.add(44);
728 try queue.add(11);
729 try queue.update(55, 5);
730 try queue.update(44, 4);
731 try queue.update(11, 1);
732 expectEqual(@as(u32, 1), queue.removeMin());
733 expectEqual(@as(u32, 4), queue.removeMin());
734 expectEqual(@as(u32, 5), queue.removeMin());
735}
736
737test "std.PriorityDequeue: update same min queue" {
738 var queue = PDQ.init(testing.allocator, lessThanComparison);
739 defer queue.deinit();
740
741 try queue.add(1);
742 try queue.add(1);
743 try queue.add(2);
744 try queue.add(2);
745 try queue.update(1, 5);
746 try queue.update(2, 4);
747 expectEqual(@as(u32, 1), queue.removeMin());
748 expectEqual(@as(u32, 2), queue.removeMin());
749 expectEqual(@as(u32, 4), queue.removeMin());
750 expectEqual(@as(u32, 5), queue.removeMin());
751}
752
753test "std.PriorityDequeue: update max queue" {
754 var queue = PDQ.init(testing.allocator, lessThanComparison);
755 defer queue.deinit();
756
757 try queue.add(55);
758 try queue.add(44);
759 try queue.add(11);
760 try queue.update(55, 5);
761 try queue.update(44, 1);
762 try queue.update(11, 4);
763
764 expectEqual(@as(u32, 5), queue.removeMax());
765 expectEqual(@as(u32, 4), queue.removeMax());
766 expectEqual(@as(u32, 1), queue.removeMax());
767}
768
769test "std.PriorityDequeue: update same max queue" {
770 var queue = PDQ.init(testing.allocator, lessThanComparison);
771 defer queue.deinit();
772
773 try queue.add(1);
774 try queue.add(1);
775 try queue.add(2);
776 try queue.add(2);
777 try queue.update(1, 5);
778 try queue.update(2, 4);
779 expectEqual(@as(u32, 5), queue.removeMax());
780 expectEqual(@as(u32, 4), queue.removeMax());
781 expectEqual(@as(u32, 2), queue.removeMax());
782 expectEqual(@as(u32, 1), queue.removeMax());
783}
784
785test "std.PriorityDequeue: iterator" {
786 var queue = PDQ.init(testing.allocator, lessThanComparison);
787 var map = std.AutoHashMap(u32, void).init(testing.allocator);
788 defer {
789 queue.deinit();
790 map.deinit();
791 }
792
793 const items = [_]u32{ 54, 12, 7, 23, 25, 13 };
794 for (items) |e| {
795 _ = try queue.add(e);
796 _ = try map.put(e, {});
797 }
798
799 var it = queue.iterator();
800 while (it.next()) |e| {
801 _ = map.remove(e);
802 }
803
804 expectEqual(@as(usize, 0), map.count());
805}
806
807test "std.PriorityDequeue: remove at index" {
808 var queue = PDQ.init(testing.allocator, lessThanComparison);
809 defer queue.deinit();
810
811 try queue.add(3);
812 try queue.add(2);
813 try queue.add(1);
814
815 var it = queue.iterator();
816 var elem = it.next();
817 var idx: usize = 0;
818 const two_idx = while (elem != null) : (elem = it.next()) {
819 if (elem.? == 2)
820 break idx;
821 idx += 1;
822 } else unreachable;
823
824 expectEqual(queue.removeIndex(two_idx), 2);
825 expectEqual(queue.removeMin(), 1);
826 expectEqual(queue.removeMin(), 3);
827 expectEqual(queue.removeMinOrNull(), null);
828}
829
830test "std.PriorityDequeue: iterator while empty" {
831 var queue = PDQ.init(testing.allocator, lessThanComparison);
832 defer queue.deinit();
833
834 var it = queue.iterator();
835
836 expectEqual(it.next(), null);
837}
838
839test "std.PriorityDequeue: shrinkRetainingCapacity and shrinkAndFree" {
840 var queue = PDQ.init(testing.allocator, lessThanComparison);
841 defer queue.deinit();
842
843 try queue.ensureCapacity(4);
844 expect(queue.capacity() >= 4);
845
846 try queue.add(1);
847 try queue.add(2);
848 try queue.add(3);
849 expect(queue.capacity() >= 4);
850 expectEqual(@as(usize, 3), queue.len);
851
852 queue.shrinkRetainingCapacity(3);
853 expect(queue.capacity() >= 4);
854 expectEqual(@as(usize, 3), queue.len);
855
856 queue.shrinkAndFree(3);
857 expectEqual(@as(usize, 3), queue.capacity());
858 expectEqual(@as(usize, 3), queue.len);
859
860 expectEqual(@as(u32, 3), queue.removeMax());
861 expectEqual(@as(u32, 2), queue.removeMax());
862 expectEqual(@as(u32, 1), queue.removeMax());
863 expect(queue.removeMaxOrNull() == null);
864}
865
866test "std.PriorityDequeue: fuzz testing min" {
867 var prng = std.rand.DefaultPrng.init(0x12345678);
868
869 const test_case_count = 100;
870 const queue_size = 1_000;
871
872 var i: usize = 0;
873 while (i < test_case_count) : (i += 1) {
874 try fuzzTestMin(&prng.random, queue_size);
875 }
876}
877
878fn fuzzTestMin(rng: *std.rand.Random, comptime queue_size: usize) !void {
879 const allocator = testing.allocator;
880 const items = try generateRandomSlice(allocator, rng, queue_size);
881
882 var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items);
883 defer queue.deinit();
884
885 var last_removed: ?u32 = null;
886 while (queue.removeMinOrNull()) |next| {
887 if (last_removed) |last| {
888 expect(last <= next);
889 }
890 last_removed = next;
891 }
892}
893
894test "std.PriorityDequeue: fuzz testing max" {
895 var prng = std.rand.DefaultPrng.init(0x87654321);
896
897 const test_case_count = 100;
898 const queue_size = 1_000;
899
900 var i: usize = 0;
901 while (i < test_case_count) : (i += 1) {
902 try fuzzTestMax(&prng.random, queue_size);
903 }
904}
905
906fn fuzzTestMax(rng: *std.rand.Random, queue_size: usize) !void {
907 const allocator = testing.allocator;
908 const items = try generateRandomSlice(allocator, rng, queue_size);
909
910 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, items);
911 defer queue.deinit();
912
913 var last_removed: ?u32 = null;
914 while (queue.removeMaxOrNull()) |next| {
915 if (last_removed) |last| {
916 expect(last >= next);
917 }
918 last_removed = next;
919 }
920}
921
922test "std.PriorityDequeue: fuzz testing min and max" {
923 var prng = std.rand.DefaultPrng.init(0x87654321);
924
925 const test_case_count = 100;
926 const queue_size = 1_000;
927
928 var i: usize = 0;
929 while (i < test_case_count) : (i += 1) {
930 try fuzzTestMinMax(&prng.random, queue_size);
931 }
932}
933
934fn fuzzTestMinMax(rng: *std.rand.Random, queue_size: usize) !void {
935 const allocator = testing.allocator;
936 const items = try generateRandomSlice(allocator, rng, queue_size);
937
938 var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items);
939 defer queue.deinit();
940
941 var last_min: ?u32 = null;
942 var last_max: ?u32 = null;
943 var i: usize = 0;
944 while (i < queue_size) : (i += 1) {
945 if (i % 2 == 0) {
946 const next = queue.removeMin();
947 if (last_min) |last| {
948 expect(last <= next);
949 }
950 last_min = next;
951 } else {
952 const next = queue.removeMax();
953 if (last_max) |last| {
954 expect(last >= next);
955 }
956 last_max = next;
957 }
958 }
959}
960
961fn generateRandomSlice(allocator: *std.mem.Allocator, rng: *std.rand.Random, size: usize) ![]u32 {
962 var array = std.ArrayList(u32).init(allocator);
963 try array.ensureCapacity(size);
964
965 var i: usize = 0;
966 while (i < size) : (i += 1) {
967 const elem = rng.int(u32);
968 try array.append(elem);
969 }
970
971 return array.toOwnedSlice();
972}
lib/std/priority_queue.zig+95-26
...@@ -6,6 +6,8 @@...@@ -6,6 +6,8 @@
6const std = @import("std.zig");6const std = @import("std.zig");
7const Allocator = std.mem.Allocator;7const Allocator = std.mem.Allocator;
8const assert = std.debug.assert;8const assert = std.debug.assert;
9const warn = std.debug.warn;
10const Order = std.math.Order;
9const testing = std.testing;11const testing = std.testing;
10const expect = testing.expect;12const expect = testing.expect;
11const expectEqual = testing.expectEqual;13const expectEqual = testing.expectEqual;
...@@ -19,15 +21,17 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -19,15 +21,17 @@ pub fn PriorityQueue(comptime T: type) type {
19 items: []T,21 items: []T,
20 len: usize,22 len: usize,
21 allocator: *Allocator,23 allocator: *Allocator,
22 compareFn: fn (a: T, b: T) bool,24 compareFn: fn (a: T, b: T) Order,
2325
24 /// Initialize and return a priority queue. Provide26 /// Initialize and return a priority queue. Provide `compareFn`
25 /// `compareFn` that returns `true` when its first argument27 /// that returns `Order.lt` when its first argument should
26 /// should get popped before its second argument. For example,28 /// get popped before its second argument, `Order.eq` if the
27 /// to make `pop` return the minimum value, provide29 /// arguments are of equal priority, or `Order.gt` if the second
30 /// argument should be popped first. For example, to make `pop`
31 /// return the smallest number, provide
28 ///32 ///
29 /// `fn lessThan(a: T, b: T) bool { return a < b; }`33 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
30 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) bool) Self {34 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) Order) Self {
31 return Self{35 return Self{
32 .items = &[_]T{},36 .items = &[_]T{},
33 .len = 0,37 .len = 0,
...@@ -60,7 +64,7 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -60,7 +64,7 @@ pub fn PriorityQueue(comptime T: type) type {
60 const child = self.items[child_index];64 const child = self.items[child_index];
61 const parent = self.items[parent_index];65 const parent = self.items[parent_index];
6266
63 if (!self.compareFn(child, parent)) break;67 if (self.compareFn(child, parent) != .lt) break;
6468
65 self.items[parent_index] = child;69 self.items[parent_index] = child;
66 self.items[child_index] = parent;70 self.items[child_index] = parent;
...@@ -132,14 +136,14 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -132,14 +136,14 @@ pub fn PriorityQueue(comptime T: type) type {
132 var smallest = self.items[index];136 var smallest = self.items[index];
133137
134 if (left) |e| {138 if (left) |e| {
135 if (self.compareFn(e, smallest)) {139 if (self.compareFn(e, smallest) == .lt) {
136 smallest_index = left_index;140 smallest_index = left_index;
137 smallest = e;141 smallest = e;
138 }142 }
139 }143 }
140144
141 if (right) |e| {145 if (right) |e| {
142 if (self.compareFn(e, smallest)) {146 if (self.compareFn(e, smallest) == .lt) {
143 smallest_index = right_index;147 smallest_index = right_index;
144 smallest = e;148 smallest = e;
145 }149 }
...@@ -158,13 +162,16 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -158,13 +162,16 @@ pub fn PriorityQueue(comptime T: type) type {
158 /// PriorityQueue takes ownership of the passed in slice. The slice must have been162 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
159 /// allocated with `allocator`.163 /// allocated with `allocator`.
160 /// Deinitialize with `deinit`.164 /// Deinitialize with `deinit`.
161 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) bool, items: []T) Self {165 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) Order, items: []T) Self {
162 var queue = Self{166 var queue = Self{
163 .items = items,167 .items = items,
164 .len = items.len,168 .len = items.len,
165 .allocator = allocator,169 .allocator = allocator,
166 .compareFn = compareFn,170 .compareFn = compareFn,
167 };171 };
172
173 if (queue.len <= 1) return queue;
174
168 const half = (queue.len >> 1) - 1;175 const half = (queue.len >> 1) - 1;
169 var i: usize = 0;176 var i: usize = 0;
170 while (i <= half) : (i += 1) {177 while (i <= half) : (i += 1) {
...@@ -183,25 +190,40 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -183,25 +190,40 @@ pub fn PriorityQueue(comptime T: type) type {
183 self.items = try self.allocator.realloc(self.items, better_capacity);190 self.items = try self.allocator.realloc(self.items, better_capacity);
184 }191 }
185192
186 pub fn resize(self: *Self, new_len: usize) !void {193 /// Reduce allocated capacity to `new_len`.
187 try self.ensureCapacity(new_len);194 pub fn shrinkAndFree(self: *Self, new_len: usize) void {
195 assert(new_len <= self.items.len);
196
197 // Cannot shrink to smaller than the current queue size without invalidating the heap property
198 assert(new_len >= self.len);
199
200 self.items = self.allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
201 error.OutOfMemory => { // no problem, capacity is still correct then.
202 self.items.len = new_len;
203 return;
204 },
205 };
188 self.len = new_len;206 self.len = new_len;
189 }207 }
190208
191 pub fn shrink(self: *Self, new_len: usize) void {209 /// Reduce length to `new_len`.
192 // TODO take advantage of the new realloc semantics210 pub fn shrinkRetainingCapacity(self: *Self, new_len: usize) void {
193 assert(new_len <= self.len);211 assert(new_len <= self.items.len);
212
213 // Cannot shrink to smaller than the current queue size without invalidating the heap property
214 assert(new_len >= self.len);
215
194 self.len = new_len;216 self.len = new_len;
195 }217 }
196218
197 pub fn update(self: *Self, elem: T, new_elem: T) !void {219 pub fn update(self: *Self, elem: T, new_elem: T) !void {
198 var update_index: usize = std.mem.indexOfScalar(T, self.items, elem) orelse return error.ElementNotFound;220 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
199 const old_elem: T = self.items[update_index];221 const old_elem: T = self.items[update_index];
200 self.items[update_index] = new_elem;222 self.items[update_index] = new_elem;
201 if (self.compareFn(new_elem, old_elem)) {223 switch (self.compareFn(new_elem, old_elem)) {
202 siftUp(self, update_index);224 .lt => siftUp(self, update_index),
203 } else {225 .gt => siftDown(self, update_index),
204 siftDown(self, update_index);226 .eq => {}, // Nothing to do as the items have equal priority
205 }227 }
206 }228 }
207229
...@@ -248,12 +270,12 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -248,12 +270,12 @@ pub fn PriorityQueue(comptime T: type) type {
248 };270 };
249}271}
250272
251fn lessThan(a: u32, b: u32) bool {273fn lessThan(a: u32, b: u32) Order {
252 return a < b;274 return std.math.order(a, b);
253}275}
254276
255fn greaterThan(a: u32, b: u32) bool {277fn greaterThan(a: u32, b: u32) Order {
256 return a > b;278 return lessThan(a, b).invert();
257}279}
258280
259const PQ = PriorityQueue(u32);281const PQ = PriorityQueue(u32);
...@@ -351,6 +373,26 @@ test "std.PriorityQueue: addSlice" {...@@ -351,6 +373,26 @@ test "std.PriorityQueue: addSlice" {
351 }373 }
352}374}
353375
376test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
377 const items = [0]u32{};
378 const queue_items = try testing.allocator.dupe(u32, &items);
379 var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]);
380 defer queue.deinit();
381 expectEqual(@as(usize, 0), queue.len);
382 expect(queue.removeOrNull() == null);
383}
384
385test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
386 const items = [1]u32{1};
387 const queue_items = try testing.allocator.dupe(u32, &items);
388 var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]);
389 defer queue.deinit();
390
391 expectEqual(@as(usize, 1), queue.len);
392 expectEqual(items[0], queue.remove());
393 expect(queue.removeOrNull() == null);
394}
395
354test "std.PriorityQueue: fromOwnedSlice" {396test "std.PriorityQueue: fromOwnedSlice" {
355 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };397 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
356 const heap_items = try testing.allocator.dupe(u32, items[0..]);398 const heap_items = try testing.allocator.dupe(u32, items[0..]);
...@@ -453,6 +495,33 @@ test "std.PriorityQueue: iterator while empty" {...@@ -453,6 +495,33 @@ test "std.PriorityQueue: iterator while empty" {
453 expectEqual(it.next(), null);495 expectEqual(it.next(), null);
454}496}
455497
498test "std.PriorityQueue: shrinkRetainingCapacity and shrinkAndFree" {
499 var queue = PQ.init(testing.allocator, lessThan);
500 defer queue.deinit();
501
502 try queue.ensureCapacity(4);
503 expect(queue.capacity() >= 4);
504
505 try queue.add(1);
506 try queue.add(2);
507 try queue.add(3);
508 expect(queue.capacity() >= 4);
509 expectEqual(@as(usize, 3), queue.len);
510
511 queue.shrinkRetainingCapacity(3);
512 expect(queue.capacity() >= 4);
513 expectEqual(@as(usize, 3), queue.len);
514
515 queue.shrinkAndFree(3);
516 expectEqual(@as(usize, 3), queue.capacity());
517 expectEqual(@as(usize, 3), queue.len);
518
519 expectEqual(@as(u32, 1), queue.remove());
520 expectEqual(@as(u32, 2), queue.remove());
521 expectEqual(@as(u32, 3), queue.remove());
522 expect(queue.removeOrNull() == null);
523}
524
456test "std.PriorityQueue: update min heap" {525test "std.PriorityQueue: update min heap" {
457 var queue = PQ.init(testing.allocator, lessThan);526 var queue = PQ.init(testing.allocator, lessThan);
458 defer queue.deinit();527 defer queue.deinit();
lib/std/std.zig+1
...@@ -31,6 +31,7 @@ pub const PackedIntArrayEndian = @import("packed_int_array.zig").PackedIntArrayE...@@ -31,6 +31,7 @@ pub const PackedIntArrayEndian = @import("packed_int_array.zig").PackedIntArrayE
31pub const PackedIntSlice = @import("packed_int_array.zig").PackedIntSlice;31pub const PackedIntSlice = @import("packed_int_array.zig").PackedIntSlice;
32pub const PackedIntSliceEndian = @import("packed_int_array.zig").PackedIntSliceEndian;32pub const PackedIntSliceEndian = @import("packed_int_array.zig").PackedIntSliceEndian;
33pub const PriorityQueue = @import("priority_queue.zig").PriorityQueue;33pub const PriorityQueue = @import("priority_queue.zig").PriorityQueue;
34pub const PriorityDequeue = @import("priority_dequeue.zig").PriorityDequeue;
34pub const Progress = @import("Progress.zig");35pub const Progress = @import("Progress.zig");
35pub const SemanticVersion = @import("SemanticVersion.zig");36pub const SemanticVersion = @import("SemanticVersion.zig");
36pub const SinglyLinkedList = @import("linked_list.zig").SinglyLinkedList;37pub const SinglyLinkedList = @import("linked_list.zig").SinglyLinkedList;