1const std = @import("std.zig");
2const Allocator = std.mem.Allocator;
3const assert = std.debug.assert;
4const Order = std.math.Order;
5const testing = std.testing;
6const expect = testing.expect;
7const expectEqual = testing.expectEqual;
8const expectError = testing.expectError;
9
10/// Priority Dequeue for storing generic data. Initialize with `init`.
11/// Provide `compareFn` that returns `Order.lt` when its second
12/// argument should get min-popped before its third argument,
13/// `Order.eq` if the arguments are of equal priority, or `Order.gt`
14/// if the third argument should be min-popped second.
15/// Popping the max element works in reverse. For example,
16/// to make `popMin` return the smallest number, provide
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 {
19 return struct {
20 const Self = @This();
21
22 items: []T,
23 len: usize,
24 context: Context,
25
26 /// A priority dequeue containing no elements.
27 pub const empty: Self = .{
28 .items = &.{},
29 .len = 0,
30 .context = undefined,
31 };
32
33 /// Initialize and return a new priority dequeue with context.
34 pub fn initContext(context: Context) Self {
35 return Self{
36 .items = &.{},
37 .len = 0,
38 .context = context,
39 };
40 }
41
42 /// Free memory used by the dequeue.
43 pub fn deinit(self: *Self, allocator: Allocator) void {
44 allocator.free(self.items);
45 self.* = undefined;
46 }
47
48 /// Insert a new element, maintaining priority.
49 pub fn push(self: *Self, allocator: Allocator, elem: T) !void {
50 try self.ensureUnusedCapacity(allocator, 1);
51 pushUnchecked(self, elem);
52 }
53
54 /// Add each element in `items` to the dequeue.
55 pub fn pushSlice(self: *Self, allocator: Allocator, items: []const T) !void {
56 try self.ensureUnusedCapacity(allocator, items.len);
57 for (items) |e| {
58 self.pushUnchecked(e);
59 }
60 }
61
62 fn pushUnchecked(self: *Self, elem: T) void {
63 self.items[self.len] = elem;
64
65 if (self.len > 0) {
66 const start = self.getStartForSiftUp(elem, self.len);
67 self.siftUp(start);
68 }
69
70 self.len += 1;
71 }
72
73 fn isMinLayer(index: usize) bool {
74 // In the min-max heap structure:
75 // The first element is on a min layer;
76 // next two are on a max layer;
77 // next four are on a min layer, and so on.
78 return 1 == @clz(index +% 1) & 1;
79 }
80
81 fn nextIsMinLayer(self: *const Self) bool {
82 return isMinLayer(self.len);
83 }
84
85 const StartIndexAndLayer = struct {
86 index: usize,
87 min_layer: bool,
88 };
89
90 fn getStartForSiftUp(self: *const Self, child: T, index: usize) StartIndexAndLayer {
91 const child_index = index;
92 const parent_index = parentIndex(child_index);
93 const parent = self.items[parent_index];
94
95 const min_layer = self.nextIsMinLayer();
96 const order = compareFn(self.context, child, parent);
97 if ((min_layer and order == .gt) or (!min_layer and order == .lt)) {
98 // We must swap the item with it's parent if it is on the "wrong" layer
99 self.items[parent_index] = child;
100 self.items[child_index] = parent;
101 return .{
102 .index = parent_index,
103 .min_layer = !min_layer,
104 };
105 } else {
106 return .{
107 .index = child_index,
108 .min_layer = min_layer,
109 };
110 }
111 }
112
113 fn siftUp(self: *Self, start: StartIndexAndLayer) void {
114 if (start.min_layer) {
115 doSiftUp(self, start.index, .lt);
116 } else {
117 doSiftUp(self, start.index, .gt);
118 }
119 }
120
121 fn doSiftUp(self: *Self, start_index: usize, target_order: Order) void {
122 var child_index = start_index;
123 while (child_index > 2) {
124 const grandparent_index = grandparentIndex(child_index);
125 const child = self.items[child_index];
126 const grandparent = self.items[grandparent_index];
127
128 // If the grandparent is already better or equal, we have gone as far as we need to
129 if (compareFn(self.context, child, grandparent) != target_order) break;
130
131 // Otherwise swap the item with it's grandparent
132 self.items[grandparent_index] = child;
133 self.items[child_index] = grandparent;
134 child_index = grandparent_index;
135 }
136 }
137
138 /// Look at the smallest element in the dequeue. Returns
139 /// `null` if empty.
140 pub fn peekMin(self: *const Self) ?T {
141 return if (self.len > 0) self.items[0] else null;
142 }
143
144 /// Look at the largest element in the dequeue. Returns
145 /// `null` if empty.
146 pub fn peekMax(self: *const Self) ?T {
147 if (self.len == 0) return null;
148 if (self.len == 1) return self.items[0];
149 if (self.len == 2) return self.items[1];
150 return self.bestItemAtIndices(1, 2, .gt).item;
151 }
152
153 fn maxIndex(self: *const Self) ?usize {
154 if (self.len == 0) return null;
155 if (self.len == 1) return 0;
156 if (self.len == 2) return 1;
157 return self.bestItemAtIndices(1, 2, .gt).index;
158 }
159
160 /// Remove and return the smallest element from the dequeue, or `null` if empty
161 pub fn popMin(self: *Self) ?T {
162 return if (self.len > 0) self.popIndex(0) else null;
163 }
164
165 /// Remove and return the largest element from the dequeue, or `null` if empty
166 pub fn popMax(self: *Self) ?T {
167 return if (self.len > 0) self.popIndex(self.maxIndex().?) else null;
168 }
169
170 /// Remove and return element at index. Indices are in the
171 /// same order as iterator, which is not necessarily priority
172 /// order.
173 pub fn popIndex(self: *Self, index: usize) T {
174 assert(self.len > index);
175 const item = self.items[index];
176 const last = self.items[self.len - 1];
177
178 self.items[index] = last;
179 self.len -= 1;
180 siftDown(self, index);
181
182 return item;
183 }
184
185 fn siftDown(self: *Self, index: usize) void {
186 if (isMinLayer(index)) {
187 self.doSiftDown(index, .lt);
188 } else {
189 self.doSiftDown(index, .gt);
190 }
191 }
192
193 fn doSiftDown(self: *Self, start_index: usize, target_order: Order) void {
194 var index = start_index;
195 const half = self.len >> 1;
196 while (true) {
197 const first_grandchild_index = firstGrandchildIndex(index);
198 const last_grandchild_index = first_grandchild_index + 3;
199
200 const elem = self.items[index];
201
202 if (last_grandchild_index < self.len) {
203 // All four grandchildren exist
204 const index2 = first_grandchild_index + 1;
205 const index3 = index2 + 1;
206
207 // Find the best grandchild
208 const best_left = self.bestItemAtIndices(first_grandchild_index, index2, target_order);
209 const best_right = self.bestItemAtIndices(index3, last_grandchild_index, target_order);
210 const best_grandchild = self.bestItem(best_left, best_right, target_order);
211
212 // If the item is better than or equal to its best grandchild, we are done
213 if (compareFn(self.context, best_grandchild.item, elem) != target_order) return;
214
215 // Otherwise, swap them
216 self.items[best_grandchild.index] = elem;
217 self.items[index] = best_grandchild.item;
218 index = best_grandchild.index;
219
220 // We might need to swap the element with it's parent
221 self.swapIfParentIsBetter(elem, index, target_order);
222 } else {
223 // The children or grandchildren are the last layer
224 const first_child_index = firstChildIndex(index);
225 if (first_child_index >= self.len) return;
226
227 const best_descendent = self.bestDescendent(first_child_index, first_grandchild_index, target_order);
228
229 // If the item is better than or equal to its best descendant, we are done
230 if (compareFn(self.context, best_descendent.item, elem) != target_order) return;
231
232 // Otherwise swap them
233 self.items[best_descendent.index] = elem;
234 self.items[index] = best_descendent.item;
235 index = best_descendent.index;
236
237 // If we didn't swap a grandchild, we are done
238 if (index < first_grandchild_index) return;
239
240 // We might need to swap the element with it's parent
241 self.swapIfParentIsBetter(elem, index, target_order);
242 return;
243 }
244
245 // If we are now in the last layer, we are done
246 if (index >= half) return;
247 }
248 }
249
250 fn swapIfParentIsBetter(self: *Self, child: T, child_index: usize, target_order: Order) void {
251 const parent_index = parentIndex(child_index);
252 const parent = self.items[parent_index];
253
254 if (compareFn(self.context, parent, child) == target_order) {
255 self.items[parent_index] = child;
256 self.items[child_index] = parent;
257 }
258 }
259
260 const ItemAndIndex = struct {
261 item: T,
262 index: usize,
263 };
264
265 fn getItem(self: *const Self, index: usize) ItemAndIndex {
266 return .{
267 .item = self.items[index],
268 .index = index,
269 };
270 }
271
272 fn bestItem(self: *const Self, item1: ItemAndIndex, item2: ItemAndIndex, target_order: Order) ItemAndIndex {
273 if (compareFn(self.context, item1.item, item2.item) == target_order) {
274 return item1;
275 } else {
276 return item2;
277 }
278 }
279
280 fn bestItemAtIndices(self: *const Self, index1: usize, index2: usize, target_order: Order) ItemAndIndex {
281 const item1 = self.getItem(index1);
282 const item2 = self.getItem(index2);
283 return self.bestItem(item1, item2, target_order);
284 }
285
286 fn bestDescendent(self: *const Self, first_child_index: usize, first_grandchild_index: usize, target_order: Order) ItemAndIndex {
287 const second_child_index = first_child_index + 1;
288 if (first_grandchild_index >= self.len) {
289 // No grandchildren, find the best child (second may not exist)
290 if (second_child_index >= self.len) {
291 return .{
292 .item = self.items[first_child_index],
293 .index = first_child_index,
294 };
295 } else {
296 return self.bestItemAtIndices(first_child_index, second_child_index, target_order);
297 }
298 }
299
300 const second_grandchild_index = first_grandchild_index + 1;
301 if (second_grandchild_index >= self.len) {
302 // One grandchild, so we know there is a second child. Compare first grandchild and second child
303 return self.bestItemAtIndices(first_grandchild_index, second_child_index, target_order);
304 }
305
306 const best_left_grandchild_index = self.bestItemAtIndices(first_grandchild_index, second_grandchild_index, target_order).index;
307 const third_grandchild_index = second_grandchild_index + 1;
308 if (third_grandchild_index >= self.len) {
309 // Two grandchildren, and we know the best. Compare this to second child.
310 return self.bestItemAtIndices(best_left_grandchild_index, second_child_index, target_order);
311 } else {
312 // Three grandchildren, compare the min of the first two with the third
313 return self.bestItemAtIndices(best_left_grandchild_index, third_grandchild_index, target_order);
314 }
315 }
316
317 /// Return the number of elements remaining in the dequeue
318 pub fn count(self: *const Self) usize {
319 return self.len;
320 }
321
322 /// Return the number of elements that can be added to the
323 /// dequeue before more memory is allocated.
324 pub fn capacity(self: *const Self) usize {
325 return self.items.len;
326 }
327
328 /// Dequeue takes ownership of the passed in slice. The slice must be de-initialize
329 /// with `deinit`.
330 pub fn fromOwnedSlice(items: []T, context: Context) Self {
331 var queue = Self{
332 .items = items,
333 .len = items.len,
334 .context = context,
335 };
336
337 if (queue.len <= 1) return queue;
338
339 const half = (queue.len >> 1) - 1;
340 var i: usize = 0;
341 while (i <= half) : (i += 1) {
342 const index = half - i;
343 queue.siftDown(index);
344 }
345 return queue;
346 }
347
348 /// Ensure that the dequeue can fit at least `new_capacity` items.
349 pub fn ensureTotalCapacity(self: *Self, allocator: Allocator, new_capacity: usize) !void {
350 var better_capacity = self.capacity();
351 if (better_capacity >= new_capacity) return;
352 while (true) {
353 better_capacity += better_capacity / 2 + 8;
354 if (better_capacity >= new_capacity) break;
355 }
356 self.items = try allocator.realloc(self.items, better_capacity);
357 }
358
359 /// Ensure that the dequeue can fit at least `additional_count` **more** items.
360 pub fn ensureUnusedCapacity(self: *Self, allocator: Allocator, additional_count: usize) !void {
361 return self.ensureTotalCapacity(allocator, self.len + additional_count);
362 }
363
364 /// Reduce allocated capacity to `new_len`.
365 pub fn shrinkAndFree(self: *Self, allocator: Allocator, new_len: usize) void {
366 assert(new_len <= self.items.len);
367
368 // Cannot shrink to smaller than the current queue size without invalidating the heap property
369 assert(new_len >= self.len);
370
371 self.items = allocator.realloc(self.items[0..], new_len) catch |e| switch (e) {
372 error.OutOfMemory => { // no problem, capacity is still correct then.
373 self.items.len = new_len;
374 return;
375 },
376 };
377 }
378
379 pub fn update(self: *Self, elem: T, new_elem: T) !void {
380 const old_index = blk: {
381 var idx: usize = 0;
382 while (idx < self.len) : (idx += 1) {
383 const item = self.items[idx];
384 if (compareFn(self.context, item, elem) == .eq) break :blk idx;
385 }
386 return error.ElementNotFound;
387 };
388 _ = self.popIndex(old_index);
389 self.pushUnchecked(new_elem);
390 }
391
392 pub const Iterator = struct {
393 queue: *PriorityDequeue(T, Context, compareFn),
394 count: usize,
395
396 pub fn next(it: *Iterator) ?T {
397 if (it.count >= it.queue.len) return null;
398 const out = it.count;
399 it.count += 1;
400 return it.queue.items[out];
401 }
402
403 pub fn reset(it: *Iterator) void {
404 it.count = 0;
405 }
406 };
407
408 /// Return an iterator that walks the queue without consuming
409 /// it. The iteration order may differ from the priority order.
410 /// Invalidated if the queue is modified.
411 pub fn iterator(self: *Self) Iterator {
412 return Iterator{
413 .queue = self,
414 .count = 0,
415 };
416 }
417
418 fn dump(self: *Self) void {
419 const print = std.debug.print;
420 print("{{ ", .{});
421 print("items: ", .{});
422 for (self.items, 0..) |e, i| {
423 if (i >= self.len) break;
424 print("{}, ", .{e});
425 }
426 print("array: ", .{});
427 for (self.items) |e| {
428 print("{}, ", .{e});
429 }
430 print("len: {} ", .{self.len});
431 print("capacity: {}", .{self.capacity()});
432 print(" }}\n", .{});
433 }
434
435 fn parentIndex(index: usize) usize {
436 return (index - 1) >> 1;
437 }
438
439 fn grandparentIndex(index: usize) usize {
440 return parentIndex(parentIndex(index));
441 }
442
443 fn firstChildIndex(index: usize) usize {
444 return (index << 1) + 1;
445 }
446
447 fn firstGrandchildIndex(index: usize) usize {
448 return firstChildIndex(firstChildIndex(index));
449 }
450 };
451}
452
453/// If a min heap is constructed from slice `{5, 8, 2, 9, 7, 1, 4, 4}` using this
454/// method, then the elements will be in order: {1, 2, 4, 4, 5, 7, 8, 9}
455fn lessThanComparison(context: void, a: u32, b: u32) Order {
456 _ = context;
457 return std.math.order(a, b);
458}
459
460/// Elements with lower priority will be removed first
461const MinHeap = PriorityDequeue(u32, void, lessThanComparison);
462
463test "push and pop min in min heap" {
464 const gpa = std.testing.allocator;
465
466 var queue: MinHeap = .empty;
467 defer queue.deinit(gpa);
468
469 try queue.push(gpa, 54);
470 try queue.push(gpa, 12);
471 try queue.push(gpa, 7);
472 try queue.push(gpa, 23);
473 try queue.push(gpa, 25);
474 try queue.push(gpa, 13);
475
476 try expectEqual(@as(u32, 7), queue.popMin());
477 try expectEqual(@as(u32, 12), queue.popMin());
478 try expectEqual(@as(u32, 13), queue.popMin());
479 try expectEqual(@as(u32, 23), queue.popMin());
480 try expectEqual(@as(u32, 25), queue.popMin());
481 try expectEqual(@as(u32, 54), queue.popMin());
482}
483
484test "push and pop min structs" {
485 const gpa = std.testing.allocator;
486
487 const S = struct {
488 size: u32,
489 };
490 var queue = PriorityDequeue(S, void, struct {
491 fn order(context: void, a: S, b: S) Order {
492 _ = context;
493 return std.math.order(a.size, b.size);
494 }
495 }.order).initContext({});
496 defer queue.deinit(gpa);
497
498 try queue.push(gpa, .{ .size = 54 });
499 try queue.push(gpa, .{ .size = 12 });
500 try queue.push(gpa, .{ .size = 7 });
501 try queue.push(gpa, .{ .size = 23 });
502 try queue.push(gpa, .{ .size = 25 });
503 try queue.push(gpa, .{ .size = 13 });
504
505 try expectEqual(@as(u32, 7), queue.popMin().?.size);
506 try expectEqual(@as(u32, 12), queue.popMin().?.size);
507 try expectEqual(@as(u32, 13), queue.popMin().?.size);
508 try expectEqual(@as(u32, 23), queue.popMin().?.size);
509 try expectEqual(@as(u32, 25), queue.popMin().?.size);
510 try expectEqual(@as(u32, 54), queue.popMin().?.size);
511}
512
513test "push and pop max in min heap" {
514 const gpa = std.testing.allocator;
515
516 var queue: MinHeap = .empty;
517 defer queue.deinit(gpa);
518
519 try queue.push(gpa, 54);
520 try queue.push(gpa, 12);
521 try queue.push(gpa, 7);
522 try queue.push(gpa, 23);
523 try queue.push(gpa, 25);
524 try queue.push(gpa, 13);
525
526 try expectEqual(@as(u32, 54), queue.popMax());
527 try expectEqual(@as(u32, 25), queue.popMax());
528 try expectEqual(@as(u32, 23), queue.popMax());
529 try expectEqual(@as(u32, 13), queue.popMax());
530 try expectEqual(@as(u32, 12), queue.popMax());
531 try expectEqual(@as(u32, 7), queue.popMax());
532}
533
534test "push and pop same min in min heap" {
535 const gpa = std.testing.allocator;
536
537 var queue: MinHeap = .empty;
538 defer queue.deinit(gpa);
539
540 try queue.push(gpa, 1);
541 try queue.push(gpa, 1);
542 try queue.push(gpa, 2);
543 try queue.push(gpa, 2);
544 try queue.push(gpa, 1);
545 try queue.push(gpa, 1);
546
547 try expectEqual(@as(u32, 1), queue.popMin());
548 try expectEqual(@as(u32, 1), queue.popMin());
549 try expectEqual(@as(u32, 1), queue.popMin());
550 try expectEqual(@as(u32, 1), queue.popMin());
551 try expectEqual(@as(u32, 2), queue.popMin());
552 try expectEqual(@as(u32, 2), queue.popMin());
553}
554
555test "push and pop same max in min heap" {
556 const gpa = std.testing.allocator;
557
558 var queue: MinHeap = .empty;
559 defer queue.deinit(gpa);
560
561 try queue.push(gpa, 1);
562 try queue.push(gpa, 1);
563 try queue.push(gpa, 2);
564 try queue.push(gpa, 2);
565 try queue.push(gpa, 1);
566 try queue.push(gpa, 1);
567
568 try expectEqual(@as(u32, 2), queue.popMax());
569 try expectEqual(@as(u32, 2), queue.popMax());
570 try expectEqual(@as(u32, 1), queue.popMax());
571 try expectEqual(@as(u32, 1), queue.popMax());
572 try expectEqual(@as(u32, 1), queue.popMax());
573 try expectEqual(@as(u32, 1), queue.popMax());
574}
575
576test "pop empty in min heap" {
577 const gpa = std.testing.allocator;
578
579 var queue: MinHeap = .empty;
580 defer queue.deinit(gpa);
581
582 try expect(queue.popMin() == null);
583 try expect(queue.popMax() == null);
584}
585
586test "edge case 3 elements popMin in min heap" {
587 const gpa = std.testing.allocator;
588
589 var queue: MinHeap = .empty;
590 defer queue.deinit(gpa);
591
592 try queue.push(gpa, 9);
593 try queue.push(gpa, 3);
594 try queue.push(gpa, 2);
595
596 try expectEqual(@as(u32, 2), queue.popMin());
597 try expectEqual(@as(u32, 3), queue.popMin());
598 try expectEqual(@as(u32, 9), queue.popMin());
599}
600
601test "edge case 3 elements popmax in min heap" {
602 const gpa = std.testing.allocator;
603
604 var queue: MinHeap = .empty;
605 defer queue.deinit(gpa);
606
607 try queue.push(gpa, 9);
608 try queue.push(gpa, 3);
609 try queue.push(gpa, 2);
610
611 try expectEqual(@as(u32, 9), queue.popMax());
612 try expectEqual(@as(u32, 3), queue.popMax());
613 try expectEqual(@as(u32, 2), queue.popMax());
614}
615
616test "peekMin in min heap" {
617 const gpa = std.testing.allocator;
618
619 var queue: MinHeap = .empty;
620 defer queue.deinit(gpa);
621
622 try expect(queue.peekMin() == null);
623
624 try queue.push(gpa, 9);
625 try queue.push(gpa, 3);
626 try queue.push(gpa, 2);
627
628 try expect(queue.peekMin().? == 2);
629 try expect(queue.peekMin().? == 2);
630}
631
632test "peekMax in min heap" {
633 const gpa = std.testing.allocator;
634
635 var queue: MinHeap = .empty;
636 defer queue.deinit(gpa);
637
638 try expect(queue.peekMin() == null);
639
640 try queue.push(gpa, 9);
641 try queue.push(gpa, 3);
642 try queue.push(gpa, 2);
643
644 try expect(queue.peekMax().? == 9);
645 try expect(queue.peekMax().? == 9);
646}
647
648test "sift up with odd indices and popMin in min heap" {
649 const gpa = std.testing.allocator;
650
651 var queue: MinHeap = .empty;
652 defer queue.deinit(gpa);
653
654 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
655 for (items) |e| {
656 try queue.push(gpa, e);
657 }
658
659 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
660 for (sorted_items) |e| {
661 try expectEqual(e, queue.popMin());
662 }
663}
664
665test "sift up with odd indices and popMax in min heap" {
666 const gpa = std.testing.allocator;
667
668 var queue: MinHeap = .empty;
669 defer queue.deinit(gpa);
670
671 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
672 for (items) |e| {
673 try queue.push(gpa, e);
674 }
675
676 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
677 for (sorted_items) |e| {
678 try expectEqual(e, queue.popMax());
679 }
680}
681
682test "pushSlice in min heap and popMin" {
683 const gpa = std.testing.allocator;
684
685 var queue: MinHeap = .empty;
686 defer queue.deinit(gpa);
687
688 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
689 try queue.pushSlice(gpa, items[0..]);
690
691 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
692 for (sorted_items) |e| {
693 try expectEqual(e, queue.popMin());
694 }
695}
696
697test "pushSlice in min heap and popMax" {
698 const gpa = std.testing.allocator;
699
700 var queue: MinHeap = .empty;
701 defer queue.deinit(gpa);
702
703 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
704 try queue.pushSlice(gpa, items[0..]);
705
706 const sorted_items = [_]u32{ 25, 24, 24, 22, 21, 16, 15, 15, 14, 13, 12, 11, 7, 7, 6, 5, 2, 1 };
707 for (sorted_items) |e| {
708 try expectEqual(e, queue.popMax());
709 }
710}
711
712test "fromOwnedSlice trivial case 0 min heap" {
713 const gpa = std.testing.allocator;
714
715 const items = [0]u32{};
716 const queue_items = try gpa.dupe(u32, &items);
717
718 var queue: MinHeap = .fromOwnedSlice(queue_items[0..], {});
719 defer queue.deinit(gpa);
720
721 try expectEqual(@as(usize, 0), queue.len);
722 try expect(queue.popMin() == null);
723}
724
725test "fromOwnedSlice trivial case 1 min heap" {
726 const gpa = std.testing.allocator;
727
728 const items = [1]u32{1};
729 const queue_items = try gpa.dupe(u32, &items);
730
731 var queue: MinHeap = .fromOwnedSlice(queue_items[0..], {});
732 defer queue.deinit(gpa);
733
734 try expectEqual(@as(usize, 1), queue.len);
735 try expectEqual(items[0], queue.popMin());
736 try expect(queue.popMin() == null);
737}
738
739test "fromOwnedSlice min heap" {
740 const gpa = std.testing.allocator;
741
742 const items = [_]u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
743 const queue_items = try gpa.dupe(u32, items[0..]);
744
745 var queue: MinHeap = .fromOwnedSlice(queue_items[0..], {});
746 defer queue.deinit(gpa);
747
748 const sorted_items = [_]u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
749 for (sorted_items) |e| {
750 try expectEqual(e, queue.popMin());
751 }
752}
753
754test "update and popMin in min heap" {
755 const gpa = std.testing.allocator;
756
757 var queue: MinHeap = .empty;
758 defer queue.deinit(gpa);
759
760 try queue.push(gpa, 55);
761 try queue.push(gpa, 44);
762 try queue.push(gpa, 11);
763 try queue.update(55, 5);
764 try queue.update(44, 4);
765 try queue.update(11, 1);
766 try expectEqual(@as(u32, 1), queue.popMin());
767 try expectEqual(@as(u32, 4), queue.popMin());
768 try expectEqual(@as(u32, 5), queue.popMin());
769}
770
771test "update same element and popMin in min heap" {
772 const gpa = std.testing.allocator;
773
774 var queue: MinHeap = .empty;
775 defer queue.deinit(gpa);
776
777 try queue.push(gpa, 1);
778 try queue.push(gpa, 1);
779 try queue.push(gpa, 2);
780 try queue.push(gpa, 2);
781 try queue.update(1, 5);
782 try queue.update(2, 4);
783 try expectEqual(@as(u32, 1), queue.popMin());
784 try expectEqual(@as(u32, 2), queue.popMin());
785 try expectEqual(@as(u32, 4), queue.popMin());
786 try expectEqual(@as(u32, 5), queue.popMin());
787}
788
789test "update and popMax in min heap" {
790 const gpa = std.testing.allocator;
791
792 var queue: MinHeap = .empty;
793 defer queue.deinit(gpa);
794
795 try queue.push(gpa, 55);
796 try queue.push(gpa, 44);
797 try queue.push(gpa, 11);
798 try queue.update(55, 5);
799 try queue.update(44, 1);
800 try queue.update(11, 4);
801
802 try expectEqual(@as(u32, 5), queue.popMax());
803 try expectEqual(@as(u32, 4), queue.popMax());
804 try expectEqual(@as(u32, 1), queue.popMax());
805}
806
807test "update same element and popMax in min heap" {
808 const gpa = std.testing.allocator;
809
810 var queue: MinHeap = .empty;
811 defer queue.deinit(gpa);
812
813 try queue.push(gpa, 1);
814 try queue.push(gpa, 1);
815 try queue.push(gpa, 2);
816 try queue.push(gpa, 2);
817 try queue.update(1, 5);
818 try queue.update(2, 4);
819 try expectEqual(@as(u32, 5), queue.popMax());
820 try expectEqual(@as(u32, 4), queue.popMax());
821 try expectEqual(@as(u32, 2), queue.popMax());
822 try expectEqual(@as(u32, 1), queue.popMax());
823}
824
825test "update after pop in min heap" {
826 const gpa = std.testing.allocator;
827
828 var queue: MinHeap = .empty;
829 defer queue.deinit(gpa);
830
831 try queue.push(gpa, 1);
832 try expectEqual(@as(u32, 1), queue.popMin());
833 try expectError(error.ElementNotFound, queue.update(1, 1));
834}
835
836test "min heap iterator" {
837 const gpa = std.testing.allocator;
838
839 var queue: MinHeap = .empty;
840 var map = std.AutoHashMap(u32, void).init(testing.allocator);
841 defer {
842 queue.deinit(gpa);
843 map.deinit();
844 }
845
846 const items = [_]u32{ 54, 12, 7, 23, 25, 13 };
847 for (items) |e| {
848 _ = try queue.push(gpa, e);
849 _ = try map.put(e, {});
850 }
851
852 var it = queue.iterator();
853 while (it.next()) |e| {
854 _ = map.remove(e);
855 }
856
857 try expectEqual(@as(usize, 0), map.count());
858}
859
860test "pop at index in min heap" {
861 const gpa = std.testing.allocator;
862
863 var queue: MinHeap = .empty;
864 defer queue.deinit(gpa);
865
866 try queue.push(gpa, 3);
867 try queue.push(gpa, 2);
868 try queue.push(gpa, 1);
869
870 var it = queue.iterator();
871 var elem = it.next();
872 var idx: usize = 0;
873 const two_idx = while (elem != null) : (elem = it.next()) {
874 if (elem.? == 2)
875 break idx;
876 idx += 1;
877 } else unreachable;
878
879 try expectEqual(queue.popIndex(two_idx), 2);
880 try expectEqual(queue.popMin(), 1);
881 try expectEqual(queue.popMin(), 3);
882 try expectEqual(queue.popMin(), null);
883}
884
885test "min heap iterator while empty" {
886 const gpa = std.testing.allocator;
887
888 var queue: MinHeap = .empty;
889 defer queue.deinit(gpa);
890
891 var it = queue.iterator();
892
893 try expectEqual(it.next(), null);
894}
895
896test "min heap shrinkAndFree" {
897 const gpa = std.testing.allocator;
898
899 var queue: MinHeap = .empty;
900 defer queue.deinit(gpa);
901
902 try queue.ensureTotalCapacity(gpa, 4);
903 try expect(queue.capacity() >= 4);
904
905 try queue.push(gpa, 1);
906 try queue.push(gpa, 2);
907 try queue.push(gpa, 3);
908 try expect(queue.capacity() >= 4);
909 try expectEqual(@as(usize, 3), queue.len);
910
911 queue.shrinkAndFree(gpa, 3);
912 try expectEqual(@as(usize, 3), queue.capacity());
913 try expectEqual(@as(usize, 3), queue.len);
914
915 try expectEqual(@as(u32, 3), queue.popMax());
916 try expectEqual(@as(u32, 2), queue.popMax());
917 try expectEqual(@as(u32, 1), queue.popMax());
918 try expect(queue.popMax() == null);
919}
920
921test "fuzz testing min" {
922 var prng = std.Random.DefaultPrng.init(std.testing.random_seed);
923 const random = prng.random();
924
925 const test_case_count = 100;
926 const queue_size = 1_000;
927
928 var i: usize = 0;
929 while (i < test_case_count) : (i += 1) {
930 try fuzzTestMin(random, queue_size);
931 }
932}
933
934fn fuzzTestMin(rng: std.Random, comptime queue_size: usize) !void {
935 const gpa = std.testing.allocator;
936
937 const items = try generateRandomSlice(gpa, rng, queue_size);
938
939 var queue: MinHeap = .fromOwnedSlice(items, {});
940 defer queue.deinit(gpa);
941
942 var last_removed: ?u32 = null;
943 while (queue.popMin()) |next| {
944 if (last_removed) |last| {
945 try expect(last <= next);
946 }
947 last_removed = next;
948 }
949}
950
951test "fuzz testing max" {
952 var prng = std.Random.DefaultPrng.init(std.testing.random_seed);
953 const random = prng.random();
954
955 const test_case_count = 100;
956 const queue_size = 1_000;
957
958 var i: usize = 0;
959 while (i < test_case_count) : (i += 1) {
960 try fuzzTestMax(random, queue_size);
961 }
962}
963
964fn fuzzTestMax(rng: std.Random, queue_size: usize) !void {
965 const gpa = std.testing.allocator;
966
967 const items = try generateRandomSlice(gpa, rng, queue_size);
968
969 var queue: MinHeap = .fromOwnedSlice(items, {});
970 defer queue.deinit(gpa);
971
972 var last_removed: ?u32 = null;
973 while (queue.popMax()) |next| {
974 if (last_removed) |last| {
975 try expect(last >= next);
976 }
977 last_removed = next;
978 }
979}
980
981test "fuzz testing min and max" {
982 var prng = std.Random.DefaultPrng.init(std.testing.random_seed);
983 const random = prng.random();
984
985 const test_case_count = 100;
986 const queue_size = 1_000;
987
988 var i: usize = 0;
989 while (i < test_case_count) : (i += 1) {
990 try fuzzTestMinMax(random, queue_size);
991 }
992}
993
994fn fuzzTestMinMax(rng: std.Random, queue_size: usize) !void {
995 const gpa = std.testing.allocator;
996
997 const items = try generateRandomSlice(gpa, rng, queue_size);
998
999 var queue: MinHeap = .fromOwnedSlice(items, {});
1000 defer queue.deinit(gpa);
1001
1002 var last_min: ?u32 = null;
1003 var last_max: ?u32 = null;
1004 var i: usize = 0;
1005 while (i < queue_size) : (i += 1) {
1006 if (i % 2 == 0) {
1007 const next = queue.popMin().?;
1008 if (last_min) |last| {
1009 try expect(last <= next);
1010 }
1011 last_min = next;
1012 } else {
1013 const next = queue.popMax().?;
1014 if (last_max) |last| {
1015 try expect(last >= next);
1016 }
1017 last_max = next;
1018 }
1019 }
1020}
1021
1022fn generateRandomSlice(allocator: std.mem.Allocator, rng: std.Random, size: usize) ![]u32 {
1023 var array = std.array_list.Managed(u32).init(allocator);
1024 try array.ensureTotalCapacity(size);
1025
1026 var i: usize = 0;
1027 while (i < size) : (i += 1) {
1028 const elem = rng.int(u32);
1029 try array.append(elem);
1030 }
1031
1032 return array.toOwnedSlice();
1033}
1034
1035fn contextLessThanComparison(context: []const u32, a: usize, b: usize) Order {
1036 return std.math.order(context[a], context[b]);
1037}
1038
1039const MinHeapWithContext = PriorityDequeue(usize, []const u32, contextLessThanComparison);
1040
1041test "push and pop" {
1042 const gpa = std.testing.allocator;
1043
1044 const context = [_]u32{ 5, 3, 4, 2, 2, 8, 0 };
1045
1046 var queue: MinHeapWithContext = .initContext(context[0..]);
1047 defer queue.deinit(gpa);
1048
1049 try queue.push(gpa, 0);
1050 try queue.push(gpa, 1);
1051 try queue.push(gpa, 2);
1052 try queue.push(gpa, 3);
1053 try queue.push(gpa, 4);
1054 try queue.push(gpa, 5);
1055 try queue.push(gpa, 6);
1056 try expectEqual(@as(usize, 6), queue.popMin());
1057 try expectEqual(@as(usize, 5), queue.popMax());
1058 try expectEqual(@as(usize, 3), queue.popMin());
1059 try expectEqual(@as(usize, 0), queue.popMax());
1060 try expectEqual(@as(usize, 4), queue.popMin());
1061 try expectEqual(@as(usize, 2), queue.popMax());
1062 try expectEqual(@as(usize, 1), queue.popMin());
1063}
1064
1065var all_cmps_unique = true;
1066
1067test "don't compare a value to a copy of itself" {
1068 const gpa = std.testing.allocator;
1069
1070 var depq = PriorityDequeue(u32, void, struct {
1071 fn uniqueLessThan(_: void, a: u32, b: u32) Order {
1072 all_cmps_unique = all_cmps_unique and (a != b);
1073 return std.math.order(a, b);
1074 }
1075 }.uniqueLessThan).initContext({});
1076 defer depq.deinit(gpa);
1077
1078 try depq.push(gpa, 1);
1079 try depq.push(gpa, 2);
1080 try depq.push(gpa, 3);
1081 try depq.push(gpa, 4);
1082 try depq.push(gpa, 5);
1083 try depq.push(gpa, 6);
1084
1085 _ = depq.popIndex(2);
1086 try expectEqual(all_cmps_unique, true);
1087}