authorgravatar for jeru.sheng@gmail.comCheng Sheng <jeru.sheng@gmail.com> 2024-07-29 04:47:55+02:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-07-28 19:47:55-07:00
log4a77c7f2588ce42c10df53fc1242f684ff506720
treef2f4e9cca8b4ac34939a2490c17029a9db4800d7
parentd1d95294fd657f771657ea671a6984b860347fb0
signaturebadge-check Signed by PGP key B5690EEEBB952194

Condense and extend std.Treap's traversal functionalities. (#20002)

The core functionalities are now in two general functions `extremeInSubtreeOnDirection()` and `nextOnDirection()` so all the other traversing functions (`getMin()`, `getMax()`, and `InorderIterator`) are all just trivial calls to these core functions. The added two functions `Node.next()` and `Node.prev()` are also just trivial calls to these. * std.Treap traversal direction: use u1 instead of usize. * Treap: fix getMin() and getMax(), and add tests for them.

2 files changed, 277 insertions(+), 47 deletions(-)

lib/compiler/resinator/source_mapping.zig+1-4
......@@ -571,10 +571,7 @@ pub const SourceMappings = struct {
571571
572572 // now subtract the span diff from the start line number of all of
573573 // the following nodes in order
574 var it = Sources.InorderIterator{
575 .current = node,
576 .previous = node.children[0],
577 };
574 var it = Sources.InorderIterator{ .current = node };
578575 // skip past current, but store it
579576 var prev = it.next().?;
580577 while (it.next()) |inorder_node| {
lib/std/treap.zig+276-43
......@@ -51,26 +51,53 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
5151 priority: usize,
5252 parent: ?*Node,
5353 children: [2]?*Node,
54
55 pub fn next(node: *Node) ?*Node {
56 return nextOnDirection(node, 1);
57 }
58 pub fn prev(node: *Node) ?*Node {
59 return nextOnDirection(node, 0);
60 }
5461 };
5562
63 fn extremeInSubtreeOnDirection(node: *Node, direction: u1) *Node {
64 var cur = node;
65 while (cur.children[direction]) |next| cur = next;
66 return cur;
67 }
68
69 fn nextOnDirection(node: *Node, direction: u1) ?*Node {
70 if (node.children[direction]) |child| {
71 return extremeInSubtreeOnDirection(child, direction ^ 1);
72 }
73 var cur = node;
74 // Traversing upward until we find `parent` to `cur` is NOT on
75 // `direction`, or equivalently, `cur` to `parent` IS on
76 // `direction` thus `parent` is the next.
77 while (true) {
78 if (cur.parent) |parent| {
79 // If `parent -> node` is NOT on `direction`, then
80 // `node -> parent` IS on `direction`
81 if (parent.children[direction] != cur) return parent;
82 cur = parent;
83 } else {
84 return null;
85 }
86 }
87 }
88
5689 /// Returns the smallest Node by key in the treap if there is one.
5790 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
5891 pub fn getMin(self: Self) ?*Node {
59 var node = self.root;
60 while (node) |current| {
61 node = current.children[0] orelse break;
62 }
63 return node;
92 if (self.root) |root| return extremeInSubtreeOnDirection(root, 0);
93 return null;
6494 }
6595
6696 /// Returns the largest Node by key in the treap if there is one.
6797 /// Use `getEntryForExisting()` to replace/remove this Node from the treap.
6898 pub fn getMax(self: Self) ?*Node {
69 var node = self.root;
70 while (node) |current| {
71 node = current.children[1] orelse break;
72 }
73 return node;
99 if (self.root) |root| return extremeInSubtreeOnDirection(root, 1);
100 return null;
74101 }
75102
76103 /// Lookup the Entry for the given key in the treap.
......@@ -117,7 +144,8 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
117144 removed,
118145 },
119146
120 /// Update's the Node at this Entry in the treap with the new node.
147 /// Update's the Node at this Entry in the treap with the new node (null for deleting). `new_node`
148 /// can have `undefind` content because the value will be initialized internally.
121149 pub fn set(self: *Entry, new_node: ?*Node) void {
122150 // Update the entry's node reference after updating the treap below.
123151 defer self.node = new_node;
......@@ -257,46 +285,26 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
257285 link.* = target;
258286 }
259287
288 /// Usage example:
289 /// var iter = treap.inorderIterator();
290 /// while (iter.next()) |node| {
291 /// ...
292 /// }
260293 pub const InorderIterator = struct {
261294 current: ?*Node,
262 previous: ?*Node = null,
263295
264296 pub fn next(it: *InorderIterator) ?*Node {
265 while (true) {
266 if (it.current) |current| {
267 const previous = it.previous;
268 it.previous = current;
269 if (previous == current.parent) {
270 if (current.children[0]) |left_child| {
271 it.current = left_child;
272 } else {
273 if (current.children[1]) |right_child| {
274 it.current = right_child;
275 } else {
276 it.current = current.parent;
277 }
278 return current;
279 }
280 } else if (previous == current.children[0]) {
281 if (current.children[1]) |right_child| {
282 it.current = right_child;
283 } else {
284 it.current = current.parent;
285 }
286 return current;
287 } else {
288 std.debug.assert(previous == current.children[1]);
289 it.current = current.parent;
290 }
291 } else {
292 return null;
293 }
294 }
297 const current = it.current;
298 it.current = if (current) |cur|
299 cur.next()
300 else
301 null;
302 return current;
295303 }
296304 };
297305
298306 pub fn inorderIterator(self: *Self) InorderIterator {
299 return .{ .current = self.root };
307 return .{ .current = self.getMin() };
300308 }
301309 };
302310}
......@@ -447,3 +455,228 @@ test "insert, find, replace, remove" {
447455 try testing.expectEqual(entry.node, treap.getEntryFor(key).node);
448456 }
449457}
458
459test "inorderIterator" {
460 var treap = TestTreap{};
461 var nodes: [10]TestNode = undefined;
462
463 // Build the tree.
464 var i: usize = 0;
465 while (i < 10) : (i += 1) {
466 const key = @as(u64, i);
467 var entry = treap.getEntryFor(key);
468 entry.set(&nodes[i]);
469 }
470
471 // Test the iterator.
472 var iter = treap.inorderIterator();
473 i = 0;
474 while (iter.next()) |node| {
475 const key = @as(u64, i);
476 try testing.expectEqual(key, node.key);
477 i += 1;
478 }
479}
480
481test "getMin, getMax, simple" {
482 var treap = TestTreap{};
483 var nodes: [3]TestNode = undefined;
484
485 try testing.expectEqual(null, treap.getMin());
486 try testing.expectEqual(null, treap.getMax());
487 { // nodes[1]
488 var entry = treap.getEntryFor(1);
489 entry.set(&nodes[1]);
490 try testing.expectEqual(&nodes[1], treap.getMin());
491 try testing.expectEqual(&nodes[1], treap.getMax());
492 }
493 { // nodes[0]
494 var entry = treap.getEntryFor(0);
495 entry.set(&nodes[0]);
496 try testing.expectEqual(&nodes[0], treap.getMin());
497 try testing.expectEqual(&nodes[1], treap.getMax());
498 }
499 { // nodes[2]
500 var entry = treap.getEntryFor(2);
501 entry.set(&nodes[2]);
502 try testing.expectEqual(&nodes[0], treap.getMin());
503 try testing.expectEqual(&nodes[2], treap.getMax());
504 }
505}
506
507test "getMin, getMax, random" {
508 var nodes: [100]TestNode = undefined;
509 var prng = std.Random.DefaultPrng.init(0xdeadbeef);
510 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
511
512 var treap = TestTreap{};
513 var min: u64 = std.math.maxInt(u64);
514 var max: u64 = 0;
515
516 try testing.expectEqual(null, treap.getMin());
517 try testing.expectEqual(null, treap.getMax());
518
519 // Insert and check min/max after each insertion.
520 iter.reset();
521 while (iter.next()) |node| {
522 const key = prng.random().int(u64);
523
524 // Insert into `treap`.
525 var entry = treap.getEntryFor(key);
526 entry.set(node);
527
528 if (key < min) min = key;
529 if (key > max) max = key;
530
531 const min_node = treap.getMin().?;
532 try std.testing.expectEqual(null, min_node.prev());
533 try std.testing.expectEqual(min, min_node.key);
534
535 const max_node = treap.getMax().?;
536 try std.testing.expectEqual(null, max_node.next());
537 try std.testing.expectEqual(max, max_node.key);
538 }
539}
540
541test "node.{prev(),next()} with sequential insertion and deletion" {
542 // Insert order: 50, 0, 1, 2, ..., 49, 51, 52, ..., 99.
543 // Delete order: 0, 1, 2, ..., 49, 51, 52, ..., 99.
544 // Check 50's neighbors.
545 var treap = TestTreap{};
546 var nodes: [100]TestNode = undefined;
547 {
548 var entry = treap.getEntryFor(50);
549 entry.set(&nodes[50]);
550 try testing.expectEqual(50, nodes[50].key);
551 try testing.expectEqual(null, nodes[50].prev());
552 try testing.expectEqual(null, nodes[50].next());
553 }
554 // Insert others.
555 var i: usize = 0;
556 while (i < 50) : (i += 1) {
557 const key = @as(u64, i);
558 const node = &nodes[i];
559 var entry = treap.getEntryFor(key);
560 entry.set(node);
561 try testing.expectEqual(key, node.key);
562 try testing.expectEqual(node, nodes[50].prev());
563 try testing.expectEqual(null, nodes[50].next());
564 }
565 i = 51;
566 while (i < 100) : (i += 1) {
567 const key = @as(u64, i);
568 const node = &nodes[i];
569 var entry = treap.getEntryFor(key);
570 entry.set(node);
571 try testing.expectEqual(key, node.key);
572 try testing.expectEqual(&nodes[49], nodes[50].prev());
573 try testing.expectEqual(&nodes[51], nodes[50].next());
574 }
575 // Remove others.
576 i = 0;
577 while (i < 49) : (i += 1) {
578 const key = @as(u64, i);
579 var entry = treap.getEntryFor(key);
580 entry.set(null);
581 try testing.expectEqual(&nodes[49], nodes[50].prev());
582 try testing.expectEqual(&nodes[51], nodes[50].next());
583 }
584 { // i = 49.
585 const key = @as(u64, i);
586 var entry = treap.getEntryFor(key);
587 entry.set(null);
588 try testing.expectEqual(null, nodes[50].prev());
589 try testing.expectEqual(&nodes[51], nodes[50].next());
590 }
591 i = 51;
592 while (i < 99) : (i += 1) {
593 const key = @as(u64, i);
594 var entry = treap.getEntryFor(key);
595 entry.set(null);
596 try testing.expectEqual(null, nodes[50].prev());
597 try testing.expectEqual(&nodes[i + 1], nodes[50].next());
598 }
599 { // i = 99.
600 const key = @as(u64, i);
601 var entry = treap.getEntryFor(key);
602 entry.set(null);
603 try testing.expectEqual(null, nodes[50].prev());
604 try testing.expectEqual(null, nodes[50].next());
605 }
606}
607
608fn findFirstGreaterOrEqual(array: []u64, value: u64) usize {
609 var i: usize = 0;
610 while (i < array.len and array[i] < value) i += 1;
611 return i;
612}
613
614fn testOrderedArrayAndTreapConsistency(array: []u64, treap: *TestTreap) !void {
615 var i: usize = 0;
616 while (i < array.len) : (i += 1) {
617 const value = array[i];
618
619 const entry = treap.getEntryFor(value);
620 try testing.expect(entry.node != null);
621 const node = entry.node.?;
622 try testing.expectEqual(value, node.key);
623
624 if (i == 0) {
625 try testing.expectEqual(node.prev(), null);
626 } else {
627 try testing.expectEqual(node.prev(), treap.getEntryFor(array[i - 1]).node);
628 }
629 if (i + 1 == array.len) {
630 try testing.expectEqual(node.next(), null);
631 } else {
632 try testing.expectEqual(node.next(), treap.getEntryFor(array[i + 1]).node);
633 }
634 }
635}
636
637test "node.{prev(),next()} with random data" {
638 var nodes: [100]TestNode = undefined;
639 var prng = std.Random.DefaultPrng.init(0xdeadbeef);
640 var iter = SliceIterRandomOrder(TestNode).init(&nodes, prng.random());
641
642 var treap = TestTreap{};
643 // A slow, stupid but correct reference. Ordered.
644 var golden = std.ArrayList(u64).init(std.testing.allocator);
645 defer golden.deinit();
646
647 // Insert.
648 iter.reset();
649 while (iter.next()) |node| {
650 const key = prng.random().int(u64);
651
652 // Insert into `golden`.
653 const i = findFirstGreaterOrEqual(golden.items, key);
654 // Ensure not found. If found: `prng`'s fault.
655 try testing.expect(i == golden.items.len or golden.items[i] > key);
656 try golden.insert(i, key);
657
658 // Insert into `treap`.
659 var entry = treap.getEntryFor(key);
660 entry.set(node);
661
662 try testOrderedArrayAndTreapConsistency(golden.items, &treap);
663 }
664
665 // Delete.
666 iter.reset();
667 while (iter.next()) |node| {
668 const key = node.key;
669
670 // Delete from `golden`.
671 const i = findFirstGreaterOrEqual(golden.items, key);
672 try testing.expect(i < golden.items.len);
673 _ = golden.orderedRemove(i);
674
675 // Delete from `treap`.
676 var entry = treap.getEntryFor(key);
677 try testing.expect(entry.node != null);
678 entry.set(null);
679
680 try testOrderedArrayAndTreapConsistency(golden.items, &treap);
681 }
682}