authorgravatar for hollmannm1@gmail.comMax Hollmann <hollmannm1@gmail.com> 2021-10-17 01:10:01+02:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2021-10-17 17:47:43-04:00
log53b87fa78a3ee2c261f3bdd6a71d5733fe17ffd5
treec3ff20f7501d461adf58ccb817f10566c9192cd1
parent22beaf5afcfb4c9c29016f19f363a24f159b5b82

Move `compareFn` from init to type constructor in `PriorityQueue` and `PriorityDequeue`.

This change significantly improves performance for simple compare functions and modifies the API to be more consistent with e.g. `HashMap`.

2 files changed, 90 insertions(+), 97 deletions(-)

lib/std/priority_dequeue.zig+52-56
...@@ -9,30 +9,27 @@ const expectEqual = testing.expectEqual;...@@ -9,30 +9,27 @@ const expectEqual = testing.expectEqual;
9const expectError = testing.expectError;9const expectError = testing.expectError;
1010
11/// Priority Dequeue for storing generic data. Initialize with `init`.11/// Priority Dequeue for storing generic data. Initialize with `init`.
12pub fn PriorityDequeue(comptime T: type) type {12/// Provide `compareFn` that returns `Order.lt` when its first
13/// argument should get min-popped before its second argument,
14/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
15/// if the second argument should be min-popped first.
16/// Popping the max element works in reverse. For example,
17/// to make `popMin` return the smallest number, provide
18/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
19pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) type {
13 return struct {20 return struct {
14 const Self = @This();21 const Self = @This();
1522
16 items: []T,23 items: []T,
17 len: usize,24 len: usize,
18 allocator: *Allocator,25 allocator: *Allocator,
19 compareFn: fn (a: T, b: T) Order,26
2027 /// Initialize and return a new priority dequeue.
21 /// Initialize and return a new priority dequeue. Provide `compareFn`28 pub fn init(allocator: *Allocator) Self {
22 /// that returns `Order.lt` when its first argument should
23 /// get min-popped before its second argument, `Order.eq` if the
24 /// arguments are of equal priority, or `Order.gt` if the second
25 /// argument should be min-popped first. Popping the max element works
26 /// in reverse. For example, to make `popMin` return the smallest
27 /// number, provide
28 ///
29 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
30 pub fn init(allocator: *Allocator, compareFn: fn (T, T) Order) Self {
31 return Self{29 return Self{
32 .items = &[_]T{},30 .items = &[_]T{},
33 .len = 0,31 .len = 0,
34 .allocator = allocator,32 .allocator = allocator,
35 .compareFn = compareFn,
36 };33 };
37 }34 }
3835
...@@ -91,7 +88,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -91,7 +88,7 @@ pub fn PriorityDequeue(comptime T: type) type {
91 const parent = self.items[parent_index];88 const parent = self.items[parent_index];
9289
93 const min_layer = self.nextIsMinLayer();90 const min_layer = self.nextIsMinLayer();
94 const order = self.compareFn(child, parent);91 const order = compareFn(child, parent);
95 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {92 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
96 // We must swap the item with it's parent if it is on the "wrong" layer93 // We must swap the item with it's parent if it is on the "wrong" layer
97 self.items[parent_index] = child;94 self.items[parent_index] = child;
...@@ -124,7 +121,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -124,7 +121,7 @@ pub fn PriorityDequeue(comptime T: type) type {
124 const grandparent = self.items[grandparent_index];121 const grandparent = self.items[grandparent_index];
125122
126 // If the grandparent is already better or equal, we have gone as far as we need to123 // If the grandparent is already better or equal, we have gone as far as we need to
127 if (self.compareFn(child, grandparent) != target_order) break;124 if (compareFn(child, grandparent) != target_order) break;
128125
129 // Otherwise swap the item with it's grandparent126 // Otherwise swap the item with it's grandparent
130 self.items[grandparent_index] = child;127 self.items[grandparent_index] = child;
...@@ -219,10 +216,10 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -219,10 +216,10 @@ pub fn PriorityDequeue(comptime T: type) type {
219 // Find the best grandchild216 // Find the best grandchild
220 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);217 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
221 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);218 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
222 const best_grandchild = self.bestItem(best_left, best_right, target_order);219 const best_grandchild = Self.bestItem(best_left, best_right, target_order);
223220
224 // If the item is better than or equal to its best grandchild, we are done221 // If the item is better than or equal to its best grandchild, we are done
225 if (self.compareFn(best_grandchild.item, elem) != target_order) return;222 if (compareFn(best_grandchild.item, elem) != target_order) return;
226223
227 // Otherwise, swap them224 // Otherwise, swap them
228 self.items[best_grandchild.index] = elem;225 self.items[best_grandchild.index] = elem;
...@@ -239,7 +236,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -239,7 +236,7 @@ pub fn PriorityDequeue(comptime T: type) type {
239 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);236 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
240237
241 // If the item is better than or equal to its best descendant, we are done238 // If the item is better than or equal to its best descendant, we are done
242 if (self.compareFn(best_descendent.item, elem) != target_order) return;239 if (compareFn(best_descendent.item, elem) != target_order) return;
243240
244 // Otherwise swap them241 // Otherwise swap them
245 self.items[best_descendent.index] = elem;242 self.items[best_descendent.index] = elem;
...@@ -263,7 +260,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -263,7 +260,7 @@ pub fn PriorityDequeue(comptime T: type) type {
263 const parent_index = parentIndex(child_index);260 const parent_index = parentIndex(child_index);
264 const parent = self.items[parent_index];261 const parent = self.items[parent_index];
265262
266 if (self.compareFn(parent, child) == target_order) {263 if (compareFn(parent, child) == target_order) {
267 self.items[parent_index] = child;264 self.items[parent_index] = child;
268 self.items[child_index] = parent;265 self.items[child_index] = parent;
269 }266 }
...@@ -281,8 +278,8 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -281,8 +278,8 @@ pub fn PriorityDequeue(comptime T: type) type {
281 };278 };
282 }279 }
283280
284 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {281 fn bestItem(item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
285 if (self.compareFn(item1.item, item2.item) == target_order) {282 if (compareFn(item1.item, item2.item) == target_order) {
286 return item1;283 return item1;
287 } else {284 } else {
288 return item2;285 return item2;
...@@ -292,7 +289,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -292,7 +289,7 @@ pub fn PriorityDequeue(comptime T: type) type {
292 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {289 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
293 var item1 = self.getItem(index1);290 var item1 = self.getItem(index1);
294 var item2 = self.getItem(index2);291 var item2 = self.getItem(index2);
295 return self.bestItem(item1, item2, target_order);292 return Self.bestItem(item1, item2, target_order);
296 }293 }
297294
298 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {295 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
...@@ -340,12 +337,11 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -340,12 +337,11 @@ pub fn PriorityDequeue(comptime T: type) type {
340 /// Dequeue takes ownership of the passed in slice. The slice must have been337 /// Dequeue takes ownership of the passed in slice. The slice must have been
341 /// allocated with `allocator`.338 /// allocated with `allocator`.
342 /// De-initialize with `deinit`.339 /// De-initialize with `deinit`.
343 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (T, T) Order, items: []T) Self {340 pub fn fromOwnedSlice(allocator: *Allocator, items: []T) Self {
344 var queue = Self{341 var queue = Self{
345 .items = items,342 .items = items,
346 .len = items.len,343 .len = items.len,
347 .allocator = allocator,344 .allocator = allocator,
348 .compareFn = compareFn,
349 };345 };
350346
351 if (queue.len <= 1) return queue;347 if (queue.len <= 1) return queue;
...@@ -400,7 +396,7 @@ pub fn PriorityDequeue(comptime T: type) type {...@@ -400,7 +396,7 @@ pub fn PriorityDequeue(comptime T: type) type {
400 }396 }
401397
402 pub const Iterator = struct {398 pub const Iterator = struct {
403 queue: *PriorityDequeue(T),399 queue: *PriorityDequeue(T, compareFn),
404 count: usize,400 count: usize,
405401
406 pub fn next(it: *Iterator) ?T {402 pub fn next(it: *Iterator) ?T {
...@@ -462,10 +458,10 @@ fn lessThanComparison(a: u32, b: u32) Order {...@@ -462,10 +458,10 @@ fn lessThanComparison(a: u32, b: u32) Order {
462 return std.math.order(a, b);458 return std.math.order(a, b);
463}459}
464460
465const PDQ = PriorityDequeue(u32);461const PDQ = PriorityDequeue(u32, lessThanComparison);
466462
467test "std.PriorityDequeue: add and remove min" {463test "std.PriorityDequeue: add and remove min" {
468 var queue = PDQ.init(testing.allocator, lessThanComparison);464 var queue = PDQ.init(testing.allocator);
469 defer queue.deinit();465 defer queue.deinit();
470466
471 try queue.add(54);467 try queue.add(54);
...@@ -487,11 +483,11 @@ test "std.PriorityDequeue: add and remove min structs" {...@@ -487,11 +483,11 @@ test "std.PriorityDequeue: add and remove min structs" {
487 const S = struct {483 const S = struct {
488 size: u32,484 size: u32,
489 };485 };
490 var queue = PriorityDequeue(S).init(testing.allocator, struct {486 var queue = PriorityDequeue(S, struct {
491 fn order(a: S, b: S) Order {487 fn order(a: S, b: S) Order {
492 return std.math.order(a.size, b.size);488 return std.math.order(a.size, b.size);
493 }489 }
494 }.order);490 }.order).init(testing.allocator);
495 defer queue.deinit();491 defer queue.deinit();
496492
497 try queue.add(.{ .size = 54 });493 try queue.add(.{ .size = 54 });
...@@ -510,7 +506,7 @@ test "std.PriorityDequeue: add and remove min structs" {...@@ -510,7 +506,7 @@ test "std.PriorityDequeue: add and remove min structs" {
510}506}
511507
512test "std.PriorityDequeue: add and remove max" {508test "std.PriorityDequeue: add and remove max" {
513 var queue = PDQ.init(testing.allocator, lessThanComparison);509 var queue = PDQ.init(testing.allocator);
514 defer queue.deinit();510 defer queue.deinit();
515511
516 try queue.add(54);512 try queue.add(54);
...@@ -529,7 +525,7 @@ test "std.PriorityDequeue: add and remove max" {...@@ -529,7 +525,7 @@ test "std.PriorityDequeue: add and remove max" {
529}525}
530526
531test "std.PriorityDequeue: add and remove same min" {527test "std.PriorityDequeue: add and remove same min" {
532 var queue = PDQ.init(testing.allocator, lessThanComparison);528 var queue = PDQ.init(testing.allocator);
533 defer queue.deinit();529 defer queue.deinit();
534530
535 try queue.add(1);531 try queue.add(1);
...@@ -548,7 +544,7 @@ test "std.PriorityDequeue: add and remove same min" {...@@ -548,7 +544,7 @@ test "std.PriorityDequeue: add and remove same min" {
548}544}
549545
550test "std.PriorityDequeue: add and remove same max" {546test "std.PriorityDequeue: add and remove same max" {
551 var queue = PDQ.init(testing.allocator, lessThanComparison);547 var queue = PDQ.init(testing.allocator);
552 defer queue.deinit();548 defer queue.deinit();
553549
554 try queue.add(1);550 try queue.add(1);
...@@ -567,7 +563,7 @@ test "std.PriorityDequeue: add and remove same max" {...@@ -567,7 +563,7 @@ test "std.PriorityDequeue: add and remove same max" {
567}563}
568564
569test "std.PriorityDequeue: removeOrNull empty" {565test "std.PriorityDequeue: removeOrNull empty" {
570 var queue = PDQ.init(testing.allocator, lessThanComparison);566 var queue = PDQ.init(testing.allocator);
571 defer queue.deinit();567 defer queue.deinit();
572568
573 try expect(queue.removeMinOrNull() == null);569 try expect(queue.removeMinOrNull() == null);
...@@ -575,7 +571,7 @@ test "std.PriorityDequeue: removeOrNull empty" {...@@ -575,7 +571,7 @@ test "std.PriorityDequeue: removeOrNull empty" {
575}571}
576572
577test "std.PriorityDequeue: edge case 3 elements" {573test "std.PriorityDequeue: edge case 3 elements" {
578 var queue = PDQ.init(testing.allocator, lessThanComparison);574 var queue = PDQ.init(testing.allocator);
579 defer queue.deinit();575 defer queue.deinit();
580576
581 try queue.add(9);577 try queue.add(9);
...@@ -588,7 +584,7 @@ test "std.PriorityDequeue: edge case 3 elements" {...@@ -588,7 +584,7 @@ test "std.PriorityDequeue: edge case 3 elements" {
588}584}
589585
590test "std.PriorityDequeue: edge case 3 elements max" {586test "std.PriorityDequeue: edge case 3 elements max" {
591 var queue = PDQ.init(testing.allocator, lessThanComparison);587 var queue = PDQ.init(testing.allocator);
592 defer queue.deinit();588 defer queue.deinit();
593589
594 try queue.add(9);590 try queue.add(9);
...@@ -601,7 +597,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {...@@ -601,7 +597,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {
601}597}
602598
603test "std.PriorityDequeue: peekMin" {599test "std.PriorityDequeue: peekMin" {
604 var queue = PDQ.init(testing.allocator, lessThanComparison);600 var queue = PDQ.init(testing.allocator);
605 defer queue.deinit();601 defer queue.deinit();
606602
607 try expect(queue.peekMin() == null);603 try expect(queue.peekMin() == null);
...@@ -615,7 +611,7 @@ test "std.PriorityDequeue: peekMin" {...@@ -615,7 +611,7 @@ test "std.PriorityDequeue: peekMin" {
615}611}
616612
617test "std.PriorityDequeue: peekMax" {613test "std.PriorityDequeue: peekMax" {
618 var queue = PDQ.init(testing.allocator, lessThanComparison);614 var queue = PDQ.init(testing.allocator);
619 defer queue.deinit();615 defer queue.deinit();
620616
621 try expect(queue.peekMin() == null);617 try expect(queue.peekMin() == null);
...@@ -629,7 +625,7 @@ test "std.PriorityDequeue: peekMax" {...@@ -629,7 +625,7 @@ test "std.PriorityDequeue: peekMax" {
629}625}
630626
631test "std.PriorityDequeue: sift up with odd indices" {627test "std.PriorityDequeue: sift up with odd indices" {
632 var queue = PDQ.init(testing.allocator, lessThanComparison);628 var queue = PDQ.init(testing.allocator);
633 defer queue.deinit();629 defer queue.deinit();
634 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };630 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
635 for (items) |e| {631 for (items) |e| {
...@@ -643,7 +639,7 @@ test "std.PriorityDequeue: sift up with odd indices" {...@@ -643,7 +639,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
643}639}
644640
645test "std.PriorityDequeue: sift up with odd indices" {641test "std.PriorityDequeue: sift up with odd indices" {
646 var queue = PDQ.init(testing.allocator, lessThanComparison);642 var queue = PDQ.init(testing.allocator);
647 defer queue.deinit();643 defer queue.deinit();
648 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };644 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
649 for (items) |e| {645 for (items) |e| {
...@@ -657,7 +653,7 @@ test "std.PriorityDequeue: sift up with odd indices" {...@@ -657,7 +653,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
657}653}
658654
659test "std.PriorityDequeue: addSlice min" {655test "std.PriorityDequeue: addSlice min" {
660 var queue = PDQ.init(testing.allocator, lessThanComparison);656 var queue = PDQ.init(testing.allocator);
661 defer queue.deinit();657 defer queue.deinit();
662 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };658 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
663 try queue.addSlice(items[0..]);659 try queue.addSlice(items[0..]);
...@@ -669,7 +665,7 @@ test "std.PriorityDequeue: addSlice min" {...@@ -669,7 +665,7 @@ test "std.PriorityDequeue: addSlice min" {
669}665}
670666
671test "std.PriorityDequeue: addSlice max" {667test "std.PriorityDequeue: addSlice max" {
672 var queue = PDQ.init(testing.allocator, lessThanComparison);668 var queue = PDQ.init(testing.allocator);
673 defer queue.deinit();669 defer queue.deinit();
674 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };670 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
675 try queue.addSlice(items[0..]);671 try queue.addSlice(items[0..]);
...@@ -683,7 +679,7 @@ test "std.PriorityDequeue: addSlice max" {...@@ -683,7 +679,7 @@ test "std.PriorityDequeue: addSlice max" {
683test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {679test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
684 const items = [0]u32{};680 const items = [0]u32{};
685 const queue_items = try testing.allocator.dupe(u32, &items);681 const queue_items = try testing.allocator.dupe(u32, &items);
686 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);682 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
687 defer queue.deinit();683 defer queue.deinit();
688 try expectEqual(@as(usize, 0), queue.len);684 try expectEqual(@as(usize, 0), queue.len);
689 try expect(queue.removeMinOrNull() == null);685 try expect(queue.removeMinOrNull() == null);
...@@ -692,7 +688,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {...@@ -692,7 +688,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
692test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {688test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
693 const items = [1]u32{1};689 const items = [1]u32{1};
694 const queue_items = try testing.allocator.dupe(u32, &items);690 const queue_items = try testing.allocator.dupe(u32, &items);
695 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);691 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
696 defer queue.deinit();692 defer queue.deinit();
697693
698 try expectEqual(@as(usize, 1), queue.len);694 try expectEqual(@as(usize, 1), queue.len);
...@@ -703,7 +699,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {...@@ -703,7 +699,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
703test "std.PriorityDequeue: fromOwnedSlice" {699test "std.PriorityDequeue: fromOwnedSlice" {
704 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };700 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
705 const queue_items = try testing.allocator.dupe(u32, items[0..]);701 const queue_items = try testing.allocator.dupe(u32, items[0..]);
706 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, queue_items[0..]);702 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
707 defer queue.deinit();703 defer queue.deinit();
708704
709 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };705 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
...@@ -713,7 +709,7 @@ test "std.PriorityDequeue: fromOwnedSlice" {...@@ -713,7 +709,7 @@ test "std.PriorityDequeue: fromOwnedSlice" {
713}709}
714710
715test "std.PriorityDequeue: update min queue" {711test "std.PriorityDequeue: update min queue" {
716 var queue = PDQ.init(testing.allocator, lessThanComparison);712 var queue = PDQ.init(testing.allocator);
717 defer queue.deinit();713 defer queue.deinit();
718714
719 try queue.add(55);715 try queue.add(55);
...@@ -728,7 +724,7 @@ test "std.PriorityDequeue: update min queue" {...@@ -728,7 +724,7 @@ test "std.PriorityDequeue: update min queue" {
728}724}
729725
730test "std.PriorityDequeue: update same min queue" {726test "std.PriorityDequeue: update same min queue" {
731 var queue = PDQ.init(testing.allocator, lessThanComparison);727 var queue = PDQ.init(testing.allocator);
732 defer queue.deinit();728 defer queue.deinit();
733729
734 try queue.add(1);730 try queue.add(1);
...@@ -744,7 +740,7 @@ test "std.PriorityDequeue: update same min queue" {...@@ -744,7 +740,7 @@ test "std.PriorityDequeue: update same min queue" {
744}740}
745741
746test "std.PriorityDequeue: update max queue" {742test "std.PriorityDequeue: update max queue" {
747 var queue = PDQ.init(testing.allocator, lessThanComparison);743 var queue = PDQ.init(testing.allocator);
748 defer queue.deinit();744 defer queue.deinit();
749745
750 try queue.add(55);746 try queue.add(55);
...@@ -760,7 +756,7 @@ test "std.PriorityDequeue: update max queue" {...@@ -760,7 +756,7 @@ test "std.PriorityDequeue: update max queue" {
760}756}
761757
762test "std.PriorityDequeue: update same max queue" {758test "std.PriorityDequeue: update same max queue" {
763 var queue = PDQ.init(testing.allocator, lessThanComparison);759 var queue = PDQ.init(testing.allocator);
764 defer queue.deinit();760 defer queue.deinit();
765761
766 try queue.add(1);762 try queue.add(1);
...@@ -776,7 +772,7 @@ test "std.PriorityDequeue: update same max queue" {...@@ -776,7 +772,7 @@ test "std.PriorityDequeue: update same max queue" {
776}772}
777773
778test "std.PriorityDequeue: iterator" {774test "std.PriorityDequeue: iterator" {
779 var queue = PDQ.init(testing.allocator, lessThanComparison);775 var queue = PDQ.init(testing.allocator);
780 var map = std.AutoHashMap(u32, void).init(testing.allocator);776 var map = std.AutoHashMap(u32, void).init(testing.allocator);
781 defer {777 defer {
782 queue.deinit();778 queue.deinit();
...@@ -798,7 +794,7 @@ test "std.PriorityDequeue: iterator" {...@@ -798,7 +794,7 @@ test "std.PriorityDequeue: iterator" {
798}794}
799795
800test "std.PriorityDequeue: remove at index" {796test "std.PriorityDequeue: remove at index" {
801 var queue = PDQ.init(testing.allocator, lessThanComparison);797 var queue = PDQ.init(testing.allocator);
802 defer queue.deinit();798 defer queue.deinit();
803799
804 try queue.add(3);800 try queue.add(3);
...@@ -821,7 +817,7 @@ test "std.PriorityDequeue: remove at index" {...@@ -821,7 +817,7 @@ test "std.PriorityDequeue: remove at index" {
821}817}
822818
823test "std.PriorityDequeue: iterator while empty" {819test "std.PriorityDequeue: iterator while empty" {
824 var queue = PDQ.init(testing.allocator, lessThanComparison);820 var queue = PDQ.init(testing.allocator);
825 defer queue.deinit();821 defer queue.deinit();
826822
827 var it = queue.iterator();823 var it = queue.iterator();
...@@ -830,7 +826,7 @@ test "std.PriorityDequeue: iterator while empty" {...@@ -830,7 +826,7 @@ test "std.PriorityDequeue: iterator while empty" {
830}826}
831827
832test "std.PriorityDequeue: shrinkAndFree" {828test "std.PriorityDequeue: shrinkAndFree" {
833 var queue = PDQ.init(testing.allocator, lessThanComparison);829 var queue = PDQ.init(testing.allocator);
834 defer queue.deinit();830 defer queue.deinit();
835831
836 try queue.ensureTotalCapacity(4);832 try queue.ensureTotalCapacity(4);
...@@ -868,7 +864,7 @@ fn fuzzTestMin(rng: *std.rand.Random, comptime queue_size: usize) !void {...@@ -868,7 +864,7 @@ fn fuzzTestMin(rng: *std.rand.Random, comptime queue_size: usize) !void {
868 const allocator = testing.allocator;864 const allocator = testing.allocator;
869 const items = try generateRandomSlice(allocator, rng, queue_size);865 const items = try generateRandomSlice(allocator, rng, queue_size);
870866
871 var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items);867 var queue = PDQ.fromOwnedSlice(allocator, items);
872 defer queue.deinit();868 defer queue.deinit();
873869
874 var last_removed: ?u32 = null;870 var last_removed: ?u32 = null;
...@@ -896,7 +892,7 @@ fn fuzzTestMax(rng: *std.rand.Random, queue_size: usize) !void {...@@ -896,7 +892,7 @@ fn fuzzTestMax(rng: *std.rand.Random, queue_size: usize) !void {
896 const allocator = testing.allocator;892 const allocator = testing.allocator;
897 const items = try generateRandomSlice(allocator, rng, queue_size);893 const items = try generateRandomSlice(allocator, rng, queue_size);
898894
899 var queue = PDQ.fromOwnedSlice(testing.allocator, lessThanComparison, items);895 var queue = PDQ.fromOwnedSlice(testing.allocator, items);
900 defer queue.deinit();896 defer queue.deinit();
901897
902 var last_removed: ?u32 = null;898 var last_removed: ?u32 = null;
...@@ -924,7 +920,7 @@ fn fuzzTestMinMax(rng: *std.rand.Random, queue_size: usize) !void {...@@ -924,7 +920,7 @@ fn fuzzTestMinMax(rng: *std.rand.Random, queue_size: usize) !void {
924 const allocator = testing.allocator;920 const allocator = testing.allocator;
925 const items = try generateRandomSlice(allocator, rng, queue_size);921 const items = try generateRandomSlice(allocator, rng, queue_size);
926922
927 var queue = PDQ.fromOwnedSlice(allocator, lessThanComparison, items);923 var queue = PDQ.fromOwnedSlice(allocator, items);
928 defer queue.deinit();924 defer queue.deinit();
929925
930 var last_min: ?u32 = null;926 var last_min: ?u32 = null;
lib/std/priority_queue.zig+38-41
...@@ -9,29 +9,26 @@ const expectEqual = testing.expectEqual;...@@ -9,29 +9,26 @@ const expectEqual = testing.expectEqual;
9const expectError = testing.expectError;9const expectError = testing.expectError;
1010
11/// Priority queue for storing generic data. Initialize with `init`.11/// Priority queue for storing generic data. Initialize with `init`.
12pub fn PriorityQueue(comptime T: type) type {12/// Provide `compareFn` that returns `Order.lt` when its first
13/// argument should get popped before its second argument,
14/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
15/// if the second argument should be popped first.
16/// For example, to make `pop` return the smallest number, provide
17/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
18pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order) type {
13 return struct {19 return struct {
14 const Self = @This();20 const Self = @This();
1521
16 items: []T,22 items: []T,
17 len: usize,23 len: usize,
18 allocator: *Allocator,24 allocator: *Allocator,
19 compareFn: fn (a: T, b: T) Order,25
2026 /// Initialize and return a priority queue.
21 /// Initialize and return a priority queue. Provide `compareFn`27 pub fn init(allocator: *Allocator) Self {
22 /// that returns `Order.lt` when its first argument should
23 /// get popped before its second argument, `Order.eq` if the
24 /// arguments are of equal priority, or `Order.gt` if the second
25 /// argument should be popped first. For example, to make `pop`
26 /// return the smallest number, provide
27 ///
28 /// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
29 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) Order) Self {
30 return Self{28 return Self{
31 .items = &[_]T{},29 .items = &[_]T{},
32 .len = 0,30 .len = 0,
33 .allocator = allocator,31 .allocator = allocator,
34 .compareFn = compareFn,
35 };32 };
36 }33 }
3734
...@@ -59,7 +56,7 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -59,7 +56,7 @@ pub fn PriorityQueue(comptime T: type) type {
59 const child = self.items[child_index];56 const child = self.items[child_index];
60 const parent = self.items[parent_index];57 const parent = self.items[parent_index];
6158
62 if (self.compareFn(child, parent) != .lt) break;59 if (compareFn(child, parent) != .lt) break;
6360
64 self.items[parent_index] = child;61 self.items[parent_index] = child;
65 self.items[child_index] = parent;62 self.items[child_index] = parent;
...@@ -131,14 +128,14 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -131,14 +128,14 @@ pub fn PriorityQueue(comptime T: type) type {
131 var smallest = self.items[index];128 var smallest = self.items[index];
132129
133 if (left) |e| {130 if (left) |e| {
134 if (self.compareFn(e, smallest) == .lt) {131 if (compareFn(e, smallest) == .lt) {
135 smallest_index = left_index;132 smallest_index = left_index;
136 smallest = e;133 smallest = e;
137 }134 }
138 }135 }
139136
140 if (right) |e| {137 if (right) |e| {
141 if (self.compareFn(e, smallest) == .lt) {138 if (compareFn(e, smallest) == .lt) {
142 smallest_index = right_index;139 smallest_index = right_index;
143 smallest = e;140 smallest = e;
144 }141 }
...@@ -157,12 +154,11 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -157,12 +154,11 @@ pub fn PriorityQueue(comptime T: type) type {
157 /// PriorityQueue takes ownership of the passed in slice. The slice must have been154 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
158 /// allocated with `allocator`.155 /// allocated with `allocator`.
159 /// Deinitialize with `deinit`.156 /// Deinitialize with `deinit`.
160 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) Order, items: []T) Self {157 pub fn fromOwnedSlice(allocator: *Allocator, items: []T) Self {
161 var queue = Self{158 var queue = Self{
162 .items = items,159 .items = items,
163 .len = items.len,160 .len = items.len,
164 .allocator = allocator,161 .allocator = allocator,
165 .compareFn = compareFn,
166 };162 };
167163
168 if (queue.len <= 1) return queue;164 if (queue.len <= 1) return queue;
...@@ -213,7 +209,7 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -213,7 +209,7 @@ pub fn PriorityQueue(comptime T: type) type {
213 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;209 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
214 const old_elem: T = self.items[update_index];210 const old_elem: T = self.items[update_index];
215 self.items[update_index] = new_elem;211 self.items[update_index] = new_elem;
216 switch (self.compareFn(new_elem, old_elem)) {212 switch (compareFn(new_elem, old_elem)) {
217 .lt => siftUp(self, update_index),213 .lt => siftUp(self, update_index),
218 .gt => siftDown(self, update_index),214 .gt => siftDown(self, update_index),
219 .eq => {}, // Nothing to do as the items have equal priority215 .eq => {}, // Nothing to do as the items have equal priority
...@@ -221,7 +217,7 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -221,7 +217,7 @@ pub fn PriorityQueue(comptime T: type) type {
221 }217 }
222218
223 pub const Iterator = struct {219 pub const Iterator = struct {
224 queue: *PriorityQueue(T),220 queue: *PriorityQueue(T, compareFn),
225 count: usize,221 count: usize,
226222
227 pub fn next(it: *Iterator) ?T {223 pub fn next(it: *Iterator) ?T {
...@@ -271,10 +267,11 @@ fn greaterThan(a: u32, b: u32) Order {...@@ -271,10 +267,11 @@ fn greaterThan(a: u32, b: u32) Order {
271 return lessThan(a, b).invert();267 return lessThan(a, b).invert();
272}268}
273269
274const PQ = PriorityQueue(u32);270const PQlt = PriorityQueue(u32, lessThan);
271const PQgt = PriorityQueue(u32, greaterThan);
275272
276test "std.PriorityQueue: add and remove min heap" {273test "std.PriorityQueue: add and remove min heap" {
277 var queue = PQ.init(testing.allocator, lessThan);274 var queue = PQlt.init(testing.allocator);
278 defer queue.deinit();275 defer queue.deinit();
279276
280 try queue.add(54);277 try queue.add(54);
...@@ -292,7 +289,7 @@ test "std.PriorityQueue: add and remove min heap" {...@@ -292,7 +289,7 @@ test "std.PriorityQueue: add and remove min heap" {
292}289}
293290
294test "std.PriorityQueue: add and remove same min heap" {291test "std.PriorityQueue: add and remove same min heap" {
295 var queue = PQ.init(testing.allocator, lessThan);292 var queue = PQlt.init(testing.allocator);
296 defer queue.deinit();293 defer queue.deinit();
297294
298 try queue.add(1);295 try queue.add(1);
...@@ -310,14 +307,14 @@ test "std.PriorityQueue: add and remove same min heap" {...@@ -310,14 +307,14 @@ test "std.PriorityQueue: add and remove same min heap" {
310}307}
311308
312test "std.PriorityQueue: removeOrNull on empty" {309test "std.PriorityQueue: removeOrNull on empty" {
313 var queue = PQ.init(testing.allocator, lessThan);310 var queue = PQlt.init(testing.allocator);
314 defer queue.deinit();311 defer queue.deinit();
315312
316 try expect(queue.removeOrNull() == null);313 try expect(queue.removeOrNull() == null);
317}314}
318315
319test "std.PriorityQueue: edge case 3 elements" {316test "std.PriorityQueue: edge case 3 elements" {
320 var queue = PQ.init(testing.allocator, lessThan);317 var queue = PQlt.init(testing.allocator);
321 defer queue.deinit();318 defer queue.deinit();
322319
323 try queue.add(9);320 try queue.add(9);
...@@ -329,7 +326,7 @@ test "std.PriorityQueue: edge case 3 elements" {...@@ -329,7 +326,7 @@ test "std.PriorityQueue: edge case 3 elements" {
329}326}
330327
331test "std.PriorityQueue: peek" {328test "std.PriorityQueue: peek" {
332 var queue = PQ.init(testing.allocator, lessThan);329 var queue = PQlt.init(testing.allocator);
333 defer queue.deinit();330 defer queue.deinit();
334331
335 try expect(queue.peek() == null);332 try expect(queue.peek() == null);
...@@ -341,7 +338,7 @@ test "std.PriorityQueue: peek" {...@@ -341,7 +338,7 @@ test "std.PriorityQueue: peek" {
341}338}
342339
343test "std.PriorityQueue: sift up with odd indices" {340test "std.PriorityQueue: sift up with odd indices" {
344 var queue = PQ.init(testing.allocator, lessThan);341 var queue = PQlt.init(testing.allocator);
345 defer queue.deinit();342 defer queue.deinit();
346 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };343 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
347 for (items) |e| {344 for (items) |e| {
...@@ -355,7 +352,7 @@ test "std.PriorityQueue: sift up with odd indices" {...@@ -355,7 +352,7 @@ test "std.PriorityQueue: sift up with odd indices" {
355}352}
356353
357test "std.PriorityQueue: addSlice" {354test "std.PriorityQueue: addSlice" {
358 var queue = PQ.init(testing.allocator, lessThan);355 var queue = PQlt.init(testing.allocator);
359 defer queue.deinit();356 defer queue.deinit();
360 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };357 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
361 try queue.addSlice(items[0..]);358 try queue.addSlice(items[0..]);
...@@ -369,7 +366,7 @@ test "std.PriorityQueue: addSlice" {...@@ -369,7 +366,7 @@ test "std.PriorityQueue: addSlice" {
369test "std.PriorityQueue: fromOwnedSlice trivial case 0" {366test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
370 const items = [0]u32{};367 const items = [0]u32{};
371 const queue_items = try testing.allocator.dupe(u32, &items);368 const queue_items = try testing.allocator.dupe(u32, &items);
372 var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]);369 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
373 defer queue.deinit();370 defer queue.deinit();
374 try expectEqual(@as(usize, 0), queue.len);371 try expectEqual(@as(usize, 0), queue.len);
375 try expect(queue.removeOrNull() == null);372 try expect(queue.removeOrNull() == null);
...@@ -378,7 +375,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {...@@ -378,7 +375,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
378test "std.PriorityQueue: fromOwnedSlice trivial case 1" {375test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
379 const items = [1]u32{1};376 const items = [1]u32{1};
380 const queue_items = try testing.allocator.dupe(u32, &items);377 const queue_items = try testing.allocator.dupe(u32, &items);
381 var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, queue_items[0..]);378 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
382 defer queue.deinit();379 defer queue.deinit();
383380
384 try expectEqual(@as(usize, 1), queue.len);381 try expectEqual(@as(usize, 1), queue.len);
...@@ -389,7 +386,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {...@@ -389,7 +386,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
389test "std.PriorityQueue: fromOwnedSlice" {386test "std.PriorityQueue: fromOwnedSlice" {
390 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };387 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
391 const heap_items = try testing.allocator.dupe(u32, items[0..]);388 const heap_items = try testing.allocator.dupe(u32, items[0..]);
392 var queue = PQ.fromOwnedSlice(testing.allocator, lessThan, heap_items[0..]);389 var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..]);
393 defer queue.deinit();390 defer queue.deinit();
394391
395 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };392 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
...@@ -399,7 +396,7 @@ test "std.PriorityQueue: fromOwnedSlice" {...@@ -399,7 +396,7 @@ test "std.PriorityQueue: fromOwnedSlice" {
399}396}
400397
401test "std.PriorityQueue: add and remove max heap" {398test "std.PriorityQueue: add and remove max heap" {
402 var queue = PQ.init(testing.allocator, greaterThan);399 var queue = PQgt.init(testing.allocator);
403 defer queue.deinit();400 defer queue.deinit();
404401
405 try queue.add(54);402 try queue.add(54);
...@@ -417,7 +414,7 @@ test "std.PriorityQueue: add and remove max heap" {...@@ -417,7 +414,7 @@ test "std.PriorityQueue: add and remove max heap" {
417}414}
418415
419test "std.PriorityQueue: add and remove same max heap" {416test "std.PriorityQueue: add and remove same max heap" {
420 var queue = PQ.init(testing.allocator, greaterThan);417 var queue = PQgt.init(testing.allocator);
421 defer queue.deinit();418 defer queue.deinit();
422419
423 try queue.add(1);420 try queue.add(1);
...@@ -435,7 +432,7 @@ test "std.PriorityQueue: add and remove same max heap" {...@@ -435,7 +432,7 @@ test "std.PriorityQueue: add and remove same max heap" {
435}432}
436433
437test "std.PriorityQueue: iterator" {434test "std.PriorityQueue: iterator" {
438 var queue = PQ.init(testing.allocator, lessThan);435 var queue = PQlt.init(testing.allocator);
439 var map = std.AutoHashMap(u32, void).init(testing.allocator);436 var map = std.AutoHashMap(u32, void).init(testing.allocator);
440 defer {437 defer {
441 queue.deinit();438 queue.deinit();
...@@ -457,7 +454,7 @@ test "std.PriorityQueue: iterator" {...@@ -457,7 +454,7 @@ test "std.PriorityQueue: iterator" {
457}454}
458455
459test "std.PriorityQueue: remove at index" {456test "std.PriorityQueue: remove at index" {
460 var queue = PQ.init(testing.allocator, lessThan);457 var queue = PQlt.init(testing.allocator);
461 defer queue.deinit();458 defer queue.deinit();
462459
463 try queue.add(3);460 try queue.add(3);
...@@ -480,7 +477,7 @@ test "std.PriorityQueue: remove at index" {...@@ -480,7 +477,7 @@ test "std.PriorityQueue: remove at index" {
480}477}
481478
482test "std.PriorityQueue: iterator while empty" {479test "std.PriorityQueue: iterator while empty" {
483 var queue = PQ.init(testing.allocator, lessThan);480 var queue = PQlt.init(testing.allocator);
484 defer queue.deinit();481 defer queue.deinit();
485482
486 var it = queue.iterator();483 var it = queue.iterator();
...@@ -489,7 +486,7 @@ test "std.PriorityQueue: iterator while empty" {...@@ -489,7 +486,7 @@ test "std.PriorityQueue: iterator while empty" {
489}486}
490487
491test "std.PriorityQueue: shrinkAndFree" {488test "std.PriorityQueue: shrinkAndFree" {
492 var queue = PQ.init(testing.allocator, lessThan);489 var queue = PQlt.init(testing.allocator);
493 defer queue.deinit();490 defer queue.deinit();
494491
495 try queue.ensureTotalCapacity(4);492 try queue.ensureTotalCapacity(4);
...@@ -512,7 +509,7 @@ test "std.PriorityQueue: shrinkAndFree" {...@@ -512,7 +509,7 @@ test "std.PriorityQueue: shrinkAndFree" {
512}509}
513510
514test "std.PriorityQueue: update min heap" {511test "std.PriorityQueue: update min heap" {
515 var queue = PQ.init(testing.allocator, lessThan);512 var queue = PQlt.init(testing.allocator);
516 defer queue.deinit();513 defer queue.deinit();
517514
518 try queue.add(55);515 try queue.add(55);
...@@ -527,7 +524,7 @@ test "std.PriorityQueue: update min heap" {...@@ -527,7 +524,7 @@ test "std.PriorityQueue: update min heap" {
527}524}
528525
529test "std.PriorityQueue: update same min heap" {526test "std.PriorityQueue: update same min heap" {
530 var queue = PQ.init(testing.allocator, lessThan);527 var queue = PQlt.init(testing.allocator);
531 defer queue.deinit();528 defer queue.deinit();
532529
533 try queue.add(1);530 try queue.add(1);
...@@ -543,7 +540,7 @@ test "std.PriorityQueue: update same min heap" {...@@ -543,7 +540,7 @@ test "std.PriorityQueue: update same min heap" {
543}540}
544541
545test "std.PriorityQueue: update max heap" {542test "std.PriorityQueue: update max heap" {
546 var queue = PQ.init(testing.allocator, greaterThan);543 var queue = PQgt.init(testing.allocator);
547 defer queue.deinit();544 defer queue.deinit();
548545
549 try queue.add(55);546 try queue.add(55);
...@@ -558,7 +555,7 @@ test "std.PriorityQueue: update max heap" {...@@ -558,7 +555,7 @@ test "std.PriorityQueue: update max heap" {
558}555}
559556
560test "std.PriorityQueue: update same max heap" {557test "std.PriorityQueue: update same max heap" {
561 var queue = PQ.init(testing.allocator, greaterThan);558 var queue = PQgt.init(testing.allocator);
562 defer queue.deinit();559 defer queue.deinit();
563560
564 try queue.add(1);561 try queue.add(1);