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:
7 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 * Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2009).
18 * Introduction to algorithms. MIT press, 3rd Edition, Ch. 13
21 #include <vppinfra/rbtree.h>
24 rb_tree_rotate_left (rb_tree_t * rt, rb_node_t * x)
26 rb_node_t *y, *tmp, *xp;
27 rb_node_index_t xi, yi;
29 xi = rb_node_index (rt, x);
31 y = rb_node_right (rt, x);
33 if (y->left != RBTREE_TNIL_INDEX)
35 tmp = rb_node_left (rt, y);
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)
51 rb_tree_rotate_right (rb_tree_t * rt, rb_node_t * y)
53 rb_node_index_t yi, xi;
56 yi = rb_node_index (rt, y);
58 x = rb_node_left (rt, y);
60 if (x->right != RBTREE_TNIL_INDEX)
62 rb_node_t *tmp = rb_node_right (rt, x);
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)
78 rb_tree_fixup_inline (rb_tree_t * rt, rb_node_t * y, rb_node_t * z)
83 while (y->color == RBTREE_RED)
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)
90 y = rb_node_right (rt, zpp);
91 if (y->color == RBTREE_RED)
93 zp->color = RBTREE_BLACK;
94 y->color = RBTREE_BLACK;
95 zpp->color = RBTREE_RED;
103 rb_tree_rotate_left (rt, z);
104 zp = rb_node (rt, z->parent);
105 zpp = rb_node (rt, zp->parent);
107 zp->color = RBTREE_BLACK;
108 zpp->color = RBTREE_RED;
109 rb_tree_rotate_right (rt, zpp);
114 y = rb_node (rt, zpp->left);
115 if (y->color == RBTREE_RED)
117 zp->color = RBTREE_BLACK;
118 y->color = RBTREE_BLACK;
119 zpp->color = RBTREE_RED;
127 rb_tree_rotate_right (rt, z);
128 zp = rb_node (rt, z->parent);
129 zpp = rb_node (rt, zp->parent);
131 zp->color = RBTREE_BLACK;
132 zpp->color = RBTREE_RED;
133 rb_tree_rotate_left (rt, zpp);
137 rb_node (rt, rt->root)->color = RBTREE_BLACK;
141 rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
143 rb_node_index_t yi = 0, xi = rt->root;
146 y = rb_node (rt, RBTREE_TNIL_INDEX);
147 while (xi != RBTREE_TNIL_INDEX)
149 x = rb_node (rt, xi);
156 yi = rb_node_index (rt, y);
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);
163 y->right = rb_node_index (rt, z);
165 /* Tree fixup stage */
166 rb_tree_fixup_inline (rt, y, z);
169 __clib_export rb_node_index_t
170 rb_tree_add (rb_tree_t * rt, u32 key)
174 pool_get_zero (rt->nodes, n);
176 n->color = RBTREE_RED;
177 rb_tree_insert (rt, n);
178 return rb_node_index (rt, n);
181 __clib_export rb_node_index_t
182 rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque)
186 pool_get_zero (rt->nodes, n);
188 n->color = RBTREE_RED;
190 rb_tree_insert (rt, n);
191 return rb_node_index (rt, n);
194 __clib_export rb_node_index_t
195 rb_tree_add_custom (rb_tree_t * rt, u32 key, uword opaque, rb_tree_lt_fn ltfn)
197 rb_node_index_t yi = 0, xi = rt->root;
198 rb_node_t *z, *y, *x;
200 pool_get_zero (rt->nodes, z);
202 z->color = RBTREE_RED;
205 y = rb_node (rt, RBTREE_TNIL_INDEX);
206 while (xi != RBTREE_TNIL_INDEX)
208 x = rb_node (rt, xi);
210 ASSERT (z->key != x->key);
211 if (ltfn (z->key, x->key))
216 yi = rb_node_index (rt, y);
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);
223 y->right = rb_node_index (rt, z);
225 rb_tree_fixup_inline (rt, y, z);
227 return rb_node_index (rt, z);
230 __clib_export rb_node_t *
231 rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
233 while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key)
235 x = rb_node_left (rt, x);
237 x = rb_node_right (rt, x);
241 __clib_export rb_node_t *
242 rb_tree_search_subtree_custom (rb_tree_t * rt, rb_node_t * x, u32 key,
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);
249 x = rb_node_right (rt, x);
253 __clib_export rb_node_t *
254 rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x)
256 while (x->left != RBTREE_TNIL_INDEX)
257 x = rb_node_left (rt, x);
261 __clib_export rb_node_t *
262 rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x)
264 while (x->right != RBTREE_TNIL_INDEX)
265 x = rb_node_right (rt, x);
269 __clib_export rb_node_t *
270 rb_tree_successor (rb_tree_t * rt, rb_node_t * x)
274 if (x->right != RBTREE_TNIL_INDEX)
275 return rb_tree_min_subtree (rt, rb_node_right (rt, x));
277 y = rb_node_parent (rt, x);
278 while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x))
281 y = rb_node_parent (rt, y);
286 __clib_export rb_node_t *
287 rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x)
291 if (x->left != RBTREE_TNIL_INDEX)
292 return rb_tree_max_subtree (rt, rb_node_left (rt, x));
294 y = rb_node_parent (rt, x);
295 while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x))
298 y = rb_node_parent (rt, y);
304 rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v)
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);
314 up->right = rb_node_index (rt, v);
315 v->parent = u->parent;
319 rb_tree_del_node_internal (rb_tree_t * rt, rb_node_t * z)
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;
326 y_original_color = y->color;
328 if (z->left == RBTREE_TNIL_INDEX)
330 x = rb_node_right (rt, z);
331 rb_tree_transplant (rt, z, x);
333 else if (z->right == RBTREE_TNIL_INDEX)
335 x = rb_node_left (rt, z);
336 rb_tree_transplant (rt, z, x);
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))
348 rb_tree_transplant (rt, y, x);
350 yr = rb_node_right (rt, y);
353 rb_tree_transplant (rt, z, y);
355 yl = rb_node_left (rt, y);
360 if (y_original_color == RBTREE_RED)
363 /* Tree fixup needed */
365 xi = rb_node_index (rt, x);
366 while (xi != rt->root && x->color == RBTREE_BLACK)
368 xp = rb_node_parent (rt, x);
371 w = rb_node_right (rt, xp);
372 if (w->color == RBTREE_RED)
374 w->color = RBTREE_BLACK;
375 xp->color = RBTREE_RED;
376 rb_tree_rotate_left (rt, xp);
377 w = rb_node_right (rt, xp);
379 wl = rb_node_left (rt, w);
380 wr = rb_node_right (rt, w);
381 if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
383 if (!rb_node_is_tnil (rt, w))
384 w->color = RBTREE_RED;
389 if (wr->color == RBTREE_BLACK)
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);
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);
406 w = rb_node_left (rt, xp);
407 if (w->color == RBTREE_RED)
409 w->color = RBTREE_BLACK;
410 xp->color = RBTREE_RED;
411 rb_tree_rotate_right (rt, xp);
412 w = rb_node_left (rt, xp);
414 wl = rb_node_left (rt, w);
415 wr = rb_node_right (rt, w);
416 if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK)
418 if (!rb_node_is_tnil (rt, w))
419 w->color = RBTREE_RED;
424 if (wl->color == RBTREE_BLACK)
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);
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);
439 xi = rb_node_index (rt, x);
441 x->color = RBTREE_BLACK;
445 rb_tree_del_node (rb_tree_t * rt, rb_node_t * z)
447 rb_tree_del_node_internal (rt, z);
448 pool_put (rt->nodes, z);
452 rb_tree_del (rb_tree_t * rt, u32 key)
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);
461 rb_tree_del_custom (rb_tree_t * rt, u32 key, rb_tree_lt_fn ltfn)
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);
470 rb_tree_n_nodes (rb_tree_t * rt)
472 return pool_elts (rt->nodes);
476 rb_tree_free_nodes (rb_tree_t * rt)
478 pool_free (rt->nodes);
479 rt->root = RBTREE_TNIL_INDEX;
483 rb_tree_init (rb_tree_t * rt)
488 rt->root = RBTREE_TNIL_INDEX;
490 /* By convention first node, index 0, is the T.nil sentinel */
491 pool_get_zero (rt->nodes, tnil);
492 tnil->color = RBTREE_BLACK;
496 rb_tree_is_init (rb_tree_t * rt)
498 if (pool_elts (rt->nodes) == 0)
504 * fd.io coding-style-patch-verification: ON
507 * eval: (c-set-style "gnu")