mirror of
https://github.com/jart/cosmopolitan.git
synced 2025-03-03 07:29:23 +00:00
Implement new red-black tree
This commit is contained in:
Justine Tunney
parent
fc65422660
commit
3756870635
3 changed files with 494 additions and 0 deletions
312
libc/intrin/rbtree.c
Normal file
312
libc/intrin/rbtree.c
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
// Copyright 2024 Justine Alexandra Roberts Tunney
|
||||
//
|
||||
// Permission to use, copy, modify, and/or distribute this software for
|
||||
// any purpose with or without fee is hereby granted, provided that the
|
||||
// above copyright notice and this permission notice appear in all copies.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
|
||||
// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
|
||||
// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
|
||||
// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
|
||||
// DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
|
||||
// PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
|
||||
// TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
// PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
#include "rbtree.h"
|
||||
#include "libc/dce.h"
|
||||
#include "libc/str/str.h"
|
||||
|
||||
#define RBTREE_DEBUG
|
||||
|
||||
struct rbtree *rbtree_next(struct rbtree *node) {
|
||||
if (!node)
|
||||
return 0;
|
||||
if (node->right)
|
||||
return rbtree_first(node->right);
|
||||
struct rbtree *parent = node->parent;
|
||||
while (parent && node == parent->right) {
|
||||
node = parent;
|
||||
parent = parent->parent;
|
||||
}
|
||||
return parent;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_prev(struct rbtree *node) {
|
||||
if (!node)
|
||||
return 0;
|
||||
if (rbtree_get_left(node))
|
||||
return rbtree_last(rbtree_get_left(node));
|
||||
struct rbtree *parent = node->parent;
|
||||
while (parent && node == rbtree_get_left(parent)) {
|
||||
node = parent;
|
||||
parent = parent->parent;
|
||||
}
|
||||
return parent;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_first(struct rbtree *node) {
|
||||
while (node && rbtree_get_left(node))
|
||||
node = rbtree_get_left(node);
|
||||
return node;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_last(struct rbtree *node) {
|
||||
while (node && node->right)
|
||||
node = node->right;
|
||||
return node;
|
||||
}
|
||||
|
||||
static void rbtree_rotate_left(struct rbtree **root, struct rbtree *x) {
|
||||
struct rbtree *y = x->right;
|
||||
x->right = rbtree_get_left(y);
|
||||
if (rbtree_get_left(y))
|
||||
rbtree_get_left(y)->parent = x;
|
||||
y->parent = x->parent;
|
||||
if (!x->parent) {
|
||||
*root = y;
|
||||
} else if (x == rbtree_get_left(x->parent)) {
|
||||
rbtree_set_left(x->parent, y);
|
||||
} else {
|
||||
x->parent->right = y;
|
||||
}
|
||||
rbtree_set_left(y, x);
|
||||
x->parent = y;
|
||||
}
|
||||
|
||||
static void rbtree_rotate_right(struct rbtree **root, struct rbtree *y) {
|
||||
struct rbtree *x = rbtree_get_left(y);
|
||||
rbtree_set_left(y, x->right);
|
||||
if (x->right)
|
||||
x->right->parent = y;
|
||||
x->parent = y->parent;
|
||||
if (!y->parent) {
|
||||
*root = x;
|
||||
} else if (y == y->parent->right) {
|
||||
y->parent->right = x;
|
||||
} else {
|
||||
rbtree_set_left(y->parent, x);
|
||||
}
|
||||
x->right = y;
|
||||
y->parent = x;
|
||||
}
|
||||
|
||||
static void rbtree_insert_fixup(struct rbtree **root, struct rbtree *node) {
|
||||
rbtree_set_red(node, 1);
|
||||
while (node != *root && rbtree_get_red(node->parent)) {
|
||||
if (node->parent == rbtree_get_left(node->parent->parent)) {
|
||||
struct rbtree *uncle = node->parent->parent->right;
|
||||
if (uncle && rbtree_get_red(uncle)) {
|
||||
rbtree_set_red(node->parent, 0);
|
||||
rbtree_set_red(uncle, 0);
|
||||
rbtree_set_red(node->parent->parent, 1);
|
||||
node = node->parent->parent;
|
||||
} else {
|
||||
if (node == node->parent->right) {
|
||||
node = node->parent;
|
||||
rbtree_rotate_left(root, node);
|
||||
}
|
||||
rbtree_set_red(node->parent, 0);
|
||||
rbtree_set_red(node->parent->parent, 1);
|
||||
rbtree_rotate_right(root, node->parent->parent);
|
||||
}
|
||||
} else {
|
||||
struct rbtree *uncle = rbtree_get_left(node->parent->parent);
|
||||
if (uncle && rbtree_get_red(uncle)) {
|
||||
rbtree_set_red(node->parent, 0);
|
||||
rbtree_set_red(uncle, 0);
|
||||
rbtree_set_red(node->parent->parent, 1);
|
||||
node = node->parent->parent;
|
||||
} else {
|
||||
if (node == rbtree_get_left(node->parent)) {
|
||||
node = node->parent;
|
||||
rbtree_rotate_right(root, node);
|
||||
}
|
||||
rbtree_set_red(node->parent, 0);
|
||||
rbtree_set_red(node->parent->parent, 1);
|
||||
rbtree_rotate_left(root, node->parent->parent);
|
||||
}
|
||||
}
|
||||
}
|
||||
rbtree_set_red(*root, 0);
|
||||
}
|
||||
|
||||
void rbtree_insert(struct rbtree **root, struct rbtree *node,
|
||||
rbtree_cmp_f *cmp) {
|
||||
bzero(node, sizeof(*node));
|
||||
if (!*root) {
|
||||
*root = node;
|
||||
} else {
|
||||
struct rbtree *search = *root;
|
||||
struct rbtree *parent = 0;
|
||||
do {
|
||||
parent = search;
|
||||
if (cmp(node, search) < 0) {
|
||||
search = rbtree_get_left(search);
|
||||
} else {
|
||||
search = search->right;
|
||||
}
|
||||
} while (search);
|
||||
if (cmp(node, parent) < 0) {
|
||||
rbtree_set_left(parent, node);
|
||||
} else {
|
||||
parent->right = node;
|
||||
}
|
||||
node->parent = parent;
|
||||
rbtree_insert_fixup(root, node);
|
||||
}
|
||||
}
|
||||
|
||||
static void rbtree_transplant(struct rbtree **root, struct rbtree *u,
|
||||
struct rbtree *v) {
|
||||
if (!u->parent) {
|
||||
*root = v;
|
||||
} else if (u == rbtree_get_left(u->parent)) {
|
||||
rbtree_set_left(u->parent, v);
|
||||
} else {
|
||||
u->parent->right = v;
|
||||
}
|
||||
if (v)
|
||||
v->parent = u->parent;
|
||||
}
|
||||
|
||||
static void rbtree_remove_fixup(struct rbtree **root, struct rbtree *node,
|
||||
struct rbtree *parent) {
|
||||
while (node != *root && (!node || !rbtree_get_red(node))) {
|
||||
if (node == rbtree_get_left(parent)) {
|
||||
struct rbtree *sibling = parent->right;
|
||||
if (rbtree_get_red(sibling)) {
|
||||
rbtree_set_red(sibling, 0);
|
||||
rbtree_set_red(parent, 1);
|
||||
rbtree_rotate_left(root, parent);
|
||||
sibling = parent->right;
|
||||
}
|
||||
if ((!rbtree_get_left(sibling) ||
|
||||
!rbtree_get_red(rbtree_get_left(sibling))) &&
|
||||
(!sibling->right || !rbtree_get_red(sibling->right))) {
|
||||
rbtree_set_red(sibling, 1);
|
||||
node = parent;
|
||||
parent = node->parent;
|
||||
} else {
|
||||
if (!sibling->right || !rbtree_get_red(sibling->right)) {
|
||||
rbtree_set_red(rbtree_get_left(sibling), 0);
|
||||
rbtree_set_red(sibling, 1);
|
||||
rbtree_rotate_right(root, sibling);
|
||||
sibling = parent->right;
|
||||
}
|
||||
rbtree_set_red(sibling, rbtree_get_red(parent));
|
||||
rbtree_set_red(parent, 0);
|
||||
rbtree_set_red(sibling->right, 0);
|
||||
rbtree_rotate_left(root, parent);
|
||||
node = *root;
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
struct rbtree *sibling = rbtree_get_left(parent);
|
||||
if (rbtree_get_red(sibling)) {
|
||||
rbtree_set_red(sibling, 0);
|
||||
rbtree_set_red(parent, 1);
|
||||
rbtree_rotate_right(root, parent);
|
||||
sibling = rbtree_get_left(parent);
|
||||
}
|
||||
if ((!sibling->right || !rbtree_get_red(sibling->right)) &&
|
||||
(!rbtree_get_left(sibling) ||
|
||||
!rbtree_get_red(rbtree_get_left(sibling)))) {
|
||||
rbtree_set_red(sibling, 1);
|
||||
node = parent;
|
||||
parent = node->parent;
|
||||
} else {
|
||||
if (!rbtree_get_left(sibling) ||
|
||||
!rbtree_get_red(rbtree_get_left(sibling))) {
|
||||
rbtree_set_red(sibling->right, 0);
|
||||
rbtree_set_red(sibling, 1);
|
||||
rbtree_rotate_left(root, sibling);
|
||||
sibling = rbtree_get_left(parent);
|
||||
}
|
||||
rbtree_set_red(sibling, rbtree_get_red(parent));
|
||||
rbtree_set_red(parent, 0);
|
||||
rbtree_set_red(rbtree_get_left(sibling), 0);
|
||||
rbtree_rotate_right(root, parent);
|
||||
node = *root;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (node)
|
||||
rbtree_set_red(node, 0);
|
||||
}
|
||||
|
||||
void rbtree_remove(struct rbtree **root, struct rbtree *node) {
|
||||
struct rbtree *y = node;
|
||||
struct rbtree *x = 0;
|
||||
struct rbtree *x_parent = 0;
|
||||
int y_original_color = rbtree_get_red(y);
|
||||
if (!rbtree_get_left(node)) {
|
||||
x = node->right;
|
||||
rbtree_transplant(root, node, node->right);
|
||||
x_parent = node->parent;
|
||||
} else if (!node->right) {
|
||||
x = rbtree_get_left(node);
|
||||
rbtree_transplant(root, node, rbtree_get_left(node));
|
||||
x_parent = node->parent;
|
||||
} else {
|
||||
y = rbtree_first(node->right);
|
||||
y_original_color = rbtree_get_red(y);
|
||||
x = y->right;
|
||||
if (y->parent == node) {
|
||||
if (x)
|
||||
x->parent = y;
|
||||
x_parent = y;
|
||||
} else {
|
||||
rbtree_transplant(root, y, y->right);
|
||||
y->right = node->right;
|
||||
y->right->parent = y;
|
||||
x_parent = y->parent;
|
||||
}
|
||||
rbtree_transplant(root, node, y);
|
||||
rbtree_set_left(y, rbtree_get_left(node));
|
||||
rbtree_get_left(y)->parent = y;
|
||||
rbtree_set_red(y, rbtree_get_red(node));
|
||||
}
|
||||
if (!y_original_color)
|
||||
rbtree_remove_fixup(root, x, x_parent);
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_get(const struct rbtree *node, const struct rbtree *key,
|
||||
rbtree_cmp_f *cmp) {
|
||||
while (node) {
|
||||
int c = cmp(key, node);
|
||||
if (c < 0) {
|
||||
node = rbtree_get_left(node);
|
||||
} else if (c > 0) {
|
||||
node = node->right;
|
||||
} else {
|
||||
return (struct rbtree *)node;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_floor(const struct rbtree *node, const struct rbtree *key,
|
||||
rbtree_cmp_f *cmp) {
|
||||
while (node) {
|
||||
if (cmp(key, node) < 0) {
|
||||
node = rbtree_get_left(node);
|
||||
} else {
|
||||
node = node->right;
|
||||
}
|
||||
}
|
||||
return (struct rbtree *)node;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_ceil(const struct rbtree *node, const struct rbtree *key,
|
||||
rbtree_cmp_f *cmp) {
|
||||
while (node) {
|
||||
if (cmp(node, key) < 0) {
|
||||
node = rbtree_get_left(node);
|
||||
} else {
|
||||
node = node->right;
|
||||
}
|
||||
}
|
||||
return (struct rbtree *)node;
|
||||
}
|
||||
57
libc/intrin/rbtree.h
Normal file
57
libc/intrin/rbtree.h
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
#ifdef _COSMO_SOURCE
|
||||
#ifndef COSMOPOLITAN_RBTREE_H_
|
||||
#define COSMOPOLITAN_RBTREE_H_
|
||||
#define rbtree_ceil __rbtree_ceil
|
||||
#define rbtree_first __rbtree_first
|
||||
#define rbtree_floor __rbtree_floor
|
||||
#define rbtree_get __rbtree_get
|
||||
#define rbtree_insert __rbtree_insert
|
||||
#define rbtree_last __rbtree_last
|
||||
#define rbtree_next __rbtree_next
|
||||
#define rbtree_prev __rbtree_prev
|
||||
#define rbtree_remove __rbtree_remove
|
||||
COSMOPOLITAN_C_START_
|
||||
|
||||
#define RBTREE_CONTAINER(t, f, p) ((t *)(((char *)(p)) - offsetof(t, f)))
|
||||
|
||||
struct rbtree {
|
||||
uintptr_t word;
|
||||
struct rbtree *right;
|
||||
struct rbtree *parent;
|
||||
};
|
||||
|
||||
typedef int rbtree_cmp_f(const struct rbtree *, const struct rbtree *);
|
||||
|
||||
static inline struct rbtree *rbtree_get_left(const struct rbtree *node) {
|
||||
return (struct rbtree *)(node->word & -2);
|
||||
}
|
||||
|
||||
static inline void rbtree_set_left(struct rbtree *node, struct rbtree *left) {
|
||||
node->word = (uintptr_t)left | (node->word & 1);
|
||||
}
|
||||
|
||||
static inline int rbtree_get_red(const struct rbtree *node) {
|
||||
return node->word & 1;
|
||||
}
|
||||
|
||||
static inline void rbtree_set_red(struct rbtree *node, int red) {
|
||||
node->word &= -2;
|
||||
node->word |= red;
|
||||
}
|
||||
|
||||
struct rbtree *rbtree_next(struct rbtree *) libcesque;
|
||||
struct rbtree *rbtree_prev(struct rbtree *) libcesque;
|
||||
struct rbtree *rbtree_first(struct rbtree *) libcesque;
|
||||
struct rbtree *rbtree_last(struct rbtree *) libcesque;
|
||||
void rbtree_remove(struct rbtree **, struct rbtree *) libcesque;
|
||||
void rbtree_insert(struct rbtree **, struct rbtree *, rbtree_cmp_f *) libcesque;
|
||||
struct rbtree *rbtree_get(const struct rbtree *, const struct rbtree *,
|
||||
rbtree_cmp_f *) libcesque;
|
||||
struct rbtree *rbtree_ceil(const struct rbtree *, const struct rbtree *,
|
||||
rbtree_cmp_f *) libcesque;
|
||||
struct rbtree *rbtree_floor(const struct rbtree *, const struct rbtree *,
|
||||
rbtree_cmp_f *) libcesque;
|
||||
|
||||
COSMOPOLITAN_C_END_
|
||||
#endif /* COSMOPOLITAN_RBTREE_H_ */
|
||||
#endif /* _COSMO_SOURCE */
|
||||
125
test/libc/intrin/rbtree_test.c
Normal file
125
test/libc/intrin/rbtree_test.c
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
// Copyright 2024 Justine Alexandra Roberts Tunney
|
||||
//
|
||||
// Permission to use, copy, modify, and/or distribute this software for
|
||||
// any purpose with or without fee is hereby granted, provided that the
|
||||
// above copyright notice and this permission notice appear in all copies.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
|
||||
// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
|
||||
// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
|
||||
// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
|
||||
// DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR
|
||||
// PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
|
||||
// TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
|
||||
// PERFORMANCE OF THIS SOFTWARE.
|
||||
|
||||
#include "libc/intrin/rbtree.h"
|
||||
#include "libc/intrin/kprintf.h"
|
||||
#include "libc/macros.internal.h"
|
||||
#include "libc/runtime/runtime.h"
|
||||
#include "libc/stdio/rand.h"
|
||||
|
||||
#define NUMBER_CONTAINER(e) RBTREE_CONTAINER(struct number, elem, e)
|
||||
|
||||
void rbtree_checker(const struct rbtree *node, const struct rbtree *parent,
|
||||
int black_count, int *black_height, rbtree_cmp_f *cmp) {
|
||||
if (!node) {
|
||||
// Leaf nodes are considered black
|
||||
if (*black_height == -1) {
|
||||
*black_height = black_count;
|
||||
} else if (black_count != *black_height) {
|
||||
// ILLEGAL TREE: Black height mismatch
|
||||
__builtin_trap();
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (node->parent != parent)
|
||||
// ILLEGAL TREE: Parent link is incorrect
|
||||
__builtin_trap();
|
||||
if (parent) {
|
||||
if (rbtree_get_left(parent) == node && cmp(parent, node) < 0)
|
||||
// ILLEGAL TREE: Binary search property violated on left child
|
||||
__builtin_trap();
|
||||
if (parent->right == node && cmp(node, parent) < 0)
|
||||
// ILLEGAL TREE: Binary search property violated on right child
|
||||
__builtin_trap();
|
||||
}
|
||||
if (!rbtree_get_red(node)) {
|
||||
black_count++;
|
||||
} else if (parent && rbtree_get_red(parent)) {
|
||||
// ILLEGAL TREE: Red node has red child
|
||||
__builtin_trap();
|
||||
}
|
||||
rbtree_checker(rbtree_get_left(node), node, black_count, black_height, cmp);
|
||||
rbtree_checker(node->right, node, black_count, black_height, cmp);
|
||||
}
|
||||
|
||||
void rbtree_check(struct rbtree *root, rbtree_cmp_f *cmp) {
|
||||
if (root) {
|
||||
if (rbtree_get_red(root))
|
||||
// ILLEGAL TREE: root node must be black
|
||||
__builtin_trap();
|
||||
int black_height = -1;
|
||||
rbtree_checker(root, 0, 0, &black_height, cmp);
|
||||
}
|
||||
}
|
||||
|
||||
struct number {
|
||||
int number;
|
||||
struct rbtree elem;
|
||||
};
|
||||
|
||||
int number_compare(const struct rbtree *ra, const struct rbtree *rb) {
|
||||
const struct number *a = NUMBER_CONTAINER(ra);
|
||||
const struct number *b = NUMBER_CONTAINER(rb);
|
||||
return (a->number > b->number) - (a->number < b->number);
|
||||
}
|
||||
|
||||
struct number *number_new(int number) {
|
||||
static int used;
|
||||
static struct number heap[8192];
|
||||
if (used == ARRAYLEN(heap))
|
||||
return 0;
|
||||
struct number *res = &heap[used++];
|
||||
res->number = number;
|
||||
return res;
|
||||
}
|
||||
|
||||
struct rbtree *tree = 0;
|
||||
|
||||
void print(void) {
|
||||
for (struct rbtree *e = rbtree_first(tree); e; e = rbtree_next(e))
|
||||
kprintf("%3d", NUMBER_CONTAINER(e)->number);
|
||||
kprintf("\n");
|
||||
}
|
||||
|
||||
void print_reversed(void) {
|
||||
for (struct rbtree *e = rbtree_last(tree); e; e = rbtree_prev(e))
|
||||
kprintf("%3d", NUMBER_CONTAINER(e)->number);
|
||||
kprintf("\n");
|
||||
}
|
||||
|
||||
void simple_test(void) {
|
||||
static const int kNumba[] = {74, 53, 96, 70, 34, 95, 30, 2, 89, 46,
|
||||
23, 2, 52, 0, 34, 12, 90, 95, 32, 65};
|
||||
for (int i = 0; i < 20; ++i) {
|
||||
int number = kNumba[i];
|
||||
kprintf("%3d", number);
|
||||
rbtree_insert(&tree, &number_new(number)->elem, number_compare);
|
||||
rbtree_check(tree, number_compare);
|
||||
}
|
||||
kprintf("\n");
|
||||
print();
|
||||
print_reversed();
|
||||
for (int i = 0; i < 20; ++i) {
|
||||
rbtree_remove(&tree, rbtree_get(tree, &(&(struct number){kNumba[i]})->elem,
|
||||
number_compare));
|
||||
rbtree_check(tree, number_compare);
|
||||
print();
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
ShowCrashReports();
|
||||
simple_test();
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue