2 files changed, 99 insertions, 24 deletions

diff --git a/src/vppinfra/rbtree.c b/src/vppinfra/rbtree.c
index 5bb2e875e86..df759cb2344 100644
--- a/src/vppinfra/rbtree.c
+++ b/src/vppinfra/rbtree.c
@@ -74,32 +74,12 @@ rb_tree_rotate_right (rb_tree_t * rt, rb_node_t * y)
   y->parent = xi;
 }
 
-static void
-rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
+static inline void
+rb_tree_fixup_inline (rb_tree_t * rt, rb_node_t * y, rb_node_t * z)
 {
-  rb_node_index_t yi = 0, xi = rt->root, zi;
-  rb_node_t *y, *zpp, *x, *zp;
+  rb_node_t *zpp, *zp;
+  rb_node_index_t zi;
 
-  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 */
   while (y->color == RBTREE_RED)
     {
       zi = rb_node_index (rt, z);
@@ -157,6 +137,35 @@ rb_tree_insert (rb_tree_t * rt, rb_node_t * z)
   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)
 {
@@ -182,6 +191,41 @@ rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque)
   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)
 {
@@ -194,6 +238,18 @@ rb_tree_search_subtree (rb_tree_t * rt, rb_node_t * x, u32 key)
 }
 
 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)
@@ -392,6 +448,18 @@ rb_tree_del (rb_tree_t * rt, u32 key)
     }
 }
 
+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)
 {
diff --git a/src/vppinfra/rbtree.h b/src/vppinfra/rbtree.h
index 79437cdd719..65580584b6d 100644
--- a/src/vppinfra/rbtree.h
+++ b/src/vppinfra/rbtree.h
@@ -45,15 +45,22 @@ typedef struct rb_tree_
   rb_node_index_t root;		/**< root index */
 } rb_tree_t;
 
+typedef int (*rb_tree_lt_fn) (u32 a, u32 b);
+
 void rb_tree_init (rb_tree_t * rt);
 rb_node_index_t rb_tree_add (rb_tree_t * rt, u32 key);
 rb_node_index_t rb_tree_add2 (rb_tree_t * rt, u32 key, uword opaque);
+rb_node_index_t rb_tree_add_custom (rb_tree_t * rt, u32 key, uword opaque,
+				    rb_tree_lt_fn ltfn);
 void rb_tree_del (rb_tree_t * rt, u32 key);
+void rb_tree_del_custom (rb_tree_t * rt, u32 key, rb_tree_lt_fn ltfn);
 void rb_tree_free_nodes (rb_tree_t * rt);
 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_search_subtree_custom (rb_tree_t * rt, rb_node_t * x,
+					  u32 key, rb_tree_lt_fn ltfn);
 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);