699 lines
18 KiB
C++
699 lines
18 KiB
C++
#pragma once
|
|
#include "concept.h"
|
|
#include "defines.h"
|
|
|
|
#include <array>
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
|
|
namespace snmalloc
|
|
{
|
|
#ifdef __cpp_concepts
|
|
template<typename Rep>
|
|
concept RBRepTypes = requires()
|
|
{
|
|
typename Rep::Holder;
|
|
typename Rep::Contents;
|
|
};
|
|
|
|
template<typename Rep>
|
|
concept RBRepMethods =
|
|
requires(typename Rep::Holder* hp, typename Rep::Contents k, bool b)
|
|
{
|
|
{
|
|
Rep::get(hp)
|
|
}
|
|
->ConceptSame<typename Rep::Contents>;
|
|
{
|
|
Rep::set(hp, k)
|
|
}
|
|
->ConceptSame<void>;
|
|
{
|
|
Rep::is_red(k)
|
|
}
|
|
->ConceptSame<bool>;
|
|
{
|
|
Rep::set_red(k, b)
|
|
}
|
|
->ConceptSame<void>;
|
|
{
|
|
Rep::ref(b, k)
|
|
}
|
|
->ConceptSame<typename Rep::Holder&>;
|
|
};
|
|
|
|
template<typename Rep>
|
|
concept RBRep = //
|
|
RBRepTypes<Rep> //
|
|
&& RBRepMethods<Rep> //
|
|
&& ConceptSame<
|
|
decltype(Rep::null),
|
|
std::add_const_t<typename Rep::Contents>>;
|
|
#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<RBStep, 128> 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.
|
|
* <p> <s>
|
|
* / \ / \
|
|
* 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.
|
|
* <p> <rn>
|
|
* / \ / \
|
|
* 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);
|
|
}
|
|
};
|
|
}
|