authorgravatar for me@arnavion.devArnavion <me@arnavion.dev> 2021-12-15 17:46:04-08:00
committergravatar for me@arnavion.devArnavion <me@arnavion.dev> 2021-12-15 17:46:04-08:00
log82e893065607cffc1d7ee9d228b8ccdb2bbc61ea
treeb3cbff6ff92f7a881c4d9898605c266f50038dab
parent877a1f2a2986e02269c64a05456dff521da27ac1

std.priority_queue: allow comparator to take a context parameter

1 files changed, 69 insertions(+), 37 deletions(-)

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}