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       ASSERT (z->key != x->key);
 211       if (ltfn (z->key, x->key))
 212         xi = x->left;
 213       else
 214         xi = x->right;
 215     }
 216   yi = rb_node_index (rt, y);
 217   z->parent = yi;
 218   if (yi == RBTREE_TNIL_INDEX)
 219     rt->root = rb_node_index (rt, z);
 220   else if (ltfn (z->key, y->key))
 221     y->left = rb_node_index (rt, z);
 222   else
 223     y->right = rb_node_index (rt, z);
 224
 225   rb_tree_fixup_inline (rt, y, z);
 226
 227   return rb_node_index (rt, z);
 228 }
 229
 230 rb_node_t *
 231 rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
 232 {
 233   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 234     if (key < x->key)
 235       x = rb_node_left (rt, x);
 236     else
 237       x = rb_node_right (rt, x);
 238   return x;
 239 }
 240
 241 rb_node_t *
 242 rb_tree_search_subtree_custom (rb_tree_t * rt, rb_node_t * x, u32 key,
 243                                rb_tree_lt_fn ltfn)
 244 {
 245   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 246     if (ltfn (key, x->key))
 247       x = rb_node_left (rt, x);
 248     else
 249       x = rb_node_right (rt, x);
 250   return x;
 251 }
 252
 253 rb_node_t *
 254 rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x)
 255 {
 256   while (x->left != RBTREE_TNIL_INDEX)
 257     x = rb_node_left (rt, x);
 258   return x;
 259 }
 260
 261 rb_node_t *
 262 rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x)
 263 {
 264   while (x->right != RBTREE_TNIL_INDEX)
 265     x = rb_node_right (rt, x);
 266   return x;
 267 }
 268
 269 rb_node_t *
 270 rb_tree_successor (rb_tree_t * rt, rb_node_t * x)
 271 {
 272   rb_node_t *y;
 273
 274   if (x->right != RBTREE_TNIL_INDEX)
 275     return rb_tree_min_subtree (rt, rb_node_right (rt, x));
 276
 277   y = rb_node_parent (rt, x);
 278   while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x))
 279     {
 280       x = y;
 281       y = rb_node_parent (rt, y);
 282     }
 283   return y;
 284 }
 285
 286 rb_node_t *
 287 rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x)
 288 {
 289   rb_node_t *y;
 290
 291   if (x->left != RBTREE_TNIL_INDEX)
 292     return rb_tree_max_subtree (rt, rb_node_left (rt, x));
 293
 294   y = rb_node_parent (rt, x);
 295   while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x))
 296     {
 297       x = y;
 298       y = rb_node_parent (rt, y);
 299     }
 300   return y;
 301 }
 302
 303 static inline void
 304 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
 305 {
 306   rb_node_t *up;
 307
 308   up = rb_node_parent (rt, u);
 309   if (u->parent == RBTREE_TNIL_INDEX)
 310     rt->root = rb_node_index (rt, v);
 311   else if (rb_node_index (rt, u) == up->left)
 312     up->left = rb_node_index (rt, v);
 313   else
 314     up->right = rb_node_index (rt, v);
 315   v->parent = u->parent;
 316 }
 317
 318 static void
 319 rb_tree_del_node_internal (rb_tree_t * rt, rb_node_t * z)
 320 {
 321   rb_node_color_t y_original_color;
 322   rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr;
 323   rb_node_index_t xi, yi;
 324
 325   y = z;
 326   y_original_color = y->color;
 327
 328   if (z->left == RBTREE_TNIL_INDEX)
 329     {
 330       x = rb_node_right (rt, z);
 331       rb_tree_transplant (rt, z, x);
 332     }
 333   else if (z->right == RBTREE_TNIL_INDEX)
 334     {
 335       x = rb_node_left (rt, z);
 336       rb_tree_transplant (rt, z, x);
 337     }
 338   else
 339     {
 340       y = rb_tree_min_subtree (rt, rb_node_right (rt, z));
 341       y_original_color = y->color;
 342       x = rb_node_right (rt, y);
 343       yi = rb_node_index (rt, y);
 344       if (y->parent == rb_node_index (rt, z))
 345         x->parent = yi;
 346       else
 347         {
 348           rb_tree_transplant (rt, y, x);
 349           y->right = z->right;
 350           yr = rb_node_right (rt, y);
 351           yr->parent = yi;
 352         }
 353       rb_tree_transplant (rt, z, y);
 354       y->left = z->left;
 355       yl = rb_node_left (rt, y);
 356       yl->parent = yi;
 357       y->color = z->color;
 358     }
 359
 360   if (y_original_color == RBTREE_RED)
 361     return;
 362
 363   /* Tree fixup needed */
 364
 365   xi = rb_node_index (rt, x);
 366   while (xi != rt->root && x->color == RBTREE_BLACK)
 367     {
 368       xp = rb_node_parent (rt, x);
 369       if (xi == xp->left)
 370         {
 371           w = rb_node_right (rt, xp);
 372           if (w->color == RBTREE_RED)
 373             {
 374               w->color = RBTREE_BLACK;
 375               xp->color = RBTREE_RED;
 376               rb_tree_rotate_left (rt, xp);
 377               w = rb_node_right (rt, xp);
 378             }
 379           wl = rb_node_left (rt, w);
 380           wr = rb_node_right (rt, w);
 381           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 382             {
 383               if (!rb_node_is_tnil (rt, w))
 384                 w->color = RBTREE_RED;
 385               x = xp;
 386             }
 387           else
 388             {
 389               if (wr->color == RBTREE_BLACK)
 390                 {
 391                   wl->color = RBTREE_BLACK;
 392                   w->color = RBTREE_RED;
 393                   rb_tree_rotate_right (rt, w);
 394                   w = rb_node_right (rt, xp);
 395                   wr = rb_node_right (rt, w);
 396                 }
 397               w->color = xp->color;
 398               xp->color = RBTREE_BLACK;
 399               wr->color = RBTREE_BLACK;
 400               rb_tree_rotate_left (rt, xp);
 401               x = rb_node (rt, rt->root);
 402             }
 403         }
 404       else
 405         {
 406           w = rb_node_left (rt, xp);
 407           if (w->color == RBTREE_RED)
 408             {
 409               w->color = RBTREE_BLACK;
 410               xp->color = RBTREE_RED;
 411               rb_tree_rotate_right (rt, xp);
 412               w = rb_node_left (rt, xp);
 413             }
 414           wl = rb_node_left (rt, w);
 415           wr = rb_node_right (rt, w);
 416           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 417             {
 418               if (!rb_node_is_tnil (rt, w))
 419                 w->color = RBTREE_RED;
 420               x = xp;
 421             }
 422           else
 423             {
 424               if (wl->color == RBTREE_BLACK)
 425                 {
 426                   wr->color = RBTREE_BLACK;
 427                   w->color = RBTREE_RED;
 428                   rb_tree_rotate_left (rt, w);
 429                   w = rb_node_left (rt, xp);
 430                   wl = rb_node_left (rt, w);
 431                 }
 432               w->color = xp->color;
 433               xp->color = RBTREE_BLACK;
 434               wl->color = RBTREE_BLACK;
 435               rb_tree_rotate_right (rt, xp);
 436               x = rb_node (rt, rt->root);
 437             }
 438         }
 439       xi = rb_node_index (rt, x);
 440     }
 441   x->color = RBTREE_BLACK;
 442 }
 443
 444 void
 445 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
 446 {
 447   rb_tree_del_node_internal (rt, z);
 448   pool_put (rt->nodes, z);
 449 }
 450
 451 void
 452 rb_tree_del (rb_tree_t * rt, u32 key)
 453 {
 454   rb_node_t *n;
 455   n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key);
 456   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 457     rb_tree_del_node (rt, n);
 458 }
 459
 460 void
 461 rb_tree_del_custom (rb_tree_t * rt, u32 key, rb_tree_lt_fn ltfn)
 462 {
 463   rb_node_t *n;
 464   n = rb_tree_search_subtree_custom (rt, rb_node (rt, rt->root), key, ltfn);
 465   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 466     rb_tree_del_node (rt, n);
 467 }
 468
 469 u32
 470 rb_tree_n_nodes (rb_tree_t * rt)
 471 {
 472   return pool_elts (rt->nodes);
 473 }
 474
 475 void
 476 rb_tree_free_nodes (rb_tree_t * rt)
 477 {
 478   pool_free (rt->nodes);
 479   rt->root = RBTREE_TNIL_INDEX;
 480 }
 481
 482 void
 483 rb_tree_init (rb_tree_t * rt)
 484 {
 485   rb_node_t *tnil;
 486
 487   rt->nodes = 0;
 488   rt->root = RBTREE_TNIL_INDEX;
 489
 490   /* By convention first node, index 0, is the T.nil sentinel */
 491   pool_get_zero (rt->nodes, tnil);
 492   tnil->color = RBTREE_BLACK;
 493 }
 494
 495 int
 496 rb_tree_is_init (rb_tree_t * rt)
 497 {
 498   if (pool_elts (rt->nodes) == 0)
 499     return 0;
 500   return 1;
 501 }
 502
 503 /*
 504  * fd.io coding-style-patch-verification: ON
 505  *
 506  * Local Variables:
 507  * eval: (c-set-style "gnu")
 508  * End:
 509  */