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;...@@ -8,27 +8,29 @@ const expectEqual = testing.expectEqual;
8const expectError = testing.expectError;8const expectError = testing.expectError;
99
10/// Priority Dequeue for storing generic data. Initialize with `init`.10/// Priority Dequeue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its first11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get min-popped before its second argument,12/// argument should get min-popped before its third argument,
13/// `Order.eq` if the arguments are of equal priority, or `Order.gt`13/// `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.
15/// Popping the max element works in reverse. For example,15/// Popping the max element works in reverse. For example,
16/// to make `popMin` return the smallest number, provide16/// to make `popMin` return the smallest number, provide
17/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`17/// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
18pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) type {18pub fn PriorityDequeue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
19 return struct {19 return struct {
20 const Self = @This();20 const Self = @This();
2121
22 items: []T,22 items: []T,
23 len: usize,23 len: usize,
24 allocator: Allocator,24 allocator: Allocator,
25 context: Context,
2526
26 /// Initialize and return a new priority dequeue.27 /// Initialize and return a new priority dequeue.
27 pub fn init(allocator: Allocator) Self {28 pub fn init(allocator: Allocator, context: Context) Self {
28 return Self{29 return Self{
29 .items = &[_]T{},30 .items = &[_]T{},
30 .len = 0,31 .len = 0,
31 .allocator = allocator,32 .allocator = allocator,
33 .context = context,
32 };34 };
33 }35 }
3436
...@@ -87,7 +89,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -87,7 +89,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
87 const parent = self.items[parent_index];89 const parent = self.items[parent_index];
8890
89 const min_layer = self.nextIsMinLayer();91 const min_layer = self.nextIsMinLayer();
90 const order = compareFn(child, parent);92 const order = compareFn(self.context, child, parent);
91 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {93 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
92 // We must swap the item with it's parent if it is on the "wrong" layer94 // We must swap the item with it's parent if it is on the "wrong" layer
93 self.items[parent_index] = child;95 self.items[parent_index] = child;
...@@ -120,7 +122,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -120,7 +122,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
120 const grandparent = self.items[grandparent_index];122 const grandparent = self.items[grandparent_index];
121123
122 // If the grandparent is already better or equal, we have gone as far as we need to124 // 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
125 // Otherwise swap the item with it's grandparent127 // Otherwise swap the item with it's grandparent
126 self.items[grandparent_index] = child;128 self.items[grandparent_index] = child;
...@@ -215,10 +217,10 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -215,10 +217,10 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
215 // Find the best grandchild217 // Find the best grandchild
216 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);218 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
217 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);219 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
220 // If the item is better than or equal to its best grandchild, we are done222 // 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
223 // Otherwise, swap them225 // Otherwise, swap them
224 self.items[best_grandchild.index] = elem;226 self.items[best_grandchild.index] = elem;
...@@ -235,7 +237,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -235,7 +237,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
235 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);237 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
236238
237 // If the item is better than or equal to its best descendant, we are done239 // 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
240 // Otherwise swap them242 // Otherwise swap them
241 self.items[best_descendent.index] = elem;243 self.items[best_descendent.index] = elem;
...@@ -259,7 +261,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -259,7 +261,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
259 const parent_index = parentIndex(child_index);261 const parent_index = parentIndex(child_index);
260 const parent = self.items[parent_index];262 const parent = self.items[parent_index];
261263
262 if (compareFn(parent, child) == target_order) {264 if (compareFn(self.context, parent, child) == target_order) {
263 self.items[parent_index] = child;265 self.items[parent_index] = child;
264 self.items[child_index] = parent;266 self.items[child_index] = parent;
265 }267 }
...@@ -277,8 +279,8 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -277,8 +279,8 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
277 };279 };
278 }280 }
279281
280 fn bestItem(item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {282 fn bestItem(self: Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
281 if (compareFn(item1.item, item2.item) == target_order) {283 if (compareFn(self.context, item1.item, item2.item) == target_order) {
282 return item1;284 return item1;
283 } else {285 } else {
284 return item2;286 return item2;
...@@ -288,7 +290,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -288,7 +290,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
288 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {290 fn bestItemAtIndices(self: Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
289 var item1 = self.getItem(index1);291 var item1 = self.getItem(index1);
290 var item2 = self.getItem(index2);292 var item2 = self.getItem(index2);
291 return Self.bestItem(item1, item2, target_order);293 return self.bestItem(item1, item2, target_order);
292 }294 }
293295
294 fn bestDescendent(self: Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {296 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...@@ -336,11 +338,12 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
336 /// Dequeue takes ownership of the passed in slice. The slice must have been338 /// Dequeue takes ownership of the passed in slice. The slice must have been
337 /// allocated with `allocator`.339 /// allocated with `allocator`.
338 /// De-initialize with `deinit`.340 /// De-initialize with `deinit`.
339 pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self {341 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
340 var queue = Self{342 var queue = Self{
341 .items = items,343 .items = items,
342 .len = items.len,344 .len = items.len,
343 .allocator = allocator,345 .allocator = allocator,
346 .context = context,
344 };347 };
345348
346 if (queue.len <= 1) return queue;349 if (queue.len <= 1) return queue;
...@@ -394,7 +397,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -394,7 +397,7 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
394 }397 }
395398
396 pub const Iterator = struct {399 pub const Iterator = struct {
397 queue: *PriorityDequeue(T, compareFn),400 queue: *PriorityDequeue(T, Context, compareFn),
398 count: usize,401 count: usize,
399402
400 pub fn next(it: *Iterator) ?T {403 pub fn next(it: *Iterator) ?T {
...@@ -453,14 +456,15 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty...@@ -453,14 +456,15 @@ pub fn PriorityDequeue(comptime T: type, comptime compareFn: fn (T, T) Order) ty
453 };456 };
454}457}
455458
456fn lessThanComparison(a: u32, b: u32) Order {459fn lessThanComparison(context: void, a: u32, b: u32) Order {
460 _ = context;
457 return std.math.order(a, b);461 return std.math.order(a, b);
458}462}
459463
460const PDQ = PriorityDequeue(u32, lessThanComparison);464const PDQ = PriorityDequeue(u32, void, lessThanComparison);
461465
462test "std.PriorityDequeue: add and remove min" {466test "std.PriorityDequeue: add and remove min" {
463 var queue = PDQ.init(testing.allocator);467 var queue = PDQ.init(testing.allocator, {});
464 defer queue.deinit();468 defer queue.deinit();
465469
466 try queue.add(54);470 try queue.add(54);
...@@ -482,11 +486,12 @@ test "std.PriorityDequeue: add and remove min structs" {...@@ -482,11 +486,12 @@ test "std.PriorityDequeue: add and remove min structs" {
482 const S = struct {486 const S = struct {
483 size: u32,487 size: u32,
484 };488 };
485 var queue = PriorityDequeue(S, struct {489 var queue = PriorityDequeue(S, void, struct {
486 fn order(a: S, b: S) Order {490 fn order(context: void, a: S, b: S) Order {
491 _ = context;
487 return std.math.order(a.size, b.size);492 return std.math.order(a.size, b.size);
488 }493 }
489 }.order).init(testing.allocator);494 }.order).init(testing.allocator, {});
490 defer queue.deinit();495 defer queue.deinit();
491496
492 try queue.add(.{ .size = 54 });497 try queue.add(.{ .size = 54 });
...@@ -505,7 +510,7 @@ test "std.PriorityDequeue: add and remove min structs" {...@@ -505,7 +510,7 @@ test "std.PriorityDequeue: add and remove min structs" {
505}510}
506511
507test "std.PriorityDequeue: add and remove max" {512test "std.PriorityDequeue: add and remove max" {
508 var queue = PDQ.init(testing.allocator);513 var queue = PDQ.init(testing.allocator, {});
509 defer queue.deinit();514 defer queue.deinit();
510515
511 try queue.add(54);516 try queue.add(54);
...@@ -524,7 +529,7 @@ test "std.PriorityDequeue: add and remove max" {...@@ -524,7 +529,7 @@ test "std.PriorityDequeue: add and remove max" {
524}529}
525530
526test "std.PriorityDequeue: add and remove same min" {531test "std.PriorityDequeue: add and remove same min" {
527 var queue = PDQ.init(testing.allocator);532 var queue = PDQ.init(testing.allocator, {});
528 defer queue.deinit();533 defer queue.deinit();
529534
530 try queue.add(1);535 try queue.add(1);
...@@ -543,7 +548,7 @@ test "std.PriorityDequeue: add and remove same min" {...@@ -543,7 +548,7 @@ test "std.PriorityDequeue: add and remove same min" {
543}548}
544549
545test "std.PriorityDequeue: add and remove same max" {550test "std.PriorityDequeue: add and remove same max" {
546 var queue = PDQ.init(testing.allocator);551 var queue = PDQ.init(testing.allocator, {});
547 defer queue.deinit();552 defer queue.deinit();
548553
549 try queue.add(1);554 try queue.add(1);
...@@ -562,7 +567,7 @@ test "std.PriorityDequeue: add and remove same max" {...@@ -562,7 +567,7 @@ test "std.PriorityDequeue: add and remove same max" {
562}567}
563568
564test "std.PriorityDequeue: removeOrNull empty" {569test "std.PriorityDequeue: removeOrNull empty" {
565 var queue = PDQ.init(testing.allocator);570 var queue = PDQ.init(testing.allocator, {});
566 defer queue.deinit();571 defer queue.deinit();
567572
568 try expect(queue.removeMinOrNull() == null);573 try expect(queue.removeMinOrNull() == null);
...@@ -570,7 +575,7 @@ test "std.PriorityDequeue: removeOrNull empty" {...@@ -570,7 +575,7 @@ test "std.PriorityDequeue: removeOrNull empty" {
570}575}
571576
572test "std.PriorityDequeue: edge case 3 elements" {577test "std.PriorityDequeue: edge case 3 elements" {
573 var queue = PDQ.init(testing.allocator);578 var queue = PDQ.init(testing.allocator, {});
574 defer queue.deinit();579 defer queue.deinit();
575580
576 try queue.add(9);581 try queue.add(9);
...@@ -583,7 +588,7 @@ test "std.PriorityDequeue: edge case 3 elements" {...@@ -583,7 +588,7 @@ test "std.PriorityDequeue: edge case 3 elements" {
583}588}
584589
585test "std.PriorityDequeue: edge case 3 elements max" {590test "std.PriorityDequeue: edge case 3 elements max" {
586 var queue = PDQ.init(testing.allocator);591 var queue = PDQ.init(testing.allocator, {});
587 defer queue.deinit();592 defer queue.deinit();
588593
589 try queue.add(9);594 try queue.add(9);
...@@ -596,7 +601,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {...@@ -596,7 +601,7 @@ test "std.PriorityDequeue: edge case 3 elements max" {
596}601}
597602
598test "std.PriorityDequeue: peekMin" {603test "std.PriorityDequeue: peekMin" {
599 var queue = PDQ.init(testing.allocator);604 var queue = PDQ.init(testing.allocator, {});
600 defer queue.deinit();605 defer queue.deinit();
601606
602 try expect(queue.peekMin() == null);607 try expect(queue.peekMin() == null);
...@@ -610,7 +615,7 @@ test "std.PriorityDequeue: peekMin" {...@@ -610,7 +615,7 @@ test "std.PriorityDequeue: peekMin" {
610}615}
611616
612test "std.PriorityDequeue: peekMax" {617test "std.PriorityDequeue: peekMax" {
613 var queue = PDQ.init(testing.allocator);618 var queue = PDQ.init(testing.allocator, {});
614 defer queue.deinit();619 defer queue.deinit();
615620
616 try expect(queue.peekMin() == null);621 try expect(queue.peekMin() == null);
...@@ -624,7 +629,7 @@ test "std.PriorityDequeue: peekMax" {...@@ -624,7 +629,7 @@ test "std.PriorityDequeue: peekMax" {
624}629}
625630
626test "std.PriorityDequeue: sift up with odd indices" {631test "std.PriorityDequeue: sift up with odd indices" {
627 var queue = PDQ.init(testing.allocator);632 var queue = PDQ.init(testing.allocator, {});
628 defer queue.deinit();633 defer queue.deinit();
629 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };634 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
630 for (items) |e| {635 for (items) |e| {
...@@ -638,7 +643,7 @@ test "std.PriorityDequeue: sift up with odd indices" {...@@ -638,7 +643,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
638}643}
639644
640test "std.PriorityDequeue: sift up with odd indices" {645test "std.PriorityDequeue: sift up with odd indices" {
641 var queue = PDQ.init(testing.allocator);646 var queue = PDQ.init(testing.allocator, {});
642 defer queue.deinit();647 defer queue.deinit();
643 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };648 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
644 for (items) |e| {649 for (items) |e| {
...@@ -652,7 +657,7 @@ test "std.PriorityDequeue: sift up with odd indices" {...@@ -652,7 +657,7 @@ test "std.PriorityDequeue: sift up with odd indices" {
652}657}
653658
654test "std.PriorityDequeue: addSlice min" {659test "std.PriorityDequeue: addSlice min" {
655 var queue = PDQ.init(testing.allocator);660 var queue = PDQ.init(testing.allocator, {});
656 defer queue.deinit();661 defer queue.deinit();
657 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };662 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
658 try queue.addSlice(items[0..]);663 try queue.addSlice(items[0..]);
...@@ -664,7 +669,7 @@ test "std.PriorityDequeue: addSlice min" {...@@ -664,7 +669,7 @@ test "std.PriorityDequeue: addSlice min" {
664}669}
665670
666test "std.PriorityDequeue: addSlice max" {671test "std.PriorityDequeue: addSlice max" {
667 var queue = PDQ.init(testing.allocator);672 var queue = PDQ.init(testing.allocator, {});
668 defer queue.deinit();673 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 };674 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..]);675 try queue.addSlice(items[0..]);
...@@ -678,7 +683,7 @@ test "std.PriorityDequeue: addSlice max" {...@@ -678,7 +683,7 @@ test "std.PriorityDequeue: addSlice max" {
678test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {683test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
679 const items = [0]u32{};684 const items = [0]u32{};
680 const queue_items = try testing.allocator.dupe(u32, &items);685 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..], {});
682 defer queue.deinit();687 defer queue.deinit();
683 try expectEqual(@as(usize, 0), queue.len);688 try expectEqual(@as(usize, 0), queue.len);
684 try expect(queue.removeMinOrNull() == null);689 try expect(queue.removeMinOrNull() == null);
...@@ -687,7 +692,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {...@@ -687,7 +692,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 0" {
687test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {692test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
688 const items = [1]u32{1};693 const items = [1]u32{1};
689 const queue_items = try testing.allocator.dupe(u32, &items);694 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..], {});
691 defer queue.deinit();696 defer queue.deinit();
692697
693 try expectEqual(@as(usize, 1), queue.len);698 try expectEqual(@as(usize, 1), queue.len);
...@@ -698,7 +703,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {...@@ -698,7 +703,7 @@ test "std.PriorityDequeue: fromOwnedSlice trivial case 1" {
698test "std.PriorityDequeue: fromOwnedSlice" {703test "std.PriorityDequeue: fromOwnedSlice" {
699 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };704 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
700 const queue_items = try testing.allocator.dupe(u32, items[0..]);705 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..], {});
702 defer queue.deinit();707 defer queue.deinit();
703708
704 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };709 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" {...@@ -708,7 +713,7 @@ test "std.PriorityDequeue: fromOwnedSlice" {
708}713}
709714
710test "std.PriorityDequeue: update min queue" {715test "std.PriorityDequeue: update min queue" {
711 var queue = PDQ.init(testing.allocator);716 var queue = PDQ.init(testing.allocator, {});
712 defer queue.deinit();717 defer queue.deinit();
713718
714 try queue.add(55);719 try queue.add(55);
...@@ -723,7 +728,7 @@ test "std.PriorityDequeue: update min queue" {...@@ -723,7 +728,7 @@ test "std.PriorityDequeue: update min queue" {
723}728}
724729
725test "std.PriorityDequeue: update same min queue" {730test "std.PriorityDequeue: update same min queue" {
726 var queue = PDQ.init(testing.allocator);731 var queue = PDQ.init(testing.allocator, {});
727 defer queue.deinit();732 defer queue.deinit();
728733
729 try queue.add(1);734 try queue.add(1);
...@@ -739,7 +744,7 @@ test "std.PriorityDequeue: update same min queue" {...@@ -739,7 +744,7 @@ test "std.PriorityDequeue: update same min queue" {
739}744}
740745
741test "std.PriorityDequeue: update max queue" {746test "std.PriorityDequeue: update max queue" {
742 var queue = PDQ.init(testing.allocator);747 var queue = PDQ.init(testing.allocator, {});
743 defer queue.deinit();748 defer queue.deinit();
744749
745 try queue.add(55);750 try queue.add(55);
...@@ -755,7 +760,7 @@ test "std.PriorityDequeue: update max queue" {...@@ -755,7 +760,7 @@ test "std.PriorityDequeue: update max queue" {
755}760}
756761
757test "std.PriorityDequeue: update same max queue" {762test "std.PriorityDequeue: update same max queue" {
758 var queue = PDQ.init(testing.allocator);763 var queue = PDQ.init(testing.allocator, {});
759 defer queue.deinit();764 defer queue.deinit();
760765
761 try queue.add(1);766 try queue.add(1);
...@@ -771,7 +776,7 @@ test "std.PriorityDequeue: update same max queue" {...@@ -771,7 +776,7 @@ test "std.PriorityDequeue: update same max queue" {
771}776}
772777
773test "std.PriorityDequeue: iterator" {778test "std.PriorityDequeue: iterator" {
774 var queue = PDQ.init(testing.allocator);779 var queue = PDQ.init(testing.allocator, {});
775 var map = std.AutoHashMap(u32, void).init(testing.allocator);780 var map = std.AutoHashMap(u32, void).init(testing.allocator);
776 defer {781 defer {
777 queue.deinit();782 queue.deinit();
...@@ -793,7 +798,7 @@ test "std.PriorityDequeue: iterator" {...@@ -793,7 +798,7 @@ test "std.PriorityDequeue: iterator" {
793}798}
794799
795test "std.PriorityDequeue: remove at index" {800test "std.PriorityDequeue: remove at index" {
796 var queue = PDQ.init(testing.allocator);801 var queue = PDQ.init(testing.allocator, {});
797 defer queue.deinit();802 defer queue.deinit();
798803
799 try queue.add(3);804 try queue.add(3);
...@@ -816,7 +821,7 @@ test "std.PriorityDequeue: remove at index" {...@@ -816,7 +821,7 @@ test "std.PriorityDequeue: remove at index" {
816}821}
817822
818test "std.PriorityDequeue: iterator while empty" {823test "std.PriorityDequeue: iterator while empty" {
819 var queue = PDQ.init(testing.allocator);824 var queue = PDQ.init(testing.allocator, {});
820 defer queue.deinit();825 defer queue.deinit();
821826
822 var it = queue.iterator();827 var it = queue.iterator();
...@@ -825,7 +830,7 @@ test "std.PriorityDequeue: iterator while empty" {...@@ -825,7 +830,7 @@ test "std.PriorityDequeue: iterator while empty" {
825}830}
826831
827test "std.PriorityDequeue: shrinkAndFree" {832test "std.PriorityDequeue: shrinkAndFree" {
828 var queue = PDQ.init(testing.allocator);833 var queue = PDQ.init(testing.allocator, {});
829 defer queue.deinit();834 defer queue.deinit();
830835
831 try queue.ensureTotalCapacity(4);836 try queue.ensureTotalCapacity(4);
...@@ -864,7 +869,7 @@ fn fuzzTestMin(rng: std.rand.Random, comptime queue_size: usize) !void {...@@ -864,7 +869,7 @@ fn fuzzTestMin(rng: std.rand.Random, comptime queue_size: usize) !void {
864 const allocator = testing.allocator;869 const allocator = testing.allocator;
865 const items = try generateRandomSlice(allocator, rng, queue_size);870 const items = try generateRandomSlice(allocator, rng, queue_size);
866871
867 var queue = PDQ.fromOwnedSlice(allocator, items);872 var queue = PDQ.fromOwnedSlice(allocator, items, {});
868 defer queue.deinit();873 defer queue.deinit();
869874
870 var last_removed: ?u32 = null;875 var last_removed: ?u32 = null;
...@@ -893,7 +898,7 @@ fn fuzzTestMax(rng: std.rand.Random, queue_size: usize) !void {...@@ -893,7 +898,7 @@ fn fuzzTestMax(rng: std.rand.Random, queue_size: usize) !void {
893 const allocator = testing.allocator;898 const allocator = testing.allocator;
894 const items = try generateRandomSlice(allocator, rng, queue_size);899 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, {});
897 defer queue.deinit();902 defer queue.deinit();
898903
899 var last_removed: ?u32 = null;904 var last_removed: ?u32 = null;
...@@ -922,7 +927,7 @@ fn fuzzTestMinMax(rng: std.rand.Random, queue_size: usize) !void {...@@ -922,7 +927,7 @@ fn fuzzTestMinMax(rng: std.rand.Random, queue_size: usize) !void {
922 const allocator = testing.allocator;927 const allocator = testing.allocator;
923 const items = try generateRandomSlice(allocator, rng, queue_size);928 const items = try generateRandomSlice(allocator, rng, queue_size);
924929
925 var queue = PDQ.fromOwnedSlice(allocator, items);930 var queue = PDQ.fromOwnedSlice(allocator, items, {});
926 defer queue.deinit();931 defer queue.deinit();
927932
928 var last_min: ?u32 = null;933 var last_min: ?u32 = null;
...@@ -957,3 +962,31 @@ fn generateRandomSlice(allocator: std.mem.Allocator, rng: std.rand.Random, size:...@@ -957,3 +962,31 @@ fn generateRandomSlice(allocator: std.mem.Allocator, rng: std.rand.Random, size:
957962
958 return array.toOwnedSlice();963 return array.toOwnedSlice();
959}964}
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;...@@ -8,26 +8,28 @@ const expectEqual = testing.expectEqual;
8const expectError = testing.expectError;8const expectError = testing.expectError;
99
10/// Priority queue for storing generic data. Initialize with `init`.10/// Priority queue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its first11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get popped before its second argument,12/// argument should get popped before its third argument,
13/// `Order.eq` if the arguments are of equal priority, or `Order.gt`13/// `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.
15/// For example, to make `pop` return the smallest number, provide15/// For example, to make `pop` return the smallest number, provide
16/// `fn lessThan(a: T, b: T) Order { return std.math.order(a, b); }`16/// `fn lessThan(context: void, a: T, b: T) Order { _ = context; return std.math.order(a, b); }`
17pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order) type {17pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareFn: fn (context: Context, a: T, b: T) Order) type {
18 return struct {18 return struct {
19 const Self = @This();19 const Self = @This();
2020
21 items: []T,21 items: []T,
22 len: usize,22 len: usize,
23 allocator: Allocator,23 allocator: Allocator,
24 context: Context,
2425
25 /// Initialize and return a priority queue.26 /// Initialize and return a priority queue.
26 pub fn init(allocator: Allocator) Self {27 pub fn init(allocator: Allocator, context: Context) Self {
27 return Self{28 return Self{
28 .items = &[_]T{},29 .items = &[_]T{},
29 .len = 0,30 .len = 0,
30 .allocator = allocator,31 .allocator = allocator,
32 .context = context,
31 };33 };
32 }34 }
3335
...@@ -55,7 +57,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order...@@ -55,7 +57,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
55 const child = self.items[child_index];57 const child = self.items[child_index];
56 const parent = self.items[parent_index];58 const parent = self.items[parent_index];
5759
58 if (compareFn(child, parent) != .lt) break;60 if (compareFn(self.context, child, parent) != .lt) break;
5961
60 self.items[parent_index] = child;62 self.items[parent_index] = child;
61 self.items[child_index] = parent;63 self.items[child_index] = parent;
...@@ -127,14 +129,14 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order...@@ -127,14 +129,14 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
127 var smallest = self.items[index];129 var smallest = self.items[index];
128130
129 if (left) |e| {131 if (left) |e| {
130 if (compareFn(e, smallest) == .lt) {132 if (compareFn(self.context, e, smallest) == .lt) {
131 smallest_index = left_index;133 smallest_index = left_index;
132 smallest = e;134 smallest = e;
133 }135 }
134 }136 }
135137
136 if (right) |e| {138 if (right) |e| {
137 if (compareFn(e, smallest) == .lt) {139 if (compareFn(self.context, e, smallest) == .lt) {
138 smallest_index = right_index;140 smallest_index = right_index;
139 smallest = e;141 smallest = e;
140 }142 }
...@@ -153,11 +155,12 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order...@@ -153,11 +155,12 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
153 /// PriorityQueue takes ownership of the passed in slice. The slice must have been155 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
154 /// allocated with `allocator`.156 /// allocated with `allocator`.
155 /// Deinitialize with `deinit`.157 /// Deinitialize with `deinit`.
156 pub fn fromOwnedSlice(allocator: Allocator, items: []T) Self {158 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
157 var queue = Self{159 var queue = Self{
158 .items = items,160 .items = items,
159 .len = items.len,161 .len = items.len,
160 .allocator = allocator,162 .allocator = allocator,
163 .context = context,
161 };164 };
162165
163 if (queue.len <= 1) return queue;166 if (queue.len <= 1) return queue;
...@@ -207,7 +210,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order...@@ -207,7 +210,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
207 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;210 var update_index: usize = std.mem.indexOfScalar(T, self.items[0..self.len], elem) orelse return error.ElementNotFound;
208 const old_elem: T = self.items[update_index];211 const old_elem: T = self.items[update_index];
209 self.items[update_index] = new_elem;212 self.items[update_index] = new_elem;
210 switch (compareFn(new_elem, old_elem)) {213 switch (compareFn(self.context, new_elem, old_elem)) {
211 .lt => siftUp(self, update_index),214 .lt => siftUp(self, update_index),
212 .gt => siftDown(self, update_index),215 .gt => siftDown(self, update_index),
213 .eq => {}, // Nothing to do as the items have equal priority216 .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...@@ -215,7 +218,7 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
215 }218 }
216219
217 pub const Iterator = struct {220 pub const Iterator = struct {
218 queue: *PriorityQueue(T, compareFn),221 queue: *PriorityQueue(T, Context, compareFn),
219 count: usize,222 count: usize,
220223
221 pub fn next(it: *Iterator) ?T {224 pub fn next(it: *Iterator) ?T {
...@@ -258,19 +261,20 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order...@@ -258,19 +261,20 @@ pub fn PriorityQueue(comptime T: type, comptime compareFn: fn (a: T, b: T) Order
258 };261 };
259}262}
260263
261fn lessThan(a: u32, b: u32) Order {264fn lessThan(context: void, a: u32, b: u32) Order {
265 _ = context;
262 return std.math.order(a, b);266 return std.math.order(a, b);
263}267}
264268
265fn greaterThan(a: u32, b: u32) Order {269fn greaterThan(context: void, a: u32, b: u32) Order {
266 return lessThan(a, b).invert();270 return lessThan(context, a, b).invert();
267}271}
268272
269const PQlt = PriorityQueue(u32, lessThan);273const PQlt = PriorityQueue(u32, void, lessThan);
270const PQgt = PriorityQueue(u32, greaterThan);274const PQgt = PriorityQueue(u32, void, greaterThan);
271275
272test "std.PriorityQueue: add and remove min heap" {276test "std.PriorityQueue: add and remove min heap" {
273 var queue = PQlt.init(testing.allocator);277 var queue = PQlt.init(testing.allocator, {});
274 defer queue.deinit();278 defer queue.deinit();
275279
276 try queue.add(54);280 try queue.add(54);
...@@ -288,7 +292,7 @@ test "std.PriorityQueue: add and remove min heap" {...@@ -288,7 +292,7 @@ test "std.PriorityQueue: add and remove min heap" {
288}292}
289293
290test "std.PriorityQueue: add and remove same min heap" {294test "std.PriorityQueue: add and remove same min heap" {
291 var queue = PQlt.init(testing.allocator);295 var queue = PQlt.init(testing.allocator, {});
292 defer queue.deinit();296 defer queue.deinit();
293297
294 try queue.add(1);298 try queue.add(1);
...@@ -306,14 +310,14 @@ test "std.PriorityQueue: add and remove same min heap" {...@@ -306,14 +310,14 @@ test "std.PriorityQueue: add and remove same min heap" {
306}310}
307311
308test "std.PriorityQueue: removeOrNull on empty" {312test "std.PriorityQueue: removeOrNull on empty" {
309 var queue = PQlt.init(testing.allocator);313 var queue = PQlt.init(testing.allocator, {});
310 defer queue.deinit();314 defer queue.deinit();
311315
312 try expect(queue.removeOrNull() == null);316 try expect(queue.removeOrNull() == null);
313}317}
314318
315test "std.PriorityQueue: edge case 3 elements" {319test "std.PriorityQueue: edge case 3 elements" {
316 var queue = PQlt.init(testing.allocator);320 var queue = PQlt.init(testing.allocator, {});
317 defer queue.deinit();321 defer queue.deinit();
318322
319 try queue.add(9);323 try queue.add(9);
...@@ -325,7 +329,7 @@ test "std.PriorityQueue: edge case 3 elements" {...@@ -325,7 +329,7 @@ test "std.PriorityQueue: edge case 3 elements" {
325}329}
326330
327test "std.PriorityQueue: peek" {331test "std.PriorityQueue: peek" {
328 var queue = PQlt.init(testing.allocator);332 var queue = PQlt.init(testing.allocator, {});
329 defer queue.deinit();333 defer queue.deinit();
330334
331 try expect(queue.peek() == null);335 try expect(queue.peek() == null);
...@@ -337,7 +341,7 @@ test "std.PriorityQueue: peek" {...@@ -337,7 +341,7 @@ test "std.PriorityQueue: peek" {
337}341}
338342
339test "std.PriorityQueue: sift up with odd indices" {343test "std.PriorityQueue: sift up with odd indices" {
340 var queue = PQlt.init(testing.allocator);344 var queue = PQlt.init(testing.allocator, {});
341 defer queue.deinit();345 defer queue.deinit();
342 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };346 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
343 for (items) |e| {347 for (items) |e| {
...@@ -351,7 +355,7 @@ test "std.PriorityQueue: sift up with odd indices" {...@@ -351,7 +355,7 @@ test "std.PriorityQueue: sift up with odd indices" {
351}355}
352356
353test "std.PriorityQueue: addSlice" {357test "std.PriorityQueue: addSlice" {
354 var queue = PQlt.init(testing.allocator);358 var queue = PQlt.init(testing.allocator, {});
355 defer queue.deinit();359 defer queue.deinit();
356 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };360 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
357 try queue.addSlice(items[0..]);361 try queue.addSlice(items[0..]);
...@@ -365,7 +369,7 @@ test "std.PriorityQueue: addSlice" {...@@ -365,7 +369,7 @@ test "std.PriorityQueue: addSlice" {
365test "std.PriorityQueue: fromOwnedSlice trivial case 0" {369test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
366 const items = [0]u32{};370 const items = [0]u32{};
367 const queue_items = try testing.allocator.dupe(u32, &items);371 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..], {});
369 defer queue.deinit();373 defer queue.deinit();
370 try expectEqual(@as(usize, 0), queue.len);374 try expectEqual(@as(usize, 0), queue.len);
371 try expect(queue.removeOrNull() == null);375 try expect(queue.removeOrNull() == null);
...@@ -374,7 +378,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {...@@ -374,7 +378,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 0" {
374test "std.PriorityQueue: fromOwnedSlice trivial case 1" {378test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
375 const items = [1]u32{1};379 const items = [1]u32{1};
376 const queue_items = try testing.allocator.dupe(u32, &items);380 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..], {});
378 defer queue.deinit();382 defer queue.deinit();
379383
380 try expectEqual(@as(usize, 1), queue.len);384 try expectEqual(@as(usize, 1), queue.len);
...@@ -385,7 +389,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {...@@ -385,7 +389,7 @@ test "std.PriorityQueue: fromOwnedSlice trivial case 1" {
385test "std.PriorityQueue: fromOwnedSlice" {389test "std.PriorityQueue: fromOwnedSlice" {
386 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };390 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
387 const heap_items = try testing.allocator.dupe(u32, items[0..]);391 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..], {});
389 defer queue.deinit();393 defer queue.deinit();
390394
391 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };395 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" {...@@ -395,7 +399,7 @@ test "std.PriorityQueue: fromOwnedSlice" {
395}399}
396400
397test "std.PriorityQueue: add and remove max heap" {401test "std.PriorityQueue: add and remove max heap" {
398 var queue = PQgt.init(testing.allocator);402 var queue = PQgt.init(testing.allocator, {});
399 defer queue.deinit();403 defer queue.deinit();
400404
401 try queue.add(54);405 try queue.add(54);
...@@ -413,7 +417,7 @@ test "std.PriorityQueue: add and remove max heap" {...@@ -413,7 +417,7 @@ test "std.PriorityQueue: add and remove max heap" {
413}417}
414418
415test "std.PriorityQueue: add and remove same max heap" {419test "std.PriorityQueue: add and remove same max heap" {
416 var queue = PQgt.init(testing.allocator);420 var queue = PQgt.init(testing.allocator, {});
417 defer queue.deinit();421 defer queue.deinit();
418422
419 try queue.add(1);423 try queue.add(1);
...@@ -431,7 +435,7 @@ test "std.PriorityQueue: add and remove same max heap" {...@@ -431,7 +435,7 @@ test "std.PriorityQueue: add and remove same max heap" {
431}435}
432436
433test "std.PriorityQueue: iterator" {437test "std.PriorityQueue: iterator" {
434 var queue = PQlt.init(testing.allocator);438 var queue = PQlt.init(testing.allocator, {});
435 var map = std.AutoHashMap(u32, void).init(testing.allocator);439 var map = std.AutoHashMap(u32, void).init(testing.allocator);
436 defer {440 defer {
437 queue.deinit();441 queue.deinit();
...@@ -453,7 +457,7 @@ test "std.PriorityQueue: iterator" {...@@ -453,7 +457,7 @@ test "std.PriorityQueue: iterator" {
453}457}
454458
455test "std.PriorityQueue: remove at index" {459test "std.PriorityQueue: remove at index" {
456 var queue = PQlt.init(testing.allocator);460 var queue = PQlt.init(testing.allocator, {});
457 defer queue.deinit();461 defer queue.deinit();
458462
459 try queue.add(3);463 try queue.add(3);
...@@ -476,7 +480,7 @@ test "std.PriorityQueue: remove at index" {...@@ -476,7 +480,7 @@ test "std.PriorityQueue: remove at index" {
476}480}
477481
478test "std.PriorityQueue: iterator while empty" {482test "std.PriorityQueue: iterator while empty" {
479 var queue = PQlt.init(testing.allocator);483 var queue = PQlt.init(testing.allocator, {});
480 defer queue.deinit();484 defer queue.deinit();
481485
482 var it = queue.iterator();486 var it = queue.iterator();
...@@ -485,7 +489,7 @@ test "std.PriorityQueue: iterator while empty" {...@@ -485,7 +489,7 @@ test "std.PriorityQueue: iterator while empty" {
485}489}
486490
487test "std.PriorityQueue: shrinkAndFree" {491test "std.PriorityQueue: shrinkAndFree" {
488 var queue = PQlt.init(testing.allocator);492 var queue = PQlt.init(testing.allocator, {});
489 defer queue.deinit();493 defer queue.deinit();
490494
491 try queue.ensureTotalCapacity(4);495 try queue.ensureTotalCapacity(4);
...@@ -508,7 +512,7 @@ test "std.PriorityQueue: shrinkAndFree" {...@@ -508,7 +512,7 @@ test "std.PriorityQueue: shrinkAndFree" {
508}512}
509513
510test "std.PriorityQueue: update min heap" {514test "std.PriorityQueue: update min heap" {
511 var queue = PQlt.init(testing.allocator);515 var queue = PQlt.init(testing.allocator, {});
512 defer queue.deinit();516 defer queue.deinit();
513517
514 try queue.add(55);518 try queue.add(55);
...@@ -523,7 +527,7 @@ test "std.PriorityQueue: update min heap" {...@@ -523,7 +527,7 @@ test "std.PriorityQueue: update min heap" {
523}527}
524528
525test "std.PriorityQueue: update same min heap" {529test "std.PriorityQueue: update same min heap" {
526 var queue = PQlt.init(testing.allocator);530 var queue = PQlt.init(testing.allocator, {});
527 defer queue.deinit();531 defer queue.deinit();
528532
529 try queue.add(1);533 try queue.add(1);
...@@ -539,7 +543,7 @@ test "std.PriorityQueue: update same min heap" {...@@ -539,7 +543,7 @@ test "std.PriorityQueue: update same min heap" {
539}543}
540544
541test "std.PriorityQueue: update max heap" {545test "std.PriorityQueue: update max heap" {
542 var queue = PQgt.init(testing.allocator);546 var queue = PQgt.init(testing.allocator, {});
543 defer queue.deinit();547 defer queue.deinit();
544548
545 try queue.add(55);549 try queue.add(55);
...@@ -554,7 +558,7 @@ test "std.PriorityQueue: update max heap" {...@@ -554,7 +558,7 @@ test "std.PriorityQueue: update max heap" {
554}558}
555559
556test "std.PriorityQueue: update same max heap" {560test "std.PriorityQueue: update same max heap" {
557 var queue = PQgt.init(testing.allocator);561 var queue = PQgt.init(testing.allocator, {});
558 defer queue.deinit();562 defer queue.deinit();
559563
560 try queue.add(1);564 try queue.add(1);
...@@ -568,3 +572,31 @@ test "std.PriorityQueue: update same max heap" {...@@ -568,3 +572,31 @@ test "std.PriorityQueue: update same max heap" {
568 try expectEqual(@as(u32, 2), queue.remove());572 try expectEqual(@as(u32, 2), queue.remove());
569 try expectEqual(@as(u32, 1), queue.remove());573 try expectEqual(@as(u32, 1), queue.remove());
570}574}
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}