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