lib/interval_tree: Add an interval tree implementation

Add an interval tree implementation based on lib/binary_search_tree's
augmented capabilities.

Signed-off-by: Pedro Falcato <[email protected]>

kernel/Makefile

@@ -63,6 +63,7 @@ LDFLAGS:=$(LDFLAGS) -fkeep-inline-functions -Wl,--build-id=none
 
 include lib/compiler-rt/builtins/Makefile
 include lib/binary_search_tree/Makefile
+include lib/interval_tree/Makefile
 
 LIBS:=$(LIBS) -nostdlib
 ARCHDIR:=arch/$(HOSTARCH)

kernel/include/onyx/interval_tree.h

@@ -0,0 +1,144 @@
+/*
+ * Copyright (c) 2023 Pedro Falcato
+ * This file is part of Onyx, and is released under the terms of the MIT License
+ * check LICENSE at the root directory for more information
+ *
+ * SPDX-License-Identifier: MIT
+ */
+#ifndef _ONYX_INTERVAL_TREE_H
+#define _ONYX_INTERVAL_TREE_H
+
+#include <lib/binary_search_tree.h>
+
+#include <onyx/assert.h>
+#include <onyx/compiler.h>
+
+#ifdef __cplusplus
+#define CONSTEXPR constexpr
+#else
+#define CONSTEXPR
+#endif
+
+struct interval_tree_root
+{
+    struct bst_root root;
+};
+
+struct interval_tree_node
+{
+    struct bst_node node;
+    unsigned long start;
+    unsigned long end;
+    unsigned long max_end;
+};
+
+CONSTEXPR
+static inline void interval_tree_root_init(struct interval_tree_root *root)
+{
+    bst_root_initialize(&root->root);
+}
+
+CONSTEXPR
+static inline void interval_tree_node_init(struct interval_tree_node *node, unsigned long start,
+                                           unsigned long end)
+{
+    bst_node_initialize(&node->node);
+    node->start = start;
+    node->end = end;
+}
+
+__BEGIN_CDECLS
+
+/**
+ * @brief Insert an interval onto the interval tree
+ *
+ * @param root Tree to insert on
+ * @param node Node to insert
+ */
+void interval_tree_insert(struct interval_tree_root *root, struct interval_tree_node *node);
+
+/**
+ * @brief Remove an interval from the interval tree
+ *
+ * @param root Interval tree root
+ * @param node Interval to remove
+ */
+void interval_tree_remove(struct interval_tree_root *root, struct interval_tree_node *node);
+
+static inline struct interval_tree_node *__interval_tree_search(
+    const struct interval_tree_root *root, unsigned long start, unsigned long end)
+{
+    /* We may not use bst_search because it won't quite work with our interval tree that /may/
+     * contain duplicates.
+     */
+    DCHECK(root);
+
+    struct bst_node *tree_node = root->root.root;
+
+    /* We go through the tree and find the smallest node that matches our range in a normal-ish bst
+     * fashion.
+     */
+    while (tree_node)
+    {
+        struct interval_tree_node *node = container_of(tree_node, struct interval_tree_node, node);
+        struct interval_tree_node *left =
+            containerof_null_safe(node->node.child[0], struct interval_tree_node, node);
+        int cmp = -2;
+
+        if (left)
+        {
+            /* Check if left has any nodes in our range */
+            if (left->max_end >= start)
+                cmp = -1; /* If so, go left */
+        }
+
+        if (cmp == -2)
+        {
+            /* Do a standard bst cmp */
+            if (node->start <= end && start <= node->end)
+                cmp = 0;
+            else
+                cmp = start < node->start ? -1 : 1;
+        }
+
+        if (!cmp)
+        {
+            /* Note: smaller nodes are already covered by the if (left) up there */
+            return node;
+        }
+
+        tree_node = tree_node->child[cmp > 0];
+    }
+
+    return NULL;
+}
+
+static inline struct interval_tree_node *__interval_tree_next(struct interval_tree_root *root,
+                                                              struct interval_tree_node *node,
+                                                              unsigned long start,
+                                                              unsigned long end)
+{
+    DCHECK(root);
+    DCHECK(node);
+    struct bst_node *next = bst_next(&root->root, &node->node);
+    if (!next)
+        return NULL;
+
+    struct interval_tree_node *next_node = container_of(next, struct interval_tree_node, node);
+    if (next_node->start <= end && start <= next_node->end)
+        return next_node;
+    return NULL;
+}
+
+#define __for_intervals_in_range(root, entry, start, end)                              \
+    for (struct interval_tree_node *__node = __interval_tree_search(root, start, end); \
+         ((entry) = __node) != NULL; __node = __interval_tree_next(root, __node, start, end))
+
+#define for_intervals_in_range(root, entry, type, member, start, end)                  \
+    for (struct interval_tree_node *__node = __interval_tree_search(root, start, end); \
+         __node && ((entry) = container_of(__node, type, member));                     \
+         __node = __interval_tree_next(root, __node, start, end))
+
+__END_CDECLS
+
+#endif

kernel/include/onyx/utils.h

@@ -17,7 +17,8 @@ void *copy_page(void *vaddr, void *p2);
 void *copy_page_to_page(void *p1, void *p2);
 
 #define container_of(ptr, type, member) ((type *) ((char *) ptr - offsetof(type, member)))
-
+#define containerof_null_safe(ptr, type, member) \
+    ((ptr) == NULL ? NULL : container_of(ptr, type, member))
 #ifndef __cplusplus
 #define min(x, y) (x < y ? x : y)
 #else

kernel/lib/interval_tree/Makefile

@@ -0,0 +1,4 @@
+interval-tree-y:= interval_tree.o
+interval-tree-$(CONFIG_KUNIT)+= interval_tree_test.o
+
+obj-y+= $(patsubst %, lib/interval_tree/%, $(interval-tree-y))

kernel/lib/interval_tree/interval_tree.c

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2023 Pedro Falcato
+ * This file is part of Onyx, and is released under the terms of the MIT License
+ * check LICENSE at the root directory for more information
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <onyx/interval_tree.h>
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+
+static unsigned long interval_tree_max_compute(struct interval_tree_node *node)
+{
+    unsigned long max = node->end;
+    const struct interval_tree_node *left =
+        containerof_null_safe(node->node.child[0], struct interval_tree_node, node);
+    const struct interval_tree_node *right =
+        containerof_null_safe(node->node.child[1], struct interval_tree_node, node);
+
+    if (left)
+        max = MAX(max, left->max_end);
+    if (right)
+        max = MAX(max, right->max_end);
+    return max;
+}
+
+BST_AUGMENTED(interval_tree, struct interval_tree_node, node, unsigned long, max_end,
+              interval_tree_max_compute);
+
+/**
+ * @brief Insert an interval onto the interval tree
+ *
+ * @param root Tree to insert on
+ * @param node Node to insert
+ */
+void interval_tree_insert(struct interval_tree_root *root, struct interval_tree_node *node)
+{
+    struct bst_node **nodep, *parent, *cur;
+
+    cur = root->root.root;
+    nodep = &root->root.root;
+    parent = NULL;
+
+    while (cur)
+    {
+        struct interval_tree_node *__node = container_of(cur, struct interval_tree_node, node);
+        int res = node->start < __node->start ? -1 : 1;
+
+        parent = cur;
+        nodep = &cur->child[res > 0];
+        cur = *nodep;
+    }
+
+    bst_link(nodep, parent, &node->node);
+    interval_tree_propagate(&node->node, NULL);
+    bst_update_rank_insert(&root->root, &node->node, &interval_tree);
+}
+
+/**
+ * @brief Remove an interval from the interval tree
+ *
+ * @param root Interval tree root
+ * @param node Interval to remove
+ */
+void interval_tree_remove(struct interval_tree_root *root, struct interval_tree_node *node)
+{
+    bst_delete_augmented(&root->root, &node->node, &interval_tree);
+}

kernel/lib/interval_tree/interval_tree_test.cpp

@@ -0,0 +1,129 @@
+/*
+ * Copyright (c) 2023 Pedro Falcato
+ * This file is part of Onyx, and is released under the terms of the MIT License
+ * check LICENSE at the root directory for more information
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <stdio.h>
+
+#include <onyx/interval_tree.h>
+#include <onyx/kunit.h>
+
+static void add_range(struct interval_tree_root *root, struct interval_tree_node *node,
+                      unsigned long start, unsigned long end)
+{
+    interval_tree_node_init(node, start, end);
+    interval_tree_insert(root, node);
+}
+
+TEST(interval_tree, insert)
+{
+#define IT_SEEN 0xDEADUL
+    struct interval_tree_node nodes[4];
+    interval_tree_root tree;
+    interval_tree_root_init(&tree);
+    add_range(&tree, &nodes[0], 0, 10);
+    add_range(&tree, &nodes[1], 1, 3);
+    add_range(&tree, &nodes[2], 32, 1000);
+    add_range(&tree, &nodes[3], 8, 15);
+
+    struct interval_tree_node *entry;
+    int i = 0;
+    __for_intervals_in_range(&tree, entry, 1, 8)
+    {
+        entry->start = IT_SEEN + i++;
+    }
+
+    EXPECT_EQ(nodes[0].start, IT_SEEN);
+    EXPECT_EQ(nodes[1].start, IT_SEEN + 1);
+    EXPECT_NE(nodes[2].start, IT_SEEN + 2);
+    EXPECT_EQ(nodes[3].start, IT_SEEN + 2);
+#undef IT_SEEN
+}
+
+TEST(interval_tree, insert2)
+{
+#define IT_SEEN 0xDEADUL
+    struct interval_tree_node nodes[4];
+    interval_tree_root tree;
+    interval_tree_root_init(&tree);
+    add_range(&tree, &nodes[0], 4, 10);
+    add_range(&tree, &nodes[1], 1, 3);
+    add_range(&tree, &nodes[2], 1, 4);
+    add_range(&tree, &nodes[3], 8, 15);
+
+    struct interval_tree_node *entry;
+    int i = 0;
+    __for_intervals_in_range(&tree, entry, 0, 8)
+    {
+        entry->start = IT_SEEN + i++;
+    }
+
+    EXPECT_GE(nodes[0].start, IT_SEEN);
+    EXPECT_GE(nodes[1].start, IT_SEEN);
+    EXPECT_GE(nodes[2].start, IT_SEEN);
+    EXPECT_GE(nodes[3].start, IT_SEEN);
+#undef IT_SEEN
+}
+
+TEST(interval_tree, iterate_empty)
+{
+    struct interval_tree_node nodes[4];
+    interval_tree_root tree;
+    interval_tree_root_init(&tree);
+    add_range(&tree, &nodes[0], 4, 5);
+    add_range(&tree, &nodes[1], 1, 3);
+    add_range(&tree, &nodes[2], 1, 4);
+    add_range(&tree, &nodes[3], 8, 15);
+
+    int i = 0;
+    struct interval_tree_node *entry;
+
+    /* Iterate on an empty range (to the right) */
+    __for_intervals_in_range(&tree, entry, 16, 20)
+    {
+        i++;
+    }
+
+    ASSERT_EQ(0, i);
+
+    /* Iterate on an empty range (in the middle of the tree) */
+
+    __for_intervals_in_range(&tree, entry, 6, 7)
+    {
+        i++;
+    }
+
+    ASSERT_EQ(0, i);
+}
+
+TEST(interval_tree, remove_works)
+{
+    struct interval_tree_node nodes[4];
+    interval_tree_root tree;
+    interval_tree_root_init(&tree);
+    add_range(&tree, &nodes[0], 4, 10);
+    add_range(&tree, &nodes[1], 1, 3);
+    add_range(&tree, &nodes[2], 1, 4);
+    add_range(&tree, &nodes[3], 9, 15);
+
+    struct interval_tree_node *entry;
+    int i = 0;
+    __for_intervals_in_range(&tree, entry, 0, 8)
+    {
+        i++;
+    }
+
+    EXPECT_EQ(3, i);
+    i = 0;
+    interval_tree_remove(&tree, &nodes[1]);
+    interval_tree_remove(&tree, &nodes[0]);
+    __for_intervals_in_range(&tree, entry, 0, 8)
+    {
+        i++;
+    }
+
+    EXPECT_EQ(1, i);
+}