Merge pull request #17383 from squeek502/gpa-optim-treap

GeneralPurposeAllocator: Considerably improve worst case performance

1 files changed, 52 insertions, 1 deletions

diff --git a/lib/std/treap.zig b/lib/std/treap.zig
index 1c4541f96f..383dc1802a 100644
--- a/lib/std/treap.zig
+++ b/lib/std/treap.zig
@@ -225,7 +225,6 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
             link.* = null;
 
             // clean up after ourselves
-            node.key = undefined;
             node.priority = 0;
             node.parent = null;
             node.children = [_]?*Node{ null, null };
@@ -257,6 +256,48 @@ pub fn Treap(comptime Key: type, comptime compareFn: anytype) type {
             assert(link.* == node);
             link.* = target;
         }
+
+        pub const InorderIterator = struct {
+            current: ?*Node,
+            previous: ?*Node = null,
+
+            pub fn next(it: *InorderIterator) ?*Node {
+                while (true) {
+                    if (it.current) |current| {
+                        const previous = it.previous;
+                        it.previous = current;
+                        if (previous == current.parent) {
+                            if (current.children[0]) |left_child| {
+                                it.current = left_child;
+                            } else {
+                                if (current.children[1]) |right_child| {
+                                    it.current = right_child;
+                                } else {
+                                    it.current = current.parent;
+                                }
+                                return current;
+                            }
+                        } else if (previous == current.children[0]) {
+                            if (current.children[1]) |right_child| {
+                                it.current = right_child;
+                            } else {
+                                it.current = current.parent;
+                            }
+                            return current;
+                        } else {
+                            std.debug.assert(previous == current.children[1]);
+                            it.current = current.parent;
+                        }
+                    } else {
+                        return null;
+                    }
+                }
+            }
+        };
+
+        pub fn inorderIterator(self: *Self) InorderIterator {
+            return .{ .current = self.root };
+        }
     };
 }
 
@@ -344,6 +385,16 @@ test "std.Treap: insert, find, replace, remove" {
         try testing.expectEqual(entry.node, treap.getEntryForExisting(node).node);
     }
 
+    // in-order iterator check
+    {
+        var it = treap.inorderIterator();
+        var last_key: u64 = 0;
+        while (it.next()) |node| {
+            try std.testing.expect(node.key >= last_key);
+            last_key = node.key;
+        }
+    }
+
     // replace check
     iter.reset();
     while (iter.next()) |node| {