rbtree: add successor and predecessor functions

Change-Id: I6934beaf5c08bae2d4f0bd3a6bb811810407c1f9
Signed-off-by: Florin Coras <fcoras@cisco.com>

3 files changed, 90 insertions, 4 deletions

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)