authorgravatar for shawn@git.icuShawn Landden <shawn@git.icu> 2020-01-19 22:10:21+04:00
committergravatar for shawn@git.icuShawn Landden <shawn@git.icu> 2020-01-19 22:10:21+04:00
log4ab9678b9553d77bc0683ee69a8936524f6a7808
treea85c01160e6ae980c77b6e6a54645679cf21a1e1
parente19008292222d37b3a08a4e6a9d41553ecfcfd7e

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(-)

lib/std/rb.zig+23
...@@ -134,6 +134,20 @@ pub const Tree = struct {...@@ -134,6 +134,20 @@ pub const Tree = struct {
134 root: ?*Node,134 root: ?*Node,
135 compareFn: fn (*Node, *Node, *Tree) Order,135 compareFn: fn (*Node, *Node, *Tree) Order,
136136
137 /// Re-sorts a tree with a new compare function
138 pub fn sort(tree: *Tree, newCompareFn: fn (*Node, *Node, *Tree) Order) !void {
139 var newTree = Tree.init(newCompareFn);
140 var node: *Node = undefined;
141 while (true) {
142 node = tree.first() orelse break;
143 tree.remove(node);
144 if (newTree.insert(node) != null) {
145 return error.NotUnique; // EEXISTS
146 }
147 }
148 tree.* = newTree;
149 }
150
137 /// If you have a need for a version that caches this, please file a bug.151 /// If you have a need for a version that caches this, please file a bug.
138 pub fn first(tree: *Tree) ?*Node {152 pub fn first(tree: *Tree) ?*Node {
139 var node: *Node = tree.root orelse return null;153 var node: *Node = tree.root orelse return null;
...@@ -244,6 +258,7 @@ pub const Tree = struct {...@@ -244,6 +258,7 @@ pub const Tree = struct {
244 return doLookup(key, tree, &parent, &is_left);258 return doLookup(key, tree, &parent, &is_left);
245 }259 }
246260
261 /// If node is not part of tree, behavior is undefined.
247 pub fn remove(tree: *Tree, nodeconst: *Node) void {262 pub fn remove(tree: *Tree, nodeconst: *Node) void {
248 var node = nodeconst;263 var node = nodeconst;
249 // as this has the same value as node, it is unsafe to access node after newnode264 // 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 {...@@ -514,6 +529,10 @@ fn testCompare(l: *Node, r: *Node, contextIgnored: *Tree) Order {
514 unreachable;529 unreachable;
515}530}
516531
532fn testCompareReverse(l: *Node, r: *Node, contextIgnored: *Tree) Order {
533 return testCompare(r, l, contextIgnored);
534}
535
517test "rb" {536test "rb" {
518 if (@import("builtin").arch == .aarch64) {537 if (@import("builtin").arch == .aarch64) {
519 // TODO https://github.com/ziglang/zig/issues/3288538 // TODO https://github.com/ziglang/zig/issues/3288
...@@ -600,4 +619,8 @@ test "multiple inserts, followed by calling first and last" {...@@ -600,4 +619,8 @@ test "multiple inserts, followed by calling first and last" {
600 var lookupNode: testNumber = undefined;619 var lookupNode: testNumber = undefined;
601 lookupNode.value = 3;620 lookupNode.value = 3;
602 assert(tree.lookup(&lookupNode.node) == &third.node);621 assert(tree.lookup(&lookupNode.node) == &third.node);
622 tree.sort(testCompareReverse) catch unreachable;
623 assert(testGetNumber(tree.first().?).value == 3);
624 assert(testGetNumber(tree.last().?).value == 0);
625 assert(tree.lookup(&lookupNode.node) == &third.node);
603}626}