1/* @title: Red black tree */
2#pragma once
3#include <container_of.h>
4#include <enforce.h>
5#include <math/align.h>
6#include <stdbool.h>
7#include <stddef.h>
8#include <stdint.h>
9
10struct rbt_parent {
11 uintptr_t word; /* Lowest bit is color */
12};
13
14enum rbt_node_color { TREE_NODE_RED = 0, TREE_NODE_BLACK = 1 };
15
16struct rbt_node {
17 struct rbt_node *left;
18 struct rbt_node *right;
19 struct rbt_parent parent_data;
20};
21
22typedef int32_t (*rbt_compare)(const struct rbt_node *a,
23 const struct rbt_node *b);
24typedef size_t (*rbt_get_data)(struct rbt_node *);
25
26struct rbt { /* TODO: stop using get_data. for now it works
27 * but in the future we may want to allow for rb-trees
28 * that are "backwards" or sorted by some other rule
29 * beyond integer field comparison */
30 rbt_get_data get_data;
31 rbt_compare compare;
32 struct rbt_node *root;
33};
34
35#define RBT_NODE_INIT \
36 (struct rbt_node) { \
37 .parent_data.word = TREE_NODE_BLACK, .left = NULL, .right = NULL \
38 }
39
40static inline struct rbt_node *rbt_parent(const struct rbt_node *node) {
41 return (struct rbt_node *) (node->parent_data.word & ~1ULL);
42}
43
44static inline enum rbt_node_color rbt_color(const struct rbt_node *node) {
45 return (enum rbt_node_color)(node->parent_data.word & 1ULL);
46}
47
48static inline void rbt_set_color(struct rbt_node *node, enum rbt_node_color c) {
49 enforce(c == TREE_NODE_BLACK || c == TREE_NODE_RED);
50 node->parent_data.word &= ~1ULL;
51 node->parent_data.word |= c;
52}
53
54static inline void rbt_set_parent(struct rbt_node *node,
55 struct rbt_node *parent) {
56 /* Lowest bit cannot be set */
57 enforce(((uintptr_t) parent & 1ULL) == 0);
58
59 node->parent_data.word &= 1ULL;
60 node->parent_data.word |= (uintptr_t) parent;
61}
62
63static inline struct rbt_node *rbt_last(const struct rbt *root) {
64 struct rbt_node *node = root->root;
65 if (!node)
66 return NULL;
67 while (node->right)
68 node = node->right;
69 return node;
70}
71
72static inline struct rbt_node *rbt_prev(struct rbt_node *node) {
73 if (node->left) {
74 /* predecessor is rightmost node of left subtree */
75 node = node->left;
76 while (node->right)
77 node = node->right;
78 return node;
79 }
80
81 /* climb up until we come from the right */
82
83 struct rbt_node *parent = rbt_parent(node);
84 while (parent && node == parent->left) {
85 node = parent;
86 parent = rbt_parent(node: parent);
87 }
88 return parent;
89}
90
91static inline void rbt_init_node(struct rbt_node *n) {
92 n->parent_data.word = TREE_NODE_BLACK;
93 n->left = n->right = NULL;
94}
95
96static inline struct rbt_node *rbt_first(const struct rbt *root) {
97 struct rbt_node *node = root->root;
98 if (!node)
99 return NULL;
100 while (node->left)
101 node = node->left;
102 return node;
103}
104
105void rbt_link_node(struct rbt_node *node, struct rbt_node *parent,
106 struct rbt_node **link);
107void rbt_insert_color(struct rbt *tree, struct rbt_node *node);
108struct rbt *rbt_init(struct rbt *t, rbt_get_data get_data, rbt_compare compare);
109struct rbt *rbt_create(rbt_get_data get, rbt_compare compare);
110struct rbt_node *rbt_find_min(struct rbt_node *node);
111struct rbt_node *rbt_find_max(struct rbt_node *node);
112
113/* `z` must be linked into `tree` */
114void rbt_delete(struct rbt *tree, struct rbt_node *z);
115struct rbt_node *rbt_search(struct rbt *tree, uint64_t data);
116void rbt_remove(struct rbt *tree, uint64_t data);
117void rbt_insert(struct rbt *tree, struct rbt_node *new_node);
118struct rbt_node *rbt_min(struct rbt *tree);
119struct rbt_node *rbt_max(struct rbt *tree);
120struct rbt_node *rbt_next(struct rbt_node *node);
121struct rbt_node *rbt_find_predecessor(struct rbt *tree, uint64_t data);
122struct rbt_node *rbt_find_successor(struct rbt *tree, uint64_t data);
123
124static inline bool rbt_empty(struct rbt *tree) {
125 return !tree->root;
126}
127
128bool rbt_has_node(struct rbt *tree, struct rbt_node *node);
129
130/* Re-position nodes whose sort key changes, with checks
131 * to first verify it needs to move at all, returning true if moved */
132bool rbt_reinsert(struct rbt *tree, struct rbt_node *node);
133
134/* True when ordering is correct */
135bool rbt_position_valid(struct rbt *tree, struct rbt_node *node);
136
137void rbt_move(struct rbt *from, struct rbt *to, struct rbt_node *node);
138
139#define rbt_for_each_safe(pos, tmp, root) \
140 for (pos = rbt_first(root), tmp = rbt_next(pos); pos != NULL; \
141 pos = tmp, tmp = rbt_next(pos))
142
143#define rbt_for_each_entry_safe(pos, tmp, tree, member) \
144 for (typeof(rbt_first(tree)) __n = rbt_first(tree), __nx = NULL; \
145 __n != NULL && \
146 (((pos) = rbt_entry(__n, typeof(*(pos)), member)), \
147 (__nx = rbt_next(__n)), \
148 ((tmp) = __nx ? rbt_entry(__nx, typeof(*(tmp)), member) : NULL), \
149 true); \
150 __n = __nx)
151
152#define rbt_for_each_safe_reverse(pos, tmp, root) \
153 for (pos = rbt_last(root), tmp = rbt_prev(pos); pos != NULL; \
154 pos = tmp, tmp = rbt_prev(pos))
155
156#define rbt_for_each_entry_safe_reverse(pos, tmp, tree, member) \
157 for (typeof(rbt_last(tree)) __n = rbt_last(tree), __nx = NULL; \
158 __n != NULL && \
159 (((pos) = rbt_entry(__n, typeof(*(pos)), member)), \
160 (__nx = rbt_prev(__n)), \
161 ((tmp) = __nx ? rbt_entry(__nx, typeof(*(tmp)), member) : NULL), \
162 true); \
163 __n = __nx)
164
165#define rbt_for_each(pos, root) \
166 for (pos = rbt_first(root); pos != NULL; pos = rbt_next(pos))
167
168#define rbt_for_each_entry(pos, tree, member) \
169 for (typeof(rbt_first(tree)) __n = rbt_first(tree); \
170 __n != NULL && \
171 (((pos) = rbt_entry(__n, typeof(*(pos)), member)), true); \
172 __n = rbt_next(__n))
173
174#define rbt_for_each_reverse(pos, root) \
175 for (pos = rbt_last(root); pos != NULL; pos = rbt_prev(pos))
176
177#define rbt_for_each_entry_reverse(pos, tree, member) \
178 for (typeof(rbt_last(tree)) __n = rbt_last(tree); \
179 __n != NULL && \
180 (((pos) = rbt_entry(__n, typeof(*(pos)), member)), true); \
181 __n = rbt_prev(__n))
182
183#define rbt_entry(ptr, type, member) container_of(ptr, type, member)
184