authorgravatar for me@arnavion.devArnavion <me@arnavion.dev> 2021-12-15 17:46:10-08:00
committergravatar for me@arnavion.devArnavion <me@arnavion.dev> 2021-12-15 17:46:10-08:00
log22e1d92d3edad13f6a3371cb11931f9f66c910c6
treeb1cfd4d5998516cb79ed783a756316ef613f6aa3
parent82e893065607cffc1d7ee9d228b8ccdb2bbc61ea

std.priority_dequeue: allow comparator to take a context parameter

1 files changed, 82 insertions(+), 49 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}