/* * Copyright (c) 2019 Cisco and/or its affiliates. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at: * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ /* * Algorithm from: * Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2009). * Introduction to algorithms. MIT press, 3rd Edition, Ch. 13 */ #include <vppinfra/rbtree.h> static inline void rb_tree_rotate_left (rb_tree_t * rt, rb_node_t * x) { rb_node_t *y, *tmp, *xp; rb_node_index_t xi, yi; xi = rb_node_index (rt, x); yi = x->right; y = rb_node_right (rt, x); x->right = y->left; if (y->left != RBTREE_TNIL_INDEX) { tmp = rb_node_left (rt, y); tmp->parent = xi; } xp = rb_node_parent (rt, x); y->parent = x->parent; if (x->parent == RBTREE_TNIL_INDEX) rt->root = rb_node_index (rt, y); else if (xp->left == xi) xp->left = yi; else xp->right = yi; y->left = xi; x->parent = yi; } static inline void rb_tree_rotate_right (rb_tree_t * rt, rb_node_t * y) { rb_node_index_t yi, xi; rb_node_t *x, *yp; yi = rb_node_index (rt, y); xi = y->left; x = rb_node_left (rt, y); y->left = x->right; if (x->right != RBTREE_TNIL_INDEX) { rb_node_t *tmp = rb_node_right (rt, x); tmp->parent = yi; } yp = rb_node_parent (rt, y); x->parent = y->parent; if (y->parent == RBTREE_TNIL_INDEX) rt->root = rb_node_index (rt, x); else if (yp->right == yi) yp->right = xi; else yp->left = xi; x->right = yi; y->parent = xi; } static inline void rb_tree_fixup_inline (rb_tree_t * rt, rb_node_t * y, rb_node_t * z) { rb_node_t *zpp, *zp; rb_node_index_t zi; while (y->color == RBTREE_RED) { zi = rb_node_index (rt, z); zp = rb_node_parent (rt, z); zpp = rb_node_parent (rt, zp); if (z->parent == zpp->left) { y = rb_node_right (rt, zpp); if (y->color == RBTREE_RED) { zp->color = RBTREE_BLACK; y->color = RBTREE_BLACK; zpp->color = RBTREE_RED; z = zpp; } else { if (zi == zp->right) { z = zp; rb_tree_rotate_left (rt, z); zp = rb_node (rt, z->parent); zpp = rb_node (rt, zp->parent); } zp->color = RBTREE_BLACK; zpp->color = RBTREE_RED; rb_tree_rotate_right (rt, zpp); } } else { y = rb_node (rt, zpp->left); if (y->color == RBTREE_RED) { zp->color = RBTREE_BLACK; y->color = RBTREE_BLACK; zpp->color = RBTREE_RED; z = zpp; } else { if (zi == zp->left) { z = zp; rb_tree_rotate_right (rt, z); zp = rb_node (rt, z->parent); zpp = rb_node (rt, zp->parent); } zp->color = RBTREE_BLACK; zpp->color = RBTREE_RED; rb_tree_rotate_left (rt, zpp); } } } rb_node (rt, rt->root)->color = RBTREE_BLACK; } static void rb_tree_insert (rb_tree_t * rt, rb_node_t * z) { rb_node_index_t yi = 0, xi = rt->root; rb_node_t *y, *x; y = rb_node (rt, RBTREE_TNIL_INDEX); while (xi != RBTREE_TNIL_INDEX) { x = rb_node (rt, xi); y = x; if (z->key < x->key) xi = x->left; else xi = x->right; } yi = rb_node_index (rt, y); z->parent = yi; if (yi == RBTREE_TNIL_INDEX) rt->root = rb_node_index (rt, z); else if (z->key < y->key) y->left = rb_node_index (rt, z); else y->right = rb_node_index (rt, z); /* Tree fixup stage */ rb_tree_fixup_inline (rt, y, z); } rb_node_index_t rb_tree_add (rb_tree_t * rt, u32 key) { rb_node_t *n; pool_get_zero (rt->nodes, n); n->key = key; n->color = RBTREE_RED; rb_tree_insert (rt, n); return rb_node_index (rt, n); } rb_node_index_t rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque) { rb_node_t *n; pool_get_zero (rt->nodes, n); n->key = key; n->color = RBTREE_RED; n->opaque = opaque; rb_tree_insert (rt, n); return rb_node_index (rt, n); } rb_node_index_t rb_tree_add_custom (rb_tree_t * rt, u32 key, uword opaque, rb_tree_lt_fn ltfn) { rb_node_index_t yi = 0, xi = rt->root; rb_node_t *z, *y, *x; pool_get_zero (rt->nodes, z); z->key = key; z->color = RBTREE_RED; z->opaque = opaque; y = rb_node (rt, RBTREE_TNIL_INDEX); while (xi != RBTREE_TNIL_INDEX) { x = rb_node (rt, xi); y = x; if (ltfn (z->key, x->key)) xi = x->left; else xi = x->right; } yi = rb_node_index (rt, y); z->parent = yi; if (yi == RBTREE_TNIL_INDEX) rt->root = rb_node_index (rt, z); else if (ltfn (z->key, y->key)) y->left = rb_node_index (rt, z); else y->right = rb_node_index (rt, z); rb_tree_fixup_inline (rt, y, z); return rb_node_index (rt, z); } rb_node_t * rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key) { while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key) if (key < x->key) x = rb_node_left (rt, x); else x = rb_node_right (rt, x); return x; } rb_node_t * rb_tree_search_subtree_custom (rb_tree_t * rt, rb_node_t * x, u32 key, rb_tree_lt_fn ltfn) { while (rb_node_index (rt, x) != RBTREE_TNIL_INDEX && key != x->key) if (ltfn (key, x->key)) x = rb_node_left (rt, x); else x = rb_node_right (rt, x); return x; } rb_node_t * rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x) { while (x->left != RBTREE_TNIL_INDEX) x = rb_node_left (rt, x); return x; } rb_node_t * rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x) { while (x->right != RBTREE_TNIL_INDEX) x = rb_node_right (rt, x); return x; } rb_node_t * rb_tree_successor (rb_tree_t * rt, rb_node_t * x) { rb_node_t *y; if (x->right != RBTREE_TNIL_INDEX) return rb_tree_min_subtree (rt, rb_node_right (rt, x)); y = rb_node_parent (rt, x); while (!rb_node_is_tnil (rt, y) && y->right == rb_node_index (rt, x)) { x = y; y = rb_node_parent (rt, y); } return y; } rb_node_t * rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x) { rb_node_t *y; if (x->left != RBTREE_TNIL_INDEX) return rb_tree_max_subtree (rt, rb_node_left (rt, x)); y = rb_node_parent (rt, x); while (!rb_node_is_tnil (rt, y) && y->left == rb_node_index (rt, x)) { x = y; y = rb_node_parent (rt, y); } return y; } static inline void rb_tree_transplant (rb_tree_t * rt, rb_node_t * u, rb_node_t * v) { rb_node_t *up; up = rb_node_parent (rt, u); if (u->parent == RBTREE_TNIL_INDEX) rt->root = rb_node_index (rt, v); else if (rb_node_index (rt, u) == up->left) up->left = rb_node_index (rt, v); else up->right = rb_node_index (rt, v); v->parent = u->parent; } void rb_tree_del_node (rb_tree_t * rt, rb_node_t * z) { rb_node_color_t y_original_color; rb_node_t *x, *y, *yr, *yl, *xp, *w, *wl, *wr; rb_node_index_t xi, yi; y = z; y_original_color = y->color; if (z->left == RBTREE_TNIL_INDEX) { x = rb_node_right (rt, z); rb_tree_transplant (rt, z, x); } else if (z->right == RBTREE_TNIL_INDEX) { x = rb_node_left (rt, z); rb_tree_transplant (rt, z, x); } else { y = rb_tree_min_subtree (rt, rb_node_right (rt, z)); y_original_color = y->color; x = rb_node_right (rt, y); yi = rb_node_index (rt, y); if (y->parent == rb_node_index (rt, z)) x->parent = yi; else { rb_tree_transplant (rt, y, x); y->right = z->right; yr = rb_node_right (rt, y); yr->parent = yi; } rb_tree_transplant (rt, z, y); y->left = z->left; yl = rb_node_left (rt, y); yl->parent = yi; y->color = z->color; } if (y_original_color == RBTREE_RED) return; /* Tree fixup needed */ xi = rb_node_index (rt, x); while (xi != rt->root && x->color == RBTREE_BLACK) { xp = rb_node_parent (rt, x); if (xi == xp->left) { w = rb_node_right (rt, xp); if (w->color == RBTREE_RED) { w->color = RBTREE_BLACK; xp->color = RBTREE_RED; rb_tree_rotate_left (rt, xp); w = rb_node_right (rt, xp); } wl = rb_node_left (rt, w); wr = rb_node_right (rt, w); if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK) { if (!rb_node_is_tnil (rt, w)) w->color = RBTREE_RED; x = xp; } else { if (wr->color == RBTREE_BLACK) { wl->color = RBTREE_BLACK; w->color = RBTREE_RED; rb_tree_rotate_right (rt, w); w = rb_node_right (rt, xp); wr = rb_node_right (rt, w); } w->color = xp->color; xp->color = RBTREE_BLACK; wr->color = RBTREE_BLACK; rb_tree_rotate_left (rt, xp); x = rb_node (rt, rt->root); } } else { w = rb_node_left (rt, xp); if (w->color == RBTREE_RED) { w->color = RBTREE_BLACK; xp->color = RBTREE_RED; rb_tree_rotate_right (rt, xp); w = rb_node_left (rt, xp); } wl = rb_node_left (rt, w); wr = rb_node_right (rt, w); if (wl->color == RBTREE_BLACK && wr->color == RBTREE_BLACK) { if (!rb_node_is_tnil (rt, w)) w->color = RBTREE_RED; x = xp; } else { if (wl->color == RBTREE_BLACK) { wr->color = RBTREE_BLACK; w->color = RBTREE_RED; rb_tree_rotate_left (rt, w); w = rb_node_left (rt, xp); wl = rb_node_left (rt, w); } w->color = xp->color; xp->color = RBTREE_BLACK; wl->color = RBTREE_BLACK; rb_tree_rotate_right (rt, xp); x = rb_node (rt, rt->root); } } xi = rb_node_index (rt, x); } x->color = RBTREE_BLACK; } void rb_tree_del (rb_tree_t * rt, u32 key) { rb_node_t *n; n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), key); if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX) { rb_tree_del_node (rt, n); pool_put (rt->nodes, n); } } void rb_tree_del_custom (rb_tree_t * rt, u32 key, rb_tree_lt_fn ltfn) { rb_node_t *n; n = rb_tree_search_subtree_custom (rt, rb_node (rt, rt->root), key, ltfn); if (rb_node_index (rt, n) != RBTREE_TNIL_INDEX) { rb_tree_del_node (rt, n); pool_put (rt->nodes, n); } } u32 rb_tree_n_nodes (rb_tree_t * rt) { return pool_elts (rt->nodes); } void rb_tree_free_nodes (rb_tree_t * rt) { pool_free (rt->nodes); rt->root = RBTREE_TNIL_INDEX; } void rb_tree_init (rb_tree_t * rt) { rb_node_t *tnil; rt->nodes = 0; rt->root = RBTREE_TNIL_INDEX; /* By convention first node, index 0, is the T.nil sentinel */ pool_get_zero (rt->nodes, tnil); tnil->color = RBTREE_BLACK; } /* * fd.io coding-style-patch-verification: ON * * Local Variables: * eval: (c-set-style "gnu") * End: */