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); }