1#include <console/printf.h>
2#include <kassert.h>
3#include <mem/alloc.h>
4#include <stddef.h>
5#include <stdint.h>
6#include <structures/rbt.h>
7
8struct rbt *rbt_create(rbt_get_data get, rbt_compare cmp) {
9 struct rbt *tree = kmalloc(sizeof(struct rbt));
10 if (!tree)
11 return NULL;
12
13 tree->compare = cmp;
14 tree->get_data = get;
15 tree->root = NULL;
16 return tree;
17}
18
19struct rbt_node *rbt_find_min(struct rbt_node *node) {
20 while (node && node->left != NULL)
21 node = node->left;
22
23 return node;
24}
25
26struct rbt_node *rbt_find_max(struct rbt_node *node) {
27 while (node && node->right != NULL) {
28 node = node->right;
29 }
30
31 return node;
32}
33
34struct rbt_node *rbt_max(struct rbt *tree) {
35 return rbt_find_max(node: tree->root);
36}
37
38struct rbt_node *rbt_min(struct rbt *tree) {
39 return rbt_find_min(node: tree->root);
40}
41
42struct rbt_node *rbt_next(struct rbt_node *node) {
43 if (!node)
44 return NULL;
45
46 if (node->right)
47 return rbt_find_min(node: node->right);
48
49 struct rbt_node *parent = node->parent;
50 while (parent && node == parent->right) {
51 node = parent;
52 parent = parent->parent;
53 }
54 return parent;
55}
56
57struct rbt_node *rbt_find_predecessor(struct rbt *tree, uint64_t data) {
58 struct rbt_node *curr = tree->root;
59 struct rbt_node *pred = NULL;
60
61 while (curr) {
62 uint64_t curr_data = tree->get_data(curr);
63 if (curr_data < data) {
64 pred = curr;
65 curr = curr->right;
66 } else {
67 curr = curr->left;
68 }
69 }
70 return pred;
71}
72
73struct rbt_node *rbt_find_successor(struct rbt *tree, uint64_t data) {
74 struct rbt_node *curr = tree->root;
75 struct rbt_node *succ = NULL;
76
77 while (curr) {
78 uint64_t curr_data = tree->get_data(curr);
79 if (curr_data > data) {
80 succ = curr;
81 curr = curr->left;
82 } else {
83 curr = curr->right;
84 }
85 }
86 return succ;
87}
88
89static void rb_transplant(struct rbt *tree, struct rbt_node *u,
90 struct rbt_node *v) {
91 if (u->parent == NULL)
92 tree->root = v;
93 else if (u == u->parent->left)
94 u->parent->left = v;
95 else
96 u->parent->right = v;
97
98 if (v)
99 v->parent = u->parent;
100}
101
102static void left_rotate(struct rbt *tree, struct rbt_node *x) {
103 struct rbt_node *y = x->right;
104 x->right = y->left;
105
106 if (y->left)
107 y->left->parent = x;
108
109 y->parent = x->parent;
110 if (!x->parent)
111 tree->root = y;
112 else if (x == x->parent->left)
113 x->parent->left = y;
114 else
115 x->parent->right = y;
116
117 y->left = x;
118 x->parent = y;
119}
120
121static void right_rotate(struct rbt *tree, struct rbt_node *y) {
122 struct rbt_node *x = y->left;
123 y->left = x->right;
124 if (x->right)
125 x->right->parent = y;
126
127 x->parent = y->parent;
128 if (!y->parent)
129 tree->root = x;
130 else if (y == y->parent->right)
131 y->parent->right = x;
132 else
133 y->parent->left = x;
134
135 x->right = y;
136 y->parent = x;
137}
138
139static void fix_deletion(struct rbt *tree, struct rbt_node *x) {
140 while (x != tree->root && (!x || x->color == TREE_NODE_BLACK)) {
141 if (!x || !x->parent)
142 break;
143
144 struct rbt_node *sibling;
145
146 if (x == x->parent->left) {
147 sibling = x->parent->right;
148
149 if (sibling && sibling->color == TREE_NODE_RED) {
150 sibling->color = TREE_NODE_BLACK;
151 x->parent->color = TREE_NODE_RED;
152 left_rotate(tree, x: x->parent);
153 sibling = x->parent->right;
154 }
155
156 if (!sibling ||
157 ((!sibling->left || sibling->left->color == TREE_NODE_BLACK) &&
158 (!sibling->right ||
159 sibling->right->color == TREE_NODE_BLACK))) {
160 if (sibling)
161 sibling->color = TREE_NODE_RED;
162 x = x->parent;
163 } else {
164 if (!sibling->right ||
165 sibling->right->color == TREE_NODE_BLACK) {
166 if (sibling->left)
167 sibling->left->color = TREE_NODE_BLACK;
168 if (sibling) {
169 sibling->color = TREE_NODE_RED;
170 right_rotate(tree, y: sibling);
171 sibling = x->parent->right;
172 }
173 }
174
175 if (sibling) {
176 sibling->color = x->parent->color;
177 x->parent->color = TREE_NODE_BLACK;
178 if (sibling->right)
179 sibling->right->color = TREE_NODE_BLACK;
180 left_rotate(tree, x: x->parent);
181 }
182 x = tree->root;
183 }
184 } else {
185 sibling = x->parent->left;
186
187 if (sibling && sibling->color == TREE_NODE_RED) {
188 sibling->color = TREE_NODE_BLACK;
189 x->parent->color = TREE_NODE_RED;
190 right_rotate(tree, y: x->parent);
191 sibling = x->parent->left;
192 }
193
194 if (!sibling ||
195 ((!sibling->left || sibling->left->color == TREE_NODE_BLACK) &&
196 (!sibling->right ||
197 sibling->right->color == TREE_NODE_BLACK))) {
198 if (sibling)
199 sibling->color = TREE_NODE_RED;
200 x = x->parent;
201 } else {
202 if (!sibling->left || sibling->left->color == TREE_NODE_BLACK) {
203 if (sibling->right)
204 sibling->right->color = TREE_NODE_BLACK;
205 if (sibling) {
206 sibling->color = TREE_NODE_RED;
207 left_rotate(tree, x: sibling);
208 sibling = x->parent->left;
209 }
210 }
211
212 if (sibling) {
213 sibling->color = x->parent->color;
214 x->parent->color = TREE_NODE_BLACK;
215 if (sibling->left)
216 sibling->left->color = TREE_NODE_BLACK;
217 right_rotate(tree, y: x->parent);
218 }
219 x = tree->root;
220 }
221 }
222 }
223
224 if (x)
225 x->color = TREE_NODE_BLACK;
226}
227
228void rbt_delete(struct rbt *tree, struct rbt_node *z) {
229 struct rbt_node *y = z;
230 struct rbt_node *x = NULL;
231 enum rbt_node_color y_original_color = y->color;
232
233 if (z->left == NULL) {
234 x = z->right;
235 rb_transplant(tree, u: z, v: z->right);
236 } else if (z->right == NULL) {
237 x = z->left;
238 rb_transplant(tree, u: z, v: z->left);
239 } else {
240 y = rbt_find_min(node: z->right);
241 y_original_color = y->color;
242 x = y->right;
243
244 if (y->parent != z) {
245 rb_transplant(tree, u: y, v: y->right);
246 y->right = z->right;
247 if (y->right)
248 y->right->parent = y;
249 }
250
251 rb_transplant(tree, u: z, v: y);
252 y->left = z->left;
253 if (y->left)
254 y->left->parent = y;
255 y->color = z->color;
256 }
257
258 if (y_original_color == TREE_NODE_BLACK) {
259 fix_deletion(tree, x);
260 }
261
262 z->left = NULL;
263 z->right = NULL;
264 z->parent = NULL;
265}
266
267static struct rbt_node *
268rbt_search_internal(struct rbt *tree, struct rbt_node *root, uint64_t data) {
269 while (root && tree->get_data(root) != data) {
270 if (data < tree->get_data(root))
271 root = root->left;
272 else
273 root = root->right;
274 }
275 return root;
276}
277
278struct rbt_node *rbt_search(struct rbt *tree, uint64_t data) {
279 struct rbt_node *root = tree->root;
280
281 while (root && tree->get_data(root) != data) {
282 if (data < tree->get_data(root))
283 root = root->left;
284 else
285 root = root->right;
286 }
287
288 return (root && tree->get_data(root) == data) ? root : NULL;
289}
290
291bool rbt_has_node(struct rbt *tree, struct rbt_node *node) {
292 struct rbt_node *iter;
293 rbt_for_each(iter, tree) {
294 if (iter == node)
295 return true;
296 }
297 return false;
298}
299
300void rbt_remove(struct rbt *tree, uint64_t data) {
301 struct rbt_node *node = rbt_search_internal(tree, root: tree->root, data);
302 if (node)
303 rbt_delete(tree, z: node);
304}
305
306static void fix_insertion(struct rbt *tree, struct rbt_node *node) {
307 while (node != tree->root && node->parent->color == TREE_NODE_RED) {
308 struct rbt_node *parent = node->parent;
309 struct rbt_node *grandparent = parent->parent;
310
311 if (parent == grandparent->left) {
312 struct rbt_node *uncle = grandparent->right;
313 if (uncle && uncle->color == TREE_NODE_RED) {
314 parent->color = TREE_NODE_BLACK;
315 uncle->color = TREE_NODE_BLACK;
316 grandparent->color = TREE_NODE_RED;
317 node = grandparent;
318 } else {
319 if (node == parent->right) {
320 node = parent;
321 left_rotate(tree, x: node);
322 parent = node->parent;
323 }
324 parent->color = TREE_NODE_BLACK;
325 grandparent->color = TREE_NODE_RED;
326 right_rotate(tree, y: grandparent);
327 }
328 } else {
329 struct rbt_node *uncle = grandparent->left;
330 if (uncle && uncle->color == TREE_NODE_RED) {
331 parent->color = TREE_NODE_BLACK;
332 uncle->color = TREE_NODE_BLACK;
333 grandparent->color = TREE_NODE_RED;
334 node = grandparent;
335 } else {
336 if (node == parent->left) {
337 node = parent;
338 right_rotate(tree, y: node);
339 parent = node->parent;
340 }
341 parent->color = TREE_NODE_BLACK;
342 grandparent->color = TREE_NODE_RED;
343 left_rotate(tree, x: grandparent);
344 }
345 }
346 }
347 tree->root->color = TREE_NODE_BLACK;
348}
349
350void rbt_insert(struct rbt *tree, struct rbt_node *new_node) {
351 new_node->left = NULL;
352 new_node->right = NULL;
353 new_node->color = TREE_NODE_RED;
354
355 if (tree->root == NULL) {
356 new_node->color = TREE_NODE_BLACK;
357 new_node->parent = NULL;
358 tree->root = new_node;
359 return;
360 }
361
362 struct rbt_node *current = tree->root;
363 struct rbt_node *parent = NULL;
364 while (current != NULL) {
365 parent = current;
366 if (tree->compare(new_node, current) < 0)
367 current = current->left;
368 else
369 current = current->right;
370 }
371
372 new_node->parent = parent;
373 if (tree->compare(new_node, parent) < 0)
374 parent->left = new_node;
375 else
376 parent->right = new_node;
377
378 fix_insertion(tree, node: new_node);
379
380 if (parent)
381 kassert(!(parent->color == TREE_NODE_RED &&
382 new_node->color == TREE_NODE_RED));
383}
384
385void rbt_link_node(struct rbt_node *node, struct rbt_node *parent,
386 struct rbt_node **link) {
387 node->parent = parent;
388 node->left = NULL;
389 node->right = NULL;
390 node->color = TREE_NODE_RED;
391 *link = node;
392}
393
394void rbt_insert_color(struct rbt *tree, struct rbt_node *node) {
395 fix_insertion(tree, node);
396}
397
398struct rbt *rbt_init(struct rbt *tree, rbt_get_data get_data, rbt_compare cmp) {
399 tree->root = NULL;
400 tree->compare = cmp;
401 tree->get_data = get_data;
402 return tree;
403}
404