rb: add sort() that re-sorts tree with new compare function

You can also specify the same compare function, but after updating the
context that the function will use (connected to the rb.Tree) before.

1 files changed, 23 insertions, 0 deletions

diff --git a/lib/std/rb.zig b/lib/std/rb.zig
index 0fefa033a9..78e1e3aedd 100644
--- a/lib/std/rb.zig
+++ b/lib/std/rb.zig
@@ -134,6 +134,20 @@ pub const Tree = struct {
     root: ?*Node,
     compareFn: fn (*Node, *Node, *Tree) Order,
 
+    /// Re-sorts a tree with a new compare function
+    pub fn sort(tree: *Tree, newCompareFn: fn (*Node, *Node, *Tree) Order) !void {
+        var newTree = Tree.init(newCompareFn);
+        var node: *Node = undefined;
+        while (true) {
+            node = tree.first() orelse break;
+            tree.remove(node);
+            if (newTree.insert(node) != null) {
+                return error.NotUnique; // EEXISTS
+            }
+        }
+        tree.* = newTree;
+    }
+
     /// If you have a need for a version that caches this, please file a bug.
     pub fn first(tree: *Tree) ?*Node {
         var node: *Node = tree.root orelse return null;
@@ -244,6 +258,7 @@ pub const Tree = struct {
         return doLookup(key, tree, &parent, &is_left);
     }
 
+    /// If node is not part of tree, behavior is undefined.
     pub fn remove(tree: *Tree, nodeconst: *Node) void {
         var node = nodeconst;
         // as this has the same value as node, it is unsafe to access node after newnode
@@ -514,6 +529,10 @@ fn testCompare(l: *Node, r: *Node, contextIgnored: *Tree) Order {
     unreachable;
 }
 
+fn testCompareReverse(l: *Node, r: *Node, contextIgnored: *Tree) Order {
+    return testCompare(r, l, contextIgnored);
+}
+
 test "rb" {
     if (@import("builtin").arch == .aarch64) {
         // TODO https://github.com/ziglang/zig/issues/3288
@@ -600,4 +619,8 @@ test "multiple inserts, followed by calling first and last" {
     var lookupNode: testNumber = undefined;
     lookupNode.value = 3;
     assert(tree.lookup(&lookupNode.node) == &third.node);
+    tree.sort(testCompareReverse) catch unreachable;
+    assert(testGetNumber(tree.first().?).value == 3);
+    assert(testGetNumber(tree.last().?).value == 0);
+    assert(tree.lookup(&lookupNode.node) == &third.node);
 }