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 rb_node_t *
  24 rb_node_right (rb_tree_t * rt, rb_node_t * n)
  25 {
  26   return pool_elt_at_index (rt->nodes, n->right);
  27 }
  28
  29 static inline rb_node_t *
  30 rb_node_left (rb_tree_t * rt, rb_node_t * n)
  31 {
  32   return pool_elt_at_index (rt->nodes, n->left);
  33 }
  34
  35 static inline rb_node_t *
  36 rb_node_parent (rb_tree_t * rt, rb_node_t * n)
  37 {
  38   return pool_elt_at_index (rt->nodes, n->parent);
  39 }
  40
  41 static inline void
  42 rb_tree_rotate_left (rb_tree_t * rt, rb_node_t * x)
  43 {
  44   rb_node_t *y, *tmp, *xp;
  45   rb_node_index_t xi, yi;
  46
  47   xi = rb_node_index (rt, x);
  48   yi = x->right;
  49   y = rb_node_right (rt, x);
  50   x->right = y->left;
  51   if (y->left != RBTREE_TNIL_INDEX)
  52     {
  53       tmp = rb_node_left (rt, y);
  54       tmp->parent = xi;
  55     }
  56   xp = rb_node_parent (rt, x);
  57   y->parent = x->parent;
  58   if (x->parent == RBTREE_TNIL_INDEX)
  59     rt->root = rb_node_index (rt, y);
  60   else if (xp->left == xi)
  61     xp->left = yi;
  62   else
  63     xp->right = yi;
  64   y->left = xi;
  65   x->parent = yi;
  66 }
  67
  68 static inline void
  69 rb_tree_rotate_right (rb_tree_t * rt, rb_node_t * y)
  70 {
  71   rb_node_index_t yi, xi;
  72   rb_node_t *x, *yp;
  73
  74   yi = rb_node_index (rt, y);
  75   xi = y->left;
  76   x = rb_node_left (rt, y);
  77   y->left = x->right;
  78   if (x->right != RBTREE_TNIL_INDEX)
  79     {
  80       rb_node_t *tmp = rb_node_right (rt, x);
  81       tmp->parent = yi;
  82     }
  83   yp = rb_node_parent (rt, y);
  84   x->parent = y->parent;
  85   if (y->parent == RBTREE_TNIL_INDEX)
  86     rt->root = rb_node_index (rt, x);
  87   else if (yp->right == yi)
  88     yp->right = xi;
  89   else
  90     yp->left = xi;
  91   x->right = yi;
  92   y->parent = xi;
  93 }
  94
  95 static void
  96 rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
  97 {
  98   rb_node_index_t yi = 0, xi = rt->root, zi;
  99   rb_node_t *y, *zpp, *x, *zp;
 100
 101   y = rb_node (rt, RBTREE_TNIL_INDEX);
 102   while (xi != RBTREE_TNIL_INDEX)
 103     {
 104       x = rb_node (rt, xi);
 105       y = x;
 106       if (z->key < x->key)
 107         xi = x->left;
 108       else
 109         xi = x->right;
 110     }
 111   yi = rb_node_index (rt, y);
 112   z->parent = yi;
 113   if (yi == RBTREE_TNIL_INDEX)
 114     rt->root = rb_node_index (rt, z);
 115   else if (z->key < y->key)
 116     y->left = rb_node_index (rt, z);
 117   else
 118     y->right = rb_node_index (rt, z);
 119
 120   /* Tree fixup stage */
 121   while (y->color == RBTREE_RED)
 122     {
 123       zi = rb_node_index (rt, z);
 124       zp = rb_node_parent (rt, z);
 125       zpp = rb_node_parent (rt, zp);
 126       if (z->parent == zpp->left)
 127         {
 128           y = rb_node_right (rt, zpp);
 129           if (y->color == RBTREE_RED)
 130             {
 131               zp->color = RBTREE_BLACK;
 132               y->color = RBTREE_BLACK;
 133               zpp->color = RBTREE_RED;
 134               z = zpp;
 135             }
 136           else
 137             {
 138               if (zi == zp->right)
 139                 {
 140                   z = zp;
 141                   rb_tree_rotate_left (rt, z);
 142                   zp = rb_node (rt, z->parent);
 143                   zpp = rb_node (rt, zp->parent);
 144                 }
 145               zp->color = RBTREE_BLACK;
 146               zpp->color = RBTREE_RED;
 147               rb_tree_rotate_right (rt, zpp);
 148             }
 149         }
 150       else
 151         {
 152           y = rb_node (rt, zpp->left);
 153           if (y->color == RBTREE_RED)
 154             {
 155               zp->color = RBTREE_BLACK;
 156               y->color = RBTREE_BLACK;
 157               zpp->color = RBTREE_RED;
 158               z = zpp;
 159             }
 160           else
 161             {
 162               if (zi == zp->left)
 163                 {
 164                   z = zp;
 165                   rb_tree_rotate_right (rt, z);
 166                   zp = rb_node (rt, z->parent);
 167                   zpp = rb_node (rt, zp->parent);
 168                 }
 169               zp->color = RBTREE_BLACK;
 170               zpp->color = RBTREE_RED;
 171               rb_tree_rotate_left (rt, zpp);
 172             }
 173         }
 174     }
 175   rb_node (rt, rt->root)->color = RBTREE_BLACK;
 176 }
 177
 178 rb_node_index_t
 179 rb_tree_add (rb_tree_t * rt, u32 key)
 180 {
 181   rb_node_t *n;
 182
 183   pool_get_zero (rt->nodes, n);
 184   n->key = key;
 185   n->color = RBTREE_RED;
 186   rb_tree_insert (rt, n);
 187   return rb_node_index (rt, n);
 188 }
 189
 190 rb_node_index_t
 191 rb_tree_add2 (rb_tree_t * rt, u32 key, u32 opaque)
 192 {
 193   rb_node_t *n;
 194
 195   pool_get_zero (rt->nodes, n);
 196   n->key = key;
 197   n->color = RBTREE_RED;
 198   n->opaque = opaque;
 199   rb_tree_insert (rt, n);
 200   return rb_node_index (rt, n);
 201 }
 202
 203 rb_node_t *
 204 rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
 205 {
 206   while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
 207     if (key < x->key)
 208       x = rb_node_left (rt, x);
 209     else
 210       x = rb_node_right (rt, x);
 211   return x;
 212 }
 213
 214 rb_node_t *
 215 rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x)
 216 {
 217   while (x->left != RBTREE_TNIL_INDEX)
 218     x = rb_node_left (rt, x);
 219   return x;
 220 }
 221
 222 rb_node_t *
 223 rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x)
 224 {
 225   while (x->right != RBTREE_TNIL_INDEX)
 226     x = rb_node_right (rt, x);
 227   return x;
 228 }
 229
 230 static inline void
 231 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
 232 {
 233   rb_node_t *up;
 234
 235   up = rb_node_parent (rt, u);
 236   if (u->parent == RBTREE_TNIL_INDEX)
 237     rt->root = rb_node_index (rt, v);
 238   else if (rb_node_index (rt, u) == up->left)
 239     up->left = rb_node_index (rt, v);
 240   else
 241     up->right = rb_node_index (rt, v);
 242   v->parent = u->parent;
 243 }
 244
 245 void
 246 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
 247 {
 248   rb_node_color_t y_original_color;
 249   rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr;
 250   rb_node_index_t xi, yi;
 251
 252   y = z;
 253   y_original_color = y->color;
 254
 255   if (z->left == RBTREE_TNIL_INDEX)
 256     {
 257       x = rb_node_right (rt, z);
 258       rb_tree_transplant (rt, z, x);
 259     }
 260   else if (z->right == RBTREE_TNIL_INDEX)
 261     {
 262       x = rb_node_left (rt, z);
 263       rb_tree_transplant (rt, z, x);
 264     }
 265   else
 266     {
 267       y = rb_tree_min_subtree (rt, rb_node_right (rt, z));
 268       y_original_color = y->color;
 269       x = rb_node_right (rt, y);
 270       yi = rb_node_index (rt, y);
 271       if (y->parent == rb_node_index (rt, z))
 272         x->parent = yi;
 273       else
 274         {
 275           rb_tree_transplant (rt, y, rb_node_right (rt, y));
 276           y->right = z->right;
 277           yr = rb_node_right (rt, y);
 278           yr->parent = yi;
 279         }
 280       rb_tree_transplant (rt, z, y);
 281       y->left = z->left;
 282       yl = rb_node_left (rt, y);
 283       yl->parent = yi;
 284       y->color = z->color;
 285     }
 286
 287   if (y_original_color == RBTREE_RED)
 288     return;
 289
 290   /* Tree fixup needed */
 291
 292   xi = rb_node_index (rt, x);
 293   while (xi != rt->root && x->color == RBTREE_BLACK)
 294     {
 295       xp = rb_node_parent (rt, x);
 296       if (xi == xp->left)
 297         {
 298           w = rb_node_right (rt, xp);
 299           if (w->color == RBTREE_RED)
 300             {
 301               w->color = RBTREE_BLACK;
 302               xp->color = RBTREE_RED;
 303               rb_tree_rotate_left (rt, xp);
 304               w = rb_node_right (rt, xp);
 305             }
 306           wl = rb_node_left (rt, w);
 307           wr = rb_node_right (rt, w);
 308           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 309             {
 310               w->color = RBTREE_RED;
 311               x = xp;
 312             }
 313           else
 314             {
 315               if (wr->color == RBTREE_BLACK)
 316                 {
 317                   wl->color = RBTREE_BLACK;
 318                   w->color = RBTREE_RED;
 319                   rb_tree_rotate_right (rt, w);
 320                   w = rb_node_right (rt, xp);
 321                 }
 322               w->color = xp->color;
 323               xp->color = RBTREE_BLACK;
 324               wr->color = RBTREE_BLACK;
 325               rb_tree_rotate_left (rt, xp);
 326               x = rb_node (rt, rt->root);
 327             }
 328         }
 329       else
 330         {
 331           w = rb_node_left (rt, xp);
 332           if (w->color == RBTREE_RED)
 333             {
 334               w->color = RBTREE_BLACK;
 335               xp->color = RBTREE_RED;
 336               rb_tree_rotate_right (rt, xp);
 337               w = rb_node_left (rt, xp);
 338             }
 339           wl = rb_node_left (rt, w);
 340           wr = rb_node_right (rt, w);
 341           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 342             {
 343               w->color = RBTREE_RED;
 344               x = xp;
 345             }
 346           else
 347             {
 348               if (wl->color == RBTREE_BLACK)
 349                 {
 350                   wr->color = RBTREE_BLACK;
 351                   w->color = RBTREE_RED;
 352                   rb_tree_rotate_left (rt, w);
 353                   w = rb_node_left (rt, xp);
 354                 }
 355               w->color = xp->color;
 356               xp->color = RBTREE_BLACK;
 357               wl->color = RBTREE_BLACK;
 358               rb_tree_rotate_right (rt, xp);
 359               x = rb_node (rt, rt->root);
 360             }
 361         }
 362       xi = rb_node_index (rt, x);
 363     }
 364   x->color = RBTREE_BLACK;
 365 }
 366
 367 void
 368 rb_tree_del (rb_tree_t * rt, u32 key)
 369 {
 370   rb_node_t *n;
 371   n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key);
 372   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 373     {
 374       rb_tree_del_node (rt, n);
 375       pool_put (rt->nodes, n);
 376     }
 377 }
 378
 379 u32
 380 rb_tree_n_nodes (rb_tree_t * rt)
 381 {
 382   return pool_elts (rt->nodes);
 383 }
 384
 385 void
 386 rb_tree_free_nodes (rb_tree_t * rt)
 387 {
 388   rb_node_t *n;
 389   pool_flush (n, rt->nodes,;);
 390 }
 391
 392 void
 393 rb_tree_init (rb_tree_t * rt)
 394 {
 395   rb_node_t *tnil;
 396
 397   rt->nodes = 0;
 398   rt->root = RBTREE_TNIL_INDEX;
 399
 400   /* By convention first node, index 0, is the T.nil sentinel */
 401   pool_get_zero (rt->nodes, tnil);
 402   tnil->color = RBTREE_BLACK;
 403 }
 404
 405 /*
 406  * fd.io coding-style-patch-verification: ON
 407  *
 408  * Local Variables:
 409  * eval: (c-set-style "gnu")
 410  * End:
 411  */