rbtree.c

   1 /*
   2  * rbtree.c
   3  *
   4  * Simple implementation of a left-leaning red-black tree with 64-bit
   5  * integer keys.  The search operation will return the highest node <=
   6  * the key; only search and insert are supported, but additional
   7  * standard llrbtree operations can be coded up at will.
   8  *
   9  * See http://www.cs.princeton.edu/~rs/talks/LLRB/RedBlack.pdf for
  10  * information about left-leaning red-black trees.
  11  */
  12
  13 #include "rbtree.h"
  14
  15 struct rbtree *rb_search(struct rbtree *tree, uint64_t key)
  16 {
  17     struct rbtree *best = NULL;
  18
  19     while (tree) {
  20         if (tree->key == key)
  21             return tree;
  22         else if (tree->key > key)
  23             tree = tree->left;
  24         else {
  25             best = tree;
  26             tree = tree->right;
  27         }
  28     }
  29     return best;
  30 }
  31
  32 static bool is_red(struct rbtree *h)
  33 {
  34     return h && h->red;
  35 }
  36
  37 static struct rbtree *rotate_left(struct rbtree *h)
  38 {
  39     struct rbtree *x = h->right;
  40     h->right = x->left;
  41     x->left = h;
  42     x->red = x->left->red;
  43     x->left->red = true;
  44     return x;
  45 }
  46
  47 static struct rbtree *rotate_right(struct rbtree *h)
  48 {
  49     struct rbtree *x = h->left;
  50     h->left = x->right;
  51     x->right = h;
  52     x->red = x->right->red;
  53     x->right->red = true;
  54     return x;
  55 }
  56
  57 static void color_flip(struct rbtree *h)
  58 {
  59     h->red = !h->red;
  60     h->left->red = !h->left->red;
  61     h->right->red = !h->right->red;
  62 }
  63
  64 struct rbtree *rb_insert(struct rbtree *tree, struct rbtree *node)
  65 {
  66     if (!tree) {
  67         node->red = true;
  68         return node;
  69     }
  70
  71     if (is_red(tree->left) && is_red(tree->right))
  72         color_flip(tree);
  73
  74     if (node->key < tree->key)
  75         tree->left = rb_insert(tree->left, node);
  76     else
  77         tree->right = rb_insert(tree->right, node);
  78
  79     if (is_red(tree->right))
  80         tree = rotate_left(tree);
  81
  82     if (is_red(tree->left) && is_red(tree->left->left))
  83         tree = rotate_right(tree);
  84
  85     return tree;
  86 }