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
8static void rbt_check_cycle(struct rbt_node *node, const char *caller) {
9 struct rbt_node *tortoise = node;
10 struct rbt_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
46struct rbt *rbt_create(rbt_get_data get, rbt_compare cmp) {
47 struct rbt *tree = kmalloc(sizeof(struct rbt));
48 if (!tree)
49 return NULL;
50
51 tree->compare = cmp;
52 tree->get_data = get;
53 tree->root = NULL;
54 return tree;
55}
56
57struct rbt_node *rbt_find_min(struct rbt_node *node) {
58 rbt_check_cycle(node, caller: __func__);
59
60 while (node && node->left != NULL)
61 node = node->left;
62
63 return node;
64}
65
66struct rbt_node *rbt_find_max(struct rbt_node *node) {
67 rbt_check_cycle(node, caller: __func__);
68
69 while (node && node->right != NULL) {
70 node = node->right;
71 }
72
73 return node;
74}
75
76struct rbt_node *rbt_max(struct rbt *tree) {
77 rbt_check_cycle(node: tree->root, caller: __func__);
78 return rbt_find_max(node: tree->root);
79}
80
81struct rbt_node *rbt_min(struct rbt *tree) {
82 rbt_check_cycle(node: tree->root, caller: __func__);
83 return rbt_find_min(node: tree->root);
84}
85
86struct rbt_node *rbt_next(struct rbt_node *node) {
87 rbt_check_cycle(node, caller: __func__);
88
89 if (!node)
90 return NULL;
91
92 if (node->right)
93 return rbt_find_min(node: node->right);
94
95 struct rbt_node *parent = node->parent;
96 while (parent && node == parent->right) {
97 node = parent;
98 parent = parent->parent;
99 }
100 return parent;
101}
102
103struct rbt_node *rbt_find_predecessor(struct rbt *tree, uint64_t data) {
104 rbt_check_cycle(node: tree->root, caller: __func__);
105 struct rbt_node *curr = tree->root;
106 struct rbt_node *pred = NULL;
107
108 while (curr) {
109 uint64_t curr_data = tree->get_data(curr);
110 if (curr_data < data) {
111 pred = curr;
112 curr = curr->right;
113 } else {
114 curr = curr->left;
115 }
116 }
117 return pred;
118}
119
120struct rbt_node *rbt_find_successor(struct rbt *tree, uint64_t data) {
121 rbt_check_cycle(node: tree->root, caller: __func__);
122 struct rbt_node *curr = tree->root;
123 struct rbt_node *succ = NULL;
124
125 while (curr) {
126 uint64_t curr_data = tree->get_data(curr);
127 if (curr_data > data) {
128 succ = curr;
129 curr = curr->left;
130 } else {
131 curr = curr->right;
132 }
133 }
134 return succ;
135}
136
137static void rb_transplant(struct rbt *tree, struct rbt_node *u,
138 struct rbt_node *v) {
139 if (u->parent == NULL)
140 tree->root = v;
141 else if (u == u->parent->left)
142 u->parent->left = v;
143 else
144 u->parent->right = v;
145
146 if (v)
147 v->parent = u->parent;
148}
149
150static void left_rotate(struct rbt *tree, struct rbt_node *x) {
151 struct rbt_node *y = x->right;
152 x->right = y->left;
153
154 if (y->left)
155 y->left->parent = x;
156
157 y->parent = x->parent;
158 if (!x->parent)
159 tree->root = y;
160 else if (x == x->parent->left)
161 x->parent->left = y;
162 else
163 x->parent->right = y;
164
165 y->left = x;
166 x->parent = y;
167}
168
169static void right_rotate(struct rbt *tree, struct rbt_node *y) {
170 struct rbt_node *x = y->left;
171 y->left = x->right;
172 if (x->right)
173 x->right->parent = y;
174
175 x->parent = y->parent;
176 if (!y->parent)
177 tree->root = x;
178 else if (y == y->parent->right)
179 y->parent->right = x;
180 else
181 y->parent->left = x;
182
183 x->right = y;
184 y->parent = x;
185}
186
187static void fix_deletion(struct rbt *tree, struct rbt_node *x) {
188 while (x != tree->root && (!x || x->color == TREE_NODE_BLACK)) {
189 if (!x || !x->parent)
190 break;
191
192 struct rbt_node *sibling;
193
194 if (x == x->parent->left) {
195 sibling = x->parent->right;
196
197 if (sibling && sibling->color == TREE_NODE_RED) {
198 sibling->color = TREE_NODE_BLACK;
199 x->parent->color = TREE_NODE_RED;
200 left_rotate(tree, x: x->parent);
201 sibling = x->parent->right;
202 }
203
204 if (!sibling ||
205 ((!sibling->left || sibling->left->color == TREE_NODE_BLACK) &&
206 (!sibling->right ||
207 sibling->right->color == TREE_NODE_BLACK))) {
208 if (sibling)
209 sibling->color = TREE_NODE_RED;
210 x = x->parent;
211 } else {
212 if (!sibling->right ||
213 sibling->right->color == TREE_NODE_BLACK) {
214 if (sibling->left)
215 sibling->left->color = TREE_NODE_BLACK;
216 if (sibling) {
217 sibling->color = TREE_NODE_RED;
218 right_rotate(tree, y: sibling);
219 sibling = x->parent->right;
220 }
221 }
222
223 if (sibling) {
224 sibling->color = x->parent->color;
225 x->parent->color = TREE_NODE_BLACK;
226 if (sibling->right)
227 sibling->right->color = TREE_NODE_BLACK;
228 left_rotate(tree, x: x->parent);
229 }
230 x = tree->root;
231 }
232 } else {
233 sibling = x->parent->left;
234
235 if (sibling && sibling->color == TREE_NODE_RED) {
236 sibling->color = TREE_NODE_BLACK;
237 x->parent->color = TREE_NODE_RED;
238 right_rotate(tree, y: x->parent);
239 sibling = x->parent->left;
240 }
241
242 if (!sibling ||
243 ((!sibling->left || sibling->left->color == TREE_NODE_BLACK) &&
244 (!sibling->right ||
245 sibling->right->color == TREE_NODE_BLACK))) {
246 if (sibling)
247 sibling->color = TREE_NODE_RED;
248 x = x->parent;
249 } else {
250 if (!sibling->left || sibling->left->color == TREE_NODE_BLACK) {
251 if (sibling->right)
252 sibling->right->color = TREE_NODE_BLACK;
253 if (sibling) {
254 sibling->color = TREE_NODE_RED;
255 left_rotate(tree, x: sibling);
256 sibling = x->parent->left;
257 }
258 }
259
260 if (sibling) {
261 sibling->color = x->parent->color;
262 x->parent->color = TREE_NODE_BLACK;
263 if (sibling->left)
264 sibling->left->color = TREE_NODE_BLACK;
265 right_rotate(tree, y: x->parent);
266 }
267 x = tree->root;
268 }
269 }
270 }
271
272 if (x)
273 x->color = TREE_NODE_BLACK;
274}
275
276void rbt_delete(struct rbt *tree, struct rbt_node *z) {
277 kassert(rbt_has_node(tree, z));
278 rbt_check_cycle(node: tree->root, caller: __func__);
279 struct rbt_node *y = z;
280 struct rbt_node *x = NULL;
281 enum rbt_node_color y_original_color = y->color;
282
283 if (z->left == NULL) {
284 x = z->right;
285 rb_transplant(tree, u: z, v: z->right);
286 } else if (z->right == NULL) {
287 x = z->left;
288 rb_transplant(tree, u: z, v: z->left);
289 } else {
290 y = rbt_find_min(node: z->right);
291 y_original_color = y->color;
292 x = y->right;
293
294 if (y->parent != z) {
295 rb_transplant(tree, u: y, v: y->right);
296 y->right = z->right;
297 if (y->right)
298 y->right->parent = y;
299 }
300
301 rb_transplant(tree, u: z, v: y);
302 y->left = z->left;
303 if (y->left)
304 y->left->parent = y;
305 y->color = z->color;
306 }
307
308 if (y_original_color == TREE_NODE_BLACK) {
309 fix_deletion(tree, x);
310 }
311
312 z->left = NULL;
313 z->right = NULL;
314 z->parent = NULL;
315}
316
317static struct rbt_node *
318rbt_search_internal(struct rbt *tree, struct rbt_node *root, uint64_t data) {
319 while (root && tree->get_data(root) != data) {
320 if (data < tree->get_data(root))
321 root = root->left;
322 else
323 root = root->right;
324 }
325 return root;
326}
327
328struct rbt_node *rbt_search(struct rbt *tree, uint64_t data) {
329 rbt_check_cycle(node: tree->root, caller: __func__);
330 struct rbt_node *root = tree->root;
331
332 while (root && tree->get_data(root) != data) {
333 if (data < tree->get_data(root))
334 root = root->left;
335 else
336 root = root->right;
337 }
338
339 return (root && tree->get_data(root) == data) ? root : NULL;
340}
341
342bool rbt_has_node(struct rbt *tree, struct rbt_node *node) {
343 rbt_check_cycle(node: tree->root, caller: __func__);
344 struct rbt_node *iter;
345 rbt_for_each(iter, tree) {
346 if (iter == node)
347 return true;
348 }
349 return false;
350}
351
352void rbt_remove(struct rbt *tree, uint64_t data) {
353 rbt_check_cycle(node: tree->root, caller: __func__);
354 struct rbt_node *node = rbt_search_internal(tree, root: tree->root, data);
355 if (node)
356 rbt_delete(tree, z: node);
357}
358
359static void fix_insertion(struct rbt *tree, struct rbt_node *node) {
360 while (node != tree->root && node->parent->color == TREE_NODE_RED) {
361 struct rbt_node *parent = node->parent;
362 struct rbt_node *grandparent = parent->parent;
363
364 if (parent == grandparent->left) {
365 struct rbt_node *uncle = grandparent->right;
366 if (uncle && uncle->color == TREE_NODE_RED) {
367 parent->color = TREE_NODE_BLACK;
368 uncle->color = TREE_NODE_BLACK;
369 grandparent->color = TREE_NODE_RED;
370 node = grandparent;
371 } else {
372 if (node == parent->right) {
373 node = parent;
374 left_rotate(tree, x: node);
375 parent = node->parent;
376 }
377 parent->color = TREE_NODE_BLACK;
378 grandparent->color = TREE_NODE_RED;
379 right_rotate(tree, y: grandparent);
380 }
381 } else {
382 struct rbt_node *uncle = grandparent->left;
383 if (uncle && uncle->color == TREE_NODE_RED) {
384 parent->color = TREE_NODE_BLACK;
385 uncle->color = TREE_NODE_BLACK;
386 grandparent->color = TREE_NODE_RED;
387 node = grandparent;
388 } else {
389 if (node == parent->left) {
390 node = parent;
391 right_rotate(tree, y: node);
392 parent = node->parent;
393 }
394 parent->color = TREE_NODE_BLACK;
395 grandparent->color = TREE_NODE_RED;
396 left_rotate(tree, x: grandparent);
397 }
398 }
399 }
400 tree->root->color = TREE_NODE_BLACK;
401}
402
403void rbt_insert(struct rbt *tree, struct rbt_node *new_node) {
404 kassert(!rbt_has_node(tree, new_node));
405 rbt_check_cycle(node: tree->root, caller: __func__);
406 new_node->left = NULL;
407 new_node->right = NULL;
408 new_node->color = TREE_NODE_RED;
409
410 if (tree->root == NULL) {
411 new_node->color = TREE_NODE_BLACK;
412 new_node->parent = NULL;
413 tree->root = new_node;
414 return;
415 }
416
417 struct rbt_node *current = tree->root;
418 struct rbt_node *parent = NULL;
419 while (current != NULL) {
420 parent = current;
421 if (tree->compare(new_node, current) < 0)
422 current = current->left;
423 else
424 current = current->right;
425 }
426
427 new_node->parent = parent;
428 if (tree->compare(new_node, parent) < 0)
429 parent->left = new_node;
430 else
431 parent->right = new_node;
432
433 fix_insertion(tree, node: new_node);
434
435 if (parent)
436 kassert(!(parent->color == TREE_NODE_RED &&
437 new_node->color == TREE_NODE_RED));
438}
439
440void rbt_link_node(struct rbt_node *node, struct rbt_node *parent,
441 struct rbt_node **link) {
442 node->parent = parent;
443 node->left = NULL;
444 node->right = NULL;
445 node->color = TREE_NODE_RED;
446 *link = node;
447}
448
449void rbt_insert_color(struct rbt *tree, struct rbt_node *node) {
450 fix_insertion(tree, node);
451}
452
453struct rbt *rbt_init(struct rbt *tree, rbt_get_data get_data, rbt_compare cmp) {
454 tree->root = NULL;
455 tree->compare = cmp;
456 tree->get_data = get_data;
457 return tree;
458}
459