From 3756870635d3a4b5a12dff0f266a8b073c830dc2 Mon Sep 17 00:00:00 2001 From: Justine Tunney Date: Fri, 5 Jul 2024 12:56:03 -0700 Subject: [PATCH] Implement new red-black tree --- libc/intrin/rbtree.c | 312 +++++++++++++++++++++++++++++++++ libc/intrin/rbtree.h | 57 ++++++ test/libc/intrin/rbtree_test.c | 125 +++++++++++++ 3 files changed, 494 insertions(+) create mode 100644 libc/intrin/rbtree.c create mode 100644 libc/intrin/rbtree.h create mode 100644 test/libc/intrin/rbtree_test.c diff --git a/libc/intrin/rbtree.c b/libc/intrin/rbtree.c new file mode 100644 index 000000000..3fc3ae5cc --- /dev/null +++ b/libc/intrin/rbtree.c @@ -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; +} diff --git a/libc/intrin/rbtree.h b/libc/intrin/rbtree.h new file mode 100644 index 000000000..0c63f67dd --- /dev/null +++ b/libc/intrin/rbtree.h @@ -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 */ diff --git a/test/libc/intrin/rbtree_test.c b/test/libc/intrin/rbtree_test.c new file mode 100644 index 000000000..5bcea721c --- /dev/null +++ b/test/libc/intrin/rbtree_test.c @@ -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(); +}