src/vppinfra/rbtree.c

   1 /*
   2  * Copyright (c) 2019 Cisco and/or its affiliates.
   3  * Licensed under the Apache License, Version 2.0 (the "License");
   4  * you may not use this file except in compliance with the License.
   5  * You may obtain a copy of the License at:
   6  *
   7  *     http://www.apache.org/licenses/LICENSE-2.0
   8  *
   9  * Unless required by applicable law or agreed to in writing, software
  10  * distributed under the License is distributed on an "AS IS" BASIS,
  11  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12  * See the License for the specific language governing permissions and
  13  * limitations under the License.
  14  */
  15 /*
  16  * Algorithm from:
  17  * Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2009).
  18  * Introduction to algorithms. MIT press, 3rd Edition, Ch. 13
  19  */
  20
  21 #include <vppinfra/rbtree.h>
  22
  23 static inline void
  24 rb_tree_rotate_left (rb_tree_t * rt, rb_node_t * x)
  25 {
  26   rb_node_t *y, *tmp, *xp;
  27   rb_node_index_t xi, yi;
  28
  29   xi = rb_node_index (rt, x);
  30   yi = x->right;
  31   y = rb_node_right (rt, x);
  32   x->right = y->left;
  33   if (y->left != RBTREE_TNIL_INDEX)
  34     {
  35       tmp = rb_node_left (rt, y);
  36       tmp->parent = xi;
  37     }
  38   xp = rb_node_parent (rt, x);
  39   y->parent = x->parent;
  40   if (x->parent == RBTREE_TNIL_INDEX)
  41     rt->root = rb_node_index (rt, y);
  42   else if (xp->left == xi)
  43     xp->left = yi;
  44   else
  45     xp->right = yi;
  46   y->left = xi;
  47   x->parent = yi;
  48 }
  49
  50 static inline void
  51 rb_tree_rotate_right (rb_tree_t * rt, rb_node_t * y)
  52 {
  53   rb_node_index_t yi, xi;
  54   rb_node_t *x, *yp;
  55
  56   yi = rb_node_index (rt, y);
  57   xi = y->left;
  58   x = rb_node_left (rt, y);
  59   y->left = x->right;
  60   if (x->right != RBTREE_TNIL_INDEX)
  61     {
  62       rb_node_t *tmp = rb_node_right (rt, x);
  63       tmp->parent = yi;
  64     }
  65   yp = rb_node_parent (rt, y);
  66   x->parent = y->parent;
  67   if (y->parent == RBTREE_TNIL_INDEX)
  68     rt->root = rb_node_index (rt, x);
  69   else if (yp->right == yi)
  70     yp->right = xi;
  71   else
  72     yp->left = xi;
  73   x->right = yi;
  74   y->parent = xi;
  75 }
  76
  77 static void
  78 rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
  79 {
  80   rb_node_index_t yi = 0, xi = rt->root, zi;
  81   rb_node_t *y, *zpp, *x, *zp;
  82
  83   y = rb_node (rt, RBTREE_TNIL_INDEX);
  84   while (xi != RBTREE_TNIL_INDEX)
  85     {
  86       x = rb_node (rt, xi);
  87       y = x;
  88       if (z->key < x->key)
  89         xi = x->left;
  90       else
  91         xi = x->right;
  92     }
  93   yi = rb_node_index (rt, y);
  94   z->parent = yi;
  95   if (yi == RBTREE_TNIL_INDEX)
  96     rt->root = rb_node_index (rt, z);
  97   else if (z->key < y->key)
  98     y->left = rb_node_index (rt, z);
  99   else
 100     y->right = rb_node_index (rt, z);
 101
 102   /* Tree fixup stage */
 103   while (y->color == RBTREE_RED)
 104     {
 105       zi = rb_node_index (rt, z);
 106       zp = rb_node_parent (rt, z);
 107       zpp = rb_node_parent (rt, zp);
 108       if (z->parent == zpp->left)
 109         {
 110           y = rb_node_right (rt, zpp);
 111           if (y->color == RBTREE_RED)
 112             {
 113               zp->color = RBTREE_BLACK;
 114               y->color = RBTREE_BLACK;
 115               zpp->color = RBTREE_RED;
 116               z = zpp;
 117             }
 118           else
 119             {
 120               if (zi == zp->right)
 121                 {
 122                   z = zp;
 123                   rb_tree_rotate_left (rt, z);
 124                   zp = rb_node (rt, z->parent);
 125                   zpp = rb_node (rt, zp->parent);
 126                 }
 127               zp->color = RBTREE_BLACK;
 128               zpp->color = RBTREE_RED;
 129               rb_tree_rotate_right (rt, zpp);
 130             }
 131         }
 132       else
 133         {
 134           y = rb_node (rt, zpp->left);
 135           if (y->color == RBTREE_RED)
 136             {
 137               zp->color = RBTREE_BLACK;
 138               y->color = RBTREE_BLACK;
 139               zpp->color = RBTREE_RED;
 140               z = zpp;
 141             }
 142           else
 143             {
 144               if (zi == zp->left)
 145                 {
 146                   z = zp;
 147                   rb_tree_rotate_right (rt, z);
 148                   zp = rb_node (rt, z->parent);
 149                   zpp = rb_node (rt, zp->parent);
 150                 }
 151               zp->color = RBTREE_BLACK;
 152               zpp->color = RBTREE_RED;
 153               rb_tree_rotate_left (rt, zpp);
 154             }
 155         }
 156     }
 157   rb_node (rt, rt->root)->color = RBTREE_BLACK;
 158 }
 159
 160 rb_node_index_t
 161 rb_tree_add (rb_tree_t * rt, u32 key)
 162 {
 163   rb_node_t *n;
 164
 165   pool_get_zero (rt->nodes, n);
 166   n->key = key;
 167   n->color = RBTREE_RED;
 168   rb_tree_insert (rt, n);
 169   return rb_node_index (rt, n);
 170 }
 171
 172 rb_node_index_t
 173 rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque)
 174 {
 175   rb_node_t *n;
 176
 177   pool_get_zero (rt->nodes, n);
 178   n->key = key;
 179   n->color = RBTREE_RED;
 180   n->opaque = opaque;
 181   rb_tree_insert (rt, n);
 182   return rb_node_index (rt, n);
 183 }
 184
 185 rb_node_t *
 186 rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
 187 {
 188   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 189     if (key < x->key)
 190       x = rb_node_left (rt, x);
 191     else
 192       x = rb_node_right (rt, x);
 193   return x;
 194 }
 195
 196 rb_node_t *
 197 rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x)
 198 {
 199   while (x->left != RBTREE_TNIL_INDEX)
 200     x = rb_node_left (rt, x);
 201   return x;
 202 }
 203
 204 rb_node_t *
 205 rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x)
 206 {
 207   while (x->right != RBTREE_TNIL_INDEX)
 208     x = rb_node_right (rt, x);
 209   return x;
 210 }
 211
 212 rb_node_t *
 213 rb_tree_successor (rb_tree_t * rt, rb_node_t * x)
 214 {
 215   rb_node_t *y;
 216
 217   if (x->right != RBTREE_TNIL_INDEX)
 218     return rb_tree_min_subtree (rt, rb_node_right (rt, x));
 219
 220   y = rb_node_parent (rt, x);
 221   while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x))
 222     {
 223       x = y;
 224       y = rb_node_parent (rt, y);
 225     }
 226   return y;
 227 }
 228
 229 rb_node_t *
 230 rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x)
 231 {
 232   rb_node_t *y;
 233
 234   if (x->left != RBTREE_TNIL_INDEX)
 235     return rb_tree_max_subtree (rt, rb_node_left (rt, x));
 236
 237   y = rb_node_parent (rt, x);
 238   while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x))
 239     {
 240       x = y;
 241       y = rb_node_parent (rt, y);
 242     }
 243   return y;
 244 }
 245
 246 static inline void
 247 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
 248 {
 249   rb_node_t *up;
 250
 251   up = rb_node_parent (rt, u);
 252   if (u->parent == RBTREE_TNIL_INDEX)
 253     rt->root = rb_node_index (rt, v);
 254   else if (rb_node_index (rt, u) == up->left)
 255     up->left = rb_node_index (rt, v);
 256   else
 257     up->right = rb_node_index (rt, v);
 258   v->parent = u->parent;
 259 }
 260
 261 void
 262 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
 263 {
 264   rb_node_color_t y_original_color;
 265   rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr;
 266   rb_node_index_t xi, yi;
 267
 268   y = z;
 269   y_original_color = y->color;
 270
 271   if (z->left == RBTREE_TNIL_INDEX)
 272     {
 273       x = rb_node_right (rt, z);
 274       rb_tree_transplant (rt, z, x);
 275     }
 276   else if (z->right == RBTREE_TNIL_INDEX)
 277     {
 278       x = rb_node_left (rt, z);
 279       rb_tree_transplant (rt, z, x);
 280     }
 281   else
 282     {
 283       y = rb_tree_min_subtree (rt, rb_node_right (rt, z));
 284       y_original_color = y->color;
 285       x = rb_node_right (rt, y);
 286       yi = rb_node_index (rt, y);
 287       if (y->parent == rb_node_index (rt, z))
 288         x->parent = yi;
 289       else
 290         {
 291           rb_tree_transplant (rt, y, rb_node_right (rt, y));
 292           y->right = z->right;
 293           yr = rb_node_right (rt, y);
 294           yr->parent = yi;
 295         }
 296       rb_tree_transplant (rt, z, y);
 297       y->left = z->left;
 298       yl = rb_node_left (rt, y);
 299       yl->parent = yi;
 300       y->color = z->color;
 301     }
 302
 303   if (y_original_color == RBTREE_RED)
 304     return;
 305
 306   /* Tree fixup needed */
 307
 308   xi = rb_node_index (rt, x);
 309   while (xi != rt->root && x->color == RBTREE_BLACK)
 310     {
 311       xp = rb_node_parent (rt, x);
 312       if (xi == xp->left)
 313         {
 314           w = rb_node_right (rt, xp);
 315           if (w->color == RBTREE_RED)
 316             {
 317               w->color = RBTREE_BLACK;
 318               xp->color = RBTREE_RED;
 319               rb_tree_rotate_left (rt, xp);
 320               w = rb_node_right (rt, xp);
 321             }
 322           wl = rb_node_left (rt, w);
 323           wr = rb_node_right (rt, w);
 324           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 325             {
 326               w->color = RBTREE_RED;
 327               x = xp;
 328             }
 329           else
 330             {
 331               if (wr->color == RBTREE_BLACK)
 332                 {
 333                   wl->color = RBTREE_BLACK;
 334                   w->color = RBTREE_RED;
 335                   rb_tree_rotate_right (rt, w);
 336                   w = rb_node_right (rt, xp);
 337                 }
 338               w->color = xp->color;
 339               xp->color = RBTREE_BLACK;
 340               wr->color = RBTREE_BLACK;
 341               rb_tree_rotate_left (rt, xp);
 342               x = rb_node (rt, rt->root);
 343             }
 344         }
 345       else
 346         {
 347           w = rb_node_left (rt, xp);
 348           if (w->color == RBTREE_RED)
 349             {
 350               w->color = RBTREE_BLACK;
 351               xp->color = RBTREE_RED;
 352               rb_tree_rotate_right (rt, xp);
 353               w = rb_node_left (rt, xp);
 354             }
 355           wl = rb_node_left (rt, w);
 356           wr = rb_node_right (rt, w);
 357           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 358             {
 359               w->color = RBTREE_RED;
 360               x = xp;
 361             }
 362           else
 363             {
 364               if (wl->color == RBTREE_BLACK)
 365                 {
 366                   wr->color = RBTREE_BLACK;
 367                   w->color = RBTREE_RED;
 368                   rb_tree_rotate_left (rt, w);
 369                   w = rb_node_left (rt, xp);
 370                 }
 371               w->color = xp->color;
 372               xp->color = RBTREE_BLACK;
 373               wl->color = RBTREE_BLACK;
 374               rb_tree_rotate_right (rt, xp);
 375               x = rb_node (rt, rt->root);
 376             }
 377         }
 378       xi = rb_node_index (rt, x);
 379     }
 380   x->color = RBTREE_BLACK;
 381 }
 382
 383 void
 384 rb_tree_del (rb_tree_t * rt, u32 key)
 385 {
 386   rb_node_t *n;
 387   n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key);
 388   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 389     {
 390       rb_tree_del_node (rt, n);
 391       pool_put (rt->nodes, n);
 392     }
 393 }
 394
 395 u32
 396 rb_tree_n_nodes (rb_tree_t * rt)
 397 {
 398   return pool_elts (rt->nodes);
 399 }
 400
 401 void
 402 rb_tree_free_nodes (rb_tree_t * rt)
 403 {
 404   rb_node_t *n;
 405   pool_flush (n, rt->nodes,;);
 406 }
 407
 408 void
 409 rb_tree_init (rb_tree_t * rt)
 410 {
 411   rb_node_t *tnil;
 412
 413   rt->nodes = 0;
 414   rt->root = RBTREE_TNIL_INDEX;
 415
 416   /* By convention first node, index 0, is the T.nil sentinel */
 417   pool_get_zero (rt->nodes, tnil);
 418   tnil->color = RBTREE_BLACK;
 419 }
 420
 421 /*
 422  * fd.io coding-style-patch-verification: ON
 423  *
 424  * Local Variables:
 425  * eval: (c-set-style "gnu")
 426  * End:
 427  */