authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2021-12-17 14:34:32-08:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2021-12-17 14:34:32-08:00
loge8b39960bbd88d73cfd0d853e8ca7383e300c58c
tree3602471fe2456649c8f163f466fb4d01eb8ca435
parentc1745b289689902945bd8171b0518d57e53557ec
parent22e1d92d3edad13f6a3371cb11931f9f66c910c6
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #10342 from Arnavion/pq-comparefn-context

std.priority_{de,}queue: allow comparator to take context parameter

2 files changed, 151 insertions(+), 86 deletions(-)

lib/std/priority_dequeue.zig+82-49
......@@ -8,27 +8,29 @@ const expectEqual = testing.expectEqual;
88const expectError = testing.expectError;
99
1010/// Priority Dequeue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its first
12/// argument should get min-popped before its second argument,
11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get min-popped before its third argument,
1313/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
14/// if the second argument should be min-popped first.
14/// if the third argument should be min-popped second.
1515/// Popping the max element works in reverse. For example,
1616/// to make `popMin` return the smallest number, provide
17/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
18pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) type {
17/// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
18pub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
1919 return struct {
2020 const Self = @This();
2121
2222 items: []T,
2323 len: usize,
2424 allocator: Allocator,
25 context: Context,
2526
2627 /// Initialize and return a new priority dequeue.
27 pub fn init(allocator: Allocator) Self {
28 pub fn init(allocator: Allocator, context: Context) Self {
2829 return Self{
2930 .items = &[_]T{},
3031 .len = 0,
3132 .allocator = allocator,
33 .context = context,
3234 };
3335 }
3436
......@@ -87,7 +89,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
8789 const parent = self.items[parent_index];
8890
8991 const min_layer = self.nextIsMinLayer();
90 const order = compareFn(child, parent);
92 const order = compareFn(self.context, child, parent);
9193 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
9294 // We must swap the item with it's parent if it is on the "wrong" layer
9395 self.items[parent_index] = child;
......@@ -120,7 +122,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
120122 const grandparent = self.items[grandparent_index];
121123
122124 // If the grandparent is already better or equal, we have gone as far as we need to
123 if (compareFn(child, grandparent) != target_order) break;
125 if (compareFn(self.context, child, grandparent) != target_order) break;
124126
125127 // Otherwise swap the item with it's grandparent
126128 self.items[grandparent_index] = child;
......@@ -215,10 +217,10 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
215217 // Find the best grandchild
216218 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
217219 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
218 const best_grandchild = Self.bestItem(best_left, best_right, target_order);
220 const best_grandchild = self.bestItem(best_left, best_right, target_order);
219221
220222 // If the item is better than or equal to its best grandchild, we are done
221 if (compareFn(best_grandchild.item, elem) != target_order) return;
223 if (compareFn(self.context, best_grandchild.item, elem) != target_order) return;
222224
223225 // Otherwise, swap them
224226 self.items[best_grandchild.index] = elem;
......@@ -235,7 +237,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
235237 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
236238
237239 // If the item is better than or equal to its best descendant, we are done
238 if (compareFn(best_descendent.item, elem) != target_order) return;
240 if (compareFn(self.context, best_descendent.item, elem) != target_order) return;
239241
240242 // Otherwise swap them
241243 self.items[best_descendent.index] = elem;
......@@ -259,7 +261,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
259261 const parent_index = parentIndex(child_index);
260262 const parent = self.items[parent_index];
261263
262 if (compareFn(parent, child) == target_order) {
264 if (compareFn(self.context, parent, child) == target_order) {
263265 self.items[parent_index] = child;
264266 self.items[child_index] = parent;
265267 }
......@@ -277,8 +279,8 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
277279 };
278280 }
279281
280 fn bestItem(item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
281 if (compareFn(item1.item, item2.item) == target_order) {
282 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
283 if (compareFn(self.context, item1.item, item2.item) == target_order) {
282284 return item1;
283285 } else {
284286 return item2;
......@@ -288,7 +290,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
288290 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
289291 var item1 = self.getItem(index1);
290292 var item2 = self.getItem(index2);
291 return Self.bestItem(item1, item2, target_order);
293 return self.bestItem(item1, item2, target_order);
292294 }
293295
294296 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
......@@ -336,11 +338,12 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
336338 /// Dequeue takes ownership of the passed in slice. The slice must have been
337339 /// allocated with `allocator`.
338340 /// De-initialize with `deinit`.
339 pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self {
341 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
340342 var queue = Self{
341343 .items = items,
342344 .len = items.len,
343345 .allocator = allocator,
346 .context = context,
344347 };
345348
346349 if (queue.len <= 1) return queue;
......@@ -394,7 +397,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
394397 }
395398
396399 pub const Iterator = struct {
397 queue: *PriorityDequeue(T, compareFn),
400 queue: *PriorityDequeue(T, Context, compareFn),
398401 count: usize,
399402
400403 pub fn next(it: *Iterator) ?T {
......@@ -453,14 +456,15 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
453456 };
454457}
455458
456fn lessThanComparison(a: u32, b: u32) Order {
459fn lessThanComparison(context: void, a: u32, b: u32) Order {
460 _ = context;
457461 return std.math.order(a, b);
458462}
459463
460const PDQ = PriorityDequeue(u32, lessThanComparison);
464const PDQ = PriorityDequeue(u32, void, lessThanComparison);
461465
462466test "std.PriorityDequeue: add and remove min" {
463 var queue = PDQ.init(testing.allocator);
467 var queue = PDQ.init(testing.allocator, {});
464468 defer queue.deinit();
465469
466470 try queue.add(54);
......@@ -482,11 +486,12 @@ test "std.PriorityDequeue: add and remove min structs" {
482486 const S = struct {
483487 size: u32,
484488 };
485 var queue = PriorityDequeue(S, struct {
486 fn order(a: S, b: S) Order {
489 var queue = PriorityDequeue(S, void, struct {
490 fn order(context: void, a: S, b: S) Order {
491 _ = context;
487492 return std.math.order(a.size, b.size);
488493 }
489 }.order).init(testing.allocator);
494 }.order).init(testing.allocator, {});
490495 defer queue.deinit();
491496
492497 try queue.add(.{ .size = 54 });
......@@ -505,7 +510,7 @@ test "std.PriorityDequeue: add and remove min structs" {
505510}
506511
507512test "std.PriorityDequeue: add and remove max" {
508 var queue = PDQ.init(testing.allocator);
513 var queue = PDQ.init(testing.allocator, {});
509514 defer queue.deinit();
510515
511516 try queue.add(54);
......@@ -524,7 +529,7 @@ test "std.PriorityDequeue: add and remove max" {
524529}
525530
526531test "std.PriorityDequeue: add and remove same min" {
527 var queue = PDQ.init(testing.allocator);
532 var queue = PDQ.init(testing.allocator, {});
528533 defer queue.deinit();
529534
530535 try queue.add(1);
......@@ -543,7 +548,7 @@ test "std.PriorityDequeue: add and remove same min" {
543548}
544549
545550test "std.PriorityDequeue: add and remove same max" {
546 var queue = PDQ.init(testing.allocator);
551 var queue = PDQ.init(testing.allocator, {});
547552 defer queue.deinit();
548553
549554 try queue.add(1);
......@@ -562,7 +567,7 @@ test "std.PriorityDequeue: add and remove same max" {
562567}
563568
564569test "std.PriorityDequeue: removeOrNull empty" {
565 var queue = PDQ.init(testing.allocator);
570 var queue = PDQ.init(testing.allocator, {});
566571 defer queue.deinit();
567572
568573 try expect(queue.removeMinOrNull() == null);
......@@ -570,7 +575,7 @@ test "std.PriorityDequeue: removeOrNull empty" {
570575}
571576
572577test "std.PriorityDequeue: edge case 3 elements" {
573 var queue = PDQ.init(testing.allocator);
578 var queue = PDQ.init(testing.allocator, {});
574579 defer queue.deinit();
575580
576581 try queue.add(9);
......@@ -583,7 +588,7 @@ test "std.PriorityDequeue: edge case 3 elements" {
583588}
584589
585590test "std.PriorityDequeue: edge case 3 elements max" {
586 var queue = PDQ.init(testing.allocator);
591 var queue = PDQ.init(testing.allocator, {});
587592 defer queue.deinit();
588593
589594 try queue.add(9);
......@@ -596,7 +601,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {
596601}
597602
598603test "std.PriorityDequeue: peekMin" {
599 var queue = PDQ.init(testing.allocator);
604 var queue = PDQ.init(testing.allocator, {});
600605 defer queue.deinit();
601606
602607 try expect(queue.peekMin() == null);
......@@ -610,7 +615,7 @@ test "std.PriorityDequeue: peekMin" {
610615}
611616
612617test "std.PriorityDequeue: peekMax" {
613 var queue = PDQ.init(testing.allocator);
618 var queue = PDQ.init(testing.allocator, {});
614619 defer queue.deinit();
615620
616621 try expect(queue.peekMin() == null);
......@@ -624,7 +629,7 @@ test "std.PriorityDequeue: peekMax" {
624629}
625630
626631test "std.PriorityDequeue: sift up with odd indices" {
627 var queue = PDQ.init(testing.allocator);
632 var queue = PDQ.init(testing.allocator, {});
628633 defer queue.deinit();
629634 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
630635 for (items) |e| {
......@@ -638,7 +643,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
638643}
639644
640645test "std.PriorityDequeue: sift up with odd indices" {
641 var queue = PDQ.init(testing.allocator);
646 var queue = PDQ.init(testing.allocator, {});
642647 defer queue.deinit();
643648 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
644649 for (items) |e| {
......@@ -652,7 +657,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
652657}
653658
654659test "std.PriorityDequeue: addSlice min" {
655 var queue = PDQ.init(testing.allocator);
660 var queue = PDQ.init(testing.allocator, {});
656661 defer queue.deinit();
657662 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
658663 try queue.addSlice(items[0..]);
......@@ -664,7 +669,7 @@ test "std.PriorityDequeue: addSlice min" {
664669}
665670
666671test "std.PriorityDequeue: addSlice max" {
667 var queue = PDQ.init(testing.allocator);
672 var queue = PDQ.init(testing.allocator, {});
668673 defer queue.deinit();
669674 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
670675 try queue.addSlice(items[0..]);
......@@ -678,7 +683,7 @@ test "std.PriorityDequeue: addSlice max" {
678683test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
679684 const items = [0]u32{};
680685 const queue_items = try testing.allocator.dupe(u32, &items);
681 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
686 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
682687 defer queue.deinit();
683688 try expectEqual(@as(usize, 0), queue.len);
684689 try expect(queue.removeMinOrNull() == null);
......@@ -687,7 +692,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
687692test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
688693 const items = [1]u32{1};
689694 const queue_items = try testing.allocator.dupe(u32, &items);
690 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
695 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
691696 defer queue.deinit();
692697
693698 try expectEqual(@as(usize, 1), queue.len);
......@@ -698,7 +703,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
698703test "std.PriorityDequeue: fromOwnedSlice" {
699704 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
700705 const queue_items = try testing.allocator.dupe(u32, items[0..]);
701 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..]);
706 var queue = PDQ.fromOwnedSlice(testing.allocator, queue_items[0..], {});
702707 defer queue.deinit();
703708
704709 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
......@@ -708,7 +713,7 @@ test "std.PriorityDequeue: fromOwnedSlice" {
708713}
709714
710715test "std.PriorityDequeue: update min queue" {
711 var queue = PDQ.init(testing.allocator);
716 var queue = PDQ.init(testing.allocator, {});
712717 defer queue.deinit();
713718
714719 try queue.add(55);
......@@ -723,7 +728,7 @@ test "std.PriorityDequeue: update min queue" {
723728}
724729
725730test "std.PriorityDequeue: update same min queue" {
726 var queue = PDQ.init(testing.allocator);
731 var queue = PDQ.init(testing.allocator, {});
727732 defer queue.deinit();
728733
729734 try queue.add(1);
......@@ -739,7 +744,7 @@ test "std.PriorityDequeue: update same min queue" {
739744}
740745
741746test "std.PriorityDequeue: update max queue" {
742 var queue = PDQ.init(testing.allocator);
747 var queue = PDQ.init(testing.allocator, {});
743748 defer queue.deinit();
744749
745750 try queue.add(55);
......@@ -755,7 +760,7 @@ test "std.PriorityDequeue: update max queue" {
755760}
756761
757762test "std.PriorityDequeue: update same max queue" {
758 var queue = PDQ.init(testing.allocator);
763 var queue = PDQ.init(testing.allocator, {});
759764 defer queue.deinit();
760765
761766 try queue.add(1);
......@@ -771,7 +776,7 @@ test "std.PriorityDequeue: update same max queue" {
771776}
772777
773778test "std.PriorityDequeue: iterator" {
774 var queue = PDQ.init(testing.allocator);
779 var queue = PDQ.init(testing.allocator, {});
775780 var map = std.AutoHashMap(u32, void).init(testing.allocator);
776781 defer {
777782 queue.deinit();
......@@ -793,7 +798,7 @@ test "std.PriorityDequeue: iterator" {
793798}
794799
795800test "std.PriorityDequeue: remove at index" {
796 var queue = PDQ.init(testing.allocator);
801 var queue = PDQ.init(testing.allocator, {});
797802 defer queue.deinit();
798803
799804 try queue.add(3);
......@@ -816,7 +821,7 @@ test "std.PriorityDequeue: remove at index" {
816821}
817822
818823test "std.PriorityDequeue: iterator while empty" {
819 var queue = PDQ.init(testing.allocator);
824 var queue = PDQ.init(testing.allocator, {});
820825 defer queue.deinit();
821826
822827 var it = queue.iterator();
......@@ -825,7 +830,7 @@ test "std.PriorityDequeue: iterator while empty" {
825830}
826831
827832test "std.PriorityDequeue: shrinkAndFree" {
828 var queue = PDQ.init(testing.allocator);
833 var queue = PDQ.init(testing.allocator, {});
829834 defer queue.deinit();
830835
831836 try queue.ensureTotalCapacity(4);
......@@ -864,7 +869,7 @@ fn fuzzTestMin(rng: std.rand.Random, comptime queue_size: usize) !void {
864869 const allocator = testing.allocator;
865870 const items = try generateRandomSlice(allocator, rng, queue_size);
866871
867 var queue = PDQ.fromOwnedSlice(allocator, items);
872 var queue = PDQ.fromOwnedSlice(allocator, items, {});
868873 defer queue.deinit();
869874
870875 var last_removed: ?u32 = null;
......@@ -893,7 +898,7 @@ fn fuzzTestMax(rng: std.rand.Random, queue_size: usize) !void {
893898 const allocator = testing.allocator;
894899 const items = try generateRandomSlice(allocator, rng, queue_size);
895900
896 var queue = PDQ.fromOwnedSlice(testing.allocator, items);
901 var queue = PDQ.fromOwnedSlice(testing.allocator, items, {});
897902 defer queue.deinit();
898903
899904 var last_removed: ?u32 = null;
......@@ -922,7 +927,7 @@ fn fuzzTestMinMax(rng: std.rand.Random, queue_size: usize) !void {
922927 const allocator = testing.allocator;
923928 const items = try generateRandomSlice(allocator, rng, queue_size);
924929
925 var queue = PDQ.fromOwnedSlice(allocator, items);
930 var queue = PDQ.fromOwnedSlice(allocator, items, {});
926931 defer queue.deinit();
927932
928933 var last_min: ?u32 = null;
......@@ -957,3 +962,31 @@ fn generateRandomSlice(allocator: std.mem.Allocator, rng: std.rand.Random, size:
957962
958963 return array.toOwnedSlice();
959964}
965
966fn contextLessThanComparison(context: []const u32, a: usize, b: usize) Order {
967 return std.math.order(context[a], context[b]);
968}
969
970const CPDQ = PriorityDequeue(usize, []const u32, contextLessThanComparison);
971
972test "std.PriorityDequeue: add and remove" {
973 const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 };
974
975 var queue = CPDQ.init(testing.allocator, context[0..]);
976 defer queue.deinit();
977
978 try queue.add(0);
979 try queue.add(1);
980 try queue.add(2);
981 try queue.add(3);
982 try queue.add(4);
983 try queue.add(5);
984 try queue.add(6);
985 try expectEqual(@as(usize, 6), queue.removeMin());
986 try expectEqual(@as(usize, 5), queue.removeMax());
987 try expectEqual(@as(usize, 3), queue.removeMin());
988 try expectEqual(@as(usize, 0), queue.removeMax());
989 try expectEqual(@as(usize, 4), queue.removeMin());
990 try expectEqual(@as(usize, 2), queue.removeMax());
991 try expectEqual(@as(usize, 1), queue.removeMin());
992}
lib/std/priority_queue.zig+69-37
......@@ -8,26 +8,28 @@ const expectEqual = testing.expectEqual;
88const expectError = testing.expectError;
99
1010/// Priority queue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its first
12/// argument should get popped before its second argument,
11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get popped before its third argument,
1313/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
14/// if the second argument should be popped first.
14/// if the third argument should be popped first.
1515/// For example, to make `pop` return the smallest number, provide
16/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`
17pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order) type {
16/// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
17pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
1818 return struct {
1919 const Self = @This();
2020
2121 items: []T,
2222 len: usize,
2323 allocator: Allocator,
24 context: Context,
2425
2526 /// Initialize and return a priority queue.
26 pub fn init(allocator: Allocator) Self {
27 pub fn init(allocator: Allocator, context: Context) Self {
2728 return Self{
2829 .items = &[_]T{},
2930 .len = 0,
3031 .allocator = allocator,
32 .context = context,
3133 };
3234 }
3335
......@@ -55,7 +57,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
5557 const child = self.items[child_index];
5658 const parent = self.items[parent_index];
5759
58 if (compareFn(child, parent) != .lt) break;
60 if (compareFn(self.context, child, parent) != .lt) break;
5961
6062 self.items[parent_index] = child;
6163 self.items[child_index] = parent;
......@@ -127,14 +129,14 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
127129 var smallest = self.items[index];
128130
129131 if (left) |e| {
130 if (compareFn(e, smallest) == .lt) {
132 if (compareFn(self.context, e, smallest) == .lt) {
131133 smallest_index = left_index;
132134 smallest = e;
133135 }
134136 }
135137
136138 if (right) |e| {
137 if (compareFn(e, smallest) == .lt) {
139 if (compareFn(self.context, e, smallest) == .lt) {
138140 smallest_index = right_index;
139141 smallest = e;
140142 }
......@@ -153,11 +155,12 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
153155 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
154156 /// allocated with `allocator`.
155157 /// Deinitialize with `deinit`.
156 pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self {
158 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
157159 var queue = Self{
158160 .items = items,
159161 .len = items.len,
160162 .allocator = allocator,
163 .context = context,
161164 };
162165
163166 if (queue.len <= 1) return queue;
......@@ -207,7 +210,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
207210 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
208211 const old_elem: T = self.items[update_index];
209212 self.items[update_index] = new_elem;
210 switch (compareFn(new_elem, old_elem)) {
213 switch (compareFn(self.context, new_elem, old_elem)) {
211214 .lt => siftUp(self, update_index),
212215 .gt => siftDown(self, update_index),
213216 .eq => {}, // Nothing to do as the items have equal priority
......@@ -215,7 +218,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
215218 }
216219
217220 pub const Iterator = struct {
218 queue: *PriorityQueue(T, compareFn),
221 queue: *PriorityQueue(T, Context, compareFn),
219222 count: usize,
220223
221224 pub fn next(it: *Iterator) ?T {
......@@ -258,19 +261,20 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
258261 };
259262}
260263
261fn lessThan(a: u32, b: u32) Order {
264fn lessThan(context: void, a: u32, b: u32) Order {
265 _ = context;
262266 return std.math.order(a, b);
263267}
264268
265fn greaterThan(a: u32, b: u32) Order {
266 return lessThan(a, b).invert();
269fn greaterThan(context: void, a: u32, b: u32) Order {
270 return lessThan(context, a, b).invert();
267271}
268272
269const PQlt = PriorityQueue(u32, lessThan);
270const PQgt = PriorityQueue(u32, greaterThan);
273const PQlt = PriorityQueue(u32, void, lessThan);
274const PQgt = PriorityQueue(u32, void, greaterThan);
271275
272276test "std.PriorityQueue: add and remove min heap" {
273 var queue = PQlt.init(testing.allocator);
277 var queue = PQlt.init(testing.allocator, {});
274278 defer queue.deinit();
275279
276280 try queue.add(54);
......@@ -288,7 +292,7 @@ test "std.PriorityQueue: add and remove min heap" {
288292}
289293
290294test "std.PriorityQueue: add and remove same min heap" {
291 var queue = PQlt.init(testing.allocator);
295 var queue = PQlt.init(testing.allocator, {});
292296 defer queue.deinit();
293297
294298 try queue.add(1);
......@@ -306,14 +310,14 @@ test "std.PriorityQueue: add and remove same min heap" {
306310}
307311
308312test "std.PriorityQueue: removeOrNull on empty" {
309 var queue = PQlt.init(testing.allocator);
313 var queue = PQlt.init(testing.allocator, {});
310314 defer queue.deinit();
311315
312316 try expect(queue.removeOrNull() == null);
313317}
314318
315319test "std.PriorityQueue: edge case 3 elements" {
316 var queue = PQlt.init(testing.allocator);
320 var queue = PQlt.init(testing.allocator, {});
317321 defer queue.deinit();
318322
319323 try queue.add(9);
......@@ -325,7 +329,7 @@ test "std.PriorityQueue: edge case 3 elements" {
325329}
326330
327331test "std.PriorityQueue: peek" {
328 var queue = PQlt.init(testing.allocator);
332 var queue = PQlt.init(testing.allocator, {});
329333 defer queue.deinit();
330334
331335 try expect(queue.peek() == null);
......@@ -337,7 +341,7 @@ test "std.PriorityQueue: peek" {
337341}
338342
339343test "std.PriorityQueue: sift up with odd indices" {
340 var queue = PQlt.init(testing.allocator);
344 var queue = PQlt.init(testing.allocator, {});
341345 defer queue.deinit();
342346 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
343347 for (items) |e| {
......@@ -351,7 +355,7 @@ test "std.PriorityQueue: sift up with odd indices" {
351355}
352356
353357test "std.PriorityQueue: addSlice" {
354 var queue = PQlt.init(testing.allocator);
358 var queue = PQlt.init(testing.allocator, {});
355359 defer queue.deinit();
356360 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
357361 try queue.addSlice(items[0..]);
......@@ -365,7 +369,7 @@ test "std.PriorityQueue: addSlice" {
365369test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
366370 const items = [0]u32{};
367371 const queue_items = try testing.allocator.dupe(u32, &items);
368 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
372 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
369373 defer queue.deinit();
370374 try expectEqual(@as(usize, 0), queue.len);
371375 try expect(queue.removeOrNull() == null);
......@@ -374,7 +378,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
374378test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
375379 const items = [1]u32{1};
376380 const queue_items = try testing.allocator.dupe(u32, &items);
377 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..]);
381 var queue = PQlt.fromOwnedSlice(testing.allocator, queue_items[0..], {});
378382 defer queue.deinit();
379383
380384 try expectEqual(@as(usize, 1), queue.len);
......@@ -385,7 +389,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
385389test "std.PriorityQueue: fromOwnedSlice" {
386390 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
387391 const heap_items = try testing.allocator.dupe(u32, items[0..]);
388 var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..]);
392 var queue = PQlt.fromOwnedSlice(testing.allocator, heap_items[0..], {});
389393 defer queue.deinit();
390394
391395 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
......@@ -395,7 +399,7 @@ test "std.PriorityQueue: fromOwnedSlice" {
395399}
396400
397401test "std.PriorityQueue: add and remove max heap" {
398 var queue = PQgt.init(testing.allocator);
402 var queue = PQgt.init(testing.allocator, {});
399403 defer queue.deinit();
400404
401405 try queue.add(54);
......@@ -413,7 +417,7 @@ test "std.PriorityQueue: add and remove max heap" {
413417}
414418
415419test "std.PriorityQueue: add and remove same max heap" {
416 var queue = PQgt.init(testing.allocator);
420 var queue = PQgt.init(testing.allocator, {});
417421 defer queue.deinit();
418422
419423 try queue.add(1);
......@@ -431,7 +435,7 @@ test "std.PriorityQueue: add and remove same max heap" {
431435}
432436
433437test "std.PriorityQueue: iterator" {
434 var queue = PQlt.init(testing.allocator);
438 var queue = PQlt.init(testing.allocator, {});
435439 var map = std.AutoHashMap(u32, void).init(testing.allocator);
436440 defer {
437441 queue.deinit();
......@@ -453,7 +457,7 @@ test "std.PriorityQueue: iterator" {
453457}
454458
455459test "std.PriorityQueue: remove at index" {
456 var queue = PQlt.init(testing.allocator);
460 var queue = PQlt.init(testing.allocator, {});
457461 defer queue.deinit();
458462
459463 try queue.add(3);
......@@ -476,7 +480,7 @@ test "std.PriorityQueue: remove at index" {
476480}
477481
478482test "std.PriorityQueue: iterator while empty" {
479 var queue = PQlt.init(testing.allocator);
483 var queue = PQlt.init(testing.allocator, {});
480484 defer queue.deinit();
481485
482486 var it = queue.iterator();
......@@ -485,7 +489,7 @@ test "std.PriorityQueue: iterator while empty" {
485489}
486490
487491test "std.PriorityQueue: shrinkAndFree" {
488 var queue = PQlt.init(testing.allocator);
492 var queue = PQlt.init(testing.allocator, {});
489493 defer queue.deinit();
490494
491495 try queue.ensureTotalCapacity(4);
......@@ -508,7 +512,7 @@ test "std.PriorityQueue: shrinkAndFree" {
508512}
509513
510514test "std.PriorityQueue: update min heap" {
511 var queue = PQlt.init(testing.allocator);
515 var queue = PQlt.init(testing.allocator, {});
512516 defer queue.deinit();
513517
514518 try queue.add(55);
......@@ -523,7 +527,7 @@ test "std.PriorityQueue: update min heap" {
523527}
524528
525529test "std.PriorityQueue: update same min heap" {
526 var queue = PQlt.init(testing.allocator);
530 var queue = PQlt.init(testing.allocator, {});
527531 defer queue.deinit();
528532
529533 try queue.add(1);
......@@ -539,7 +543,7 @@ test "std.PriorityQueue: update same min heap" {
539543}
540544
541545test "std.PriorityQueue: update max heap" {
542 var queue = PQgt.init(testing.allocator);
546 var queue = PQgt.init(testing.allocator, {});
543547 defer queue.deinit();
544548
545549 try queue.add(55);
......@@ -554,7 +558,7 @@ test "std.PriorityQueue: update max heap" {
554558}
555559
556560test "std.PriorityQueue: update same max heap" {
557 var queue = PQgt.init(testing.allocator);
561 var queue = PQgt.init(testing.allocator, {});
558562 defer queue.deinit();
559563
560564 try queue.add(1);
......@@ -568,3 +572,31 @@ test "std.PriorityQueue: update same max heap" {
568572 try expectEqual(@as(u32, 2), queue.remove());
569573 try expectEqual(@as(u32, 1), queue.remove());
570574}
575
576fn contextLessThan(context: []const u32, a: usize, b: usize) Order {
577 return std.math.order(context[a], context[b]);
578}
579
580const CPQlt = PriorityQueue(usize, []const u32, contextLessThan);
581
582test "std.PriorityQueue: add and remove min heap with contextful comparator" {
583 const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 };
584
585 var queue = CPQlt.init(testing.allocator, context[0..]);
586 defer queue.deinit();
587
588 try queue.add(0);
589 try queue.add(1);
590 try queue.add(2);
591 try queue.add(3);
592 try queue.add(4);
593 try queue.add(5);
594 try queue.add(6);
595 try expectEqual(@as(usize, 6), queue.remove());
596 try expectEqual(@as(usize, 4), queue.remove());
597 try expectEqual(@as(usize, 3), queue.remove());
598 try expectEqual(@as(usize, 1), queue.remove());
599 try expectEqual(@as(usize, 2), queue.remove());
600 try expectEqual(@as(usize, 0), queue.remove());
601 try expectEqual(@as(usize, 5), queue.remove());
602}