diff --git a/src/ds/redblacktree.h b/src/ds/redblacktree.h new file mode 100644 index 0000000..7ea6e9a --- /dev/null +++ b/src/ds/redblacktree.h @@ -0,0 +1,698 @@ +#pragma once +#include "concept.h" +#include "defines.h" + +#include +#include +#include + +namespace snmalloc +{ +#ifdef __cpp_concepts + template + concept RBRepTypes = requires() + { + typename Rep::Holder; + typename Rep::Contents; + }; + + template + concept RBRepMethods = + requires(typename Rep::Holder* hp, typename Rep::Contents k, bool b) + { + { + Rep::get(hp) + } + ->ConceptSame; + { + Rep::set(hp, k) + } + ->ConceptSame; + { + Rep::is_red(k) + } + ->ConceptSame; + { + Rep::set_red(k, b) + } + ->ConceptSame; + { + Rep::ref(b, k) + } + ->ConceptSame; + }; + + template + concept RBRep = // + RBRepTypes // + && RBRepMethods // + && ConceptSame< + decltype(Rep::null), + std::add_const_t>; +#endif + + /** + * Contains a self balancing binary tree. + * + * The template parameter Rep provides the representation of the nodes as a + * collection of functions and types that are requires. See the associated + * test for an example. + * + * run_checks enables invariant checking on the tree. Enabled in Debug. + * TRACE prints all the sets of the rebalancing operations. Only enabled by + * the test when debugging a specific failure. + */ + template< + SNMALLOC_CONCEPT(RBRep) Rep, + bool run_checks = DEBUG, + bool TRACE = false> + class RBTree + { + using H = typename Rep::Holder; + using K = typename Rep::Contents; + + // Container that behaves like a C++ Ref type to enable assignment + // to treat left, right and root uniformly. + class ChildRef + { + H* ptr; + + public: + ChildRef() : ptr(nullptr) {} + + ChildRef(H& p) : ptr(&p) {} + + operator K() + { + return Rep::get(ptr); + } + + K operator=(K t) + { + // Use representations assigment, so we update the correct bits + // color and other things way also be stored in the Holder. + Rep::set(ptr, t); + return t; + } + + bool operator==(ChildRef& t) + { + return ptr == t.ptr; + } + + bool operator!=(ChildRef& t) + { + return ptr != t.ptr; + } + + H* addr() + { + return ptr; + } + }; + + // Root field of the tree + H root{}; + + static ChildRef get_dir(bool direction, K k) + { + return {Rep::ref(direction, k)}; + } + + ChildRef get_root() + { + return {root}; + } + + void invariant() + { + invariant(get_root()); + } + + /* + * Verify structural invariants. Returns the black depth of the `curr`ent + * node. + */ + int invariant(K curr, K lower = Rep::MinKey, K upper = Rep::MaxKey) + { + if constexpr (!run_checks) + { + UNUSED(curr, lower, upper); + return 0; + } + if (curr == Rep::null) + return 1; + + if (curr < lower || curr > upper) + { + if constexpr (TRACE) + { + std::cout << "Invariant failed: " << curr << " is out of bounds " + << lower << ", " << upper << std::endl; + print(); + } + snmalloc::error("Invariant failed"); + } + + if ( + Rep::is_red(curr) && + (Rep::is_red(get_dir(true, curr)) || Rep::is_red(get_dir(false, curr)))) + { + if constexpr (TRACE) + { + std::cout << "Red invariant failed: " << curr + << " is red and has red children" << std::endl; + print(); + } + snmalloc::error("Invariant failed"); + } + + int left_inv = invariant(get_dir(true, curr), lower, curr); + int right_inv = invariant(get_dir(false, curr), curr, upper); + + if (left_inv != right_inv) + { + if constexpr (TRACE) + { + std::cout << "Balance failed: " << curr + << " has different black depths on left and right" + << std::endl; + print(); + } + snmalloc::error("Invariant failed"); + } + + if (Rep::is_red(curr)) + return left_inv; + else + return left_inv + 1; + } + + struct RBStep + { + ChildRef node; + bool dir = false; + }; + + public: + // Internal representation of a path in the tree. + // Exposed to allow for some composite operations to be defined + // externally. + class RBPath + { + friend class RBTree; + + std::array path; + size_t length = 0; + + RBPath(typename Rep::Holder& root) : path{} + { + path[0] = {root, false}; + length = 1; + } + + ChildRef ith(size_t n) + { + SNMALLOC_ASSERT(length >= n); + return path[length - n - 1].node; + } + + bool ith_dir(size_t n) + { + SNMALLOC_ASSERT(length >= n); + return path[length - n - 1].dir; + } + + ChildRef curr() + { + return ith(0); + } + + bool curr_dir() + { + return ith_dir(0); + } + + ChildRef parent() + { + return ith(1); + } + + bool parent_dir() + { + return ith_dir(1); + } + + ChildRef grand_parent() + { + return ith(2); + } + + // Extend path in `direction`. + // If `direction` contains `Rep::null`, do not extend the path. + // Returns false if path is not extended. + bool move(bool direction) + { + auto next = get_dir(direction, curr()); + if (next == Rep::null) + return false; + path[length] = {next, direction}; + length++; + return true; + } + + // Extend path in `direction`. + // If `direction` contains zero, do not extend the path. + // Returns false if path is extended with null. + bool move_inc_null(bool direction) + { + auto next = get_dir(direction, curr()); + path[length] = {next, direction}; + length++; + return next != Rep::null; + } + + // Remove top element from the path. + void pop() + { + SNMALLOC_ASSERT(length > 0); + length--; + } + + // If a path is changed in place, then some references can be stale. + // This rewalks the updated path, and corrects any internal references. + // `expected` is used to run the update, or if `false` used to check + // that no update is required. + void fixup(bool expected = true) + { + if (!run_checks && !expected) + return; + + // During a splice in remove the path can be invalidated, + // this refreshs the path so that the it refers to the spliced + // nodes fields. + // TODO optimise usage to avoid traversing whole path. + for (size_t i = 1; i < length; i++) + { + auto parent = path[i - 1].node; + auto& curr = path[i].node; + auto dir = path[i].dir; + auto actual = get_dir(dir, parent); + if (actual != curr) + { + if (!expected) + { + snmalloc::error("Performed an unexpected fixup."); + } + curr = actual; + } + } + } + + void print() + { + if constexpr (TRACE) + { + for (size_t i = 0; i < length; i++) + { + std::cout << "->" << K(path[i].node) << "@" << path[i].node.addr() + << " (" << path[i].dir << ") "; + } + std::cout << std::endl; + } + } + }; + + private: + void debug_log(const char* msg, RBPath& path) + { + debug_log(msg, path, get_root()); + } + + void debug_log(const char* msg, RBPath& path, ChildRef base) + { + if constexpr (TRACE) + { + std::cout << "-------" << std::endl; + std::cout << msg << std::endl; + path.print(); + print(base); + } + else + { + UNUSED(msg, path, base); + } + } + + public: + RBTree() {} + + void print() + { + print(get_root()); + } + + void print(ChildRef curr, const char* indent = "", size_t depth = 0) + { + if constexpr (TRACE) + { + std::cout << indent << "\\_"; + + if (curr == Rep::null) + { + std::cout << "null" << std::endl; + return; + } + +#ifdef _MSC_VER + auto colour = Rep::is_red(curr) ? "R-" : "B-"; + auto reset = ""; +#else + auto colour = Rep::is_red(curr) ? "\e[1;31m" : "\e[1;34m"; + auto reset = "\e[0m"; +#endif + + std::cout << colour << curr << reset << curr.addr() << " (" << depth + << ")" << std::endl; + if ((get_dir(true, curr) != 0) || (get_dir(false, curr) != 0)) + { + auto s_indent = std::string(indent); + print(get_dir(true, curr), (s_indent + "|").c_str(), depth + 1); + print(get_dir(false, curr), (s_indent + " ").c_str(), depth + 1); + } + } + } + + bool find(RBPath& path, K value) + { + bool dir; + + if (path.curr() == Rep::null) + return false; + + do + { + if (path.curr() == value) + return true; + dir = path.curr() > value; + } while (path.move_inc_null(dir)); + + return false; + } + + bool remove_path(RBPath& path) + { + ChildRef splice = path.curr(); + SNMALLOC_ASSERT(splice != Rep::null); + + debug_log("Removing", path); + + /* + * Find immediately smaller leaf element (rightmost descendant of left + * child) to serve as the replacement for this node. We may not have a + * left subtree, so this may not move the path at all. + */ + path.move(true); + while (path.move(false)) + { + } + + K curr = path.curr(); + + { + // Locally extract right-child-less replacement, replacing it with its + // left child, if any + K child = get_dir(true, path.curr()); + // Unlink target replacing with possible child. + path.curr() = child; + } + + bool leaf_red = Rep::is_red(curr); + + if (path.curr() != splice) + { + // If we had a left child, replace ourselves with the extracted value + // from above + Rep::set_red(curr, Rep::is_red(splice)); + get_dir(true, curr) = K(get_dir(true, splice)); + get_dir(false, curr) = K(get_dir(false, splice)); + splice = curr; + path.fixup(); + } + + debug_log("Splice done", path); + + // Red leaf removal requires no rebalancing. + if (leaf_red) + return true; + + // Now in the double black case. + // End of path is considered double black, that is, one black element + // shorter than satisfies the invariant. The following algorithm moves up + // the path until it finds a close red element or the root. If we convert + // the tree to one, in which the root is double black, then the algorithm + // is complete, as there is nothing to be out of balance with. Otherwise, + // we are searching for nearby red elements so we can rotate the tree to + // rebalance. The following slides nicely cover the case analysis below + // https://www.cs.purdue.edu/homes/ayg/CS251/slides/chap13c.pdf + while (path.curr() != ChildRef(root)) + { + K parent = path.parent(); + bool cur_dir = path.curr_dir(); + K sibling = get_dir(!cur_dir, parent); + + /* Handle red sibling case. + * This performs a rotation to give a black sibling. + * + * p s(b) + * / \ / \ + * c s(r) --> p(r) m + * / \ / \ + * n m c n + * + * By invariant we know that p, n and m are all initially black. + */ + if (Rep::is_red(sibling)) + { + debug_log("Red sibling", path, path.parent()); + K nibling = get_dir(cur_dir, sibling); + get_dir(!cur_dir, parent) = nibling; + get_dir(cur_dir, sibling) = parent; + Rep::set_red(parent, true); + Rep::set_red(sibling, false); + path.parent() = sibling; + // Manually fix path. Using path.fixup would alter the complexity + // class. + path.pop(); + path.move(cur_dir); + path.move_inc_null(cur_dir); + path.fixup(false); + debug_log("Red sibling - done", path, path.parent()); + continue; + } + + /* Handle red nibling case 1. + *

+ * / \ / \ + * c s --> p rn + * / \ / \ + * on rn c on + */ + if (Rep::is_red(get_dir(!cur_dir, sibling))) + { + debug_log("Red nibling 1", path, path.parent()); + K r_nibling = get_dir(!cur_dir, sibling); + K o_nibling = get_dir(cur_dir, sibling); + get_dir(cur_dir, sibling) = parent; + get_dir(!cur_dir, parent) = o_nibling; + path.parent() = sibling; + Rep::set_red(r_nibling, false); + Rep::set_red(sibling, Rep::is_red(parent)); + Rep::set_red(parent, false); + debug_log("Red nibling 1 - done", path, path.parent()); + break; + } + + /* Handle red nibling case 2. + *

+ * / \ / \ + * c s --> p s + * / \ / \ / \ + * rn on c rno rns on + * / \ + * rno rns + */ + if (Rep::is_red(get_dir(cur_dir, sibling))) + { + debug_log("Red nibling 2", path, path.parent()); + K r_nibling = get_dir(cur_dir, sibling); + K r_nibling_same = get_dir(cur_dir, r_nibling); + K r_nibling_opp = get_dir(!cur_dir, r_nibling); + get_dir(!cur_dir, parent) = r_nibling_same; + get_dir(cur_dir, sibling) = r_nibling_opp; + get_dir(cur_dir, r_nibling) = parent; + get_dir(!cur_dir, r_nibling) = sibling; + path.parent() = r_nibling; + Rep::set_red(r_nibling, Rep::is_red(parent)); + Rep::set_red(parent, false); + debug_log("Red nibling 2 - done", path, path.parent()); + break; + } + + // Handle black sibling and niblings, and red parent. + if (Rep::is_red(parent)) + { + // std::cout << "Black sibling and red parent case" << std::endl; + Rep::set_red(parent, false); + Rep::set_red(sibling, true); + break; + } + // Handle black sibling and niblings and black parent. + debug_log( + "Black sibling, niblings and black parent case", path, path.parent()); + Rep::set_red(sibling, true); + path.pop(); + invariant(path.curr()); + debug_log( + "Black sibling, niblings and black parent case - done", + path, + path.curr()); + } + return true; + } + + // Insert an element at the given path. + void insert_path(RBPath path, K value) + { + SNMALLOC_ASSERT(path.curr() == Rep::null); + path.curr() = value; + get_dir(true, path.curr()) = Rep::null; + get_dir(false, path.curr()) = Rep::null; + Rep::set_red(value, true); + + debug_log("Insert ", path); + + // Propogate double red up to rebalance. + // These notes were particularly clear for explaining insert + // https://www.cs.cmu.edu/~fp/courses/15122-f10/lectures/17-rbtrees.pdf + while (path.curr() != get_root()) + { + SNMALLOC_ASSERT(Rep::is_red(path.curr())); + if (!Rep::is_red(path.parent())) + { + invariant(); + return; + } + bool curr_dir = path.curr_dir(); + K curr = path.curr(); + K parent = path.parent(); + K grand_parent = path.grand_parent(); + SNMALLOC_ASSERT(!Rep::is_red(grand_parent)); + if (path.parent_dir() == curr_dir) + { + debug_log("Insert - double red case 1", path, path.grand_parent()); + /* Same direction case + * G - grand parent + * P - parent + * C - current + * S - sibling + * + * G P + * / \ / \ + * A P --> G C + * / \ / \ + * S C A S + */ + K sibling = get_dir(!curr_dir, parent); + Rep::set_red(curr, false); + get_dir(curr_dir, grand_parent) = sibling; + get_dir(!curr_dir, parent) = grand_parent; + path.grand_parent() = parent; + debug_log( + "Insert - double red case 1 - done", path, path.grand_parent()); + } + else + { + debug_log("Insert - double red case 2", path, path.grand_parent()); + /* G - grand parent + * P - parent + * C - current + * Cg - Current child for grand parent + * Cp - Current child for parent + * + * G C + * / \ / \ + * A P G P + * / \ --> / \ / \ + * C B A Cg Cp B + * / \ + * Cg Cp + */ + K child_g = get_dir(curr_dir, curr); + K child_p = get_dir(!curr_dir, curr); + + Rep::set_red(parent, false); + path.grand_parent() = curr; + get_dir(curr_dir, curr) = grand_parent; + get_dir(!curr_dir, curr) = parent; + get_dir(curr_dir, parent) = child_p; + get_dir(!curr_dir, grand_parent) = child_g; + debug_log( + "Insert - double red case 2 - done", path, path.grand_parent()); + } + + // Move to what replaced grand parent. + path.pop(); + path.pop(); + invariant(path.curr()); + } + Rep::set_red(get_root(), false); + invariant(); + } + + K remove_min() + { + if (get_root() == Rep::null) + return Rep::null; + + auto path = get_root_path(); + while (path.move(true)) + { + } + + K result = path.curr(); + + remove_path(path); + return result; + } + + bool remove_elem(K value) + { + if (get_root() == Rep::null) + return false; + + auto path = get_root_path(); + if (!find(path, value)) + return false; + + remove_path(path); + return true; + } + + bool insert_elem(K value) + { + auto path = get_root_path(); + + if (find(path, value)) + return false; + + insert_path(path, value); + return true; + } + + RBPath get_root_path() + { + return RBPath(root); + } + }; +} diff --git a/src/test/func/redblack/redblack.cc b/src/test/func/redblack/redblack.cc new file mode 100644 index 0000000..185d3c1 --- /dev/null +++ b/src/test/func/redblack/redblack.cc @@ -0,0 +1,178 @@ +#include "test/opt.h" +#include "test/setup.h" +#include "test/usage.h" +#include "test/xoroshiro.h" + +#include +#include +#include +#include + +// Redblack tree needs some libraries with trace enabled. +#include "ds/redblacktree.h" +#include "snmalloc.h" + +struct Wrapper +{ + // The redblack tree is going to be used inside the pagemap, + // and the redblack tree cannot use all the bits. Applying an offset + // to the stored value ensures that we have some abstraction over + // the representation. + static constexpr size_t offset = 10000; + + size_t value = offset << 1; +}; + +// Simple representation that is like the pagemap. +// Bottom bit of left is used to store the colour. +// We shift the fields up to make room for the colour. +struct node +{ + Wrapper left; + Wrapper right; +}; + +inline static node array[2048]; + +class Rep +{ +public: + using key = size_t; + + static constexpr key null = 0; + static constexpr key MinKey = 0; + static constexpr key MaxKey = ~MinKey; + + using Holder = Wrapper; + using Contents = size_t; + + static void set(Holder* ptr, Contents r) + { + ptr->value = ((r + Wrapper::offset) << 1) + (ptr->value & 1); + } + + static Contents get(Holder* ptr) + { + return (ptr->value >> 1) - Wrapper::offset; + } + + static Holder& ref(bool direction, key k) + { + if (direction) + return array[k].left; + else + return array[k].right; + } + + static bool is_red(key k) + { + return (array[k].left.value & 1) == 1; + } + + static void set_red(key k, bool new_is_red) + { + if (new_is_red != is_red(k)) + array[k].left.value ^= 1; + } +}; + +template +void test(size_t size, unsigned int seed) +{ + /// Perform a pseudo-random series of + /// additions and removals from the tree. + + xoroshiro::p64r32 rand(seed); + snmalloc::RBTree tree; + std::vector entries; + + bool first = true; + std::cout << "size: " << size << " seed: " << seed << std::endl; + for (size_t i = 0; i < 20 * size; i++) + { + auto batch = 1 + rand.next() % (3 + (size / 2)); + auto op = rand.next() % 4; + if (op < 2 || first) + { + first = false; + for (auto j = batch; j > 0; j--) + { + auto index = 1 + rand.next() % size; + if (tree.insert_elem(index)) + { + entries.push_back(index); + } + } + } + else if (op == 3) + { + for (auto j = batch; j > 0; j--) + { + if (entries.size() == 0) + continue; + auto index = rand.next() % entries.size(); + auto elem = entries[index]; + if (!tree.remove_elem(elem)) + { + std::cout << "Failed to remove element: " << elem << std::endl; + abort(); + } + entries.erase(entries.begin() + static_cast(index)); + } + } + else + { + for (auto j = batch; j > 0; j--) + { + // print(); + auto min = tree.remove_min(); + auto s = entries.size(); + if (min == 0) + break; + + entries.erase( + std::remove(entries.begin(), entries.end(), min), entries.end()); + if (s != entries.size() + 1) + { + std::cout << "Failed to remove min: " << min << std::endl; + abort(); + } + } + } + if (entries.size() == 0) + { + break; + } + } +} + +int main(int argc, char** argv) +{ + setup(); + + opt::Opt opt(argc, argv); + + auto seed = opt.is("--seed", 0); + auto size = opt.is("--size", 0); + + if (seed == 0 && size == 0) + { + for (size = 1; size <= 300; size = size + 1 + (size >> 3)) + for (seed = 1; seed < 5 + (8 * size); seed++) + { + test(size, seed); + } + + return 0; + } + + if (seed == 0 || size == 0) + { + std::cout << "Set both --seed and --size" << std::endl; + return 1; + } + + // Trace particular example + test(size, seed); + return 0; +} \ No newline at end of file