mirror of
https://github.com/ziglang/zig.git
synced 2024-11-27 15:42:49 +00:00
603 lines
18 KiB
Zig
603 lines
18 KiB
Zig
const std = @import("std.zig");
|
|
const assert = std.debug.assert;
|
|
const testing = std.testing;
|
|
const mem = std.mem; // For mem.Compare
|
|
|
|
const Color = enum(u1) {
|
|
Black,
|
|
Red,
|
|
};
|
|
const Red = Color.Red;
|
|
const Black = Color.Black;
|
|
|
|
const ReplaceError = error{NotEqual};
|
|
|
|
/// Insert this into your struct that you want to add to a red-black tree.
|
|
/// Do not use a pointer. Turn the *rb.Node results of the functions in rb
|
|
/// (after resolving optionals) to your structure using @fieldParentPtr(). Example:
|
|
///
|
|
/// const Number = struct {
|
|
/// node: rb.Node,
|
|
/// value: i32,
|
|
/// };
|
|
/// fn number(node: *rb.Node) Number {
|
|
/// return @fieldParentPtr(Number, "node", node);
|
|
/// }
|
|
pub const Node = struct {
|
|
left: ?*Node,
|
|
right: ?*Node,
|
|
|
|
/// parent | color
|
|
parent_and_color: usize,
|
|
|
|
pub fn next(constnode: *Node) ?*Node {
|
|
var node = constnode;
|
|
|
|
if (node.right) |right| {
|
|
var n = right;
|
|
while (n.left) |left|
|
|
n = left;
|
|
return n;
|
|
}
|
|
|
|
while (true) {
|
|
var parent = node.getParent();
|
|
if (parent) |p| {
|
|
if (node != p.right)
|
|
return p;
|
|
node = p;
|
|
} else
|
|
return null;
|
|
}
|
|
}
|
|
|
|
pub fn prev(constnode: *Node) ?*Node {
|
|
var node = constnode;
|
|
|
|
if (node.left) |left| {
|
|
var n = left;
|
|
while (n.right) |right|
|
|
n = right;
|
|
return n;
|
|
}
|
|
|
|
while (true) {
|
|
var parent = node.getParent();
|
|
if (parent) |p| {
|
|
if (node != p.left)
|
|
return p;
|
|
node = p;
|
|
} else
|
|
return null;
|
|
}
|
|
}
|
|
|
|
pub fn isRoot(node: *Node) bool {
|
|
return node.getParent() == null;
|
|
}
|
|
|
|
fn isRed(node: *Node) bool {
|
|
return node.getColor() == Red;
|
|
}
|
|
|
|
fn isBlack(node: *Node) bool {
|
|
return node.getColor() == Black;
|
|
}
|
|
|
|
fn setParent(node: *Node, parent: ?*Node) void {
|
|
node.parent_and_color = @ptrToInt(parent) | (node.parent_and_color & 1);
|
|
}
|
|
|
|
fn getParent(node: *Node) ?*Node {
|
|
const mask: usize = 1;
|
|
comptime {
|
|
assert(@alignOf(*Node) >= 2);
|
|
}
|
|
const maybe_ptr = node.parent_and_color & ~mask;
|
|
return if (maybe_ptr == 0) null else @intToPtr(*Node, maybe_ptr);
|
|
}
|
|
|
|
fn setColor(node: *Node, color: Color) void {
|
|
const mask: usize = 1;
|
|
node.parent_and_color = (node.parent_and_color & ~mask) | @enumToInt(color);
|
|
}
|
|
|
|
fn getColor(node: *Node) Color {
|
|
return @intToEnum(Color, @intCast(u1, node.parent_and_color & 1));
|
|
}
|
|
|
|
fn setChild(node: *Node, child: ?*Node, is_left: bool) void {
|
|
if (is_left) {
|
|
node.left = child;
|
|
} else {
|
|
node.right = child;
|
|
}
|
|
}
|
|
|
|
fn getFirst(nodeconst: *Node) *Node {
|
|
var node = nodeconst;
|
|
while (node.left) |left| {
|
|
node = left;
|
|
}
|
|
return node;
|
|
}
|
|
|
|
fn getLast(node: *Node) *Node {
|
|
while (node.right) |right| {
|
|
node = right;
|
|
}
|
|
return node;
|
|
}
|
|
};
|
|
|
|
pub const Tree = struct {
|
|
root: ?*Node,
|
|
compareFn: fn (*Node, *Node) mem.Compare,
|
|
|
|
/// 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;
|
|
|
|
while (node.left) |left| {
|
|
node = left;
|
|
}
|
|
|
|
return node;
|
|
}
|
|
|
|
pub fn last(tree: *Tree) ?*Node {
|
|
var node: *Node = tree.root orelse return null;
|
|
|
|
while (node.right) |right| {
|
|
node = right;
|
|
}
|
|
|
|
return node;
|
|
}
|
|
|
|
/// Duplicate keys are not allowed. The item with the same key already in the
|
|
/// tree will be returned, and the item will not be inserted.
|
|
pub fn insert(tree: *Tree, node_const: *Node) ?*Node {
|
|
var node = node_const;
|
|
var maybe_key: ?*Node = undefined;
|
|
var maybe_parent: ?*Node = undefined;
|
|
var is_left: bool = undefined;
|
|
|
|
maybe_key = doLookup(node, tree, &maybe_parent, &is_left);
|
|
if (maybe_key) |key| {
|
|
return key;
|
|
}
|
|
|
|
node.left = null;
|
|
node.right = null;
|
|
node.setColor(Red);
|
|
node.setParent(maybe_parent);
|
|
|
|
if (maybe_parent) |parent| {
|
|
parent.setChild(node, is_left);
|
|
} else {
|
|
tree.root = node;
|
|
}
|
|
|
|
while (node.getParent()) |*parent| {
|
|
if (parent.*.isBlack())
|
|
break;
|
|
// the root is always black
|
|
var grandpa = parent.*.getParent() orelse unreachable;
|
|
|
|
if (parent.* == grandpa.left) {
|
|
var maybe_uncle = grandpa.right;
|
|
|
|
if (maybe_uncle) |uncle| {
|
|
if (uncle.isBlack())
|
|
break;
|
|
|
|
parent.*.setColor(Black);
|
|
uncle.setColor(Black);
|
|
grandpa.setColor(Red);
|
|
node = grandpa;
|
|
} else {
|
|
if (node == parent.*.right) {
|
|
rotateLeft(parent.*, tree);
|
|
node = parent.*;
|
|
parent.* = node.getParent().?; // Just rotated
|
|
}
|
|
parent.*.setColor(Black);
|
|
grandpa.setColor(Red);
|
|
rotateRight(grandpa, tree);
|
|
}
|
|
} else {
|
|
var maybe_uncle = grandpa.left;
|
|
|
|
if (maybe_uncle) |uncle| {
|
|
if (uncle.isBlack())
|
|
break;
|
|
|
|
parent.*.setColor(Black);
|
|
uncle.setColor(Black);
|
|
grandpa.setColor(Red);
|
|
node = grandpa;
|
|
} else {
|
|
if (node == parent.*.left) {
|
|
rotateRight(parent.*, tree);
|
|
node = parent.*;
|
|
parent.* = node.getParent().?; // Just rotated
|
|
}
|
|
parent.*.setColor(Black);
|
|
grandpa.setColor(Red);
|
|
rotateLeft(grandpa, tree);
|
|
}
|
|
}
|
|
}
|
|
// This was an insert, there is at least one node.
|
|
tree.root.?.setColor(Black);
|
|
return null;
|
|
}
|
|
|
|
/// lookup searches for the value of key, using binary search. It will
|
|
/// return a pointer to the node if it is there, otherwise it will return null.
|
|
/// Complexity guaranteed O(log n), where n is the number of nodes book-kept
|
|
/// by tree.
|
|
pub fn lookup(tree: *Tree, key: *Node) ?*Node {
|
|
var parent: ?*Node = undefined;
|
|
var is_left: bool = undefined;
|
|
return doLookup(key, tree, &parent, &is_left);
|
|
}
|
|
|
|
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
|
|
var newnode: ?*Node = nodeconst;
|
|
var maybe_parent: ?*Node = node.getParent();
|
|
var color: Color = undefined;
|
|
var next: *Node = undefined;
|
|
|
|
// This clause is to avoid optionals
|
|
if (node.left == null and node.right == null) {
|
|
if (maybe_parent) |parent| {
|
|
parent.setChild(null, parent.left == node);
|
|
} else
|
|
tree.root = null;
|
|
color = node.getColor();
|
|
newnode = null;
|
|
} else {
|
|
if (node.left == null) {
|
|
next = node.right.?; // Not both null as per above
|
|
} else if (node.right == null) {
|
|
next = node.left.?; // Not both null as per above
|
|
} else
|
|
next = node.right.?.getFirst(); // Just checked for null above
|
|
|
|
if (maybe_parent) |parent| {
|
|
parent.setChild(next, parent.left == node);
|
|
} else
|
|
tree.root = next;
|
|
|
|
if (node.left != null and node.right != null) {
|
|
const left = node.left.?;
|
|
const right = node.right.?;
|
|
|
|
color = next.getColor();
|
|
next.setColor(node.getColor());
|
|
|
|
next.left = left;
|
|
left.setParent(next);
|
|
|
|
if (next != right) {
|
|
var parent = next.getParent().?; // Was traversed via child node (right/left)
|
|
next.setParent(node.getParent());
|
|
|
|
newnode = next.right;
|
|
parent.left = node;
|
|
|
|
next.right = right;
|
|
right.setParent(next);
|
|
} else {
|
|
next.setParent(maybe_parent);
|
|
maybe_parent = next;
|
|
newnode = next.right;
|
|
}
|
|
} else {
|
|
color = node.getColor();
|
|
newnode = next;
|
|
}
|
|
}
|
|
|
|
if (newnode) |n|
|
|
n.setParent(maybe_parent);
|
|
|
|
if (color == Red)
|
|
return;
|
|
if (newnode) |n| {
|
|
n.setColor(Black);
|
|
return;
|
|
}
|
|
|
|
while (node == tree.root) {
|
|
// If not root, there must be parent
|
|
var parent = maybe_parent.?;
|
|
if (node == parent.left) {
|
|
var sibling = parent.right.?; // Same number of black nodes.
|
|
|
|
if (sibling.isRed()) {
|
|
sibling.setColor(Black);
|
|
parent.setColor(Red);
|
|
rotateLeft(parent, tree);
|
|
sibling = parent.right.?; // Just rotated
|
|
}
|
|
if ((if (sibling.left) |n| n.isBlack() else true) and
|
|
(if (sibling.right) |n| n.isBlack() else true))
|
|
{
|
|
sibling.setColor(Red);
|
|
node = parent;
|
|
maybe_parent = parent.getParent();
|
|
continue;
|
|
}
|
|
if (if (sibling.right) |n| n.isBlack() else true) {
|
|
sibling.left.?.setColor(Black); // Same number of black nodes.
|
|
sibling.setColor(Red);
|
|
rotateRight(sibling, tree);
|
|
sibling = parent.right.?; // Just rotated
|
|
}
|
|
sibling.setColor(parent.getColor());
|
|
parent.setColor(Black);
|
|
sibling.right.?.setColor(Black); // Same number of black nodes.
|
|
rotateLeft(parent, tree);
|
|
newnode = tree.root;
|
|
break;
|
|
} else {
|
|
var sibling = parent.left.?; // Same number of black nodes.
|
|
|
|
if (sibling.isRed()) {
|
|
sibling.setColor(Black);
|
|
parent.setColor(Red);
|
|
rotateRight(parent, tree);
|
|
sibling = parent.left.?; // Just rotated
|
|
}
|
|
if ((if (sibling.left) |n| n.isBlack() else true) and
|
|
(if (sibling.right) |n| n.isBlack() else true))
|
|
{
|
|
sibling.setColor(Red);
|
|
node = parent;
|
|
maybe_parent = parent.getParent();
|
|
continue;
|
|
}
|
|
if (if (sibling.left) |n| n.isBlack() else true) {
|
|
sibling.right.?.setColor(Black); // Same number of black nodes
|
|
sibling.setColor(Red);
|
|
rotateLeft(sibling, tree);
|
|
sibling = parent.left.?; // Just rotated
|
|
}
|
|
sibling.setColor(parent.getColor());
|
|
parent.setColor(Black);
|
|
sibling.left.?.setColor(Black); // Same number of black nodes
|
|
rotateRight(parent, tree);
|
|
newnode = tree.root;
|
|
break;
|
|
}
|
|
|
|
if (node.isRed())
|
|
break;
|
|
}
|
|
|
|
if (newnode) |n|
|
|
n.setColor(Black);
|
|
}
|
|
|
|
/// This is a shortcut to avoid removing and re-inserting an item with the same key.
|
|
pub fn replace(tree: *Tree, old: *Node, newconst: *Node) !void {
|
|
var new = newconst;
|
|
|
|
// I assume this can get optimized out if the caller already knows.
|
|
if (tree.compareFn(old, new) != mem.Compare.Equal) return ReplaceError.NotEqual;
|
|
|
|
if (old.getParent()) |parent| {
|
|
parent.setChild(new, parent.left == old);
|
|
} else
|
|
tree.root = new;
|
|
|
|
if (old.left) |left|
|
|
left.setParent(new);
|
|
if (old.right) |right|
|
|
right.setParent(new);
|
|
|
|
new.* = old.*;
|
|
}
|
|
|
|
pub fn init(tree: *Tree, f: fn (*Node, *Node) mem.Compare) void {
|
|
tree.root = null;
|
|
tree.compareFn = f;
|
|
}
|
|
};
|
|
|
|
fn rotateLeft(node: *Node, tree: *Tree) void {
|
|
var p: *Node = node;
|
|
var q: *Node = node.right orelse unreachable;
|
|
var parent: *Node = undefined;
|
|
|
|
if (!p.isRoot()) {
|
|
parent = p.getParent().?;
|
|
if (parent.left == p) {
|
|
parent.left = q;
|
|
} else {
|
|
parent.right = q;
|
|
}
|
|
q.setParent(parent);
|
|
} else {
|
|
tree.root = q;
|
|
q.setParent(null);
|
|
}
|
|
p.setParent(q);
|
|
|
|
p.right = q.left;
|
|
if (p.right) |right| {
|
|
right.setParent(p);
|
|
}
|
|
q.left = p;
|
|
}
|
|
|
|
fn rotateRight(node: *Node, tree: *Tree) void {
|
|
var p: *Node = node;
|
|
var q: *Node = node.left orelse unreachable;
|
|
var parent: *Node = undefined;
|
|
|
|
if (!p.isRoot()) {
|
|
parent = p.getParent().?;
|
|
if (parent.left == p) {
|
|
parent.left = q;
|
|
} else {
|
|
parent.right = q;
|
|
}
|
|
q.setParent(parent);
|
|
} else {
|
|
tree.root = q;
|
|
q.setParent(null);
|
|
}
|
|
p.setParent(q);
|
|
|
|
p.left = q.right;
|
|
if (p.left) |left| {
|
|
left.setParent(p);
|
|
}
|
|
q.right = p;
|
|
}
|
|
|
|
fn doLookup(key: *Node, tree: *Tree, pparent: *?*Node, is_left: *bool) ?*Node {
|
|
var maybe_node: ?*Node = tree.root;
|
|
|
|
pparent.* = null;
|
|
is_left.* = false;
|
|
|
|
while (maybe_node) |node| {
|
|
var res: mem.Compare = tree.compareFn(node, key);
|
|
if (res == mem.Compare.Equal) {
|
|
return node;
|
|
}
|
|
pparent.* = node;
|
|
if (res == mem.Compare.GreaterThan) {
|
|
is_left.* = true;
|
|
maybe_node = node.left;
|
|
} else if (res == mem.Compare.LessThan) {
|
|
is_left.* = false;
|
|
maybe_node = node.right;
|
|
} else {
|
|
unreachable;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
const testNumber = struct {
|
|
node: Node,
|
|
value: usize,
|
|
};
|
|
|
|
fn testGetNumber(node: *Node) *testNumber {
|
|
return @fieldParentPtr(testNumber, "node", node);
|
|
}
|
|
|
|
fn testCompare(l: *Node, r: *Node) mem.Compare {
|
|
var left = testGetNumber(l);
|
|
var right = testGetNumber(r);
|
|
|
|
if (left.value < right.value) {
|
|
return mem.Compare.LessThan;
|
|
} else if (left.value == right.value) {
|
|
return mem.Compare.Equal;
|
|
} else if (left.value > right.value) {
|
|
return mem.Compare.GreaterThan;
|
|
}
|
|
unreachable;
|
|
}
|
|
|
|
test "rb" {
|
|
if (@import("builtin").arch == .aarch64) {
|
|
// TODO https://github.com/ziglang/zig/issues/3288
|
|
return error.SkipZigTest;
|
|
}
|
|
|
|
var tree: Tree = undefined;
|
|
var ns: [10]testNumber = undefined;
|
|
ns[0].value = 42;
|
|
ns[1].value = 41;
|
|
ns[2].value = 40;
|
|
ns[3].value = 39;
|
|
ns[4].value = 38;
|
|
ns[5].value = 39;
|
|
ns[6].value = 3453;
|
|
ns[7].value = 32345;
|
|
ns[8].value = 392345;
|
|
ns[9].value = 4;
|
|
|
|
var dup: testNumber = undefined;
|
|
dup.value = 32345;
|
|
|
|
tree.init(testCompare);
|
|
_ = tree.insert(&ns[1].node);
|
|
_ = tree.insert(&ns[2].node);
|
|
_ = tree.insert(&ns[3].node);
|
|
_ = tree.insert(&ns[4].node);
|
|
_ = tree.insert(&ns[5].node);
|
|
_ = tree.insert(&ns[6].node);
|
|
_ = tree.insert(&ns[7].node);
|
|
_ = tree.insert(&ns[8].node);
|
|
_ = tree.insert(&ns[9].node);
|
|
tree.remove(&ns[3].node);
|
|
testing.expect(tree.insert(&dup.node) == &ns[7].node);
|
|
try tree.replace(&ns[7].node, &dup.node);
|
|
|
|
var num: *testNumber = undefined;
|
|
num = testGetNumber(tree.first().?);
|
|
while (num.node.next() != null) {
|
|
testing.expect(testGetNumber(num.node.next().?).value > num.value);
|
|
num = testGetNumber(num.node.next().?);
|
|
}
|
|
}
|
|
|
|
test "inserting and looking up" {
|
|
var tree: Tree = undefined;
|
|
tree.init(testCompare);
|
|
var number: testNumber = undefined;
|
|
number.value = 1000;
|
|
_ = tree.insert(&number.node);
|
|
var dup: testNumber = undefined;
|
|
//Assert that tuples with identical value fields finds the same pointer
|
|
dup.value = 1000;
|
|
assert(tree.lookup(&dup.node) == &number.node);
|
|
//Assert that tuples with identical values do not clobber when inserted.
|
|
_ = tree.insert(&dup.node);
|
|
assert(tree.lookup(&dup.node) == &number.node);
|
|
assert(tree.lookup(&number.node) != &dup.node);
|
|
assert(testGetNumber(tree.lookup(&dup.node).?).value == testGetNumber(&dup.node).value);
|
|
//Assert that if looking for a non-existing value, return null.
|
|
var non_existing_value: testNumber = undefined;
|
|
non_existing_value.value = 1234;
|
|
assert(tree.lookup(&non_existing_value.node) == null);
|
|
}
|
|
|
|
test "multiple inserts, followed by calling first and last" {
|
|
if (@import("builtin").arch == .aarch64) {
|
|
// TODO https://github.com/ziglang/zig/issues/3288
|
|
return error.SkipZigTest;
|
|
}
|
|
var tree: Tree = undefined;
|
|
tree.init(testCompare);
|
|
var zeroth: testNumber = undefined;
|
|
zeroth.value = 0;
|
|
var first: testNumber = undefined;
|
|
first.value = 1;
|
|
var second: testNumber = undefined;
|
|
second.value = 2;
|
|
var third: testNumber = undefined;
|
|
third.value = 3;
|
|
_ = tree.insert(&zeroth.node);
|
|
_ = tree.insert(&first.node);
|
|
_ = tree.insert(&second.node);
|
|
_ = tree.insert(&third.node);
|
|
assert(testGetNumber(tree.first().?).value == 0);
|
|
assert(testGetNumber(tree.last().?).value == 3);
|
|
var lookupNode: testNumber = undefined;
|
|
lookupNode.value = 3;
|
|
assert(tree.lookup(&lookupNode.node) == &third.node);
|
|
}
|