| // SPDX-License-Identifier: Unlicense |
| // |
| // Based on Julienne Walker's <http://eternallyconfuzzled.com/> rb_tree |
| // implementation. |
| // |
| // Modified by Mirek Rusin <http://github.com/mirek/rb_tree>. |
| // |
| // This is free and unencumbered software released into the public domain. |
| // |
| // Anyone is free to copy, modify, publish, use, compile, sell, or |
| // distribute this software, either in source code form or as a compiled |
| // binary, for any purpose, commercial or non-commercial, and by any |
| // means. |
| // |
| // In jurisdictions that recognize copyright laws, the author or authors |
| // of this software dedicate any and all copyright interest in the |
| // software to the public domain. We make this dedication for the benefit |
| // of the public at large and to the detriment of our heirs and |
| // successors. We intend this dedication to be an overt act of |
| // relinquishment in perpetuity of all present and future rights to this |
| // software under copyright law. |
| // |
| // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, |
| // EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF |
| // MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. |
| // IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR |
| // OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, |
| // ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR |
| // OTHER DEALINGS IN THE SOFTWARE. |
| // |
| // For more information, please refer to <http://unlicense.org/> |
| // |
| |
| #include "rb_tree.h" |
| |
| // rb_node |
| |
| struct rb_node * |
| rb_node_alloc () { |
| return malloc(sizeof(struct rb_node)); |
| } |
| |
| struct rb_node * |
| rb_node_init (struct rb_node *self, void *value) { |
| if (self) { |
| self->red = 1; |
| self->link[0] = self->link[1] = NULL; |
| self->value = value; |
| } |
| return self; |
| } |
| |
| struct rb_node * |
| rb_node_create (void *value) { |
| return rb_node_init(rb_node_alloc(), value); |
| } |
| |
| void |
| rb_node_dealloc (struct rb_node *self) { |
| if (self) { |
| free(self); |
| } |
| } |
| |
| static int |
| rb_node_is_red (const struct rb_node *self) { |
| return self ? self->red : 0; |
| } |
| |
| static struct rb_node * |
| rb_node_rotate (struct rb_node *self, int dir) { |
| struct rb_node *result = NULL; |
| if (self) { |
| result = self->link[!dir]; |
| self->link[!dir] = result->link[dir]; |
| result->link[dir] = self; |
| self->red = 1; |
| result->red = 0; |
| } |
| return result; |
| } |
| |
| static struct rb_node * |
| rb_node_rotate2 (struct rb_node *self, int dir) { |
| struct rb_node *result = NULL; |
| if (self) { |
| self->link[!dir] = rb_node_rotate(self->link[!dir], !dir); |
| result = rb_node_rotate(self, dir); |
| } |
| return result; |
| } |
| |
| // rb_tree - default callbacks |
| |
| int |
| rb_tree_node_cmp_ptr_cb (struct rb_tree *self, struct rb_node *a, struct rb_node *b) { |
| return (a->value > b->value) - (a->value < b->value); |
| } |
| |
| void |
| rb_tree_node_dealloc_cb (struct rb_tree *self, struct rb_node *node) { |
| if (self) { |
| if (node) { |
| rb_node_dealloc(node); |
| } |
| } |
| } |
| |
| // rb_tree |
| |
| struct rb_tree * |
| rb_tree_alloc () { |
| return malloc(sizeof(struct rb_tree)); |
| } |
| |
| struct rb_tree * |
| rb_tree_init (struct rb_tree *self, rb_tree_node_cmp_f node_cmp_cb) { |
| if (self) { |
| self->root = NULL; |
| self->size = 0; |
| self->cmp = node_cmp_cb ? node_cmp_cb : rb_tree_node_cmp_ptr_cb; |
| } |
| return self; |
| } |
| |
| struct rb_tree * |
| rb_tree_create (rb_tree_node_cmp_f node_cb) { |
| return rb_tree_init(rb_tree_alloc(), node_cb); |
| } |
| |
| void |
| rb_tree_dealloc (struct rb_tree *self, rb_tree_node_f node_cb) { |
| if (self) { |
| if (node_cb) { |
| struct rb_node *node = self->root; |
| struct rb_node *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(self, 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; |
| } |
| } |
| free(self); |
| } |
| } |
| |
| int |
| rb_tree_test (struct rb_tree *self, struct rb_node *root) { |
| int lh, rh; |
| |
| if ( root == NULL ) |
| return 1; |
| else { |
| struct rb_node *ln = root->link[0]; |
| struct rb_node *rn = root->link[1]; |
| |
| /* Consecutive red links */ |
| if (rb_node_is_red(root)) { |
| if (rb_node_is_red(ln) || rb_node_is_red(rn)) { |
| printf("Red violation"); |
| return 0; |
| } |
| } |
| |
| lh = rb_tree_test(self, ln); |
| rh = rb_tree_test(self, rn); |
| |
| /* Invalid binary search tree */ |
| if ( ( ln != NULL && self->cmp(self, ln, root) >= 0 ) |
| || ( rn != NULL && self->cmp(self, rn, root) <= 0)) |
| { |
| puts ( "Binary tree violation" ); |
| return 0; |
| } |
| |
| /* Black height mismatch */ |
| if ( lh != 0 && rh != 0 && lh != rh ) { |
| puts ( "Black violation" ); |
| return 0; |
| } |
| |
| /* Only count black links */ |
| if ( lh != 0 && rh != 0 ) |
| return rb_node_is_red ( root ) ? lh : lh + 1; |
| else |
| return 0; |
| } |
| } |
| |
| void * |
| rb_tree_find(struct rb_tree *self, void *value) { |
| void *result = NULL; |
| if (self) { |
| struct rb_node node = { .value = value }; |
| struct rb_node *it = self->root; |
| int cmp = 0; |
| while (it) { |
| if ((cmp = self->cmp(self, 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; |
| } |
| |
| // Creates (malloc'ates) |
| int |
| rb_tree_insert (struct rb_tree *self, void *value) { |
| return rb_tree_insert_node(self, rb_node_create(value)); |
| } |
| |
| // Returns 1 on success, 0 otherwise. |
| int |
| rb_tree_insert_node (struct rb_tree *self, struct rb_node *node) { |
| if (self && node) { |
| if (self->root == NULL) { |
| self->root = node; |
| } else { |
| struct rb_node head = { 0 }; // False tree root |
| struct rb_node *g, *t; // Grandparent & parent |
| struct rb_node *p, *q; // Iterator & parent |
| int dir = 0, last = 0; |
| |
| // Set up our helpers |
| t = &head; |
| g = p = NULL; |
| q = t->link[1] = self->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 (rb_node_is_red(q->link[0]) && rb_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 (rb_node_is_red(q) && rb_node_is_red(p)) { |
| |
| // Hard red violation: rotations necessary |
| int dir2 = t->link[1] == g; |
| if (q == p->link[last]) { |
| t->link[dir2] = rb_node_rotate(g, !last); |
| } else { |
| t->link[dir2] = rb_node_rotate2(g, !last); |
| } |
| } |
| |
| // Stop working if we inserted a node. This |
| // check also disallows duplicates in the tree |
| if (self->cmp(self, q, node) == 0) { |
| break; |
| } |
| |
| last = dir; |
| dir = self->cmp(self, 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) |
| self->root = head.link[1]; |
| } |
| |
| // Make the root black for simplified logic |
| self->root->red = 0; |
| ++self->size; |
| return 1; |
| } |
| return 0; |
| } |
| |
| // Returns 1 if the value was removed, 0 otherwise. Optional node callback |
| // can be provided to dealloc node and/or user data. Use rb_tree_node_dealloc |
| // default callback to deallocate node created by rb_tree_insert(...). |
| int |
| rb_tree_remove_with_cb (struct rb_tree *self, void *value, rb_tree_node_f node_cb) { |
| if (self->root != NULL) { |
| struct rb_node head = {0}; // False tree root |
| struct rb_node node = { .value = value }; // Value wrapper node |
| struct rb_node *q, *p, *g; // Helpers |
| struct rb_node *f = NULL; // Found item |
| int dir = 1; |
| |
| // Set up our helpers |
| q = &head; |
| g = p = NULL; |
| q->link[1] = self->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 = self->cmp(self, q, &node) < 0; |
| |
| // Save the node with matching value and keep |
| // going; we'll do removal tasks at the end |
| if (self->cmp(self, q, &node) == 0) { |
| f = q; |
| } |
| |
| // Push the red node down with rotations and color flips |
| if (!rb_node_is_red(q) && !rb_node_is_red(q->link[dir])) { |
| if (rb_node_is_red(q->link[!dir])) { |
| p = p->link[last] = rb_node_rotate(q, dir); |
| } else if (!rb_node_is_red(q->link[!dir])) { |
| struct rb_node *s = p->link[!last]; |
| if (s) { |
| if (!rb_node_is_red(s->link[!last]) && !rb_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 (rb_node_is_red(s->link[last])) { |
| g->link[dir2] = rb_node_rotate2(p, last); |
| } else if (rb_node_is_red(s->link[!last])) { |
| g->link[dir2] = rb_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(self, q); |
| } |
| q = NULL; |
| } |
| |
| // Update the root (it may be different) |
| self->root = head.link[1]; |
| |
| // Make the root black for simplified logic |
| if (self->root != NULL) { |
| self->root->red = 0; |
| } |
| |
| --self->size; |
| } |
| return 1; |
| } |
| |
| int |
| rb_tree_remove (struct rb_tree *self, void *value) { |
| int result = 0; |
| if (self) { |
| result = rb_tree_remove_with_cb(self, value, rb_tree_node_dealloc_cb); |
| } |
| return result; |
| } |
| |
| size_t |
| rb_tree_size (struct rb_tree *self) { |
| size_t result = 0; |
| if (self) { |
| result = self->size; |
| } |
| return result; |
| } |
| |
| // rb_iter |
| |
| struct rb_iter * |
| rb_iter_alloc () { |
| return malloc(sizeof(struct rb_iter)); |
| } |
| |
| struct rb_iter * |
| rb_iter_init (struct rb_iter *self) { |
| if (self) { |
| self->tree = NULL; |
| self->node = NULL; |
| self->top = 0; |
| } |
| return self; |
| } |
| |
| struct rb_iter * |
| rb_iter_create () { |
| return rb_iter_init(rb_iter_alloc()); |
| } |
| |
| void |
| rb_iter_dealloc (struct rb_iter *self) { |
| if (self) { |
| free(self); |
| } |
| } |
| |
| // Internal function, init traversal object, dir determines whether |
| // to begin traversal at the smallest or largest valued node. |
| static void * |
| rb_iter_start (struct rb_iter *self, struct rb_tree *tree, int dir) { |
| void *result = NULL; |
| if (self) { |
| self->tree = tree; |
| self->node = tree->root; |
| self->top = 0; |
| |
| // Save the path for later selfersal |
| if (self->node != NULL) { |
| while (self->node->link[dir] != NULL) { |
| self->path[self->top++] = self->node; |
| self->node = self->node->link[dir]; |
| } |
| } |
| |
| result = self->node == NULL ? NULL : self->node->value; |
| } |
| return result; |
| } |
| |
| // Traverse a red black tree in the user-specified direction (0 asc, 1 desc) |
| static void * |
| rb_iter_move (struct rb_iter *self, int dir) { |
| if (self->node->link[dir] != NULL) { |
| |
| // Continue down this branch |
| self->path[self->top++] = self->node; |
| self->node = self->node->link[dir]; |
| while ( self->node->link[!dir] != NULL ) { |
| self->path[self->top++] = self->node; |
| self->node = self->node->link[!dir]; |
| } |
| } else { |
| |
| // Move to the next branch |
| struct rb_node *last = NULL; |
| do { |
| if (self->top == 0) { |
| self->node = NULL; |
| break; |
| } |
| last = self->node; |
| self->node = self->path[--self->top]; |
| } while (last == self->node->link[dir]); |
| } |
| return self->node == NULL ? NULL : self->node->value; |
| } |
| |
| void * |
| rb_iter_first (struct rb_iter *self, struct rb_tree *tree) { |
| return rb_iter_start(self, tree, 0); |
| } |
| |
| void * |
| rb_iter_last (struct rb_iter *self, struct rb_tree *tree) { |
| return rb_iter_start(self, tree, 1); |
| } |
| |
| void * |
| rb_iter_next (struct rb_iter *self) { |
| return rb_iter_move(self, 1); |
| } |
| |
| void * |
| rb_iter_prev (struct rb_iter *self) { |
| return rb_iter_move(self, 0); |
| } |