1 files changed, 113 insertions, 0 deletions

diff --git a/epan/wmem/wmem_interval_tree.c b/epan/wmem/wmem_interval_tree.c
index 1d89a06539..3bcd4d3d35 100644
--- a/epan/wmem/wmem_interval_tree.c
+++ b/epan/wmem/wmem_interval_tree.c
@@ -136,6 +136,119 @@ wmem_itree_insert(wmem_itree_t *tree, const guint64 low, const guint64 high, voi
     update_max_edge(node);
 }
 
+struct point_search_result {
+    wmem_tree_node_t *left;
+    wmem_tree_node_t *middle;
+    wmem_tree_node_t *right;
+};
+
+static void
+wmem_itree_find_points_in_subtree(wmem_tree_node_t *node, guint64 low, guint64 pos, guint64 high, struct point_search_result *result)
+{
+    const wmem_range_t *current;
+
+    if (!node) {
+        return;
+    }
+    current = (wmem_range_t *)node->key;
+
+    /* there is no child that can possibly match */
+    if (low > current->max_edge) {
+        return;
+    }
+
+    /* TODO implement a wmem_tree_remove method such that this check can be
+     * removed. */
+    if (!node->is_removed) {
+        if (current->low <= pos && pos <= current->high) {
+            /* overlap: point x overlaps the current interval. */
+            result->middle = node;
+        } else {
+            if (low != pos && current->high == low) {
+                /* current node is an interval at x-1 */
+                result->left = node;
+            }
+            if (high != pos && current->low == high) {
+                /* current node is an interval at x+1 */
+                result->right = node;
+            }
+        }
+    }
+
+    wmem_itree_find_points_in_subtree(node->left, low, pos, high, result);
+    wmem_itree_find_points_in_subtree(node->right, low, pos, high, result);
+}
+
+
+void
+wmem_itree_insert_point(wmem_itree_t *tree, const guint64 pos, wmem_itree_data_callback callback)
+{
+    const guint64 low = pos > 0 ? pos - 1 : pos;
+    const guint64 high = pos < G_MAXUINT64 ? pos + 1 : pos;
+    struct point_search_result result = { NULL, NULL, NULL };
+    wmem_range_t *leftKey, *rightKey;
+    void *leftData = NULL, *rightData = NULL;
+
+    /*
+     * Look for intervals overlapping at x-1, x or x+1.
+     * If an interval matches "x", then the interval range remains unmodified.
+     * Else if it is adjacent to both x-1 or x+1, merge those intervals.
+     * Else if it is adjacent to either x-1 or x+1, extend interval.
+     * Otherwise add a new interval.
+     */
+    wmem_itree_find_points_in_subtree(tree->root, low, pos, high, &result);
+    if (result.middle) {
+        /* interval exists, do not replace data. */
+        /* TODO maybe allow existing data to be replaced? Something like:
+         * result.middle->data = callback(result.middle->data, result.middle->data) */
+    } else if (result.left && result.right) {
+        /* need to merge the intervals. */
+        leftData = result.left->data;
+        rightData = result.right->data;
+
+        /* remove (insert NULL for) right key, it will be merged into left node. */
+        leftKey = (wmem_range_t *)result.left->key;
+        rightKey = (wmem_range_t *)result.right->key;
+        g_assert(leftKey->high + 1 == pos);
+        g_assert(rightKey->low - 1 == pos);
+        const wmem_range_t oldRightRange = *rightKey;
+        /* TODO use wmem_tree_remove when available. */
+        wmem_tree_insert(tree, rightKey, NULL, (compare_func)wmem_tree_compare_ranges);
+        leftKey->high = oldRightRange.high;
+        leftKey->max_edge = oldRightRange.max_edge;
+        update_max_edge(result.left);
+
+        result.left->data = callback(leftData, rightData);
+    } else if (result.left) {
+        /* Interval at x-1, simply extend left node. */
+        leftData = result.left->data;
+
+        leftKey = (wmem_range_t *)result.left->key;
+        g_assert(leftKey->high + 1 == pos);
+        leftKey->high = pos;
+        leftKey->max_edge = MAX(leftKey->max_edge, pos);
+        update_max_edge(result.left);
+
+        result.left->data = callback(leftData, NULL);
+    } else if (result.right) {
+        /* Interval at x+1. As tree is keyed by low, need to remove+insert. */
+        rightData = result.right->data;
+
+        rightKey = (wmem_range_t *)result.right->key;
+        g_assert(rightKey->low - 1 == pos);
+        const guint64 rightHigh = rightKey->high;
+        /* TODO use wmem_tree_remove when available. */
+        wmem_tree_insert(tree, rightKey, NULL, (compare_func)wmem_tree_compare_ranges);
+
+        rightData = callback(NULL, rightData);
+        wmem_itree_insert(tree, pos, rightHigh, rightData);
+    } else {
+        /* This is a new interval. */
+        leftData = callback(NULL, NULL);
+        wmem_itree_insert(tree, pos, pos, leftData);
+    }
+}
+
 
 static void
 wmem_itree_find_intervals_in_subtree(wmem_tree_node_t *node, wmem_range_t requested, wmem_list_t *results)