Line data    Source code

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