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 inline void
  78 rb_tree_fixup_inline (rb_tree_t * rt, rb_node_t * y, rb_node_t * z)
  79 {
  80   rb_node_t *zpp, *zp;
  81   rb_node_index_t zi;
  82
  83   while (y->color == RBTREE_RED)
  84     {
  85       zi = rb_node_index (rt, z);
  86       zp = rb_node_parent (rt, z);
  87       zpp = rb_node_parent (rt, zp);
  88       if (z->parent == zpp->left)
  89         {
  90           y = rb_node_right (rt, zpp);
  91           if (y->color == RBTREE_RED)
  92             {
  93               zp->color = RBTREE_BLACK;
  94               y->color = RBTREE_BLACK;
  95               zpp->color = RBTREE_RED;
  96               z = zpp;
  97             }
  98           else
  99             {
 100               if (zi == zp->right)
 101                 {
 102                   z = zp;
 103                   rb_tree_rotate_left (rt, z);
 104                   zp = rb_node (rt, z->parent);
 105                   zpp = rb_node (rt, zp->parent);
 106                 }
 107               zp->color = RBTREE_BLACK;
 108               zpp->color = RBTREE_RED;
 109               rb_tree_rotate_right (rt, zpp);
 110             }
 111         }
 112       else
 113         {
 114           y = rb_node (rt, zpp->left);
 115           if (y->color == RBTREE_RED)
 116             {
 117               zp->color = RBTREE_BLACK;
 118               y->color = RBTREE_BLACK;
 119               zpp->color = RBTREE_RED;
 120               z = zpp;
 121             }
 122           else
 123             {
 124               if (zi == zp->left)
 125                 {
 126                   z = zp;
 127                   rb_tree_rotate_right (rt, z);
 128                   zp = rb_node (rt, z->parent);
 129                   zpp = rb_node (rt, zp->parent);
 130                 }
 131               zp->color = RBTREE_BLACK;
 132               zpp->color = RBTREE_RED;
 133               rb_tree_rotate_left (rt, zpp);
 134             }
 135         }
 136     }
 137   rb_node (rt, rt->root)->color = RBTREE_BLACK;
 138 }
 139
 140 static void
 141 rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
 142 {
 143   rb_node_index_t yi = 0, xi = rt->root;
 144   rb_node_t *y, *x;
 145
 146   y = rb_node (rt, RBTREE_TNIL_INDEX);
 147   while (xi != RBTREE_TNIL_INDEX)
 148     {
 149       x = rb_node (rt, xi);
 150       y = x;
 151       if (z->key < x->key)
 152         xi = x->left;
 153       else
 154         xi = x->right;
 155     }
 156   yi = rb_node_index (rt, y);
 157   z->parent = yi;
 158   if (yi == RBTREE_TNIL_INDEX)
 159     rt->root = rb_node_index (rt, z);
 160   else if (z->key < y->key)
 161     y->left = rb_node_index (rt, z);
 162   else
 163     y->right = rb_node_index (rt, z);
 164
 165   /* Tree fixup stage */
 166   rb_tree_fixup_inline (rt, y, z);
 167 }
 168
 169 rb_node_index_t
 170 rb_tree_add (rb_tree_t * rt, u32 key)
 171 {
 172   rb_node_t *n;
 173
 174   pool_get_zero (rt->nodes, n);
 175   n->key = key;
 176   n->color = RBTREE_RED;
 177   rb_tree_insert (rt, n);
 178   return rb_node_index (rt, n);
 179 }
 180
 181 rb_node_index_t
 182 rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque)
 183 {
 184   rb_node_t *n;
 185
 186   pool_get_zero (rt->nodes, n);
 187   n->key = key;
 188   n->color = RBTREE_RED;
 189   n->opaque = opaque;
 190   rb_tree_insert (rt, n);
 191   return rb_node_index (rt, n);
 192 }
 193
 194 rb_node_index_t
 195 rb_tree_add_custom (rb_tree_t * rt, u32 key, uword opaque, rb_tree_lt_fn ltfn)
 196 {
 197   rb_node_index_t yi = 0, xi = rt->root;
 198   rb_node_t *z, *y, *x;
 199
 200   pool_get_zero (rt->nodes, z);
 201   z->key = key;
 202   z->color = RBTREE_RED;
 203   z->opaque = opaque;
 204
 205   y = rb_node (rt, RBTREE_TNIL_INDEX);
 206   while (xi != RBTREE_TNIL_INDEX)
 207     {
 208       x = rb_node (rt, xi);
 209       y = x;
 210       if (ltfn (z->key, x->key))
 211         xi = x->left;
 212       else
 213         xi = x->right;
 214     }
 215   yi = rb_node_index (rt, y);
 216   z->parent = yi;
 217   if (yi == RBTREE_TNIL_INDEX)
 218     rt->root = rb_node_index (rt, z);
 219   else if (ltfn (z->key, y->key))
 220     y->left = rb_node_index (rt, z);
 221   else
 222     y->right = rb_node_index (rt, z);
 223
 224   rb_tree_fixup_inline (rt, y, z);
 225
 226   return rb_node_index (rt, z);
 227 }
 228
 229 rb_node_t *
 230 rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
 231 {
 232   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 233     if (key < x->key)
 234       x = rb_node_left (rt, x);
 235     else
 236       x = rb_node_right (rt, x);
 237   return x;
 238 }
 239
 240 rb_node_t *
 241 rb_tree_search_subtree_custom (rb_tree_t * rt, rb_node_t * x, u32 key,
 242                                rb_tree_lt_fn ltfn)
 243 {
 244   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 245     if (ltfn (key, x->key))
 246       x = rb_node_left (rt, x);
 247     else
 248       x = rb_node_right (rt, x);
 249   return x;
 250 }
 251
 252 rb_node_t *
 253 rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x)
 254 {
 255   while (x->left != RBTREE_TNIL_INDEX)
 256     x = rb_node_left (rt, x);
 257   return x;
 258 }
 259
 260 rb_node_t *
 261 rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x)
 262 {
 263   while (x->right != RBTREE_TNIL_INDEX)
 264     x = rb_node_right (rt, x);
 265   return x;
 266 }
 267
 268 rb_node_t *
 269 rb_tree_successor (rb_tree_t * rt, rb_node_t * x)
 270 {
 271   rb_node_t *y;
 272
 273   if (x->right != RBTREE_TNIL_INDEX)
 274     return rb_tree_min_subtree (rt, rb_node_right (rt, x));
 275
 276   y = rb_node_parent (rt, x);
 277   while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x))
 278     {
 279       x = y;
 280       y = rb_node_parent (rt, y);
 281     }
 282   return y;
 283 }
 284
 285 rb_node_t *
 286 rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x)
 287 {
 288   rb_node_t *y;
 289
 290   if (x->left != RBTREE_TNIL_INDEX)
 291     return rb_tree_max_subtree (rt, rb_node_left (rt, x));
 292
 293   y = rb_node_parent (rt, x);
 294   while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x))
 295     {
 296       x = y;
 297       y = rb_node_parent (rt, y);
 298     }
 299   return y;
 300 }
 301
 302 static inline void
 303 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
 304 {
 305   rb_node_t *up;
 306
 307   up = rb_node_parent (rt, u);
 308   if (u->parent == RBTREE_TNIL_INDEX)
 309     rt->root = rb_node_index (rt, v);
 310   else if (rb_node_index (rt, u) == up->left)
 311     up->left = rb_node_index (rt, v);
 312   else
 313     up->right = rb_node_index (rt, v);
 314   v->parent = u->parent;
 315 }
 316
 317 void
 318 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
 319 {
 320   rb_node_color_t y_original_color;
 321   rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr;
 322   rb_node_index_t xi, yi;
 323
 324   y = z;
 325   y_original_color = y->color;
 326
 327   if (z->left == RBTREE_TNIL_INDEX)
 328     {
 329       x = rb_node_right (rt, z);
 330       rb_tree_transplant (rt, z, x);
 331     }
 332   else if (z->right == RBTREE_TNIL_INDEX)
 333     {
 334       x = rb_node_left (rt, z);
 335       rb_tree_transplant (rt, z, x);
 336     }
 337   else
 338     {
 339       y = rb_tree_min_subtree (rt, rb_node_right (rt, z));
 340       y_original_color = y->color;
 341       x = rb_node_right (rt, y);
 342       yi = rb_node_index (rt, y);
 343       if (y->parent == rb_node_index (rt, z))
 344         x->parent = yi;
 345       else
 346         {
 347           rb_tree_transplant (rt, y, x);
 348           y->right = z->right;
 349           yr = rb_node_right (rt, y);
 350           yr->parent = yi;
 351         }
 352       rb_tree_transplant (rt, z, y);
 353       y->left = z->left;
 354       yl = rb_node_left (rt, y);
 355       yl->parent = yi;
 356       y->color = z->color;
 357     }
 358
 359   if (y_original_color == RBTREE_RED)
 360     return;
 361
 362   /* Tree fixup needed */
 363
 364   xi = rb_node_index (rt, x);
 365   while (xi != rt->root && x->color == RBTREE_BLACK)
 366     {
 367       xp = rb_node_parent (rt, x);
 368       if (xi == xp->left)
 369         {
 370           w = rb_node_right (rt, xp);
 371           if (w->color == RBTREE_RED)
 372             {
 373               w->color = RBTREE_BLACK;
 374               xp->color = RBTREE_RED;
 375               rb_tree_rotate_left (rt, xp);
 376               w = rb_node_right (rt, xp);
 377             }
 378           wl = rb_node_left (rt, w);
 379           wr = rb_node_right (rt, w);
 380           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 381             {
 382               if (!rb_node_is_tnil (rt, w))
 383                 w->color = RBTREE_RED;
 384               x = xp;
 385             }
 386           else
 387             {
 388               if (wr->color == RBTREE_BLACK)
 389                 {
 390                   wl->color = RBTREE_BLACK;
 391                   w->color = RBTREE_RED;
 392                   rb_tree_rotate_right (rt, w);
 393                   w = rb_node_right (rt, xp);
 394                   wr = rb_node_right (rt, w);
 395                 }
 396               w->color = xp->color;
 397               xp->color = RBTREE_BLACK;
 398               wr->color = RBTREE_BLACK;
 399               rb_tree_rotate_left (rt, xp);
 400               x = rb_node (rt, rt->root);
 401             }
 402         }
 403       else
 404         {
 405           w = rb_node_left (rt, xp);
 406           if (w->color == RBTREE_RED)
 407             {
 408               w->color = RBTREE_BLACK;
 409               xp->color = RBTREE_RED;
 410               rb_tree_rotate_right (rt, xp);
 411               w = rb_node_left (rt, xp);
 412             }
 413           wl = rb_node_left (rt, w);
 414           wr = rb_node_right (rt, w);
 415           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 416             {
 417               if (!rb_node_is_tnil (rt, w))
 418                 w->color = RBTREE_RED;
 419               x = xp;
 420             }
 421           else
 422             {
 423               if (wl->color == RBTREE_BLACK)
 424                 {
 425                   wr->color = RBTREE_BLACK;
 426                   w->color = RBTREE_RED;
 427                   rb_tree_rotate_left (rt, w);
 428                   w = rb_node_left (rt, xp);
 429                   wl = rb_node_left (rt, w);
 430                 }
 431               w->color = xp->color;
 432               xp->color = RBTREE_BLACK;
 433               wl->color = RBTREE_BLACK;
 434               rb_tree_rotate_right (rt, xp);
 435               x = rb_node (rt, rt->root);
 436             }
 437         }
 438       xi = rb_node_index (rt, x);
 439     }
 440   x->color = RBTREE_BLACK;
 441 }
 442
 443 void
 444 rb_tree_del (rb_tree_t * rt, u32 key)
 445 {
 446   rb_node_t *n;
 447   n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key);
 448   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 449     {
 450       rb_tree_del_node (rt, n);
 451       pool_put (rt->nodes, n);
 452     }
 453 }
 454
 455 void
 456 rb_tree_del_custom (rb_tree_t * rt, u32 key, rb_tree_lt_fn ltfn)
 457 {
 458   rb_node_t *n;
 459   n = rb_tree_search_subtree_custom (rt, rb_node (rt, rt->root), key, ltfn);
 460   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 461     {
 462       rb_tree_del_node (rt, n);
 463       pool_put (rt->nodes, n);
 464     }
 465 }
 466
 467 u32
 468 rb_tree_n_nodes (rb_tree_t * rt)
 469 {
 470   return pool_elts (rt->nodes);
 471 }
 472
 473 void
 474 rb_tree_free_nodes (rb_tree_t * rt)
 475 {
 476   pool_free (rt->nodes);
 477   rt->root = RBTREE_TNIL_INDEX;
 478 }
 479
 480 void
 481 rb_tree_init (rb_tree_t * rt)
 482 {
 483   rb_node_t *tnil;
 484
 485   rt->nodes = 0;
 486   rt->root = RBTREE_TNIL_INDEX;
 487
 488   /* By convention first node, index 0, is the T.nil sentinel */
 489   pool_get_zero (rt->nodes, tnil);
 490   tnil->color = RBTREE_BLACK;
 491 }
 492
 493 /*
 494  * fd.io coding-style-patch-verification: ON
 495  *
 496  * Local Variables:
 497  * eval: (c-set-style "gnu")
 498  * End:
 499  */