1#include <console/panic.h>
2#include <enforce.h>
3#include <math/min_max.h>
4#include <mem/alloc.h>
5#include <stddef.h>
6#include <stdint.h>
7#include <structures/rbit.h>
8
9static void rbit_check_cycle(struct rbit_node *node, const char *caller) {
10 struct rbit_node *tortoise = node;
11 struct rbit_node *hare = node;
12
13 while (hare && hare->left != NULL) {
14 tortoise = tortoise ? tortoise->left : NULL;
15 hare = hare->left->left;
16
17 if (tortoise == NULL || hare == NULL)
18 break;
19
20 if (tortoise == hare) {
21 panic(CORRUPTION,
22 "Cycle detected in %s (left spine): tortoise == hare at "
23 "node=%p",
24 caller, (void *) tortoise);
25 break;
26 }
27 }
28
29 tortoise = node;
30 hare = node;
31
32 while (hare && hare->right != NULL) {
33 tortoise = tortoise ? tortoise->right : NULL;
34 hare = hare->right->right;
35
36 if (tortoise == NULL || hare == NULL)
37 break;
38
39 if (tortoise == hare) {
40 panic(CORRUPTION,
41 "Cycle detected in %s (right spine): tortoise == hare at "
42 "node=%p",
43 caller, (void *) tortoise);
44 break;
45 }
46 }
47}
48
49static size_t node_max(struct rbit_node *n) {
50 return n ? n->max : 0;
51}
52
53static bool default_augment(struct rbit_node *n) {
54 size_t old = n->max;
55 n->max = MAX(n->interval.high, node_max(n->left), node_max(n->right));
56 return n->max != old;
57}
58
59static bool node_augment(struct rbit *tree, struct rbit_node *n) {
60 if (!n)
61 return false;
62 return tree->augment ? tree->augment(n) : default_augment(n);
63}
64
65static void augment_up(struct rbit *tree, struct rbit_node *n) {
66 while (n) {
67 node_augment(tree, n);
68 n = n->parent;
69 }
70}
71
72struct rbit *rbit_init(struct rbit *rbit) {
73 rbit->root = NULL;
74 rbit->augment = NULL;
75 return rbit;
76}
77
78struct rbit *rbit_tree_create(void) {
79 struct rbit *tree = vnew(tree);
80 if (!tree)
81 return NULL;
82
83 tree->root = NULL;
84 tree->augment = NULL;
85 return tree;
86}
87
88struct rbit_node *rbit_find_min(struct rbit_node *node) {
89 rbit_check_cycle(node, caller: __func__);
90 while (node && node->left != NULL) {
91 node = node->left;
92 }
93 return node;
94}
95
96struct rbit_node *rbit_find_max(struct rbit_node *node) {
97 rbit_check_cycle(node, caller: __func__);
98 while (node && node->right != NULL) {
99 node = node->right;
100 }
101 return node;
102}
103
104struct rbit_node *rbit_min(struct rbit *tree) {
105 rbit_check_cycle(node: tree->root, caller: __func__);
106 return rbit_find_min(node: tree->root);
107}
108
109struct rbit_node *rbit_max(struct rbit *tree) {
110 rbit_check_cycle(node: tree->root, caller: __func__);
111 return rbit_find_max(node: tree->root);
112}
113
114struct rbit_node *rbit_next(struct rbit_node *node) {
115 rbit_check_cycle(node, caller: __func__);
116
117 if (!node)
118 return NULL;
119
120 if (node->right)
121 return rbit_find_min(node: node->right);
122
123 struct rbit_node *parent = node->parent;
124 while (parent && node == parent->right) {
125 node = parent;
126 parent = parent->parent;
127 }
128 return parent;
129}
130
131struct rbit_node *rbit_find_predecessor(struct rbit *tree, size_t low) {
132 rbit_check_cycle(node: tree->root, caller: __func__);
133 struct rbit_node *curr = tree->root;
134 struct rbit_node *pred = NULL;
135
136 while (curr) {
137 if (curr->interval.low < low) {
138 pred = curr;
139 curr = curr->right;
140 } else {
141 curr = curr->left;
142 }
143 }
144 return pred;
145}
146
147struct rbit_node *rbit_find_successor(struct rbit *tree, size_t low) {
148 rbit_check_cycle(node: tree->root, caller: __func__);
149 struct rbit_node *curr = tree->root;
150 struct rbit_node *succ = NULL;
151
152 while (curr) {
153 if (curr->interval.low > low) {
154 succ = curr;
155 curr = curr->left;
156 } else {
157 curr = curr->right;
158 }
159 }
160 return succ;
161}
162
163bool rbit_has_node(struct rbit *tree, struct rbit_node *node) {
164 rbit_check_cycle(node: tree->root, caller: __func__);
165 struct rbit_node *iter;
166 rbit_for_each(iter, tree) {
167 if (iter == node)
168 return true;
169 }
170 return false;
171}
172
173static void rbit_transplant(struct rbit *tree, struct rbit_node *u,
174 struct rbit_node *v) {
175 if (u->parent == NULL)
176 tree->root = v;
177 else if (u == u->parent->left)
178 u->parent->left = v;
179 else
180 u->parent->right = v;
181
182 if (v)
183 v->parent = u->parent;
184}
185
186static void left_rotate(struct rbit *tree, struct rbit_node *x) {
187 struct rbit_node *y = x->right;
188 x->right = y->left;
189
190 if (y->left)
191 y->left->parent = x;
192
193 y->parent = x->parent;
194 if (!x->parent)
195 tree->root = y;
196 else if (x == x->parent->left)
197 x->parent->left = y;
198 else
199 x->parent->right = y;
200
201 y->left = x;
202 x->parent = y;
203
204 /* bottom-up: x (now the lower node) before y (the new apex) */
205 node_augment(tree, n: x);
206 node_augment(tree, n: y);
207}
208
209static void right_rotate(struct rbit *tree, struct rbit_node *y) {
210 struct rbit_node *x = y->left;
211 y->left = x->right;
212 if (x->right)
213 x->right->parent = y;
214
215 x->parent = y->parent;
216 if (!y->parent)
217 tree->root = x;
218 else if (y == y->parent->right)
219 y->parent->right = x;
220 else
221 y->parent->left = x;
222
223 x->right = y;
224 y->parent = x;
225
226 /* bottom-up: y (now the lower node) before x (the new apex) */
227 node_augment(tree, n: y);
228 node_augment(tree, n: x);
229}
230
231static bool is_black(struct rbit_node *n) {
232 return !n || n->color == RBIT_BLACK;
233}
234
235static void fix_deletion(struct rbit *tree, struct rbit_node *x,
236 struct rbit_node *x_parent) {
237 while (x != tree->root && is_black(n: x)) {
238 if (!x_parent)
239 break;
240
241 struct rbit_node *sibling;
242
243 if (x == x_parent->left) {
244 sibling = x_parent->right;
245
246 if (sibling && sibling->color == RBIT_RED) {
247 sibling->color = RBIT_BLACK;
248 x_parent->color = RBIT_RED;
249 left_rotate(tree, x: x_parent);
250 sibling = x_parent->right;
251 }
252
253 if (!sibling ||
254 (is_black(n: sibling->left) && is_black(n: sibling->right))) {
255 if (sibling)
256 sibling->color = RBIT_RED;
257 x = x_parent;
258 x_parent = x->parent;
259 } else {
260 if (is_black(n: sibling->right)) {
261 if (sibling->left)
262 sibling->left->color = RBIT_BLACK;
263 sibling->color = RBIT_RED;
264 right_rotate(tree, y: sibling);
265 sibling = x_parent->right;
266 }
267
268 sibling->color = x_parent->color;
269 x_parent->color = RBIT_BLACK;
270 if (sibling->right)
271 sibling->right->color = RBIT_BLACK;
272 left_rotate(tree, x: x_parent);
273 x = tree->root;
274 x_parent = NULL;
275 }
276 } else {
277 sibling = x_parent->left;
278
279 if (sibling && sibling->color == RBIT_RED) {
280 sibling->color = RBIT_BLACK;
281 x_parent->color = RBIT_RED;
282 right_rotate(tree, y: x_parent);
283 sibling = x_parent->left;
284 }
285
286 if (!sibling ||
287 (is_black(n: sibling->left) && is_black(n: sibling->right))) {
288 if (sibling)
289 sibling->color = RBIT_RED;
290 x = x_parent;
291 x_parent = x->parent;
292 } else {
293 if (is_black(n: sibling->left)) {
294 if (sibling->right)
295 sibling->right->color = RBIT_BLACK;
296 sibling->color = RBIT_RED;
297 left_rotate(tree, x: sibling);
298 sibling = x_parent->left;
299 }
300
301 sibling->color = x_parent->color;
302 x_parent->color = RBIT_BLACK;
303 if (sibling->left)
304 sibling->left->color = RBIT_BLACK;
305 right_rotate(tree, y: x_parent);
306 x = tree->root;
307 x_parent = NULL;
308 }
309 }
310 }
311
312 if (x)
313 x->color = RBIT_BLACK;
314}
315
316static size_t validate_rbit(struct rbit_node *node, size_t *black_height) {
317 if (node == NULL) {
318 *black_height = 1;
319 return 1;
320 }
321
322 if (node->color == RBIT_RED) {
323 if ((node->left && node->left->color == RBIT_RED) ||
324 (node->right && node->right->color == RBIT_RED)) {
325 panic(CORRUPTION, "Red-Red violation at node [%zu,%zu]",
326 node->interval.low, node->interval.high);
327 return 0;
328 }
329 }
330
331 size_t left_black_height = 0;
332 size_t right_black_height = 0;
333
334 if (!validate_rbit(node: node->left, black_height: &left_black_height))
335 return 0;
336 if (!validate_rbit(node: node->right, black_height: &right_black_height))
337 return 0;
338
339 if (left_black_height != right_black_height) {
340 panic(CORRUPTION,
341 "Black-height violation at node [%zu,%zu] (left=%zu, right=%zu)",
342 node->interval.low, node->interval.high, left_black_height,
343 right_black_height);
344 return 0;
345 }
346
347 size_t expected_max =
348 MAX(node->interval.high, node_max(node->left), node_max(node->right));
349 if (node->max != expected_max) {
350 panic(CORRUPTION,
351 "Max-invariant violation at node [%zu,%zu]: stored=%zu, "
352 "expected=%zu",
353 node->interval.low, node->interval.high, node->max, expected_max);
354 return 0;
355 }
356
357 *black_height = left_black_height + (node->color == RBIT_BLACK ? 1 : 0);
358 return 1;
359}
360
361void rbit_delete(struct rbit *tree, struct rbit_node *cc_nonnull_ptr z) {
362 enforce_precond(z != NULL);
363 enforce(rbit_has_node(tree, z));
364 rbit_check_cycle(node: tree->root, caller: __func__);
365 struct rbit_node *y = z;
366 struct rbit_node *x = NULL;
367 struct rbit_node *x_parent = NULL;
368 struct rbit_node *fix_from = NULL;
369 enum rbit_color y_original_color = y->color;
370
371 if (z->left == NULL) {
372 x = z->right;
373 x_parent = z->parent;
374 fix_from = z->parent;
375 rbit_transplant(tree, u: z, v: z->right);
376 } else if (z->right == NULL) {
377 x = z->left;
378 x_parent = z->parent;
379 fix_from = z->parent;
380 rbit_transplant(tree, u: z, v: z->left);
381 } else {
382 y = rbit_find_min(node: z->right);
383 y_original_color = y->color;
384 x = y->right;
385
386 if (y->parent != z) {
387 x_parent = y->parent;
388 fix_from = y->parent;
389 rbit_transplant(tree, u: y, v: y->right);
390 y->right = z->right;
391 if (y->right)
392 y->right->parent = y;
393 } else {
394 x_parent = y;
395 fix_from = y;
396 }
397
398 rbit_transplant(tree, u: z, v: y);
399 y->left = z->left;
400 if (y->left)
401 y->left->parent = y;
402 y->color = z->color;
403 node_augment(tree, n: y);
404 }
405
406 augment_up(tree, n: fix_from);
407
408 if (y_original_color == RBIT_BLACK) {
409 fix_deletion(tree, x, x_parent);
410 }
411
412 if (x)
413 augment_up(tree, n: x);
414 augment_up(tree, n: x_parent);
415
416 z->left = NULL;
417 z->right = NULL;
418 z->parent = NULL;
419
420 size_t dummy_bh = 0;
421 enforce(validate_rbit(tree->root, &dummy_bh), CORRUPTION);
422}
423
424struct rbit_node *rbit_search(struct rbit_node *root, struct interval iv) {
425 rbit_check_cycle(node: root, caller: __func__);
426 while (root &&
427 !(root->interval.low == iv.low && root->interval.high == iv.high)) {
428 if (iv.low < root->interval.low)
429 root = root->left;
430 else
431 root = root->right;
432 }
433 return root;
434}
435
436static int intervals_overlap(struct interval a, struct interval b) {
437 return a.low <= b.high && b.low <= a.high;
438}
439
440struct rbit_node *rbit_overlap_search(struct rbit_node *root,
441 struct interval iv) {
442 rbit_check_cycle(node: root, caller: __func__);
443 while (root && !intervals_overlap(a: root->interval, b: iv)) {
444 if (root->left && root->left->max >= iv.low)
445 root = root->left;
446 else
447 root = root->right;
448 }
449 return root;
450}
451
452void rbit_remove(struct rbit *tree, struct interval iv) {
453 struct rbit_node *node = rbit_search(root: tree->root, iv);
454 if (node)
455 rbit_delete(tree, z: node);
456}
457
458static void fix_insertion(struct rbit *tree, struct rbit_node *node) {
459 while (node != tree->root && node->parent->color == RBIT_RED) {
460 struct rbit_node *parent = node->parent;
461 struct rbit_node *grandparent = parent->parent;
462
463 if (parent == grandparent->left) {
464 struct rbit_node *uncle = grandparent->right;
465 if (uncle && uncle->color == RBIT_RED) {
466 parent->color = RBIT_BLACK;
467 uncle->color = RBIT_BLACK;
468 grandparent->color = RBIT_RED;
469 node = grandparent;
470 } else {
471 if (node == parent->right) {
472 node = parent;
473 left_rotate(tree, x: node);
474 parent = node->parent;
475 }
476 parent->color = RBIT_BLACK;
477 grandparent->color = RBIT_RED;
478 right_rotate(tree, y: grandparent);
479 }
480 } else {
481 struct rbit_node *uncle = grandparent->left;
482 if (uncle && uncle->color == RBIT_RED) {
483 parent->color = RBIT_BLACK;
484 uncle->color = RBIT_BLACK;
485 grandparent->color = RBIT_RED;
486 node = grandparent;
487 } else {
488 if (node == parent->left) {
489 node = parent;
490 right_rotate(tree, y: node);
491 parent = node->parent;
492 }
493 parent->color = RBIT_BLACK;
494 grandparent->color = RBIT_RED;
495 left_rotate(tree, x: grandparent);
496 }
497 }
498 }
499 tree->root->color = RBIT_BLACK;
500}
501
502void rbit_insert(struct rbit *tree, struct rbit_node *cc_nonnull_ptr new_node) {
503 enforce_precond(new_node != NULL);
504 enforce(!rbit_has_node(tree, new_node), DUPLICATE);
505 rbit_check_cycle(node: tree->root, caller: __func__);
506 new_node->left = NULL;
507 new_node->right = NULL;
508 new_node->color = RBIT_RED;
509 node_augment(tree, n: new_node);
510
511 if (tree->root == NULL) {
512 new_node->color = RBIT_BLACK;
513 new_node->parent = NULL;
514 tree->root = new_node;
515 return;
516 }
517
518 struct rbit_node *current = tree->root;
519 struct rbit_node *parent = NULL;
520 while (current != NULL) {
521 parent = current;
522 if (new_node->interval.low < current->interval.low)
523 current = current->left;
524 else
525 current = current->right;
526 }
527
528 new_node->parent = parent;
529 if (new_node->interval.low < parent->interval.low)
530 parent->left = new_node;
531 else
532 parent->right = new_node;
533
534 augment_up(tree, n: parent);
535
536 fix_insertion(tree, node: new_node);
537
538 if (parent)
539 enforce(!(parent->color == RBIT_RED && new_node->color == RBIT_RED),
540 CORRUPTION);
541
542 size_t dummy_bh = 0;
543 enforce(validate_rbit(tree->root, &dummy_bh), CORRUPTION);
544}
545