1#include <console/panic.h>
2#include <console/printf.h>
3#include <enforce.h>
4#include <mem/alloc.h>
5#include <stddef.h>
6#include <stdint.h>
7#include <structures/rbt.h>
8
9#ifdef DEBUG_RBT
10
11#define RBT_DEBUG_MAX_DEPTH 128
12
13static inline void rbt_debug_step(size_t *steps, const struct rbt_node *at,
14 const char *caller) {
15 if (cc_unlikely(++*steps > RBT_DEBUG_MAX_DEPTH))
16 panic(CORRUPTION, "%s: walk passed %d steps at node %p -- cycle?",
17 caller, RBT_DEBUG_MAX_DEPTH, at);
18}
19
20static void rbt_debug_check_linked(const struct rbt *tree,
21 const struct rbt_node *node,
22 const char *caller) {
23 const struct rbt_node *n = node;
24 size_t steps = 0;
25
26 while (rbt_parent(n)) {
27 const struct rbt_node *p = rbt_parent(n);
28 if (cc_unlikely(p->left != n && p->right != n))
29 panic(CORRUPTION,
30 "%s: node %p (from %p) claims parent %p, whose children are "
31 "%p/%p, tree %p",
32 caller, n, node, p, p->left, p->right, tree);
33
34 n = p;
35 rbt_debug_step(&steps, n, caller);
36 }
37
38 if (cc_unlikely(n != tree->root))
39 panic(CORRUPTION, "%s: node %p is under root %p, not tree %p (root %p)",
40 caller, node, n, tree, tree->root);
41}
42
43#else
44
45static inline void rbt_debug_step(size_t *steps, const struct rbt_node *at,
46 const char *caller) {
47 cc_ignore(steps, at, caller);
48}
49
50static inline void rbt_debug_check_linked(const struct rbt *tree,
51 const struct rbt_node *node,
52 const char *caller) {
53 cc_ignore(tree, node, caller);
54}
55
56#endif /* DEBUG_RBT */
57
58#ifdef DEBUG_RBT_DEEP
59
60static size_t rbt_debug_validate_subtree(const struct rbt *tree,
61 const struct rbt_node *node,
62 size_t depth, const char *caller) {
63 if (!node)
64 return 1;
65
66 if (cc_unlikely(depth > RBT_DEBUG_MAX_DEPTH))
67 panic(CORRUPTION, "%s: tree %p deeper than %d at node %p -- cycle?",
68 caller, tree, RBT_DEBUG_MAX_DEPTH, node);
69
70 if (cc_unlikely(node->left && node->left == node->right))
71 panic(CORRUPTION, "%s: node %p has %p as both children, tree %p",
72 caller, node, node->left, tree);
73
74 if (cc_unlikely(node->left && rbt_parent(node->left) != node))
75 panic(CORRUPTION, "%s: node %p left child %p has parent %p, tree %p",
76 caller, node, node->left, rbt_parent(node->left), tree);
77
78 if (cc_unlikely(node->right && rbt_parent(node->right) != node))
79 panic(CORRUPTION, "%s: node %p right child %p has parent %p, tree %p",
80 caller, node, node->right, rbt_parent(node->right), tree);
81
82 if (cc_unlikely(rbt_color(node) == TREE_NODE_RED &&
83 ((node->left && rbt_color(node->left) == TREE_NODE_RED) ||
84 (node->right && rbt_color(node->right) == TREE_NODE_RED))))
85 panic(CORRUPTION, "%s: red-red violation at node %p, tree %p", caller,
86 node, tree);
87
88 size_t lh = rbt_debug_validate_subtree(tree, node->left, depth + 1, caller);
89 size_t rh =
90 rbt_debug_validate_subtree(tree, node->right, depth + 1, caller);
91
92 if (cc_unlikely(lh != rh))
93 panic(CORRUPTION,
94 "%s: black-height violation at node %p (left %zu, right %zu), "
95 "tree %p",
96 caller, node, lh, rh, tree);
97
98 return lh + (rbt_color(node) == TREE_NODE_BLACK ? 1 : 0);
99}
100
101static void rbt_debug_validate(const struct rbt *tree, const char *caller) {
102 const struct rbt_node *root = tree->root;
103 if (!root)
104 return;
105
106 if (cc_unlikely(rbt_parent(root)))
107 panic(CORRUPTION, "%s: root %p of tree %p has parent %p", caller, root,
108 tree, rbt_parent(root));
109
110 if (cc_unlikely(rbt_color(root) != TREE_NODE_BLACK))
111 panic(CORRUPTION, "%s: root %p of tree %p is red", caller, root, tree);
112
113 rbt_debug_validate_subtree(tree, root, 0, caller);
114}
115
116#else
117
118static inline void rbt_debug_validate(const struct rbt *tree,
119 const char *caller) {
120 cc_ignore(tree, caller);
121}
122
123#endif /* DEBUG_RBT_DEEP */
124
125struct rbt *rbt_create(rbt_get_data get, rbt_compare cmp) {
126 struct rbt *tree = vnew(tree);
127 if (!tree)
128 return NULL;
129
130 tree->compare = cmp;
131 tree->get_data = get;
132 tree->root = NULL;
133 return tree;
134}
135
136struct rbt_node *rbt_find_min(struct rbt_node *node) {
137 size_t steps = 0;
138 while (node && node->left != NULL) {
139 node = node->left;
140 rbt_debug_step(steps: &steps, at: node, caller: __func__);
141 }
142
143 return node;
144}
145
146struct rbt_node *rbt_find_max(struct rbt_node *node) {
147 size_t steps = 0;
148 while (node && node->right != NULL) {
149 node = node->right;
150 rbt_debug_step(steps: &steps, at: node, caller: __func__);
151 }
152
153 return node;
154}
155
156struct rbt_node *rbt_max(struct rbt *tree) {
157 return rbt_find_max(node: tree->root);
158}
159
160struct rbt_node *rbt_min(struct rbt *tree) {
161 return rbt_find_min(node: tree->root);
162}
163
164struct rbt_node *rbt_next(struct rbt_node *node) {
165 if (!node)
166 return NULL;
167
168 if (node->right)
169 return rbt_find_min(node: node->right);
170
171 size_t steps = 0;
172 struct rbt_node *parent = rbt_parent(node);
173 while (parent && node == parent->right) {
174 node = parent;
175 parent = rbt_parent(node: parent);
176 rbt_debug_step(steps: &steps, at: node, caller: __func__);
177 }
178 return parent;
179}
180
181struct rbt_node *rbt_find_predecessor(struct rbt *tree, uint64_t data) {
182 struct rbt_node *curr = tree->root;
183 struct rbt_node *pred = NULL;
184
185 while (curr) {
186 uint64_t curr_data = tree->get_data(curr);
187 if (curr_data < data) {
188 pred = curr;
189 curr = curr->right;
190 } else {
191 curr = curr->left;
192 }
193 }
194 return pred;
195}
196
197struct rbt_node *rbt_find_successor(struct rbt *tree, uint64_t data) {
198 struct rbt_node *curr = tree->root;
199 struct rbt_node *succ = NULL;
200
201 while (curr) {
202 uint64_t curr_data = tree->get_data(curr);
203 if (curr_data > data) {
204 succ = curr;
205 curr = curr->left;
206 } else {
207 curr = curr->right;
208 }
209 }
210 return succ;
211}
212
213static void rb_transplant(struct rbt *tree, struct rbt_node *u,
214 struct rbt_node *v) {
215 if (rbt_parent(node: u) == NULL)
216 tree->root = v;
217 else if (u == rbt_parent(node: u)->left)
218 rbt_parent(node: u)->left = v;
219 else
220 rbt_parent(node: u)->right = v;
221
222 if (v)
223 rbt_set_parent(node: v, parent: rbt_parent(node: u));
224}
225
226static void left_rotate(struct rbt *tree, struct rbt_node *x) {
227 struct rbt_node *y = x->right;
228 x->right = y->left;
229
230 if (y->left)
231 rbt_set_parent(node: y->left, parent: x);
232
233 rbt_set_parent(node: y, parent: rbt_parent(node: x));
234 if (!rbt_parent(node: x))
235 tree->root = y;
236 else if (x == rbt_parent(node: x)->left)
237 rbt_parent(node: x)->left = y;
238 else
239 rbt_parent(node: x)->right = y;
240
241 y->left = x;
242 rbt_set_parent(node: x, parent: y);
243}
244
245static void right_rotate(struct rbt *tree, struct rbt_node *y) {
246 struct rbt_node *x = y->left;
247 y->left = x->right;
248 if (x->right)
249 rbt_set_parent(node: x->right, parent: y);
250
251 rbt_set_parent(node: x, parent: rbt_parent(node: y));
252 if (!rbt_parent(node: y))
253 tree->root = x;
254 else if (y == rbt_parent(node: y)->right)
255 rbt_parent(node: y)->right = x;
256 else
257 rbt_parent(node: y)->left = x;
258
259 x->right = y;
260 rbt_set_parent(node: y, parent: x);
261}
262
263static bool rbt_is_black(const struct rbt_node *node) {
264 return !node || rbt_color(node) == TREE_NODE_BLACK;
265}
266
267static void fix_deletion(struct rbt *tree, struct rbt_node *x,
268 struct rbt_node *x_parent) {
269 size_t steps = 0;
270 while (x != tree->root && (!x || rbt_color(node: x) == TREE_NODE_BLACK)) {
271 if (!x_parent)
272 break;
273
274 rbt_debug_step(steps: &steps, at: x_parent, caller: __func__);
275
276 struct rbt_node *sibling;
277
278 if (x == x_parent->left) {
279 sibling = x_parent->right;
280
281 if (sibling && rbt_color(node: sibling) == TREE_NODE_RED) {
282 rbt_set_color(node: sibling, c: TREE_NODE_BLACK);
283 rbt_set_color(node: x_parent, c: TREE_NODE_RED);
284 left_rotate(tree, x: x_parent);
285 sibling = x_parent->right;
286 }
287
288 if (!sibling ||
289 (rbt_is_black(node: sibling->left) && rbt_is_black(node: sibling->right))) {
290 if (sibling)
291 rbt_set_color(node: sibling, c: TREE_NODE_RED);
292 x = x_parent;
293 x_parent = rbt_parent(node: x);
294 } else {
295 if (rbt_is_black(node: sibling->right)) {
296 if (sibling->left)
297 rbt_set_color(node: sibling->left, c: TREE_NODE_BLACK);
298 rbt_set_color(node: sibling, c: TREE_NODE_RED);
299 right_rotate(tree, y: sibling);
300 sibling = x_parent->right;
301 }
302
303 rbt_set_color(node: sibling, c: rbt_color(node: x_parent));
304 rbt_set_color(node: x_parent, c: TREE_NODE_BLACK);
305 if (sibling->right)
306 rbt_set_color(node: sibling->right, c: TREE_NODE_BLACK);
307 left_rotate(tree, x: x_parent);
308 x = tree->root;
309 x_parent = NULL;
310 }
311 } else {
312 sibling = x_parent->left;
313
314 if (sibling && rbt_color(node: sibling) == TREE_NODE_RED) {
315 rbt_set_color(node: sibling, c: TREE_NODE_BLACK);
316 rbt_set_color(node: x_parent, c: TREE_NODE_RED);
317 right_rotate(tree, y: x_parent);
318 sibling = x_parent->left;
319 }
320
321 if (!sibling ||
322 (rbt_is_black(node: sibling->left) && rbt_is_black(node: sibling->right))) {
323 if (sibling)
324 rbt_set_color(node: sibling, c: TREE_NODE_RED);
325 x = x_parent;
326 x_parent = rbt_parent(node: x);
327 } else {
328 if (rbt_is_black(node: sibling->left)) {
329 if (sibling->right)
330 rbt_set_color(node: sibling->right, c: TREE_NODE_BLACK);
331 rbt_set_color(node: sibling, c: TREE_NODE_RED);
332 left_rotate(tree, x: sibling);
333 sibling = x_parent->left;
334 }
335
336 rbt_set_color(node: sibling, c: rbt_color(node: x_parent));
337 rbt_set_color(node: x_parent, c: TREE_NODE_BLACK);
338 if (sibling->left)
339 rbt_set_color(node: sibling->left, c: TREE_NODE_BLACK);
340 right_rotate(tree, y: x_parent);
341 x = tree->root;
342 x_parent = NULL;
343 }
344 }
345 }
346
347 if (x)
348 rbt_set_color(node: x, c: TREE_NODE_BLACK);
349}
350
351void rbt_delete(struct rbt *tree, struct rbt_node *z) {
352 enforce(rbt_parent(z) != NULL || z == tree->root,
353 "rbt_delete: node %p is not linked into tree %p (root %p)", z, tree,
354 tree->root);
355 rbt_debug_check_linked(tree, node: z, caller: __func__);
356 rbt_debug_validate(tree, caller: __func__);
357
358 struct rbt_node *y = z;
359 struct rbt_node *x = NULL;
360 struct rbt_node *x_parent = NULL;
361 enum rbt_node_color y_original_color = rbt_color(node: y);
362
363 if (z->left == NULL) {
364 x = z->right;
365 x_parent = rbt_parent(node: z);
366 rb_transplant(tree, u: z, v: z->right);
367 } else if (z->right == NULL) {
368 x = z->left;
369 x_parent = rbt_parent(node: z);
370 rb_transplant(tree, u: z, v: z->left);
371 } else {
372 y = rbt_find_min(node: z->right);
373 y_original_color = rbt_color(node: y);
374 x = y->right;
375
376 if (rbt_parent(node: y) != z) {
377 x_parent = rbt_parent(node: y);
378 rb_transplant(tree, u: y, v: y->right);
379 y->right = z->right;
380 if (y->right)
381 rbt_set_parent(node: y->right, parent: y);
382 } else {
383 x_parent = y;
384 }
385
386 rb_transplant(tree, u: z, v: y);
387 y->left = z->left;
388 if (y->left)
389 rbt_set_parent(node: y->left, parent: y);
390 rbt_set_color(node: y, c: rbt_color(node: z));
391 }
392
393 if (y_original_color == TREE_NODE_BLACK) {
394 fix_deletion(tree, x, x_parent);
395 }
396
397 z->left = NULL;
398 z->right = NULL;
399 rbt_set_parent(node: z, NULL);
400
401 rbt_debug_validate(tree, caller: __func__);
402}
403
404static struct rbt_node *
405rbt_search_internal(struct rbt *tree, struct rbt_node *root, uint64_t data) {
406 while (root && tree->get_data(root) != data) {
407 if (data < tree->get_data(root))
408 root = root->left;
409 else
410 root = root->right;
411 }
412 return root;
413}
414
415struct rbt_node *rbt_search(struct rbt *tree, uint64_t data) {
416 struct rbt_node *root = tree->root;
417
418 while (root && tree->get_data(root) != data) {
419 if (data < tree->get_data(root))
420 root = root->left;
421 else
422 root = root->right;
423 }
424
425 return (root && tree->get_data(root) == data) ? root : NULL;
426}
427
428bool rbt_has_node(struct rbt *tree, struct rbt_node *node) {
429 struct rbt_node *iter;
430 rbt_for_each(iter, tree) {
431 if (iter == node)
432 return true;
433 }
434 return false;
435}
436
437bool rbt_position_valid(struct rbt *tree, struct rbt_node *node) {
438 struct rbt_node *prev = rbt_prev(node);
439 if (prev && tree->compare(prev, node) > 0)
440 return false;
441
442 struct rbt_node *next = rbt_next(node);
443 if (next && tree->compare(node, next) > 0)
444 return false;
445
446 return true;
447}
448
449bool rbt_reinsert(struct rbt *tree, struct rbt_node *node) {
450 if (rbt_position_valid(tree, node))
451 return false;
452
453 rbt_delete(tree, z: node);
454 rbt_insert(tree, new_node: node);
455 return true;
456}
457
458void rbt_move(struct rbt *from, struct rbt *to, struct rbt_node *node) {
459 rbt_delete(tree: from, z: node);
460 rbt_insert(tree: to, new_node: node);
461}
462
463void rbt_remove(struct rbt *tree, uint64_t data) {
464 struct rbt_node *node = rbt_search_internal(tree, root: tree->root, data);
465 if (node)
466 rbt_delete(tree, z: node);
467}
468
469static void fix_insertion(struct rbt *tree, struct rbt_node *node) {
470 size_t steps = 0;
471 while (node != tree->root && rbt_color(node: rbt_parent(node)) == TREE_NODE_RED) {
472 rbt_debug_step(steps: &steps, at: node, caller: __func__);
473 struct rbt_node *parent = rbt_parent(node);
474 struct rbt_node *grandparent = rbt_parent(node: parent);
475
476 if (parent == grandparent->left) {
477 struct rbt_node *uncle = grandparent->right;
478 if (uncle && rbt_color(node: uncle) == TREE_NODE_RED) {
479 rbt_set_color(node: parent, c: TREE_NODE_BLACK);
480 rbt_set_color(node: uncle, c: TREE_NODE_BLACK);
481 rbt_set_color(node: grandparent, c: TREE_NODE_RED);
482 node = grandparent;
483 } else {
484 if (node == parent->right) {
485 node = parent;
486 left_rotate(tree, x: node);
487 parent = rbt_parent(node);
488 }
489 rbt_set_color(node: parent, c: TREE_NODE_BLACK);
490 rbt_set_color(node: grandparent, c: TREE_NODE_RED);
491 right_rotate(tree, y: grandparent);
492 }
493 } else {
494 struct rbt_node *uncle = grandparent->left;
495 if (uncle && rbt_color(node: uncle) == TREE_NODE_RED) {
496 rbt_set_color(node: parent, c: TREE_NODE_BLACK);
497 rbt_set_color(node: uncle, c: TREE_NODE_BLACK);
498 rbt_set_color(node: grandparent, c: TREE_NODE_RED);
499 node = grandparent;
500 } else {
501 if (node == parent->left) {
502 node = parent;
503 right_rotate(tree, y: node);
504 parent = rbt_parent(node);
505 }
506 rbt_set_color(node: parent, c: TREE_NODE_BLACK);
507 rbt_set_color(node: grandparent, c: TREE_NODE_RED);
508 left_rotate(tree, x: grandparent);
509 }
510 }
511 }
512 rbt_set_color(node: tree->root, c: TREE_NODE_BLACK);
513}
514
515void rbt_insert(struct rbt *tree, struct rbt_node *new_node) {
516 rbt_debug_validate(tree, caller: __func__);
517
518 new_node->left = NULL;
519 new_node->right = NULL;
520 rbt_set_color(node: new_node, c: TREE_NODE_RED);
521
522 if (tree->root == NULL) {
523 rbt_set_color(node: new_node, c: TREE_NODE_BLACK);
524 rbt_set_parent(node: new_node, NULL);
525 tree->root = new_node;
526 return;
527 }
528
529 struct rbt_node *current = tree->root;
530 struct rbt_node *parent = NULL;
531 int32_t cmp = 0;
532 size_t steps = 0;
533 while (current != NULL) {
534#ifdef DEBUG_RBT
535 /* Only finds a duplicate whose key still places it on this path */
536 if (cc_unlikely(current == new_node))
537 panic(DUPLICATE, "rbt_insert: node %p is already in tree %p",
538 new_node, tree);
539#endif
540
541 rbt_debug_step(steps: &steps, at: current, caller: __func__);
542 parent = current;
543 cmp = tree->compare(new_node, current);
544 if (cmp < 0)
545 current = current->left;
546 else
547 current = current->right;
548 }
549
550 rbt_set_parent(node: new_node, parent);
551
552#ifdef DEBUG_LOCK_CHK
553 if ((cmp < 0 && parent->left) || (cmp >= 0 && parent->right))
554 panic(CORRUPTION,
555 "rbt_insert: node %p would be attached a child of %p in tree %p, "
556 "potential subtree orophaning",
557 new_node, parent, tree);
558#endif
559
560 if (cmp < 0)
561 parent->left = new_node;
562 else
563 parent->right = new_node;
564
565 fix_insertion(tree, node: new_node);
566
567 if (parent)
568 enforce(!(rbt_color(parent) == TREE_NODE_RED &&
569 rbt_color(new_node) == TREE_NODE_RED),
570 CORRUPTION);
571
572 rbt_debug_validate(tree, caller: __func__);
573}
574
575void rbt_link_node(struct rbt_node *node, struct rbt_node *parent,
576 struct rbt_node **link) {
577 rbt_set_parent(node, parent);
578 node->left = NULL;
579 node->right = NULL;
580 rbt_set_color(node, c: TREE_NODE_RED);
581 *link = node;
582}
583
584void rbt_insert_color(struct rbt *tree, struct rbt_node *node) {
585 fix_insertion(tree, node);
586 rbt_debug_validate(tree, caller: __func__);
587}
588
589struct rbt *rbt_init(struct rbt *tree, rbt_get_data get_data, rbt_compare cmp) {
590 tree->root = NULL;
591 tree->compare = cmp;
592 tree->get_data = get_data;
593 return tree;
594}
595