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 rb_node_t *
 231 rb_tree_successor (rb_tree_t * rt, rb_node_t * x)
 232 {
 233   rb_node_t *y;
 234
 235   if (x->right != RBTREE_TNIL_INDEX)
 236     return rb_tree_min_subtree (rt, rb_node_right (rt, x));
 237
 238   y = rb_node_parent (rt, x);
 239   while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x))
 240     {
 241       x = y;
 242       y = rb_node_parent (rt, y);
 243     }
 244   return y;
 245 }
 246
 247 rb_node_t *
 248 rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x)
 249 {
 250   rb_node_t *y;
 251
 252   if (x->left != RBTREE_TNIL_INDEX)
 253     return rb_tree_max_subtree (rt, rb_node_left (rt, x));
 254
 255   y = rb_node_parent (rt, x);
 256   while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x))
 257     {
 258       x = y;
 259       y = rb_node_parent (rt, y);
 260     }
 261   return y;
 262 }
 263
 264 static inline void
 265 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
 266 {
 267   rb_node_t *up;
 268
 269   up = rb_node_parent (rt, u);
 270   if (u->parent == RBTREE_TNIL_INDEX)
 271     rt->root = rb_node_index (rt, v);
 272   else if (rb_node_index (rt, u) == up->left)
 273     up->left = rb_node_index (rt, v);
 274   else
 275     up->right = rb_node_index (rt, v);
 276   v->parent = u->parent;
 277 }
 278
 279 void
 280 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
 281 {
 282   rb_node_color_t y_original_color;
 283   rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr;
 284   rb_node_index_t xi, yi;
 285
 286   y = z;
 287   y_original_color = y->color;
 288
 289   if (z->left == RBTREE_TNIL_INDEX)
 290     {
 291       x = rb_node_right (rt, z);
 292       rb_tree_transplant (rt, z, x);
 293     }
 294   else if (z->right == RBTREE_TNIL_INDEX)
 295     {
 296       x = rb_node_left (rt, z);
 297       rb_tree_transplant (rt, z, x);
 298     }
 299   else
 300     {
 301       y = rb_tree_min_subtree (rt, rb_node_right (rt, z));
 302       y_original_color = y->color;
 303       x = rb_node_right (rt, y);
 304       yi = rb_node_index (rt, y);
 305       if (y->parent == rb_node_index (rt, z))
 306         x->parent = yi;
 307       else
 308         {
 309           rb_tree_transplant (rt, y, rb_node_right (rt, y));
 310           y->right = z->right;
 311           yr = rb_node_right (rt, y);
 312           yr->parent = yi;
 313         }
 314       rb_tree_transplant (rt, z, y);
 315       y->left = z->left;
 316       yl = rb_node_left (rt, y);
 317       yl->parent = yi;
 318       y->color = z->color;
 319     }
 320
 321   if (y_original_color == RBTREE_RED)
 322     return;
 323
 324   /* Tree fixup needed */
 325
 326   xi = rb_node_index (rt, x);
 327   while (xi != rt->root && x->color == RBTREE_BLACK)
 328     {
 329       xp = rb_node_parent (rt, x);
 330       if (xi == xp->left)
 331         {
 332           w = rb_node_right (rt, xp);
 333           if (w->color == RBTREE_RED)
 334             {
 335               w->color = RBTREE_BLACK;
 336               xp->color = RBTREE_RED;
 337               rb_tree_rotate_left (rt, xp);
 338               w = rb_node_right (rt, xp);
 339             }
 340           wl = rb_node_left (rt, w);
 341           wr = rb_node_right (rt, w);
 342           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 343             {
 344               w->color = RBTREE_RED;
 345               x = xp;
 346             }
 347           else
 348             {
 349               if (wr->color == RBTREE_BLACK)
 350                 {
 351                   wl->color = RBTREE_BLACK;
 352                   w->color = RBTREE_RED;
 353                   rb_tree_rotate_right (rt, w);
 354                   w = rb_node_right (rt, xp);
 355                 }
 356               w->color = xp->color;
 357               xp->color = RBTREE_BLACK;
 358               wr->color = RBTREE_BLACK;
 359               rb_tree_rotate_left (rt, xp);
 360               x = rb_node (rt, rt->root);
 361             }
 362         }
 363       else
 364         {
 365           w = rb_node_left (rt, xp);
 366           if (w->color == RBTREE_RED)
 367             {
 368               w->color = RBTREE_BLACK;
 369               xp->color = RBTREE_RED;
 370               rb_tree_rotate_right (rt, xp);
 371               w = rb_node_left (rt, xp);
 372             }
 373           wl = rb_node_left (rt, w);
 374           wr = rb_node_right (rt, w);
 375           if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
 376             {
 377               w->color = RBTREE_RED;
 378               x = xp;
 379             }
 380           else
 381             {
 382               if (wl->color == RBTREE_BLACK)
 383                 {
 384                   wr->color = RBTREE_BLACK;
 385                   w->color = RBTREE_RED;
 386                   rb_tree_rotate_left (rt, w);
 387                   w = rb_node_left (rt, xp);
 388                 }
 389               w->color = xp->color;
 390               xp->color = RBTREE_BLACK;
 391               wl->color = RBTREE_BLACK;
 392               rb_tree_rotate_right (rt, xp);
 393               x = rb_node (rt, rt->root);
 394             }
 395         }
 396       xi = rb_node_index (rt, x);
 397     }
 398   x->color = RBTREE_BLACK;
 399 }
 400
 401 void
 402 rb_tree_del (rb_tree_t * rt, u32 key)
 403 {
 404   rb_node_t *n;
 405   n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key);
 406   if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX)
 407     {
 408       rb_tree_del_node (rt, n);
 409       pool_put (rt->nodes, n);
 410     }
 411 }
 412
 413 u32
 414 rb_tree_n_nodes (rb_tree_t * rt)
 415 {
 416   return pool_elts (rt->nodes);
 417 }
 418
 419 void
 420 rb_tree_free_nodes (rb_tree_t * rt)
 421 {
 422   rb_node_t *n;
 423   pool_flush (n, rt->nodes,;);
 424 }
 425
 426 void
 427 rb_tree_init (rb_tree_t * rt)
 428 {
 429   rb_node_t *tnil;
 430
 431   rt->nodes = 0;
 432   rt->root = RBTREE_TNIL_INDEX;
 433
 434   /* By convention first node, index 0, is the T.nil sentinel */
 435   pool_get_zero (rt->nodes, tnil);
 436   tnil->color = RBTREE_BLACK;
 437 }
 438
 439 /*
 440  * fd.io coding-style-patch-verification: ON
 441  *
 442  * Local Variables:
 443  * eval: (c-set-style "gnu")
 444  * End:
 445  */