From f7cda7a9bc399258959084154263c42c5a91e030 Mon Sep 17 00:00:00 2001 From: Florin Coras Date: Fri, 19 Apr 2019 18:50:34 -0700 Subject: rbtree: add successor and predecessor functions Change-Id: I6934beaf5c08bae2d4f0bd3a6bb811810407c1f9 Signed-off-by: Florin Coras --- src/plugins/unittest/rbtree_test.c | 58 +++++++++++++++++++++++++++++++++++--- src/vppinfra/rbtree.c | 34 ++++++++++++++++++++++ src/vppinfra/rbtree.h | 2 ++ 3 files changed, 90 insertions(+), 4 deletions(-) (limited to 'src') diff --git a/src/plugins/unittest/rbtree_test.c b/src/plugins/unittest/rbtree_test.c index a20ac1b111b..490be9c156e 100644 --- a/src/plugins/unittest/rbtree_test.c +++ b/src/plugins/unittest/rbtree_test.c @@ -41,7 +41,7 @@ rbtree_test_basic (vlib_main_t * vm, unformat_input_t * input) int __clib_unused verbose, n_keys = 1e3, i; u32 *test_keys = 0, search_key; rb_tree_t _rt, *rt = &_rt; - rb_node_t *n; + rb_node_t *n, *aux; while (unformat_check_input (input) != UNFORMAT_END_OF_INPUT) { @@ -82,13 +82,21 @@ rbtree_test_basic (vlib_main_t * vm, unformat_input_t * input) RBTREE_TEST (n->key == n_keys / 2, "search result should be %u", n_keys / 2); + aux = rb_tree_successor (rt, n); + RBTREE_TEST (aux->key == n_keys / 2 + 1, "successor should be %u is %u", + n_keys / 2 + 1, aux->key); + + aux = rb_tree_predecessor (rt, n); + RBTREE_TEST (aux->key == n_keys / 2 - 1, "predecessor should be %u is %u", + n_keys / 2 - 1, aux->key); + n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), n_keys); RBTREE_TEST (rb_node_is_tnil (rt, n), "search result should be tnil"); /* * Delete even keys */ - for (i = 0; i < (n_keys - 1) / 2; i += 2) + for (i = 0; i < n_keys; i += 2) rb_tree_del (rt, i); n = rb_tree_max_subtree (rt, rb_node (rt, rt->root)); @@ -99,17 +107,26 @@ rbtree_test_basic (vlib_main_t * vm, unformat_input_t * input) search_key = 2 * ((n_keys - 1) / 4); n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), search_key); - RBTREE_TEST (rb_node_is_tnil (rt, n), "search result should be tnil"); + RBTREE_TEST (rb_node_is_tnil (rt, n), "search result for %u should be tnil", + search_key); search_key += 1; n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), search_key); RBTREE_TEST (n->key == search_key, "search result should be %u", search_key); + aux = rb_tree_successor (rt, n); + RBTREE_TEST (aux->key == search_key + 2, "successor should be %u is %u", + search_key + 2, aux->key); + + aux = rb_tree_predecessor (rt, n); + RBTREE_TEST (aux->key == search_key - 2, "predecessor should be %u is %u", + search_key - 2, aux->key); + /* * Re-add even keys */ - for (i = 0; i < (n_keys - 1) / 2; i += 2) + for (i = 0; i < n_keys; i += 2) rb_tree_add (rt, i); RBTREE_TEST (rb_tree_n_nodes (rt) == n_keys + 1, "number nodes %u is %u", @@ -131,6 +148,14 @@ rbtree_test_basic (vlib_main_t * vm, unformat_input_t * input) RBTREE_TEST (n->key == search_key, "search result should be %u", search_key); + aux = rb_tree_successor (rt, n); + RBTREE_TEST (aux->key == search_key + 1, "successor should be %u is %u", + search_key + 1, aux->key); + + aux = rb_tree_predecessor (rt, n); + RBTREE_TEST (aux->key == search_key - 1, "predecessor should be %u is %u", + search_key - 1, aux->key); + /* * Delete all keys */ @@ -150,6 +175,31 @@ rbtree_test_basic (vlib_main_t * vm, unformat_input_t * input) n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), search_key); RBTREE_TEST (rb_node_is_tnil (rt, n), "search result should be tnil"); + /* + * Test successor/predecessor + */ + u8 test_vals[] = { 2, 3, 4, 6, 7, 9, 13, 15, 17, 18, 20 }; + for (i = 0; i < sizeof (test_vals) / sizeof (u8); i++) + rb_tree_add (rt, test_vals[i]); + + search_key = 13; + n = rb_tree_search_subtree (rt, rb_node (rt, rt->root), search_key); + RBTREE_TEST (n->key == search_key, "search result should be %u", + search_key); + + aux = rb_tree_successor (rt, n); + RBTREE_TEST (aux->key == 15, "successor should be %u is %u", 15, aux->key); + + aux = rb_tree_predecessor (rt, n); + RBTREE_TEST (aux->key == 9, "predecessor should be %u is %u", 9, aux->key); + + n = aux; + aux = rb_tree_predecessor (rt, n); + RBTREE_TEST (aux->key == 7, "predecessor should be %u is %u", 7, aux->key); + + /* + * Cleanup + */ rb_tree_free_nodes (rt); RBTREE_TEST (rb_tree_n_nodes (rt) == 0, "number nodes %u is %u", 0, rb_tree_n_nodes (rt)); diff --git a/src/vppinfra/rbtree.c b/src/vppinfra/rbtree.c index d99367a2c6f..3770c2304ff 100644 --- a/src/vppinfra/rbtree.c +++ b/src/vppinfra/rbtree.c @@ -227,6 +227,40 @@ rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * 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) { diff --git a/src/vppinfra/rbtree.h b/src/vppinfra/rbtree.h index ace07ac5622..73180d1d54c 100644 --- a/src/vppinfra/rbtree.h +++ b/src/vppinfra/rbtree.h @@ -54,6 +54,8 @@ u32 rb_tree_n_nodes (rb_tree_t * rt); rb_node_t *rb_tree_min_subtree (rb_tree_t * rt, rb_node_t * x); rb_node_t *rb_tree_max_subtree (rb_tree_t * rt, rb_node_t * x); rb_node_t *rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key); +rb_node_t *rb_tree_successor (rb_tree_t * rt, rb_node_t * x); +rb_node_t *rb_tree_predecessor (rb_tree_t * rt, rb_node_t * x); static inline rb_node_index_t rb_node_index (rb_tree_t * rt, rb_node_t * n) -- cgit 1.2.3-korg