1 /*
2 * Copyright (c) 2018 Intel Corporation
3 *
4 * SPDX-License-Identifier: Apache-2.0
5 */
6
7 /* These assertions are very useful when debugging the tree code
8 * itself, but produce significant performance degradation as they are
9 * checked many times per operation. Leave them off unless you're
10 * working on the rbtree code itself
11 */
12 #define CHECK(n) /**/
13 /* #define CHECK(n) __ASSERT_NO_MSG(n) */
14
15 #include <zephyr/kernel.h>
16 #include <zephyr/sys/rb.h>
17 #include <stdbool.h>
18
19 enum rb_color { RED = 0U, BLACK = 1U };
20
get_child(struct rbnode * n,uint8_t side)21 static struct rbnode *get_child(struct rbnode *n, uint8_t side)
22 {
23 CHECK(n);
24 if (side != 0U) {
25 return n->children[1];
26 }
27
28 uintptr_t l = (uintptr_t) n->children[0];
29
30 l &= ~1UL;
31 return (struct rbnode *) l;
32 }
33
set_child(struct rbnode * n,uint8_t side,void * val)34 static void set_child(struct rbnode *n, uint8_t side, void *val)
35 {
36 CHECK(n);
37 if (side != 0U) {
38 n->children[1] = val;
39 } else {
40 uintptr_t old = (uintptr_t) n->children[0];
41 uintptr_t new = (uintptr_t) val;
42
43 n->children[0] = (void *) (new | (old & 1UL));
44 }
45 }
46
get_color(struct rbnode * n)47 static enum rb_color get_color(struct rbnode *n)
48 {
49 CHECK(n);
50 return ((uintptr_t)n->children[0]) & 1UL;
51 }
52
is_black(struct rbnode * n)53 static bool is_black(struct rbnode *n)
54 {
55 return get_color(n) == BLACK;
56 }
57
is_red(struct rbnode * n)58 static bool is_red(struct rbnode *n)
59 {
60 return get_color(n) == RED;
61 }
62
set_color(struct rbnode * n,enum rb_color color)63 static void set_color(struct rbnode *n, enum rb_color color)
64 {
65 CHECK(n);
66
67 uintptr_t *p = (void *) &n->children[0];
68
69 *p = (*p & ~1UL) | (uint8_t)color;
70 }
71
72 /* Searches the tree down to a node that is either identical with the
73 * "node" argument or has an empty/leaf child pointer where "node"
74 * should be, leaving all nodes found in the resulting stack. Note
75 * that tree must not be empty and that stack should be allocated to
76 * contain at least tree->max_depth entries! Returns the number of
77 * entries pushed onto the stack.
78 */
find_and_stack(struct rbtree * tree,struct rbnode * node,struct rbnode ** stack)79 static int find_and_stack(struct rbtree *tree, struct rbnode *node,
80 struct rbnode **stack)
81 {
82 int sz = 0;
83
84 stack[sz++] = tree->root;
85
86 while (stack[sz - 1] != node) {
87 uint8_t side = tree->lessthan_fn(node, stack[sz - 1]) ? 0U : 1U;
88 struct rbnode *ch = get_child(stack[sz - 1], side);
89
90 if (ch != NULL) {
91 stack[sz++] = ch;
92 } else {
93 break;
94 }
95 }
96
97 return sz;
98 }
99
z_rb_get_minmax(struct rbtree * tree,uint8_t side)100 struct rbnode *z_rb_get_minmax(struct rbtree *tree, uint8_t side)
101 {
102 struct rbnode *n;
103
104 for (n = tree->root; (n != NULL) && (get_child(n, side) != NULL);
105 n = get_child(n, side)) {
106 ;
107 }
108 return n;
109 }
110
get_side(struct rbnode * parent,struct rbnode * child)111 static uint8_t get_side(struct rbnode *parent, struct rbnode *child)
112 {
113 CHECK(get_child(parent, 0U) == child || get_child(parent, 1U) == child);
114
115 return (get_child(parent, 1U) == child) ? 1U : 0U;
116 }
117
118 /* Swaps the position of the two nodes at the top of the provided
119 * stack, modifying the stack accordingly. Does not change the color
120 * of either node. That is, it effects the following transition (or
121 * its mirror if N is on the other side of P, of course):
122 *
123 * P N
124 * N c --> a P
125 * a b b c
126 *
127 */
rotate(struct rbnode ** stack,int stacksz)128 static void rotate(struct rbnode **stack, int stacksz)
129 {
130 CHECK(stacksz >= 2);
131
132 struct rbnode *parent = stack[stacksz - 2];
133 struct rbnode *child = stack[stacksz - 1];
134 uint8_t side = get_side(parent, child);
135 struct rbnode *a = get_child(child, side);
136 struct rbnode *b = get_child(child, (side == 0U) ? 1U : 0U);
137
138 if (stacksz >= 3) {
139 struct rbnode *grandparent = stack[stacksz - 3];
140
141 set_child(grandparent, get_side(grandparent, parent), child);
142 }
143
144 set_child(child, side, a);
145 set_child(child, (side == 0U) ? 1U : 0U, parent);
146 set_child(parent, side, b);
147 stack[stacksz - 2] = child;
148 stack[stacksz - 1] = parent;
149 }
150
151 /* The node at the top of the provided stack is red, and its parent is
152 * too. Iteratively fix the tree so it becomes a valid red black tree
153 * again
154 */
fix_extra_red(struct rbnode ** stack,int stacksz)155 static void fix_extra_red(struct rbnode **stack, int stacksz)
156 {
157 while (stacksz > 1) {
158 struct rbnode *node = stack[stacksz - 1];
159 struct rbnode *parent = stack[stacksz - 2];
160
161 /* Correct child colors are a precondition of the loop */
162 CHECK((get_child(node, 0U) == NULL) ||
163 is_black(get_child(node, 0U)));
164 CHECK((get_child(node, 1U) == NULL) ||
165 is_black(get_child(node, 1U)));
166
167 if (is_black(parent)) {
168 return;
169 }
170
171 /* We are guaranteed to have a grandparent if our
172 * parent is red, as red nodes cannot be the root
173 */
174 CHECK(stacksz >= 2);
175
176 struct rbnode *grandparent = stack[stacksz - 3];
177 uint8_t side = get_side(grandparent, parent);
178 struct rbnode *aunt = get_child(grandparent,
179 (side == 0U) ? 1U : 0U);
180
181 if ((aunt != NULL) && is_red(aunt)) {
182 set_color(grandparent, RED);
183 set_color(parent, BLACK);
184 set_color(aunt, BLACK);
185
186 /* We colored the grandparent red, which might
187 * have a red parent, so continue iterating
188 * from there.
189 */
190 stacksz -= 2;
191 continue;
192 }
193
194 /* We can rotate locally to fix the whole tree. First
195 * make sure that node is on the same side of parent
196 * as parent is of grandparent.
197 */
198 uint8_t parent_side = get_side(parent, node);
199
200 if (parent_side != side) {
201 rotate(stack, stacksz);
202 }
203
204 /* Rotate the grandparent with parent, swapping colors */
205 rotate(stack, stacksz - 1);
206 set_color(stack[stacksz - 3], BLACK);
207 set_color(stack[stacksz - 2], RED);
208 return;
209 }
210
211 /* If we exit the loop, it's because our node is now the root,
212 * which must be black.
213 */
214 set_color(stack[0], BLACK);
215 }
216
rb_insert(struct rbtree * tree,struct rbnode * node)217 void rb_insert(struct rbtree *tree, struct rbnode *node)
218 {
219 set_child(node, 0U, NULL);
220 set_child(node, 1U, NULL);
221
222 if (tree->root == NULL) {
223 tree->root = node;
224 tree->max_depth = 1;
225 set_color(node, BLACK);
226 return;
227 }
228
229 #ifdef CONFIG_MISRA_SANE
230 struct rbnode **stack = &tree->iter_stack[0];
231 #else
232 struct rbnode *stack[tree->max_depth + 1];
233 #endif
234
235 int stacksz = find_and_stack(tree, node, stack);
236
237 struct rbnode *parent = stack[stacksz - 1];
238
239 uint8_t side = tree->lessthan_fn(node, parent) ? 0U : 1U;
240
241 set_child(parent, side, node);
242 set_color(node, RED);
243
244 stack[stacksz++] = node;
245 fix_extra_red(stack, stacksz);
246
247 if (stacksz > tree->max_depth) {
248 tree->max_depth = stacksz;
249 }
250
251 /* We may have rotated up into the root! */
252 tree->root = stack[0];
253 CHECK(is_black(tree->root));
254 }
255
256 /* Called for a node N (at the top of the stack) which after a
257 * deletion operation is "missing a black" in its subtree. By
258 * construction N must be black (because if it was red it would be
259 * trivially fixed by recoloring and we wouldn't be here). Fixes up
260 * the tree to preserve red/black rules. The "null_node" pointer is
261 * for situations where we are removing a childless black node. The
262 * tree munging needs a real node for simplicity, so we use it and
263 * then clean it up (replace it with a simple NULL child in the
264 * parent) when finished.
265 */
fix_missing_black(struct rbnode ** stack,int stacksz,struct rbnode * null_node)266 static void fix_missing_black(struct rbnode **stack, int stacksz,
267 struct rbnode *null_node)
268 {
269 /* Loop upward until we reach the root */
270 while (stacksz > 1) {
271 struct rbnode *c0, *c1, *inner, *outer;
272 struct rbnode *n = stack[stacksz - 1];
273 struct rbnode *parent = stack[stacksz - 2];
274 uint8_t n_side = get_side(parent, n);
275 struct rbnode *sib = get_child(parent,
276 (n_side == 0U) ? 1U : 0U);
277
278 CHECK(is_black(n));
279
280 /* Guarantee the sibling is black, rotating N down a
281 * level if needed (after rotate() our parent is the
282 * child of our previous-sibling, so N is lower in the
283 * tree)
284 */
285 if (!is_black(sib)) {
286 stack[stacksz - 1] = sib;
287 rotate(stack, stacksz);
288 set_color(parent, RED);
289 set_color(sib, BLACK);
290 stack[stacksz++] = n;
291
292 parent = stack[stacksz - 2];
293 sib = get_child(parent, (n_side == 0U) ? 1U : 0U);
294 }
295
296 CHECK(sib);
297
298 /* Cases where the sibling has only black children
299 * have simple resolutions
300 */
301 c0 = get_child(sib, 0U);
302 c1 = get_child(sib, 1U);
303 if (((c0 == NULL) || is_black(c0)) && ((c1 == NULL) ||
304 is_black(c1))) {
305 if (n == null_node) {
306 set_child(parent, n_side, NULL);
307 }
308
309 set_color(sib, RED);
310 if (is_black(parent)) {
311 /* Balance the sibling's subtree by
312 * coloring it red, then our parent
313 * has a missing black so iterate
314 * upward
315 */
316 stacksz--;
317 continue;
318 } else {
319 /* Recoloring makes the whole tree OK */
320 set_color(parent, BLACK);
321 return;
322 }
323 }
324
325 CHECK((c0 && is_red(c0)) || (c1 && is_red(c1)));
326
327 /* We know sibling has at least one red child. Fix it
328 * so that the far/outer position (i.e. on the
329 * opposite side from N) is definitely red.
330 */
331 outer = get_child(sib, (n_side == 0U) ? 1U : 0U);
332 if (!((outer != NULL) && is_red(outer))) {
333 inner = get_child(sib, n_side);
334
335 stack[stacksz - 1] = sib;
336 stack[stacksz++] = inner;
337 rotate(stack, stacksz);
338 set_color(sib, RED);
339 set_color(inner, BLACK);
340
341 /* Restore stack state to have N on the top
342 * and make sib reflect the new sibling
343 */
344 sib = stack[stacksz - 2];
345 outer = get_child(sib, (n_side == 0U) ? 1U : 0U);
346 stack[stacksz - 2] = n;
347 stacksz--;
348 }
349
350 /* Finally, the sibling must have a red child in the
351 * far/outer slot. We can rotate sib with our parent
352 * and recolor to produce a valid tree.
353 */
354 CHECK(is_red(outer));
355 set_color(sib, get_color(parent));
356 set_color(parent, BLACK);
357 set_color(outer, BLACK);
358 stack[stacksz - 1] = sib;
359 rotate(stack, stacksz);
360 if (n == null_node) {
361 set_child(parent, n_side, NULL);
362 }
363 return;
364 }
365 }
366
rb_remove(struct rbtree * tree,struct rbnode * node)367 void rb_remove(struct rbtree *tree, struct rbnode *node)
368 {
369 struct rbnode *tmp;
370 #ifdef CONFIG_MISRA_SANE
371 struct rbnode **stack = &tree->iter_stack[0];
372 #else
373 struct rbnode *stack[tree->max_depth + 1];
374 #endif
375
376 int stacksz = find_and_stack(tree, node, stack);
377
378 if (node != stack[stacksz - 1]) {
379 return;
380 }
381
382 /* We can only remove a node with zero or one child, if we
383 * have two then pick the "biggest" child of side 0 (smallest
384 * of 1 would work too) and swap our spot in the tree with
385 * that one
386 */
387 if ((get_child(node, 0U) != NULL) && (get_child(node, 1U) != NULL)) {
388 int stacksz0 = stacksz;
389 struct rbnode *hiparent, *loparent;
390 struct rbnode *node2 = get_child(node, 0U);
391
392 hiparent = (stacksz > 1) ? stack[stacksz - 2] : NULL;
393 stack[stacksz++] = node2;
394 while (get_child(node2, 1U) != NULL) {
395 node2 = get_child(node2, 1U);
396 stack[stacksz++] = node2;
397 }
398
399 loparent = stack[stacksz - 2];
400
401 /* Now swap the position of node/node2 in the tree.
402 * Design note: this is a spot where being an
403 * intrusive data structure hurts us fairly badly.
404 * The trees you see in textbooks do this by swapping
405 * the "data" pointers between the two nodes, but we
406 * have a few special cases to check. In principle
407 * this works by swapping the child pointers between
408 * the nodes and retargeting the nodes pointing to
409 * them from their parents, but: (1) the upper node
410 * may be the root of the tree and not have a parent,
411 * and (2) the lower node may be a direct child of the
412 * upper node. Remember to swap the color bits of the
413 * two nodes also. And of course we don't have parent
414 * pointers, so the stack tracking this structure
415 * needs to be swapped too!
416 */
417 if (hiparent != NULL) {
418 set_child(hiparent, get_side(hiparent, node), node2);
419 } else {
420 tree->root = node2;
421 }
422
423 if (loparent == node) {
424 set_child(node, 0U, get_child(node2, 0U));
425 set_child(node2, 0U, node);
426 } else {
427 set_child(loparent, get_side(loparent, node2), node);
428 tmp = get_child(node, 0U);
429 set_child(node, 0U, get_child(node2, 0U));
430 set_child(node2, 0U, tmp);
431 }
432
433 set_child(node2, 1U, get_child(node, 1U));
434 set_child(node, 1U, NULL);
435
436 tmp = stack[stacksz0 - 1];
437 stack[stacksz0 - 1] = stack[stacksz - 1];
438 stack[stacksz - 1] = tmp;
439
440 enum rb_color ctmp = get_color(node);
441
442 set_color(node, get_color(node2));
443 set_color(node2, ctmp);
444 }
445
446 CHECK((get_child(node, 0U) == NULL) ||
447 (get_child(node, 1U) == NULL));
448
449 struct rbnode *child = get_child(node, 0U);
450
451 if (child == NULL) {
452 child = get_child(node, 1U);
453 }
454
455 /* Removing the root */
456 if (stacksz < 2) {
457 tree->root = child;
458 if (child != NULL) {
459 set_color(child, BLACK);
460 } else {
461 tree->max_depth = 0;
462 }
463 return;
464 }
465
466 struct rbnode *parent = stack[stacksz - 2];
467
468 /* Special case: if the node to be removed is childless, then
469 * we leave it in place while we do the missing black
470 * rotations, which will replace it with a proper NULL when
471 * they isolate it.
472 */
473 if (child == NULL) {
474 if (is_black(node)) {
475 fix_missing_black(stack, stacksz, node);
476 } else {
477 /* Red childless nodes can just be dropped */
478 set_child(parent, get_side(parent, node), NULL);
479 }
480 } else {
481 set_child(parent, get_side(parent, node), child);
482
483 /* Check colors, if one was red (at least one must have been
484 * black in a valid tree), then we're done.
485 */
486 __ASSERT(is_black(node) || is_black(child), "both nodes red?!");
487 if (is_red(node) || is_red(child)) {
488 set_color(child, BLACK);
489 }
490 }
491
492 /* We may have rotated up into the root! */
493 tree->root = stack[0];
494 }
495
496 #ifndef CONFIG_MISRA_SANE
z_rb_walk(struct rbnode * node,rb_visit_t visit_fn,void * cookie)497 void z_rb_walk(struct rbnode *node, rb_visit_t visit_fn, void *cookie)
498 {
499 if (node != NULL) {
500 z_rb_walk(get_child(node, 0U), visit_fn, cookie);
501 visit_fn(node, cookie);
502 z_rb_walk(get_child(node, 1U), visit_fn, cookie);
503 }
504 }
505 #endif
506
z_rb_child(struct rbnode * node,uint8_t side)507 struct rbnode *z_rb_child(struct rbnode *node, uint8_t side)
508 {
509 return get_child(node, side);
510 }
511
z_rb_is_black(struct rbnode * node)512 int z_rb_is_black(struct rbnode *node)
513 {
514 return is_black(node);
515 }
516
rb_contains(struct rbtree * tree,struct rbnode * node)517 bool rb_contains(struct rbtree *tree, struct rbnode *node)
518 {
519 struct rbnode *n = tree->root;
520
521 while ((n != NULL) && (n != node)) {
522 n = get_child(n, tree->lessthan_fn(n, node));
523 }
524
525 return n == node;
526 }
527
528 /* Pushes the node and its chain of left-side children onto the stack
529 * in the foreach struct, returning the last node, which is the next
530 * node to iterate. By construction node will always be a right child
531 * or the root, so is_left must be false.
532 */
stack_left_limb(struct rbnode * n,struct _rb_foreach * f)533 static inline struct rbnode *stack_left_limb(struct rbnode *n,
534 struct _rb_foreach *f)
535 {
536 f->top++;
537 f->stack[f->top] = n;
538 f->is_left[f->top] = 0U;
539
540 while ((n = get_child(n, 0U)) != NULL) {
541 f->top++;
542 f->stack[f->top] = n;
543 f->is_left[f->top] = 1;
544 }
545
546 return f->stack[f->top];
547 }
548
549 /* The foreach tracking works via a dynamic stack allocated via
550 * alloca(). The current node is found in stack[top] (and its parent
551 * is thus stack[top-1]). The side of each stacked node from its
552 * parent is stored in is_left[] (i.e. if is_left[top] is true, then
553 * node/stack[top] is the left child of stack[top-1]). The special
554 * case of top == -1 indicates that the stack is uninitialized and we
555 * need to push an initial stack starting at the root.
556 */
z_rb_foreach_next(struct rbtree * tree,struct _rb_foreach * f)557 struct rbnode *z_rb_foreach_next(struct rbtree *tree, struct _rb_foreach *f)
558 {
559 struct rbnode *n;
560
561 if (tree->root == NULL) {
562 return NULL;
563 }
564
565 /* Initialization condition, pick the leftmost child of the
566 * root as our first node, initializing the stack on the way.
567 */
568 if (f->top == -1) {
569 return stack_left_limb(tree->root, f);
570 }
571
572 /* The next child from a given node is the leftmost child of
573 * it's right subtree if it has a right child
574 */
575 n = get_child(f->stack[f->top], 1U);
576 if (n != NULL) {
577 return stack_left_limb(n, f);
578 }
579
580 /* Otherwise if the node is a left child of its parent, the
581 * next node is the parent (note that the root is stacked
582 * above with is_left set to 0, so this condition still works
583 * even if node has no parent).
584 */
585 if (f->is_left[f->top] != 0U) {
586 return f->stack[--f->top];
587 }
588
589 /* If we had no left tree and are a right child then our
590 * parent was already walked, so walk up the stack looking for
591 * a left child (whose parent is unwalked, and thus next).
592 */
593 while ((f->top > 0) && (f->is_left[f->top] == 0U)) {
594 f->top--;
595 }
596
597 f->top--;
598 return (f->top >= 0) ? f->stack[f->top] : NULL;
599 }
600
