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