From 802dcab8ad6f31f266fa98e128b0a0bec3a15edf Mon Sep 17 00:00:00 2001 From: Richard Qin Date: Tue, 21 Apr 2026 23:19:39 +0800 Subject: feat: add basic implementation and enhanced operations of red-black tree from Linux kernel source --- na/include/lume/rbtree.h | 219 +++++++++++++++++++++++++++++ na/include/lume/rbtree_augmented.h | 272 +++++++++++++++++++++++++++++++++++++ na/include/lume/rbtree_types.h | 28 ++++ 3 files changed, 519 insertions(+) create mode 100644 na/include/lume/rbtree.h create mode 100644 na/include/lume/rbtree_augmented.h create mode 100644 na/include/lume/rbtree_types.h (limited to 'na/include') 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 + */ + +#pragma once + +#include + +/* ------------------------------------------------------------------ */ +/* Linux-compat macros for freestanding kernel */ +/* ------------------------------------------------------------------ */ + +#ifndef container_of +#define container_of(ptr, type, member) \ + reinterpret_cast( \ + reinterpret_cast(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 + * (C) 2002 David Woodhouse + * (C) 2012 Michel Lespinasse + */ + +#pragma once + +#include + +/* ------------------------------------------------------------------ */ +/* 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 + +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 } + -- cgit v1.2.3