feat: add basic implementation and enhanced operations of red-black tree from Linux kernel source

3 files changed, 519 insertions, 0 deletions

diff --git a/na/include/lume/rbtree.h b/na/include/lume/rbtree.h
new file mode 100644
index 0000000..331875c
--- /dev/null
+++ b/na/include/lume/rbtree.h
@@ -0,0 +1,219 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+/*
+ * Red-Black Trees — public API
+ * Adapted from Linux kernel include/linux/rbtree.h
+ *
+ * (C) 1999  Andrea Arcangeli <andrea@suse.de>
+ */
+
+#pragma once
+
+#include <lume/rbtree_types.h>
+
+/* ------------------------------------------------------------------ */
+/*  Linux-compat macros for freestanding kernel                       */
+/* ------------------------------------------------------------------ */
+
+#ifndef container_of
+#define container_of(ptr, type, member) \
+	reinterpret_cast<type *>( \
+		reinterpret_cast<char *>(ptr) - __builtin_offsetof(type, member))
+#endif
+
+#ifndef READ_ONCE
+#define READ_ONCE(x)  (*(volatile __typeof__(x) *)&(x))
+#endif
+
+/* ------------------------------------------------------------------ */
+/*  Core macros                                                       */
+/* ------------------------------------------------------------------ */
+
+#define rb_parent(r)   ((struct rb_node *)((r)->__rb_parent_color & ~3UL))
+
+#define rb_entry(ptr, type, member)  container_of(ptr, type, member)
+
+#define RB_EMPTY_ROOT(root)  (READ_ONCE((root)->rb_node) == nullptr)
+
+#define RB_EMPTY_NODE(node)  \
+	((node)->__rb_parent_color == (unsigned long)(node))
+#define RB_CLEAR_NODE(node)  \
+	((node)->__rb_parent_color = (unsigned long)(node))
+
+/* ------------------------------------------------------------------ */
+/*  Core API (implemented in lib/rbtree.cc)                           */
+/* ------------------------------------------------------------------ */
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+void rb_insert_color(struct rb_node *, struct rb_root *);
+void rb_erase(struct rb_node *, struct rb_root *);
+
+struct rb_node *rb_next(const struct rb_node *);
+struct rb_node *rb_prev(const struct rb_node *);
+struct rb_node *rb_first(const struct rb_root *);
+struct rb_node *rb_last(const struct rb_root *);
+
+struct rb_node *rb_first_postorder(const struct rb_root *);
+struct rb_node *rb_next_postorder(const struct rb_node *);
+
+void rb_replace_node(struct rb_node *victim, struct rb_node *_new,
+		     struct rb_root *root);
+
+#ifdef __cplusplus
+}
+#endif
+
+/* ------------------------------------------------------------------ */
+/*  Inline helpers                                                    */
+/* ------------------------------------------------------------------ */
+
+static inline void rb_link_node(struct rb_node *node, struct rb_node *parent,
+				struct rb_node **rb_link)
+{
+	node->__rb_parent_color = (unsigned long)parent;
+	node->rb_left = node->rb_right = nullptr;
+	*rb_link = node;
+}
+
+#define rb_entry_safe(ptr, type, member) \
+	({ __typeof__(ptr) ____ptr = (ptr); \
+	   ____ptr ? rb_entry(____ptr, type, member) : nullptr; \
+	})
+
+/* ------------------------------------------------------------------ */
+/*  Cached (leftmost) tree helpers                                    */
+/* ------------------------------------------------------------------ */
+
+#define rb_first_cached(root) (root)->rb_leftmost
+
+static inline void rb_insert_color_cached(struct rb_node *node,
+					  struct rb_root_cached *root,
+					  bool leftmost)
+{
+	if (leftmost)
+		root->rb_leftmost = node;
+	rb_insert_color(node, &root->rb_root);
+}
+
+static inline void rb_erase_cached(struct rb_node *node,
+				   struct rb_root_cached *root)
+{
+	if (root->rb_leftmost == node)
+		root->rb_leftmost = rb_next(node);
+	rb_erase(node, &root->rb_root);
+}
+
+static inline void rb_replace_node_cached(struct rb_node *victim,
+					  struct rb_node *_new,
+					  struct rb_root_cached *root)
+{
+	if (root->rb_leftmost == victim)
+		root->rb_leftmost = _new;
+	rb_replace_node(victim, _new, &root->rb_root);
+}
+
+static inline void rb_erase_init(struct rb_node *n, struct rb_root *root)
+{
+	rb_erase(n, root);
+	RB_CLEAR_NODE(n);
+}
+
+/* ------------------------------------------------------------------ */
+/*  Generic insert / find helpers (inline, using function pointers)   */
+/* ------------------------------------------------------------------ */
+
+static inline void
+rb_add(struct rb_node *node, struct rb_root *tree,
+       bool (*less)(struct rb_node *, const struct rb_node *))
+{
+	struct rb_node **link = &tree->rb_node;
+	struct rb_node *parent = nullptr;
+
+	while (*link) {
+		parent = *link;
+		if (less(node, parent))
+			link = &parent->rb_left;
+		else
+			link = &parent->rb_right;
+	}
+
+	rb_link_node(node, parent, link);
+	rb_insert_color(node, tree);
+}
+
+static inline void
+rb_add_cached(struct rb_node *node, struct rb_root_cached *tree,
+	      bool (*less)(struct rb_node *, const struct rb_node *))
+{
+	struct rb_node **link = &tree->rb_root.rb_node;
+	struct rb_node *parent = nullptr;
+	bool leftmost = true;
+
+	while (*link) {
+		parent = *link;
+		if (less(node, parent)) {
+			link = &parent->rb_left;
+		} else {
+			link = &parent->rb_right;
+			leftmost = false;
+		}
+	}
+
+	rb_link_node(node, parent, link);
+	rb_insert_color_cached(node, tree, leftmost);
+}
+
+static inline struct rb_node *
+rb_find(const void *key, const struct rb_root *tree,
+	int (*cmp)(const void *key, const struct rb_node *))
+{
+	struct rb_node *node = tree->rb_node;
+
+	while (node) {
+		int c = cmp(key, node);
+		if (c < 0)
+			node = node->rb_left;
+		else if (c > 0)
+			node = node->rb_right;
+		else
+			return node;
+	}
+	return nullptr;
+}
+
+static inline struct rb_node *
+rb_find_add(struct rb_node *node, struct rb_root *tree,
+	    int (*cmp)(struct rb_node *, const struct rb_node *))
+{
+	struct rb_node **link = &tree->rb_node;
+	struct rb_node *parent = nullptr;
+	int c;
+
+	while (*link) {
+		parent = *link;
+		c = cmp(node, parent);
+		if (c < 0)
+			link = &parent->rb_left;
+		else if (c > 0)
+			link = &parent->rb_right;
+		else
+			return parent;
+	}
+
+	rb_link_node(node, parent, link);
+	rb_insert_color(node, tree);
+	return nullptr;
+}
+
+/* ------------------------------------------------------------------ */
+/*  Iteration macros                                                  */
+/* ------------------------------------------------------------------ */
+
+#define rbtree_postorder_for_each_entry_safe(pos, n, root, field) \
+	for (pos = rb_entry_safe(rb_first_postorder(root), __typeof__(*pos), field); \
+	     pos && ({ n = rb_entry_safe(rb_next_postorder(&pos->field), \
+			__typeof__(*pos), field); 1; }); \
+	     pos = n)
+
diff --git a/na/include/lume/rbtree_augmented.h b/na/include/lume/rbtree_augmented.h
new file mode 100644
index 0000000..07f367b
--- /dev/null
+++ b/na/include/lume/rbtree_augmented.h
@@ -0,0 +1,272 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+/*
+ * Red-Black Trees — augmented operations & internal helpers
+ * Adapted from Linux kernel include/linux/rbtree_augmented.h
+ *
+ * (C) 1999  Andrea Arcangeli <andrea@suse.de>
+ * (C) 2002  David Woodhouse <dwmw2@infradead.org>
+ * (C) 2012  Michel Lespinasse <walken@google.com>
+ */
+
+#pragma once
+
+#include <lume/rbtree.h>
+
+/* ------------------------------------------------------------------ */
+/*  Linux-compat macros for freestanding kernel                       */
+/* ------------------------------------------------------------------ */
+
+#ifndef WRITE_ONCE
+#define WRITE_ONCE(x, val)  do { \
+	__asm__ volatile("" ::: "memory"); \
+	(x) = (val); \
+	__asm__ volatile("" ::: "memory"); \
+} while (0)
+#endif
+
+#ifndef unlikely
+#define unlikely(x)   __builtin_expect(!!(x), 0)
+#endif
+
+#ifndef __always_inline
+#define __always_inline  inline __attribute__((always_inline))
+#endif
+
+/* ------------------------------------------------------------------ */
+/*  Augmented callback structure                                      */
+/* ------------------------------------------------------------------ */
+
+struct rb_augment_callbacks {
+	void (*propagate)(struct rb_node *node, struct rb_node *stop);
+	void (*copy)(struct rb_node *old, struct rb_node *_new);
+	void (*rotate)(struct rb_node *old, struct rb_node *_new);
+};
+
+extern "C" void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+	void (*augment_rotate)(struct rb_node *old, struct rb_node *_new));
+
+extern "C" void __rb_erase_color(struct rb_node *parent, struct rb_root *root,
+	void (*augment_rotate)(struct rb_node *old, struct rb_node *_new));
+
+static inline void
+rb_insert_augmented(struct rb_node *node, struct rb_root *root,
+		    const struct rb_augment_callbacks *augment)
+{
+	__rb_insert_augmented(node, root, augment->rotate);
+}
+
+static inline void
+rb_insert_augmented_cached(struct rb_node *node,
+			   struct rb_root_cached *root, bool newleft,
+			   const struct rb_augment_callbacks *augment)
+{
+	if (newleft)
+		root->rb_leftmost = node;
+	rb_insert_augmented(node, &root->rb_root, augment);
+}
+
+/* ------------------------------------------------------------------ */
+/*  Color / parent encoding helpers                                   */
+/* ------------------------------------------------------------------ */
+
+#define RB_RED		0
+#define RB_BLACK	1
+
+#define __rb_parent(pc)    ((struct rb_node *)(pc & ~3UL))
+
+#define __rb_color(pc)     ((pc) & 1)
+#define __rb_is_black(pc)  __rb_color(pc)
+#define __rb_is_red(pc)    (!__rb_color(pc))
+#define rb_color(rb)       __rb_color((rb)->__rb_parent_color)
+#define rb_is_red(rb)      __rb_is_red((rb)->__rb_parent_color)
+#define rb_is_black(rb)    __rb_is_black((rb)->__rb_parent_color)
+
+static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
+{
+	rb->__rb_parent_color = rb_color(rb) + (unsigned long)p;
+}
+
+static inline void rb_set_parent_color(struct rb_node *rb,
+				       struct rb_node *p, int color)
+{
+	rb->__rb_parent_color = (unsigned long)p + color;
+}
+
+static inline void
+__rb_change_child(struct rb_node *old, struct rb_node *_new,
+		  struct rb_node *parent, struct rb_root *root)
+{
+	if (parent) {
+		if (parent->rb_left == old)
+			WRITE_ONCE(parent->rb_left, _new);
+		else
+			WRITE_ONCE(parent->rb_right, _new);
+	} else
+		WRITE_ONCE(root->rb_node, _new);
+}
+
+/* ------------------------------------------------------------------ */
+/*  Augmented erase (inline)                                          */
+/* ------------------------------------------------------------------ */
+
+static __always_inline struct rb_node *
+__rb_erase_augmented(struct rb_node *node, struct rb_root *root,
+		     const struct rb_augment_callbacks *augment)
+{
+	struct rb_node *child = node->rb_right;
+	struct rb_node *tmp = node->rb_left;
+	struct rb_node *parent, *rebalance;
+	unsigned long pc;
+
+	if (!tmp) {
+		/*
+		 * Case 1: node to erase has no more than 1 child (easy!)
+		 */
+		pc = node->__rb_parent_color;
+		parent = __rb_parent(pc);
+		__rb_change_child(node, child, parent, root);
+		if (child) {
+			child->__rb_parent_color = pc;
+			rebalance = nullptr;
+		} else
+			rebalance = __rb_is_black(pc) ? parent : nullptr;
+		tmp = parent;
+	} else if (!child) {
+		/* Still case 1, but this time the child is node->rb_left */
+		tmp->__rb_parent_color = pc = node->__rb_parent_color;
+		parent = __rb_parent(pc);
+		__rb_change_child(node, tmp, parent, root);
+		rebalance = nullptr;
+		tmp = parent;
+	} else {
+		struct rb_node *successor = child, *child2;
+
+		tmp = child->rb_left;
+		if (!tmp) {
+			/*
+			 * Case 2: node's successor is its right child
+			 */
+			parent = successor;
+			child2 = successor->rb_right;
+
+			augment->copy(node, successor);
+		} else {
+			/*
+			 * Case 3: node's successor is leftmost under
+			 * node's right child subtree
+			 */
+			do {
+				parent = successor;
+				successor = tmp;
+				tmp = tmp->rb_left;
+			} while (tmp);
+			child2 = successor->rb_right;
+			WRITE_ONCE(parent->rb_left, child2);
+			WRITE_ONCE(successor->rb_right, child);
+			rb_set_parent(child, successor);
+
+			augment->copy(node, successor);
+			augment->propagate(parent, successor);
+		}
+
+		tmp = node->rb_left;
+		WRITE_ONCE(successor->rb_left, tmp);
+		rb_set_parent(tmp, successor);
+
+		pc = node->__rb_parent_color;
+		tmp = __rb_parent(pc);
+		__rb_change_child(node, successor, tmp, root);
+
+		if (child2) {
+			rb_set_parent_color(child2, parent, RB_BLACK);
+			rebalance = nullptr;
+		} else {
+			rebalance = rb_is_black(successor) ? parent : nullptr;
+		}
+		successor->__rb_parent_color = pc;
+		tmp = successor;
+	}
+
+	augment->propagate(tmp, nullptr);
+	return rebalance;
+}
+
+static __always_inline void
+rb_erase_augmented(struct rb_node *node, struct rb_root *root,
+		   const struct rb_augment_callbacks *augment)
+{
+	struct rb_node *rebalance = __rb_erase_augmented(node, root, augment);
+	if (rebalance)
+		__rb_erase_color(rebalance, root, augment->rotate);
+}
+
+static __always_inline void
+rb_erase_augmented_cached(struct rb_node *node, struct rb_root_cached *root,
+			  const struct rb_augment_callbacks *augment)
+{
+	if (root->rb_leftmost == node)
+		root->rb_leftmost = rb_next(node);
+	rb_erase_augmented(node, &root->rb_root, augment);
+}
+
+/* ------------------------------------------------------------------ */
+/*  Augmented callback declaration macros                             */
+/* ------------------------------------------------------------------ */
+
+#define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,				\
+			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE)	\
+static inline void							\
+RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
+{									\
+	while (rb != stop) {						\
+		RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD);	\
+		if (RBCOMPUTE(node, true))				\
+			break;						\
+		rb = rb_parent(&node->RBFIELD);				\
+	}								\
+}									\
+static inline void							\
+RBNAME ## _copy(struct rb_node *rb_old, struct rb_node *rb_new)		\
+{									\
+	RBSTRUCT *old_ = rb_entry(rb_old, RBSTRUCT, RBFIELD);		\
+	RBSTRUCT *new_ = rb_entry(rb_new, RBSTRUCT, RBFIELD);		\
+	new_->RBAUGMENTED = old_->RBAUGMENTED;				\
+}									\
+static void								\
+RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
+{									\
+	RBSTRUCT *old_ = rb_entry(rb_old, RBSTRUCT, RBFIELD);		\
+	RBSTRUCT *new_ = rb_entry(rb_new, RBSTRUCT, RBFIELD);		\
+	new_->RBAUGMENTED = old_->RBAUGMENTED;				\
+	RBCOMPUTE(old_, false);						\
+}									\
+RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
+	RBNAME ## _propagate,						\
+	RBNAME ## _copy,						\
+	RBNAME ## _rotate						\
+};
+
+#define RB_DECLARE_CALLBACKS_MAX(RBSTATIC, RBNAME, RBSTRUCT, RBFIELD,	      \
+				 RBTYPE, RBAUGMENTED, RBCOMPUTE)	      \
+static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit_early)    \
+{									      \
+	RBSTRUCT *child;						      \
+	RBTYPE max = RBCOMPUTE(node);					      \
+	if (node->RBFIELD.rb_left) {					      \
+		child = rb_entry(node->RBFIELD.rb_left, RBSTRUCT, RBFIELD);   \
+		if (child->RBAUGMENTED > max)				      \
+			max = child->RBAUGMENTED;			      \
+	}								      \
+	if (node->RBFIELD.rb_right) {					      \
+		child = rb_entry(node->RBFIELD.rb_right, RBSTRUCT, RBFIELD);  \
+		if (child->RBAUGMENTED > max)				      \
+			max = child->RBAUGMENTED;			      \
+	}								      \
+	if (exit_early && node->RBAUGMENTED == max)			      \
+		return true;						      \
+	node->RBAUGMENTED = max;					      \
+	return false;							      \
+}									      \
+RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,					      \
+		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)
+
diff --git a/na/include/lume/rbtree_types.h b/na/include/lume/rbtree_types.h
new file mode 100644
index 0000000..688ade4
--- /dev/null
+++ b/na/include/lume/rbtree_types.h
@@ -0,0 +1,28 @@
+/* SPDX-License-Identifier: GPL-2.0-only */
+/*
+ * Red-Black Tree type definitions
+ * Adapted from Linux kernel include/linux/rbtree_types.h
+ */
+
+#pragma once
+
+#include <lume/types.h>
+
+struct rb_node {
+	unsigned long  __rb_parent_color;
+	struct rb_node *rb_right;
+	struct rb_node *rb_left;
+} __attribute__((aligned(sizeof(long))));
+
+struct rb_root {
+	struct rb_node *rb_node;
+};
+
+struct rb_root_cached {
+	struct rb_root rb_root;
+	struct rb_node *rb_leftmost;
+};
+
+#define RB_ROOT		(struct rb_root) { nullptr, }
+#define RB_ROOT_CACHED	(struct rb_root_cached) { {nullptr, }, nullptr }
+