Introduces augmented interval trees

Interval trees (wmem_itree_t) are implemented as an extension of wmem_tree with a
guint64-based range as the key.
This is useful for instance in MPTCP analysis, to look for packets
matching a range defined by a mapping across TCP subflows.

Change-Id: Iea706d44fe975e390a4191ad0257ef37d5c71525
Reviewed-on: https://code.wireshark.org/review/11714
Reviewed-by: Evan Huus <eapache@gmail.com>
Petri-Dish: Evan Huus <eapache@gmail.com>
Tested-by: Petri Dish Buildbot <buildbot-no-reply@wireshark.org>
Reviewed-by: Guy Harris <guy@alum.mit.edu>

9 files changed, 592 insertions, 61 deletions

diff --git a/epan/CMakeLists.txt b/epan/CMakeLists.txt
index fe468c5534..5bb903cab4 100644
--- a/epan/CMakeLists.txt
+++ b/epan/CMakeLists.txt
@@ -1545,6 +1545,7 @@ set(WMEM_FILES
 	wmem/wmem_allocator_block_fast.c
 	wmem/wmem_allocator_simple.c
 	wmem/wmem_allocator_strict.c
+	wmem/wmem_interval_tree.c
 	wmem/wmem_list.c
 	wmem/wmem_map.c
 	wmem/wmem_miscutl.c
diff --git a/epan/wmem/Makefile.common b/epan/wmem/Makefile.common
index dfd39f5af1..10d734a414 100644
--- a/epan/wmem/Makefile.common
+++ b/epan/wmem/Makefile.common
@@ -35,7 +35,8 @@ LIBWMEM_SRC =				\
 	wmem_stack.c			\
 	wmem_strbuf.c			\
 	wmem_strutl.c			\
-	wmem_tree.c			\
+	wmem_tree.c				\
+	wmem_interval_tree.c	\
 	wmem_user_cb.c
 
 LIBWMEM_INCLUDES =			\
@@ -56,7 +57,9 @@ LIBWMEM_INCLUDES =			\
 	wmem_stack.h			\
 	wmem_strbuf.h			\
 	wmem_strutl.h			\
-	wmem_tree.h			\
+	wmem_tree.h				\
+	wmem_tree-int.h	\
+	wmem_interval_tree.h	\
 	wmem_user_cb.h			\
 	wmem_user_cb_int.h
 
diff --git a/epan/wmem/wmem.h b/epan/wmem/wmem.h
index a10f369f69..3ad3853f71 100644
--- a/epan/wmem/wmem.h
+++ b/epan/wmem/wmem.h
@@ -35,6 +35,7 @@
 #include "wmem_strbuf.h"
 #include "wmem_strutl.h"
 #include "wmem_tree.h"
+#include "wmem_interval_tree.h"
 #include "wmem_user_cb.h"
 
 #endif /* __WMEM_H__ */
diff --git a/epan/wmem/wmem_interval_tree.c b/epan/wmem/wmem_interval_tree.c
new file mode 100644
index 0000000000..ca91d45ca5
--- /dev/null
+++ b/epan/wmem/wmem_interval_tree.c
@@ -0,0 +1,194 @@
+/* wmem_interval_tree.c
+ * Implements an augmented interval tree
+ * Based on the red-black tree implementation in epan/wmem.*
+ * Copyright 2015, Matthieu coudron <matthieu.coudron@lip6.fr>
+ *
+ * Wireshark - Network traffic analyzer
+ * By Gerald Combs <gerald@wireshark.org>
+ * Copyright 1998 Gerald Combs
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+ */
+
+#include "config.h"
+
+#include <string.h>
+#include <stdio.h>
+#include <glib.h>
+
+#include "wmem_core.h"
+#include "wmem_tree-int.h"
+#include "wmem_strutl.h"
+#include "wmem_interval_tree.h"
+#include "wmem_user_cb.h"
+
+
+
+static void
+print_range(const void *value)
+{
+    wmem_range_t *range = (wmem_range_t *)value;
+    if(!value) {
+        return;
+    }
+    printf("Range: low=%lu high=%lu max_edge=%lu\n", range->low, range->high, range->max_edge);
+}
+
+/**
+ * In an augmented interval tree, each node saves the maximum edge of its child subtrees
+ * This function compares the children max_edge with the current max_edge
+ * and propagates any change to the parent nodes.
+ */
+static void
+update_max_edge(wmem_tree_node_t *node)
+{
+    wmem_range_t *range;
+    wmem_range_t *range_l;
+    wmem_range_t *range_r;
+    guint64 maxEdge  = 0;
+
+    if(!node) {
+        return ;
+    }
+
+    range = (wmem_range_t *)node->key;
+
+    range_l = (node->left) ? (wmem_range_t *) (node->left->key) : NULL;
+    range_r = (node->right) ? (wmem_range_t *) (node->right->key) : NULL;
+
+    maxEdge = range->max_edge;
+
+    if(range_r) {
+        maxEdge = MAX(maxEdge, range_r->max_edge) ;
+    }
+    if(range_l) {
+        maxEdge = MAX(maxEdge, range_l->max_edge) ;
+    }
+
+    /* a change was made, update the parent nodes */
+    if(range->max_edge != maxEdge) {
+        range->max_edge = maxEdge;
+        update_max_edge(node->parent);
+    }
+}
+
+gboolean
+wmem_itree_range_overlap(const wmem_range_t *r1, const wmem_range_t *r2)
+{
+    return (r1->low <= r2->high && r2->low <= r1->high);
+}
+
+wmem_itree_t *
+wmem_itree_new(wmem_allocator_t *allocator)
+{
+    wmem_itree_t *tree      = wmem_tree_new(allocator);
+    tree->post_rotation_cb  = &update_max_edge;
+    return tree;
+}
+
+gboolean
+wmem_itree_is_empty(wmem_itree_t *tree)
+{
+    return wmem_tree_is_empty(tree);
+}
+
+static int
+wmem_tree_compare_ranges(const wmem_range_t *ra, const wmem_range_t *rb)
+{
+    if( ra->low == rb->low) {
+        return 0;
+    }
+    else if(ra->low < rb->low) {
+        return -1;
+    }
+    else {
+        return 1;
+    }
+}
+
+
+void
+wmem_itree_insert(wmem_itree_t *tree, const guint64 low, const guint64 high, void *data)
+{
+    wmem_tree_node_t *node;
+    wmem_range_t *range = (wmem_range_t *)wmem_new(tree->allocator, wmem_range_t);
+
+    g_assert(low <= high);
+    range->low = low;
+    range->high = high;
+    range->max_edge = high;
+
+    node = wmem_tree_insert(tree, range, data, (compare_func)wmem_tree_compare_ranges);
+
+    /* Even If no rotations, still a need to update max_edge */
+    update_max_edge(node);
+}
+
+
+static void
+wmem_itree_find_intervals_in_subtree(wmem_tree_node_t *node, wmem_range_t requested, wmem_list_t *results)
+{
+    const wmem_range_t* current;
+
+    if(!node) {
+        return;
+    }
+    current = (wmem_range_t*)node->key;
+
+    /* there is no child that can possibly match */
+    if(requested.low > current->max_edge) {
+        return;
+    }
+
+    if(wmem_itree_range_overlap(current, &requested)) {
+        wmem_list_prepend(results, node->data);
+    }
+
+    wmem_itree_find_intervals_in_subtree(node->left, requested, results);
+    wmem_itree_find_intervals_in_subtree(node->right, requested, results);
+}
+
+wmem_list_t *
+wmem_itree_find_intervals(wmem_itree_t *tree, wmem_allocator_t *allocator, guint64 low, guint64 high)
+{
+    wmem_list_t *results = NULL;
+    wmem_range_t requested;
+    requested.low  = low;
+    requested.high = high;
+    results = wmem_list_new(allocator);
+
+    wmem_itree_find_intervals_in_subtree(tree->root, requested, results);
+    return results;
+}
+
+
+void
+wmem_print_itree(wmem_tree_t *tree)
+{
+    wmem_print_tree(tree, &print_range, NULL);
+}
+
+/*
+ * Editor modelines  -  http://www.wireshark.org/tools/modelines.html
+ *
+ * Local variables:
+ * c-basic-offset: 4
+ * tab-width: 8
+ * indent-tabs-mode: nil
+ * End:
+ *
+ * vi: set shiftwidth=4 tabstop=8 expandtab:
+ * :indentSize=4:tabSize=8:noTabs=true:
+ */
diff --git a/epan/wmem/wmem_interval_tree.h b/epan/wmem/wmem_interval_tree.h
new file mode 100644
index 0000000000..c832cb780f
--- /dev/null
+++ b/epan/wmem/wmem_interval_tree.h
@@ -0,0 +1,110 @@
+/* wmem_interval_tree.h
+ * Definitions for the Wireshark Memory Manager Red-Black Tree
+ * Based on the red-black tree implementation in epan/emem.*
+ * Copyright 2015, Matthieu coudron <matthieu.coudron@lip6.fr>
+ *
+ * Wireshark - Network traffic analyzer
+ * By Gerald Combs <gerald@wireshark.org>
+ * Copyright 1998 Gerald Combs
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+ */
+#ifndef __WMEM_INTERVAL_TREE_H__
+#define __WMEM_INTERVAL_TREE_H__
+
+#include "wmem_core.h"
+#include "wmem_tree.h"
+#include "wmem_list.h"
+
+
+#ifdef __cplusplus
+extern "C" {
+#endif /* __cplusplus */
+
+/** @addtogroup wmem
+ *  @{
+ *    @defgroup wmem-interval-tree Interval Tree
+ *
+ * http://www.geeksforgeeks.org/interval-tree/
+ * The idea is to augment a self-balancing Binary Search Tree (BST) like Red Black Tree, AVL Tree, etc ...
+ * to maintain a set of intervals so that all operations can be done in O(Logn) time.
+ *    @{
+ * Following wikipedia's convention this is an augmented tree rather then an interval tree
+ * http://www.wikiwand.com/en/Interval_tree
+ */
+
+struct _wmem_tree_t;
+typedef struct _wmem_tree_t wmem_itree_t;
+
+struct _wmem_range_t {
+    guint64 low;        /* low is used as the key in the binary tree */
+    guint64 high;       /* Max value of the range */
+    guint64 max_edge;   /* max value among subtrees */
+};
+
+WS_DLL_PUBLIC
+wmem_itree_t *
+wmem_itree_new(wmem_allocator_t *allocator)
+G_GNUC_MALLOC;
+
+
+/** Returns true if the tree is empty (has no nodes). */
+WS_DLL_PUBLIC
+gboolean
+wmem_itree_is_empty(wmem_itree_t *tree);
+
+
+/** Inserts a range low-high indexed by "low" in O(log(n)).
+ * As in wmem_tree, if a key "low" already exists, it will be overwritten with the new data
+ *
+ */
+WS_DLL_PUBLIC
+void
+wmem_itree_insert(wmem_itree_t *tree, const guint64 low, const guint64 high, void *data);
+
+
+/*
+ * Save results in a wmem_list with the scope passed as a parameter.
+ * wmem_list_t is always allocated even if there is no result
+ */
+WS_DLL_PUBLIC
+wmem_list_t *
+wmem_itree_find_intervals(wmem_itree_t *tree, wmem_allocator_t *allocator, guint64 low, guint64 high);
+
+
+/**
+ * Print ranges along the tree
+ */
+void
+wmem_print_itree(wmem_itree_t *tree);
+
+#ifdef __cplusplus
+}
+#endif /* __cplusplus */
+
+#endif /* __WMEM_INTERVAL_TREE_H__ */
+
+/*
+ * Editor modelines  -  http://www.wireshark.org/tools/modelines.html
+ *
+ * Local variables:
+ * c-basic-offset: 4
+ * tab-width: 8
+ * indent-tabs-mode: nil
+ * End:
+ *
+ * vi: set shiftwidth=4 tabstop=8 expandtab:
+ * :indentSize=4:tabSize=8:noTabs=true:
+ */
diff --git a/epan/wmem/wmem_test.c b/epan/wmem/wmem_test.c
index 307125c115..9d95d0e66f 100644
--- a/epan/wmem/wmem_test.c
+++ b/epan/wmem/wmem_test.c
@@ -27,6 +27,7 @@
 #include <glib.h>
 
 #include "wmem.h"
+#include "wmem_tree-int.h"
 #include "wmem_allocator.h"
 #include "wmem_allocator_block.h"
 #include "wmem_allocator_block_fast.h"
@@ -982,6 +983,109 @@ wmem_test_tree(void)
     wmem_destroy_allocator(allocator);
 }
 
+
+/* to be used as userdata in the callback wmem_test_itree_check_overlap_cb*/
+typedef struct wmem_test_itree_user_data {
+    wmem_range_t range;
+    guint counter;
+} wmem_test_itree_user_data_t;
+
+
+/* increase userData counter in case the range match the userdata range */
+static gboolean
+wmem_test_itree_check_overlap_cb (const void *key, void *value _U_, void *userData)
+{
+    const wmem_range_t *ckey = (const wmem_range_t *)key;
+    struct wmem_test_itree_user_data * d = (struct wmem_test_itree_user_data *)userData;
+    g_assert(key);
+    g_assert(d);
+
+    if(wmem_itree_range_overlap(ckey, &d->range)) {
+        d->counter++;
+    }
+
+    return FALSE;
+}
+
+
+static gboolean
+wmem_test_overlap(guint64 low, guint64 high, guint64 lowbis, guint64 highbis)
+{
+    wmem_range_t r1 = {low, high, 0};
+    wmem_range_t r2 = {lowbis, highbis, 0};
+    return wmem_itree_range_overlap(&r1, &r2);
+}
+
+static void
+wmem_test_itree(void)
+{
+    wmem_allocator_t   *allocator, *extra_allocator;
+    wmem_itree_t       *tree;
+    int i = 0;
+    gint32 max_rand = 0;
+    wmem_test_itree_user_data_t userData;
+
+    allocator       = wmem_allocator_new(WMEM_ALLOCATOR_STRICT);
+    extra_allocator = wmem_allocator_new(WMEM_ALLOCATOR_STRICT);
+
+    tree = wmem_itree_new(allocator);
+    g_assert(tree);
+    g_assert(wmem_itree_is_empty(tree));
+
+    wmem_free_all(allocator);
+
+    /* make sure that wmem_test_overlap is correct (well it's no proof but...)*/
+    g_assert(wmem_test_overlap(0, 10, 0, 4));
+    g_assert(wmem_test_overlap(0, 10, 9, 14));
+    g_assert(wmem_test_overlap(5, 10, 3, 8));
+    g_assert(wmem_test_overlap(5, 10, 1, 12));
+    g_assert(!wmem_test_overlap(0, 10, 11, 12));
+
+    /* Generate a reference range, then fill an itree with random ranges,
+    then we count greedily the number of overlapping ranges and compare
+    the result with the optimized result
+     */
+
+    userData.counter = 0;
+
+    tree = wmem_itree_new(allocator);
+    wmem_range_t range, r2;
+
+    /* even though keys are uint64_t, we use G_MAXINT32 as a max because of the type returned by
+      g_test_rand_int_range.
+     */
+    max_rand = G_MAXINT32;
+    r2.max_edge = range.max_edge = 0;
+    range.low = g_test_rand_int_range(0, max_rand);
+    range.high = g_test_rand_int_range( (gint32)range.low, (gint32)max_rand);
+    userData.range = range;
+
+    for (i=0; i<CONTAINER_ITERS; i++) {
+
+        wmem_list_t *results = NULL;
+
+        /* reset the search */
+        userData.counter = 0;
+        r2.low = (guint64)g_test_rand_int_range(0, 100);
+        r2.high = (guint64)g_test_rand_int_range( (gint32)r2.low, 100);
+
+        wmem_itree_insert(tree, r2.low, r2.high, GINT_TO_POINTER(i));
+
+        /* greedy search */
+        wmem_tree_foreach(tree, wmem_test_itree_check_overlap_cb, &userData);
+
+        /* Optimized search */
+        results = wmem_itree_find_intervals(tree, allocator, range.low, range.high);
+
+        /* keep it as a loop instead of wmem_list_count in case one */
+        g_assert(wmem_list_count(results) == userData.counter);
+    }
+
+    wmem_destroy_allocator(extra_allocator);
+    wmem_destroy_allocator(allocator);
+}
+
+
 int
 main(int argc, char **argv)
 {
@@ -1007,6 +1111,7 @@ main(int argc, char **argv)
     g_test_add_func("/wmem/datastruct/stack",  wmem_test_stack);
     g_test_add_func("/wmem/datastruct/strbuf", wmem_test_strbuf);
     g_test_add_func("/wmem/datastruct/tree",   wmem_test_tree);
+    g_test_add_func("/wmem/datastruct/itree",  wmem_test_itree);
 
     ret = g_test_run();
 
diff --git a/epan/wmem/wmem_tree-int.h b/epan/wmem/wmem_tree-int.h
new file mode 100644
index 0000000000..4ced1f3796
--- /dev/null
+++ b/epan/wmem/wmem_tree-int.h
@@ -0,0 +1,99 @@
+/* wmem_tree_internals.h
+ * Definitions for the Wireshark Memory Manager Red-Black Tree
+ * Copyright 2013, Evan Huus <eapache@gmail.com>
+ *
+ * Wireshark - Network traffic analyzer
+ * By Gerald Combs <gerald@wireshark.org>
+ * Copyright 1998 Gerald Combs
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+ */
+
+#ifndef __WMEM_TREE_INT_H__
+#define __WMEM_TREE_INT_H__
+
+
+#ifdef __cplusplus
+extern "C" {
+#endif /* __cplusplus */
+
+typedef enum _wmem_node_color_t {
+    WMEM_NODE_COLOR_RED,
+    WMEM_NODE_COLOR_BLACK
+} wmem_node_color_t;
+
+
+struct _wmem_tree_node_t {
+    struct _wmem_tree_node_t *parent;
+    struct _wmem_tree_node_t *left;
+    struct _wmem_tree_node_t *right;
+
+    const void *key;
+    void *data;
+
+    wmem_node_color_t color;
+    gboolean          is_subtree;
+    gboolean          is_removed;
+
+
+};
+
+typedef struct _wmem_tree_node_t wmem_tree_node_t;
+
+
+typedef struct _wmem_itree_node_t wmem_itree_node_t;
+
+struct _wmem_tree_t {
+    wmem_allocator_t *master;
+    wmem_allocator_t *allocator;
+    wmem_tree_node_t *root;
+    guint             master_cb_id;
+    guint             slave_cb_id;
+
+    void (*post_rotation_cb)(wmem_tree_node_t *);
+};
+
+typedef struct _wmem_tree_t wmem_tree_t;
+
+
+
+typedef int (*compare_func)(const void *a, const void *b);
+
+wmem_tree_node_t *
+wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cmp);
+
+typedef struct _wmem_range_t wmem_range_t;
+
+gboolean
+wmem_itree_range_overlap(const wmem_range_t *r1, const wmem_range_t *r2);
+
+#ifdef __cplusplus
+}
+#endif /* __cplusplus */
+
+#endif /* __WMEM_TREE__INTERNALS_H__ */
+
+/*
+ * Editor modelines  -  http://www.wireshark.org/tools/modelines.html
+ *
+ * Local variables:
+ * c-basic-offset: 4
+ * tab-width: 8
+ * indent-tabs-mode: nil
+ * End:
+ *
+ * vi: set shiftwidth=4 tabstop=8 expandtab:
+ * :indentSize=4:tabSize=8:noTabs=true:
+ */
diff --git a/epan/wmem/wmem_tree.c b/epan/wmem/wmem_tree.c
index 070f951b7d..8332c45a48 100644
--- a/epan/wmem/wmem_tree.c
+++ b/epan/wmem/wmem_tree.c
@@ -31,37 +31,11 @@
 #include "wmem_core.h"
 #include "wmem_strutl.h"
 #include "wmem_tree.h"
+#include "wmem_tree-int.h"
 #include "wmem_user_cb.h"
 
-typedef enum _wmem_node_color_t {
-    WMEM_NODE_COLOR_RED,
-    WMEM_NODE_COLOR_BLACK
-} wmem_node_color_t;
 
-struct _wmem_tree_node_t {
-    struct _wmem_tree_node_t *parent;
-    struct _wmem_tree_node_t *left;
-    struct _wmem_tree_node_t *right;
 
-    const void *key;
-    void *data;
-
-    wmem_node_color_t color;
-    gboolean          is_subtree;
-    gboolean          is_removed;
-};
-
-typedef struct _wmem_tree_node_t wmem_tree_node_t;
-
-struct _wmem_tree_t {
-    wmem_allocator_t *master;
-    wmem_allocator_t *allocator;
-    wmem_tree_node_t *root;
-    guint             master_cb_id;
-    guint             slave_cb_id;
-};
-
-typedef int (*compare_func)(const void *a, const void *b);
 
 static wmem_tree_node_t *
 node_uncle(wmem_tree_node_t *node)
@@ -111,6 +85,10 @@ rotate_left(wmem_tree_t *tree, wmem_tree_node_t *node)
         node->right->parent = node;
     }
     node->parent->left = node;
+
+    if (tree->post_rotation_cb) {
+        tree->post_rotation_cb (node);
+    }
 }
 
 static void
@@ -135,6 +113,11 @@ rotate_right(wmem_tree_t *tree, wmem_tree_node_t *node)
         node->left->parent = node;
     }
     node->parent->right = node;
+
+
+    if (tree->post_rotation_cb) {
+        tree->post_rotation_cb (node);
+    }
 }
 
 static void
@@ -232,7 +215,7 @@ wmem_tree_new(wmem_allocator_t *allocator)
     tree->master    = allocator;
     tree->allocator = allocator;
     tree->root      = NULL;
-
+    tree->post_rotation_cb = NULL;
     return tree;
 }
 
@@ -272,6 +255,7 @@ wmem_tree_new_autoreset(wmem_allocator_t *master, wmem_allocator_t *slave)
     tree->master    = master;
     tree->allocator = slave;
     tree->root      = NULL;
+    tree->post_rotation_cb      = NULL;
 
     tree->master_cb_id = wmem_register_callback(master, wmem_tree_destroy_cb,
             tree);
@@ -310,8 +294,13 @@ create_node(wmem_allocator_t *allocator, wmem_tree_node_t *parent, const void *k
 }
 
 #define CREATE_DATA(TRANSFORM, DATA) ((TRANSFORM) ? (TRANSFORM)(DATA) : (DATA))
-static void *
-lookup_or_insert32(wmem_tree_t *tree, guint32 key,
+
+
+/**
+ * return inserted node
+ */
+static wmem_tree_node_t *
+lookup_or_insert32_node(wmem_tree_t *tree, guint32 key,
         void*(*func)(void*), void* data, gboolean is_subtree, gboolean replace)
 {
     wmem_tree_node_t *node     = tree->root;
@@ -322,7 +311,7 @@ lookup_or_insert32(wmem_tree_t *tree, guint32 key,
         new_node = create_node(tree->allocator, NULL, GUINT_TO_POINTER(key),
                 CREATE_DATA(func, data), WMEM_NODE_COLOR_BLACK, is_subtree);
         tree->root = new_node;
-        return new_node->data;
+        return new_node;
     }
 
     /* it was not the new root so walk the tree until we find where to
@@ -334,7 +323,7 @@ lookup_or_insert32(wmem_tree_t *tree, guint32 key,
             if (replace) {
                 node->data = CREATE_DATA(func, data);
             }
-            return node->data;
+            return node;
         }
         else if (key < GPOINTER_TO_UINT(node->key)) {
             if (node->left) {
@@ -365,7 +354,16 @@ lookup_or_insert32(wmem_tree_t *tree, guint32 key,
     /* node will now point to the newly created node */
     rb_insert_case1(tree, new_node);
 
-    return new_node->data;
+    return new_node;
+}
+
+
+static void *
+lookup_or_insert32(wmem_tree_t *tree, guint32 key,
+        void*(*func)(void*), void* data, gboolean is_subtree, gboolean replace)
+{
+    wmem_tree_node_t *node = lookup_or_insert32_node(tree, key, func, data, is_subtree, replace);
+    return node->data;
 }
 
 static void *
@@ -395,7 +393,7 @@ wmem_tree_lookup(wmem_tree_t *tree, const void *key, compare_func cmp)
     return NULL;
 }
 
-static void
+wmem_tree_node_t *
 wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cmp)
 {
     wmem_tree_node_t *node = tree->root;
@@ -405,7 +403,7 @@ wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cm
     if (!node) {
         tree->root = create_node(tree->allocator, node, key,
                 data, WMEM_NODE_COLOR_BLACK, FALSE);
-        return;
+        return tree->root;
     }
 
     /* it was not the new root so walk the tree until we find where to
@@ -416,7 +414,7 @@ wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cm
         if (result == 0) {
             node->data = data;
             node->is_removed = data ? FALSE : TRUE;
-            return;
+            return node;
         }
         else if (result < 0) {
             if (node->left) {
@@ -443,6 +441,8 @@ wmem_tree_insert(wmem_tree_t *tree, const void *key, void *data, compare_func cm
 
     /* node will now point to the newly created node */
     rb_insert_case1(tree, new_node);
+
+    return new_node;
 }
 
 void
@@ -700,59 +700,71 @@ wmem_tree_foreach(wmem_tree_t* tree, wmem_foreach_func callback,
     return wmem_tree_foreach_nodes(tree->root, callback, user_data);
 }
 
-static void wmem_print_subtree(wmem_tree_t *tree, guint32 level);
+static void wmem_print_subtree(wmem_tree_t *tree, guint32 level, wmem_printer_func key_printer, wmem_printer_func data_printer);
 
 static void
-wmem_tree_print_nodes(const char *prefix, wmem_tree_node_t *node, guint32 level)
-{
+wmem_print_indent(guint32 level) {
     guint32 i;
+    for (i=0; i<level; i++) {
+        printf("    ");
+    }
+}
 
+static void
+wmem_tree_print_nodes(const char *prefix, wmem_tree_node_t *node, guint32 level,
+    wmem_printer_func key_printer, wmem_printer_func data_printer)
+{
     if (!node)
         return;
 
-    for (i=0; i<level; i++) {
-        printf("    ");
-    }
+    wmem_print_indent(level);
 
-    printf("%sNODE:%p parent:%p left:%p right:%p colour:%s key:%u %s:%p\n",
+    printf("%sNODE:%p parent:%p left:%p right:%p colour:%s key:%p %s:%p\n",
             prefix,
             (void *)node, (void *)node->parent,
             (void *)node->left, (void *)node->right,
-            node->color?"Black":"Red", GPOINTER_TO_UINT(node->key),
+            node->color?"Black":"Red", node->key,
             node->is_subtree?"tree":"data", node->data);
+    if(key_printer) {
+        wmem_print_indent(level);
+        key_printer(node->key);
+        printf("\n");
+    }
+    if(data_printer) {
+        wmem_print_indent(level);
+        data_printer(node->data);
+        printf("\n");
+    }
+
     if (node->left)
-        wmem_tree_print_nodes("L-", node->left, level+1);
+        wmem_tree_print_nodes("L-", node->left, level+1, key_printer, data_printer);
     if (node->right)
-        wmem_tree_print_nodes("R-", node->right, level+1);
+        wmem_tree_print_nodes("R-", node->right, level+1, key_printer, data_printer);
 
     if (node->is_subtree)
-        wmem_print_subtree((wmem_tree_t *)node->data, level+1);
+        wmem_print_subtree((wmem_tree_t *)node->data, level+1, key_printer, data_printer);
 }
 
+
 static void
-wmem_print_subtree(wmem_tree_t *tree, guint32 level)
+wmem_print_subtree(wmem_tree_t *tree, guint32 level, wmem_printer_func key_printer, wmem_printer_func data_printer)
 {
-    guint32 i;
-
     if (!tree)
         return;
 
-    for (i=0; i<level; i++) {
-        printf("    ");
-    }
+    wmem_print_indent(level);
 
     printf("WMEM tree:%p root:%p\n", (void *)tree, (void *)tree->root);
     if (tree->root) {
-        wmem_tree_print_nodes("Root-", tree->root, level);
+        wmem_tree_print_nodes("Root-", tree->root, level, key_printer, data_printer);
     }
 }
 
 void
-wmem_print_tree(wmem_tree_t *tree)
+wmem_print_tree(wmem_tree_t *tree, wmem_printer_func key_printer, wmem_printer_func data_printer)
 {
-    wmem_print_subtree(tree, 0);
+    wmem_print_subtree(tree, 0, key_printer, data_printer);
 }
-
 /*
  * Editor modelines  -  http://www.wireshark.org/tools/modelines.html
  *
diff --git a/epan/wmem/wmem_tree.h b/epan/wmem/wmem_tree.h
index 33058c0174..a31de3c1e9 100644
--- a/epan/wmem/wmem_tree.h
+++ b/epan/wmem/wmem_tree.h
@@ -206,6 +206,11 @@ wmem_tree_lookup32_array_le(wmem_tree_t *tree, wmem_tree_key_t *key);
  */
 typedef gboolean (*wmem_foreach_func)(const void *key, void *value, void *userdata);
 
+
+/** Function type to print key/data of nodes in wmem_print_tree_verbose */
+typedef void (*wmem_printer_func)(const void *data);
+
+
 /** Traverse the tree and call callback(value, userdata) for each value found.
  * Returns TRUE if the traversal was ended prematurely by the callback.
  */
@@ -214,9 +219,10 @@ gboolean
 wmem_tree_foreach(wmem_tree_t* tree, wmem_foreach_func callback,
         void *user_data);
 
-/** Prints the structure of the tree to stdout. Primarily for debugging. */
+
+/* Accepts callbacks to print the key and/or data (both printers can be null) */
 void
-wmem_print_tree(wmem_tree_t *tree);
+wmem_print_tree(wmem_tree_t *tree, wmem_printer_func key_printer, wmem_printer_func data_printer);
 
 /**   @}
  *  @} */