1 files changed, 59 insertions, 3 deletions

diff --git a/epan/wmem/wmem_tree.c b/epan/wmem/wmem_tree.c
index 71815c7567..ee45c58695 100644
--- a/epan/wmem/wmem_tree.c
+++ b/epan/wmem/wmem_tree.c
@@ -206,6 +206,31 @@ rb_insert_case1(wmem_tree_t *tree, wmem_tree_node_t *node)
     }
 }
 
+static void
+rb_transplant(wmem_tree_t *tree, wmem_tree_node_t *node, wmem_tree_node_t *child)
+{
+    if (node->parent == NULL) {
+        tree->root = child;
+    } else if (node == node->parent->left) {
+        node->parent->left = child;
+    } else /* if (node == node->parent->right) */ {
+        node->parent->right = child;
+    }
+
+    if (child) {
+        child->parent = node->parent;
+    }
+}
+
+static wmem_tree_node_t *
+tree_minimum(wmem_tree_node_t *node)
+{
+    while (node->left) {
+        node = node->left;
+    }
+    return node;
+}
+
 wmem_tree_t *
 wmem_tree_new(wmem_allocator_t *allocator)
 {
@@ -423,7 +448,7 @@ lookup_or_insert32(wmem_tree_t *tree, guint32 key,
     return node->data;
 }
 
-static void *
+static wmem_tree_node_t *
 wmem_tree_lookup(wmem_tree_t *tree, const void *key, compare_func cmp)
 {
     wmem_tree_node_t *node;
@@ -437,7 +462,7 @@ wmem_tree_lookup(wmem_tree_t *tree, const void *key, compare_func cmp)
     while (node) {
         int result = cmp(key, node->key);
         if (result == 0) {
-            return node->data;
+            return node;
         }
         else if (result < 0) {
             node = node->left;
@@ -503,6 +528,35 @@ wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cm
 }
 
 void
+wmem_tree_remove(wmem_tree_t *tree, const void *key, compare_func cmp)
+{
+    wmem_tree_node_t *node, *movingNode;
+    wmem_node_color_t origColor;
+
+    node = wmem_tree_lookup(tree, key, cmp);
+    if (node == NULL) {
+        return;
+    }
+
+    /* node that will be removed (for at most one leaf node) or
+     * moved (for two leaf nodes). */
+    movingNode = node;
+    origColor = movingNode->color;
+
+    if (node->left == NULL) {
+        fixNode = node->right;
+        rb_transplant(tree, node, node->right);
+    } else if (node->right == NULL) {
+        fixNode = node->left;
+        rb_transplant(tree, node, node->left);
+    } else {
+        movingNode = tree_minimum(node->right);
+        origColor = movingNode->color;
+
+    }
+}
+
+void
 wmem_tree_insert32(wmem_tree_t *tree, guint32 key, void *data)
 {
     lookup_or_insert32(tree, key, NULL, data, FALSE, TRUE);
@@ -622,6 +676,7 @@ void *
 wmem_tree_lookup_string(wmem_tree_t* tree, const gchar* k, guint32 flags)
 {
     compare_func cmp;
+    wmem_tree_node_t *node;
 
     if (flags & WMEM_TREE_STRING_NOCASE) {
         cmp = (compare_func)g_ascii_strcasecmp;
@@ -629,7 +684,8 @@ wmem_tree_lookup_string(wmem_tree_t* tree, const gchar* k, guint32 flags)
         cmp = (compare_func)strcmp;
     }
 
-    return wmem_tree_lookup(tree, k, cmp);
+    node = wmem_tree_lookup(tree, k, cmp);
+    return node ? node->data : NULL;
 }
 
 void *