src/util/rb_tree.h

   1 /*
   2  * Copyright © 2017 Jason Ekstrand
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the "Software"),
   6  * to deal in the Software without restriction, including without limitation
   7  * the rights to use, copy, modify, merge, publish, distribute, sublicense,
   8  * and/or sell copies of the Software, and to permit persons to whom the
   9  * Software is furnished to do so, subject to the following conditions:
  10  *
  11  * The above copyright notice and this permission notice shall be included in
  12  * all copies or substantial portions of the Software.
  13  *
  14  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  15  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  16  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  17  * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  18  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  19  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  20  * DEALINGS IN THE SOFTWARE.
  21  */
  22
  23 #ifndef RB_TREE_H
  24 #define RB_TREE_H
  25
  26 #include <stdbool.h>
  27 #include <stddef.h>
  28 #include <stdint.h>
  29 #include <stdlib.h>
  30
  31 /** A red-black tree node
  32  *
  33  * This struct represents a node in the red-black tree.  This struct should
  34  * be embedded as a member in whatever structure you wish to put in the
  35  * tree.
  36  */
  37 struct rb_node {
  38     /** Parent and color of this node
  39      *
  40      * The least significant bit represents the color and is est to 1 for
  41      * black and 0 for red.  The other bits are the pointer to the parent
  42      * and that pointer can be retrieved by masking off the bottom bit and
  43      * casting to a pointer.
  44      */
  45     uintptr_t parent;
  46
  47     /** Left child of this node, null for a leaf */
  48     struct rb_node *left;
  49
  50     /** Right child of this node, null for a leaf */
  51     struct rb_node *right;
  52 };
  53
  54 /** Return the parent node of the given node or NULL if it is the root */
  55 static inline struct rb_node *
  56 rb_node_parent(struct rb_node *n)
  57 {
  58     return (struct rb_node *)(n->parent & ~(uintptr_t)1);
  59 }
  60
  61 /** A red-black tree
  62  *
  63  * This struct represents the red-black tree itself.  It is just a pointer
  64  * to the root node with no other metadata.
  65  */
  66 struct rb_tree {
  67     struct rb_node *root;
  68 };
  69
  70 /** Initialize a red-black tree */
  71 void rb_tree_init(struct rb_tree *T);
  72
  73 /** Returns true if the red-black tree is empty */
  74 static inline bool
  75 rb_tree_is_empty(const struct rb_tree *T)
  76 {
  77     return T->root == NULL;
  78 }
  79
  80 /** Retrieve the data structure containing a node
  81  *
  82  * \param   type    The type of the containing data structure
  83  *
  84  * \param   node    A pointer to a rb_node
  85  *
  86  * \param   field   The rb_node field in the containing data structure
  87  */
  88 #define rb_node_data(type, node, field) \
  89     ((type *)(((char *)(node)) - offsetof(type, field)))
  90
  91 /** Insert a node into a tree at a particular location
  92  *
  93  * This function should probably not be used directly as it relies on the
  94  * caller to ensure that the parent node is correct.  Use rb_tree_insert
  95  * instead.
  96  *
  97  * \param   T           The red-black tree into which to insert the new node
  98  *
  99  * \param   parent      The node in the tree that will be the parent of the
 100  *                      newly inserted node
 101  *
 102  * \param   node        The node to insert
 103  *
 104  * \param   insert_left If true, the new node will be the left child of
 105  *                      \p parent, otherwise it will be the right child
 106  */
 107 void rb_tree_insert_at(struct rb_tree *T, struct rb_node *parent,
 108                        struct rb_node *node, bool insert_left);
 109
 110 /** Insert a node into a tree
 111  *
 112  * \param   T       The red-black tree into which to insert the new node
 113  *
 114  * \param   node    The node to insert
 115  *
 116  * \param   cmp     A comparison function to use to order the nodes.
 117  */
 118 static inline void
 119 rb_tree_insert(struct rb_tree *T, struct rb_node *node,
 120                int (*cmp)(const struct rb_node *, const struct rb_node *))
 121 {
 122     /* This function is declared inline in the hopes that the compiler can
 123      * optimize away the comparison function pointer call.
 124      */
 125     struct rb_node *y = NULL;
 126     struct rb_node *x = T->root;
 127     bool left = false;
 128     while (x != NULL) {
 129         y = x;
 130         left = cmp(x, node) < 0;
 131         if (left)
 132             x = x->left;
 133         else
 134             x = x->right;
 135     }
 136
 137     rb_tree_insert_at(T, y, node, left);
 138 }
 139
 140 /** Remove a node from a tree
 141  *
 142  * \param   T       The red-black tree from which to remove the node
 143  *
 144  * \param   node    The node to remove
 145  */
 146 void rb_tree_remove(struct rb_tree *T, struct rb_node *z);
 147
 148 /** Search the tree for a node
 149  *
 150  * If a node with a matching key exists, the first matching node found will
 151  * be returned.  If no matching node exists, NULL is returned.
 152  *
 153  * \param   T       The red-black tree to search
 154  *
 155  * \param   key     The key to search for
 156  *
 157  * \param   cmp     A comparison function to use to order the nodes
 158  */
 159 static inline struct rb_node *
 160 rb_tree_search(struct rb_tree *T, const void *key,
 161                int (*cmp)(const struct rb_node *, const void *))
 162 {
 163     /* This function is declared inline in the hopes that the compiler can
 164      * optimize away the comparison function pointer call.
 165      */
 166     struct rb_node *x = T->root;
 167     while (x != NULL) {
 168         int c = cmp(x, key);
 169         if (c < 0)
 170             x = x->left;
 171         else if (c > 0)
 172             x = x->right;
 173         else
 174             return x;
 175     }
 176
 177     return x;
 178 }
 179
 180 /** Sloppily search the tree for a node
 181  *
 182  * This function searches the tree for a given node.  If a node with a
 183  * matching key exists, that first matching node found will be returned.
 184  * If no node with an exactly matching key exists, the node returned will
 185  * be either the right-most node comparing less than \p key or the
 186  * right-most node comparing greater than \p key.  If the tree is empty,
 187  * NULL is returned.
 188  *
 189  * \param   T       The red-black tree to search
 190  *
 191  * \param   key     The key to search for
 192  *
 193  * \param   cmp     A comparison function to use to order the nodes
 194  */
 195 static inline struct rb_node *
 196 rb_tree_search_sloppy(struct rb_tree *T, const void *key,
 197                       int (*cmp)(const struct rb_node *, const void *))
 198 {
 199     /* This function is declared inline in the hopes that the compiler can
 200      * optimize away the comparison function pointer call.
 201      */
 202     struct rb_node *y = NULL;
 203     struct rb_node *x = T->root;
 204     while (x != NULL) {
 205         y = x;
 206         int c = cmp(x, key);
 207         if (c < 0)
 208             x = x->left;
 209         else if (c > 0)
 210             x = x->right;
 211         else
 212             return x;
 213     }
 214
 215     return y;
 216 }
 217
 218 /** Get the first (left-most) node in the tree or NULL */
 219 struct rb_node *rb_tree_first(struct rb_tree *T);
 220
 221 /** Get the last (right-most) node in the tree or NULL */
 222 struct rb_node *rb_tree_last(struct rb_tree *T);
 223
 224 /** Get the next node (to the right) in the tree or NULL */
 225 struct rb_node *rb_node_next(struct rb_node *node);
 226
 227 /** Get the next previous (to the left) in the tree or NULL */
 228 struct rb_node *rb_node_prev(struct rb_node *node);
 229
 230 #define rb_node_next_or_null(n) ((n) == NULL ? NULL : rb_node_next(n))
 231 #define rb_node_prev_or_null(n) ((n) == NULL ? NULL : rb_node_prev(n))
 232
 233 /** Iterate over the nodes in the tree
 234  *
 235  * \param   type    The type of the containing data structure
 236  *
 237  * \param   node    The variable name for current node in the iteration;
 238  *                  this will be declared as a pointer to \p type
 239  *
 240  * \param   T       The red-black tree
 241  *
 242  * \param   field   The rb_node field in containing data structure
 243  */
 244 #define rb_tree_foreach(type, iter, T, field) \
 245    for (type *iter, *__node = (type *)rb_tree_first(T); \
 246         __node != NULL && \
 247         (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
 248         __node = (type *)rb_node_next((struct rb_node *)__node))
 249
 250 /** Iterate over the nodes in the tree, allowing the current node to be freed
 251  *
 252  * \param   type    The type of the containing data structure
 253  *
 254  * \param   node    The variable name for current node in the iteration;
 255  *                  this will be declared as a pointer to \p type
 256  *
 257  * \param   T       The red-black tree
 258  *
 259  * \param   field   The rb_node field in containing data structure
 260  */
 261 #define rb_tree_foreach_safe(type, iter, T, field) \
 262    for (type *iter, \
 263              *__node = (type *)rb_tree_first(T), \
 264              *__next = (type *)rb_node_next_or_null((struct rb_node *)__node); \
 265         __node != NULL && \
 266         (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
 267         __node = __next, \
 268         __next = (type *)rb_node_next_or_null((struct rb_node *)__node))
 269
 270 /** Iterate over the nodes in the tree in reverse
 271  *
 272  * \param   type    The type of the containing data structure
 273  *
 274  * \param   node    The variable name for current node in the iteration;
 275  *                  this will be declared as a pointer to \p type
 276  *
 277  * \param   T       The red-black tree
 278  *
 279  * \param   field   The rb_node field in containing data structure
 280  */
 281 #define rb_tree_foreach_rev(type, iter, T, field) \
 282    for (type *iter, *__node = (type *)rb_tree_last(T); \
 283         __node != NULL && \
 284         (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
 285         __node = (type *)rb_node_prev((struct rb_node *)__node))
 286
 287 /** Iterate over the nodes in the tree in reverse, allowing the current node to be freed
 288  *
 289  * \param   type    The type of the containing data structure
 290  *
 291  * \param   node    The variable name for current node in the iteration;
 292  *                  this will be declared as a pointer to \p type
 293  *
 294  * \param   T       The red-black tree
 295  *
 296  * \param   field   The rb_node field in containing data structure
 297  */
 298 #define rb_tree_foreach_rev_safe(type, iter, T, field) \
 299    for (type *iter, \
 300              *__node = (type *)rb_tree_last(T), \
 301              *__prev = (type *)rb_node_prev_or_null((struct rb_node *)__node); \
 302         __node != NULL && \
 303         (iter = rb_node_data(type, (struct rb_node *)__node, field), true); \
 304         __node = __prev, \
 305         __prev = (type *)rb_node_prev_or_null((struct rb_node *)__node))
 306
 307 /** Validate a red-black tree
 308  *
 309  * This function walks the tree and validates that this is a valid red-
 310  * black tree.  If anything is wrong, it will assert-fail.
 311  */
 312 void rb_tree_validate(struct rb_tree *T);
 313
 314 #endif /* RB_TREE_H */