Enrico Fraccaroli
564 lines
15 KiB
C
564 lines
15 KiB
C
/// @file rbtree.c
|
|
/// @brief Red/Black tree.
|
|
/// @copyright (c) 2014-2021 This file is distributed under the MIT License.
|
|
/// See LICENSE.md for details.
|
|
|
|
#include "klib/rbtree.h"
|
|
|
|
#include "assert.h"
|
|
#include "io/debug.h"
|
|
#include "mem/slab.h"
|
|
|
|
/// @brief Stores information of a node.
|
|
struct rbtree_node_t {
|
|
/// Color red (1), black (0)
|
|
int red;
|
|
/// Link left [0] and right [1]
|
|
rbtree_node_t *link[2];
|
|
/// User provided, used indirectly via rbtree_tree_node_cmp_f.
|
|
void *value;
|
|
};
|
|
|
|
/// @brief Stores information of a rbtree.
|
|
struct rbtree_t {
|
|
/// Root of the tree.
|
|
rbtree_node_t *root;
|
|
/// Comparison function for insertion.
|
|
rbtree_tree_node_cmp_f cmp;
|
|
/// Size of the tree.
|
|
unsigned int size;
|
|
};
|
|
|
|
/// @brief Stores information for iterating a rbtree.
|
|
struct rbtree_iter_t {
|
|
/// Pointer to the tree itself.
|
|
rbtree_t *tree;
|
|
/// Current node
|
|
rbtree_node_t *node;
|
|
/// Traversal path
|
|
rbtree_node_t *path[RBTREE_ITER_MAX_HEIGHT];
|
|
/// Top of stack
|
|
unsigned int top;
|
|
};
|
|
|
|
rbtree_node_t *rbtree_node_alloc()
|
|
{
|
|
return kmalloc(sizeof(rbtree_node_t));
|
|
}
|
|
|
|
rbtree_node_t *rbtree_node_init(rbtree_node_t *node, void *value)
|
|
{
|
|
if (node) {
|
|
node->red = 1;
|
|
node->link[0] = node->link[1] = NULL;
|
|
node->value = value;
|
|
}
|
|
return node;
|
|
}
|
|
|
|
rbtree_node_t *rbtree_node_create(void *value)
|
|
{
|
|
return rbtree_node_init(rbtree_node_alloc(), value);
|
|
}
|
|
|
|
void *rbtree_node_get_value(rbtree_node_t *node)
|
|
{
|
|
if (node)
|
|
return node->value;
|
|
return NULL;
|
|
}
|
|
|
|
void rbtree_node_dealloc(rbtree_node_t *node)
|
|
{
|
|
if (node)
|
|
kfree(node);
|
|
}
|
|
|
|
static int rbtree_node_is_red(const rbtree_node_t *node)
|
|
{
|
|
return node ? node->red : 0;
|
|
}
|
|
|
|
static rbtree_node_t *rbtree_node_rotate(rbtree_node_t *node, int dir)
|
|
{
|
|
rbtree_node_t *result = NULL;
|
|
if (node) {
|
|
result = node->link[!dir];
|
|
node->link[!dir] = result->link[dir];
|
|
result->link[dir] = node;
|
|
node->red = 1;
|
|
result->red = 0;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static rbtree_node_t *rbtree_node_rotate2(rbtree_node_t *node, int dir)
|
|
{
|
|
rbtree_node_t *result = NULL;
|
|
if (node) {
|
|
node->link[!dir] = rbtree_node_rotate(node->link[!dir], !dir);
|
|
result = rbtree_node_rotate(node, dir);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// rbtree_t - default callbacks
|
|
|
|
static int rbtree_tree_node_cmp_ptr_cb(
|
|
rbtree_t *tree,
|
|
rbtree_node_t *a,
|
|
rbtree_node_t *b)
|
|
{
|
|
(void)tree;
|
|
return (a->value > b->value) - (a->value < b->value);
|
|
}
|
|
|
|
static void rbtree_tree_node_dealloc_cb(rbtree_t *tree, rbtree_node_t *node)
|
|
{
|
|
if (tree)
|
|
if (node)
|
|
rbtree_node_dealloc(node);
|
|
}
|
|
|
|
// rbtree_t
|
|
|
|
rbtree_t *rbtree_tree_alloc()
|
|
{
|
|
return kmalloc(sizeof(rbtree_t));
|
|
}
|
|
|
|
rbtree_t *rbtree_tree_init(rbtree_t *tree, rbtree_tree_node_cmp_f node_cmp_cb)
|
|
{
|
|
if (tree) {
|
|
tree->root = NULL;
|
|
tree->size = 0;
|
|
tree->cmp = node_cmp_cb ? node_cmp_cb : rbtree_tree_node_cmp_ptr_cb;
|
|
}
|
|
return tree;
|
|
}
|
|
|
|
rbtree_t *rbtree_tree_create(rbtree_tree_node_cmp_f node_cb)
|
|
{
|
|
return rbtree_tree_init(rbtree_tree_alloc(), node_cb);
|
|
}
|
|
|
|
void rbtree_tree_dealloc(rbtree_t *tree, rbtree_tree_node_f node_cb)
|
|
{
|
|
assert(tree);
|
|
if (node_cb) {
|
|
rbtree_node_t *node = tree->root;
|
|
rbtree_node_t *save = NULL;
|
|
|
|
// Rotate away the left links so that
|
|
// we can treat this like the destruction
|
|
// of a linked list
|
|
while (node) {
|
|
if (node->link[0] == NULL) {
|
|
// No left links, just kill the node and move on
|
|
save = node->link[1];
|
|
node_cb(tree, node);
|
|
kfree(node);
|
|
node = NULL;
|
|
} else {
|
|
// Rotate away the left link and check again
|
|
save = node->link[0];
|
|
node->link[0] = save->link[1];
|
|
save->link[1] = node;
|
|
}
|
|
node = save;
|
|
}
|
|
}
|
|
kfree(tree);
|
|
}
|
|
|
|
void *rbtree_tree_find(rbtree_t *tree, void *value)
|
|
{
|
|
void *result = NULL;
|
|
if (tree) {
|
|
rbtree_node_t node = { .value = value };
|
|
rbtree_node_t *it = tree->root;
|
|
int cmp = 0;
|
|
while (it) {
|
|
if ((cmp = tree->cmp(tree, it, &node))) {
|
|
// If the tree supports duplicates, they should be
|
|
// chained to the right subtree for this to work
|
|
it = it->link[cmp < 0];
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
result = it ? it->value : NULL;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
void *rbtree_tree_find_by_value(rbtree_t *tree,
|
|
rbtree_tree_cmp_f cmp_fun,
|
|
void *value)
|
|
{
|
|
void *result = NULL;
|
|
if (tree) {
|
|
rbtree_node_t *it = tree->root;
|
|
int cmp = 0;
|
|
while (it) {
|
|
if ((cmp = cmp_fun(tree, it, value))) {
|
|
// If the tree supports duplicates, they should be
|
|
// chained to the right subtree for this to work
|
|
it = it->link[cmp < 0];
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
result = it ? it->value : NULL;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// Creates (kmalloc'ates)
|
|
int rbtree_tree_insert(rbtree_t *tree, void *value)
|
|
{
|
|
return rbtree_tree_insert_node(tree, rbtree_node_create(value));
|
|
}
|
|
|
|
// Returns 1 on success, 0 otherwise.
|
|
int rbtree_tree_insert_node(rbtree_t *tree, rbtree_node_t *node)
|
|
{
|
|
if (tree && node) {
|
|
if (tree->root == NULL) {
|
|
tree->root = node;
|
|
} else {
|
|
rbtree_node_t head = { 0 }; // False tree root
|
|
rbtree_node_t *g, *t; // Grandparent & parent
|
|
rbtree_node_t *p, *q; // Iterator & parent
|
|
int dir = 0, last = 0;
|
|
|
|
// Set up our helpers
|
|
t = &head;
|
|
g = p = NULL;
|
|
q = t->link[1] = tree->root;
|
|
|
|
// Search down the tree for a place to insert
|
|
while (1) {
|
|
if (q == NULL) {
|
|
// Insert node at the first null link.
|
|
p->link[dir] = q = node;
|
|
} else if (rbtree_node_is_red(q->link[0]) &&
|
|
rbtree_node_is_red(q->link[1])) {
|
|
// Simple red violation: color flip
|
|
q->red = 1;
|
|
q->link[0]->red = 0;
|
|
q->link[1]->red = 0;
|
|
}
|
|
|
|
if (rbtree_node_is_red(q) && rbtree_node_is_red(p)) {
|
|
// Hard red violation: rotations necessary
|
|
int dir2 = t->link[1] == g;
|
|
if (q == p->link[last]) {
|
|
t->link[dir2] = rbtree_node_rotate(g, !last);
|
|
} else {
|
|
t->link[dir2] = rbtree_node_rotate2(g, !last);
|
|
}
|
|
}
|
|
|
|
// Stop working if we inserted a node. This
|
|
// check also disallows duplicates in the tree
|
|
if (tree->cmp(tree, q, node) == 0) {
|
|
break;
|
|
}
|
|
|
|
last = dir;
|
|
dir = tree->cmp(tree, q, node) < 0;
|
|
|
|
// Move the helpers down
|
|
if (g != NULL) {
|
|
t = g;
|
|
}
|
|
|
|
g = p, p = q;
|
|
q = q->link[dir];
|
|
}
|
|
|
|
// Update the root (it may be different)
|
|
tree->root = head.link[1];
|
|
}
|
|
|
|
// Make the root black for simplified logic
|
|
tree->root->red = 0;
|
|
++tree->size;
|
|
}
|
|
|
|
return 1;
|
|
}
|
|
|
|
// Returns 1 if the value was removed, 0 otherwise. Optional node callback
|
|
// can be provided to dealloc node and/or user data. Use rbtree_tree_node_dealloc
|
|
// default callback to deallocate node created by rbtree_tree_insert(...).
|
|
int rbtree_tree_remove_with_cb(rbtree_t *tree,
|
|
void *value,
|
|
rbtree_tree_node_f node_cb)
|
|
{
|
|
if (tree->root != NULL) {
|
|
rbtree_node_t head = { 0 }; // False tree root
|
|
rbtree_node_t node = { .value = value }; // Value wrapper node
|
|
rbtree_node_t *q, *p, *g; // Helpers
|
|
rbtree_node_t *f = NULL; // Found item
|
|
int dir = 1;
|
|
|
|
// Set up our helpers
|
|
q = &head;
|
|
g = p = NULL;
|
|
q->link[1] = tree->root;
|
|
|
|
// Search and push a red node down
|
|
// to fix red violations as we go
|
|
while (q->link[dir] != NULL) {
|
|
int last = dir;
|
|
|
|
// Move the helpers down
|
|
g = p, p = q;
|
|
q = q->link[dir];
|
|
dir = tree->cmp(tree, q, &node) < 0;
|
|
|
|
// Save the node with matching value and keep
|
|
// going; we'll do removal tasks at the end
|
|
if (tree->cmp(tree, q, &node) == 0) {
|
|
f = q;
|
|
}
|
|
|
|
// Push the red node down with rotations and color flips
|
|
if (!rbtree_node_is_red(q) && !rbtree_node_is_red(q->link[dir])) {
|
|
if (rbtree_node_is_red(q->link[!dir])) {
|
|
p = p->link[last] = rbtree_node_rotate(q, dir);
|
|
} else if (!rbtree_node_is_red(q->link[!dir])) {
|
|
rbtree_node_t *s = p->link[!last];
|
|
if (s) {
|
|
if (!rbtree_node_is_red(s->link[!last]) &&
|
|
!rbtree_node_is_red(s->link[last])) {
|
|
// Color flip
|
|
p->red = 0;
|
|
s->red = 1;
|
|
q->red = 1;
|
|
} else {
|
|
int dir2 = g->link[1] == p;
|
|
if (rbtree_node_is_red(s->link[last])) {
|
|
g->link[dir2] = rbtree_node_rotate2(p, last);
|
|
} else if (rbtree_node_is_red(s->link[!last])) {
|
|
g->link[dir2] = rbtree_node_rotate(p, last);
|
|
}
|
|
|
|
// Ensure correct coloring
|
|
q->red = g->link[dir2]->red = 1;
|
|
g->link[dir2]->link[0]->red = 0;
|
|
g->link[dir2]->link[1]->red = 0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Replace and remove the saved node
|
|
if (f) {
|
|
void *tmp = f->value;
|
|
f->value = q->value;
|
|
q->value = tmp;
|
|
|
|
p->link[p->link[1] == q] = q->link[q->link[0] == NULL];
|
|
|
|
if (node_cb) {
|
|
node_cb(tree, q);
|
|
}
|
|
q = NULL;
|
|
}
|
|
|
|
// Update the root (it may be different)
|
|
tree->root = head.link[1];
|
|
|
|
// Make the root black for simplified logic
|
|
if (tree->root != NULL) {
|
|
tree->root->red = 0;
|
|
}
|
|
|
|
--tree->size;
|
|
}
|
|
return 1;
|
|
}
|
|
|
|
int rbtree_tree_remove(rbtree_t *tree, void *value)
|
|
{
|
|
int result = 0;
|
|
if (tree) {
|
|
result = rbtree_tree_remove_with_cb(
|
|
tree, value, rbtree_tree_node_dealloc_cb);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
unsigned int rbtree_tree_size(rbtree_t *tree)
|
|
{
|
|
unsigned int result = 0;
|
|
if (tree) {
|
|
result = tree->size;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// rbtree_iter_t
|
|
|
|
rbtree_iter_t *rbtree_iter_alloc()
|
|
{
|
|
return kmalloc(sizeof(rbtree_iter_t));
|
|
}
|
|
|
|
rbtree_iter_t *rbtree_iter_init(rbtree_iter_t *iter)
|
|
{
|
|
if (iter) {
|
|
iter->tree = NULL;
|
|
iter->node = NULL;
|
|
iter->top = 0;
|
|
}
|
|
return iter;
|
|
}
|
|
|
|
rbtree_iter_t *rbtree_iter_create()
|
|
{
|
|
return rbtree_iter_init(rbtree_iter_alloc());
|
|
}
|
|
|
|
void rbtree_iter_dealloc(rbtree_iter_t *iter)
|
|
{
|
|
if (iter) {
|
|
kfree(iter);
|
|
}
|
|
}
|
|
|
|
// Internal function, init traversal object, dir determines whether
|
|
// to begin traversal at the smallest or largest valued node.
|
|
static void *rbtree_iter_start(rbtree_iter_t *iter, rbtree_t *tree, int dir)
|
|
{
|
|
void *result = NULL;
|
|
if (iter) {
|
|
iter->tree = tree;
|
|
iter->node = tree->root;
|
|
iter->top = 0;
|
|
|
|
// Save the path for later selfersal
|
|
if (iter->node != NULL) {
|
|
while (iter->node->link[dir] != NULL) {
|
|
iter->path[iter->top++] = iter->node;
|
|
iter->node = iter->node->link[dir];
|
|
}
|
|
}
|
|
|
|
result = iter->node == NULL ? NULL : iter->node->value;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// Traverse a red black tree in the user-specified direction (0 asc, 1 desc)
|
|
static void *rbtree_iter_move(rbtree_iter_t *iter, int dir)
|
|
{
|
|
if (iter->node->link[dir] != NULL) {
|
|
// Continue down this branch
|
|
iter->path[iter->top++] = iter->node;
|
|
iter->node = iter->node->link[dir];
|
|
while (iter->node->link[!dir] != NULL) {
|
|
iter->path[iter->top++] = iter->node;
|
|
iter->node = iter->node->link[!dir];
|
|
}
|
|
} else {
|
|
// Move to the next branch
|
|
rbtree_node_t *last = NULL;
|
|
do {
|
|
if (iter->top == 0) {
|
|
iter->node = NULL;
|
|
break;
|
|
}
|
|
last = iter->node;
|
|
iter->node = iter->path[--iter->top];
|
|
} while (last == iter->node->link[dir]);
|
|
}
|
|
return iter->node == NULL ? NULL : iter->node->value;
|
|
}
|
|
|
|
void *rbtree_iter_first(rbtree_iter_t *iter, rbtree_t *tree)
|
|
{
|
|
return rbtree_iter_start(iter, tree, 0);
|
|
}
|
|
|
|
void *rbtree_iter_last(rbtree_iter_t *iter, rbtree_t *tree)
|
|
{
|
|
return rbtree_iter_start(iter, tree, 1);
|
|
}
|
|
|
|
void *rbtree_iter_next(rbtree_iter_t *iter)
|
|
{
|
|
return rbtree_iter_move(iter, 1);
|
|
}
|
|
|
|
void *rbtree_iter_prev(rbtree_iter_t *iter)
|
|
{
|
|
return rbtree_iter_move(iter, 0);
|
|
}
|
|
|
|
int rbtree_tree_test(rbtree_t *tree, rbtree_node_t *root)
|
|
{
|
|
int lh, rh;
|
|
|
|
if (root == NULL)
|
|
return 1;
|
|
else {
|
|
rbtree_node_t *ln = root->link[0];
|
|
rbtree_node_t *rn = root->link[1];
|
|
|
|
/* Consecutive red links */
|
|
if (rbtree_node_is_red(root)) {
|
|
if (rbtree_node_is_red(ln) || rbtree_node_is_red(rn)) {
|
|
pr_err("Red violation");
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
lh = rbtree_tree_test(tree, ln);
|
|
rh = rbtree_tree_test(tree, rn);
|
|
|
|
/* Invalid binary search tree */
|
|
if ((ln != NULL && tree->cmp(tree, ln, root) >= 0) || (rn != NULL && tree->cmp(tree, rn, root) <= 0)) {
|
|
pr_err("Binary tree violation");
|
|
return 0;
|
|
}
|
|
|
|
/* Black height mismatch */
|
|
if (lh != 0 && rh != 0 && lh != rh) {
|
|
pr_err("Black violation");
|
|
return 0;
|
|
}
|
|
|
|
/* Only count black links */
|
|
if (lh != 0 && rh != 0)
|
|
return rbtree_node_is_red(root) ? lh : lh + 1;
|
|
else
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
static void rbtree_tree_print_iter(rbtree_t *tree,
|
|
rbtree_node_t *node,
|
|
rbtree_tree_node_f fun)
|
|
{
|
|
assert(tree);
|
|
assert(node);
|
|
assert(fun);
|
|
fun(tree, node);
|
|
if (node->link[0])
|
|
rbtree_tree_print_iter(tree, node->link[0], fun);
|
|
if (node->link[1])
|
|
rbtree_tree_print_iter(tree, node->link[1], fun);
|
|
}
|
|
|
|
void rbtree_tree_print(rbtree_t *tree, rbtree_tree_node_f fun)
|
|
{
|
|
assert(tree);
|
|
assert(fun);
|
|
rbtree_tree_print_iter(tree, tree->root, fun);
|
|
}
|