create a_regress_simple_evar
[liba.git] / src / rbt.c
blobd3496280657219539a12ab8ad71a0d7ea57e2859
1 #include "a/rbt.h"
3 /*
4 red-black trees properties: https://en.wikipedia.org/wiki/Rbtree
6 1) A node is either red or black
7 2) The root is black
8 3) All leaves (null) are black
9 4) Both children of every red node are black
10 5) Every simple path from root to leaves contains the same number of black nodes.
12 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
13 consecutive red nodes in a path and every red node is therefore followed by
14 a black. So if B is the number of black nodes on every simple path (as per
15 5), then the longest possible path due to 4 is 2B.
17 We shall indicate color with case, where black nodes are uppercase and red
18 nodes will be lowercase. Unknown color nodes shall be drawn as red within
19 parentheses and have some accompanying text comment.
22 /* Replaces the child of the specified red–black tree node. */
23 static A_INLINE void a_rbt_new_child(a_rbt *root, a_rbt_node *parent, a_rbt_node *node, a_rbt_node *newnode)
25 if (parent)
27 if (parent->left == node)
29 parent->left = newnode;
31 else
33 parent->right = newnode;
36 else
38 root->node = newnode;
42 /* Sets the parent and color of the specified red–black tree node. */
43 static A_INLINE void a_rbt_set_parent_color(a_rbt_node *node, a_rbt_node *parent, unsigned int color)
45 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
46 node->parent_ = (a_uptr)parent + color;
47 #else /* !A_SIZE_POINTER */
48 node->parent = parent;
49 node->color = color;
50 #endif /* A_SIZE_POINTER */
53 /* Sets the parent of the specified red–black tree node. */
54 static A_INLINE void a_rbt_set_parent(a_rbt_node *node, a_rbt_node *parent)
56 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
57 node->parent_ = (a_uptr)parent + (node->parent_ & 1);
58 #else /* !A_SIZE_POINTER */
59 node->parent = parent;
60 #endif /* A_SIZE_POINTER */
63 /* Returns the color of the specified red–black tree node. */
64 static A_INLINE unsigned int a_rbt_color(a_rbt_node const *node)
66 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
67 return (unsigned int)(node->parent_ & 1);
68 #else /* !A_SIZE_POINTER */
69 return node->color;
70 #endif /* A_SIZE_POINTER */
73 /* Sets the black color of the specified red–black tree node. */
74 static A_INLINE void a_rbt_set_black(a_rbt_node *node)
76 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
77 node->parent_ |= 1;
78 #else /* !A_SIZE_POINTER */
79 node->color = 1;
80 #endif /* A_SIZE_POINTER */
84 Helper function for rotations:
85 - node's parent and color get assigned to new
86 - node gets assigned new as a parent and 'color' as a color.
88 static A_INLINE void a_rbt_set_parents(a_rbt *root, a_rbt_node *node, a_rbt_node *newnode, unsigned int color)
90 a_rbt_node *const parent = a_rbt_parent(node);
91 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
92 newnode->parent_ = node->parent_;
93 #else /* !A_SIZE_POINTER */
94 newnode->parent = node->parent;
95 newnode->color = node->color;
96 #endif /* A_SIZE_POINTER */
97 a_rbt_set_parent_color(node, newnode, color);
98 a_rbt_new_child(root, parent, node, newnode);
101 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
102 #define A_RBT_PARENT(node) a_cast_r(a_rbt_node *, (node)->parent_)
103 #else /* !A_SIZE_POINTER */
104 #define A_RBT_PARENT(node) (node)->parent
105 #endif /* A_SIZE_POINTER */
107 void a_rbt_insert_adjust(a_rbt *root, a_rbt_node *node)
109 for (a_rbt_node *parent = A_RBT_PARENT(node), *gparent, *tmp;;)
111 /* Loop invariant: node is red. */
112 if (A_UNLIKELY(!parent))
115 The inserted node is root. Either this is the first node,
116 or we recursed at Case 1 below and are no longer violating 4).
118 a_rbt_set_parent_color(node, A_NULL, 1);
119 break;
122 If there is a black parent, we are done. Otherwise, take some corrective action as, per 4),
123 we don't want a red root or two consecutive red nodes.
125 if (a_rbt_color(parent)) { break; }
126 gparent = A_RBT_PARENT(parent);
127 tmp = gparent->right;
128 if (parent != tmp) /* parent == gparent->left */
130 if (tmp && a_rbt_color(tmp) == 0)
133 Case 1 - node's uncle is red (color flips).
136 / \ / \
137 p u --> P U
141 However, since g's parent might be red, and 4) does not allow this, we need to recurse at g.
143 a_rbt_set_parent_color(tmp, gparent, 1);
144 a_rbt_set_parent_color(parent, gparent, 1);
145 node = gparent;
146 parent = a_rbt_parent(node);
147 a_rbt_set_parent_color(node, parent, 0);
148 continue;
150 tmp = parent->right;
151 if (node == tmp)
154 Case 2 - node's uncle is black and node is the parent's right child (left rotate at parent).
157 / \ / \
158 p U --> n U
162 This still leaves us in violation of 4), the continuation into Case 3 will fix that.
164 tmp = node->left;
165 parent->right = tmp;
166 node->left = parent;
167 if (tmp) { a_rbt_set_parent_color(tmp, parent, 1); }
168 a_rbt_set_parent_color(parent, node, 0);
169 parent = node;
170 tmp = node->right;
173 Case 3 - node's uncle is black and node is
174 the parent's left child (right rotate at gparent).
177 / \ / \
178 p U --> n g
182 gparent->left = tmp; /* == parent->right */
183 parent->right = gparent;
184 if (tmp) { a_rbt_set_parent_color(tmp, gparent, 1); }
185 a_rbt_set_parents(root, gparent, parent, 0);
186 break;
188 else
190 tmp = gparent->left;
191 if (tmp && a_rbt_color(tmp) == 0)
193 /* Case 1 - color flips */
194 a_rbt_set_parent_color(tmp, gparent, 1);
195 a_rbt_set_parent_color(parent, gparent, 1);
196 node = gparent;
197 parent = a_rbt_parent(node);
198 a_rbt_set_parent_color(node, parent, 0);
199 continue;
201 tmp = parent->left;
202 if (node == tmp)
204 /* Case 2 - right rotate at parent */
205 tmp = node->right;
206 parent->left = tmp;
207 node->right = parent;
208 if (tmp) { a_rbt_set_parent_color(tmp, parent, 1); }
209 a_rbt_set_parent_color(parent, node, 0);
210 parent = node;
211 tmp = node->left;
213 /* Case 3 - left rotate at gparent */
214 gparent->right = tmp; /* == parent->left */
215 parent->left = gparent;
216 if (tmp) { a_rbt_set_parent_color(tmp, gparent, 1); }
217 a_rbt_set_parents(root, gparent, parent, 0);
218 break;
223 static A_INLINE void a_rbt_remove_adjust(a_rbt *root, a_rbt_node *parent)
225 for (a_rbt_node *node = A_NULL, *sibling, *tmp1, *tmp2;;)
228 Loop invariants:
229 - node is black (or null on first iteration)
230 - node is not the root (parent is not null)
231 - All leaf paths going through parent and node have a
232 black node count that is 1 lower than other leaf paths.
234 if (node != parent->right) /* node == parent->left */
236 sibling = parent->right;
237 if (a_rbt_color(sibling) == 0)
240 Case 1 - left rotate at parent
243 / \ / \
244 N s --> p Sr
245 / \ / \
246 Sl Sr N Sl
248 tmp1 = sibling->left;
249 parent->right = tmp1;
250 sibling->left = parent;
251 a_rbt_set_parent_color(tmp1, parent, 1);
252 a_rbt_set_parents(root, parent, sibling, 0);
253 sibling = tmp1;
254 A_ASSUME(tmp1);
256 tmp1 = sibling->right;
257 if (!tmp1 || a_rbt_color(tmp1))
259 tmp2 = sibling->left;
260 if (!tmp2 || a_rbt_color(tmp2))
263 Case 2 - sibling color flip (p could be either color here)
265 (p) (p)
266 / \ / \
267 N S --> N s
268 / \ / \
269 Sl Sr Sl Sr
271 This leaves us violating 5) which can be fixed by flipping p to black if it was red,
272 or by recursing at p. p is red when coming from Case 1.
274 a_rbt_set_parent_color(sibling, parent, 0);
275 if (a_rbt_color(parent) == 0)
277 a_rbt_set_black(parent);
279 else
281 node = parent;
282 parent = a_rbt_parent(node);
283 if (parent) { continue; }
285 break;
288 Case 3 - right rotate at sibling (p could be either color here)
290 (p) (p)
291 / \ / \
292 N S --> N sl
293 / \ \
294 sl Sr S
298 Note: p might be red, and then both p and sl are red after rotation(which breaks property 4).
299 This is fixed in Case 4 (in a_rbt_set_parents() which set sl the color of p and set black)
301 (p) (sl)
302 / \ / \
303 N sl --> P S
304 \ / \
305 S N Sr
309 tmp1 = tmp2->right;
310 sibling->left = tmp1;
311 tmp2->right = sibling;
312 parent->right = tmp2;
313 if (tmp1) { a_rbt_set_parent_color(tmp1, sibling, 1); }
314 tmp1 = sibling;
315 sibling = tmp2;
318 Case 4 - left rotate at parent + color flips (p and sl could be either color here.
319 After rotation, p becomes black, s acquires p's color, and sl keeps its color)
321 (p) (s)
322 / \ / \
323 N S --> P Sr
324 / \ / \
325 (sl) sr N (sl)
327 tmp2 = sibling->left;
328 parent->right = tmp2;
329 sibling->left = parent;
330 a_rbt_set_parent_color(tmp1, sibling, 1);
331 if (tmp2) { a_rbt_set_parent(tmp2, parent); }
332 a_rbt_set_parents(root, parent, sibling, 1);
333 break;
335 else
337 A_ASSUME(parent->left);
338 sibling = parent->left;
339 if (a_rbt_color(sibling) == 0)
341 /* Case 1 - right rotate at parent */
342 tmp1 = sibling->right;
343 parent->left = tmp1;
344 sibling->right = parent;
345 a_rbt_set_parent_color(tmp1, parent, 1);
346 a_rbt_set_parents(root, parent, sibling, 0);
347 sibling = tmp1;
348 A_ASSUME(tmp1);
350 tmp1 = sibling->left;
351 if (!tmp1 || a_rbt_color(tmp1))
353 tmp2 = sibling->right;
354 if (!tmp2 || a_rbt_color(tmp2))
356 /* Case 2 - sibling color flip */
357 a_rbt_set_parent_color(sibling, parent, 0);
358 if (a_rbt_color(parent) == 0)
360 a_rbt_set_black(parent);
362 else
364 node = parent;
365 parent = a_rbt_parent(node);
366 if (parent) { continue; }
368 break;
370 /* Case 3 - left rotate at sibling */
371 tmp1 = tmp2->left;
372 sibling->right = tmp1;
373 tmp2->left = sibling;
374 parent->left = tmp2;
375 if (tmp1) { a_rbt_set_parent_color(tmp1, sibling, 1); }
376 tmp1 = sibling;
377 sibling = tmp2;
379 /* Case 4 - right rotate at parent + color flips */
380 tmp2 = sibling->right;
381 parent->left = tmp2;
382 sibling->right = parent;
383 a_rbt_set_parent_color(tmp1, sibling, 1);
384 if (tmp2) { a_rbt_set_parent(tmp2, parent); }
385 a_rbt_set_parents(root, parent, sibling, 1);
386 break;
391 void a_rbt_remove(a_rbt *root, a_rbt_node *node)
393 a_rbt_node *child = node->right;
394 a_rbt_node *tmp = node->left;
395 a_rbt_node *parent, *adjust;
396 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
397 a_uptr parent_;
398 #else /* !A_SIZE_POINTER */
399 unsigned int color;
400 #endif /* A_SIZE_POINTER */
402 if (!tmp)
405 Case 1: node to erase has no more than 1 child (easy!)
406 Note that if there is one child it must be red due to 5) and node must be black due to 4).
407 We adjust colors locally so as to bypass a_rbt_remove_adjust() later on.
409 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
410 parent_ = node->parent_;
411 #else /* !A_SIZE_POINTER */
412 color = node->color;
413 #endif /* A_SIZE_POINTER */
414 parent = a_rbt_parent(node);
415 a_rbt_new_child(root, parent, node, child);
416 if (child)
418 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
419 child->parent_ = parent_;
420 #else /* !A_SIZE_POINTER */
421 child->parent = parent;
422 child->color = color;
423 #endif /* A_SIZE_POINTER */
424 adjust = A_NULL;
426 else
428 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
429 adjust = (parent_ & 1) ? parent : A_NULL;
430 #else /* !A_SIZE_POINTER */
431 adjust = color ? parent : A_NULL;
432 #endif /* A_SIZE_POINTER */
435 else if (!child)
437 /* Still case 1, but this time the child is node->left */
438 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
439 tmp->parent_ = node->parent_;
440 #else /* !A_SIZE_POINTER */
441 tmp->parent = node->parent;
442 tmp->color = node->color;
443 #endif /* A_SIZE_POINTER */
444 parent = a_rbt_parent(node);
445 a_rbt_new_child(root, parent, node, tmp);
446 adjust = A_NULL;
448 else
450 a_rbt_node *successor = child, *child2;
452 tmp = child->left;
453 if (!tmp)
456 Case 2: node's successor is its right child
458 (n) (s)
459 / \ / \
460 (x) (s) -> (x) (c)
464 parent = successor;
465 child2 = successor->right;
467 else
470 Case 3: node's successor is leftmost under node's right child subtree
472 (n) (s)
473 / \ / \
474 (x) (y) -> (x) (y)
476 (p) (p)
478 (s) (c)
482 do {
483 parent = successor;
484 successor = tmp;
485 tmp = tmp->left;
486 } while (tmp);
487 child2 = successor->right;
488 parent->left = child2;
489 successor->right = child;
490 a_rbt_set_parent(child, successor);
493 tmp = node->left;
494 successor->left = tmp;
495 a_rbt_set_parent(tmp, successor);
497 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
498 parent_ = node->parent_;
499 #else /* !A_SIZE_POINTER */
500 color = node->color;
501 #endif /* A_SIZE_POINTER */
502 tmp = a_rbt_parent(node);
503 a_rbt_new_child(root, tmp, node, successor);
505 if (child2)
507 a_rbt_set_parent_color(child2, parent, 1);
508 adjust = A_NULL;
510 else
512 adjust = a_rbt_color(successor) ? parent : A_NULL;
514 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
515 successor->parent_ = parent_;
516 #else /* !A_SIZE_POINTER */
517 successor->parent = tmp;
518 successor->color = color;
519 #endif /* A_SIZE_POINTER */
522 if (adjust) { a_rbt_remove_adjust(root, adjust); }
525 a_rbt_node *a_rbt_insert(a_rbt *root, a_rbt_node *node, int (*cmp)(void const *, void const *))
527 a_rbt_node *parent = root->node;
528 a_rbt_node **link = &root->node;
529 while (*link)
531 parent = *link;
532 int const res = cmp(node, parent);
533 if (res < 0)
535 link = &parent->left;
537 else if (res > 0)
539 link = &parent->right;
541 else { return parent; }
543 *link = a_rbt_init(node, parent);
544 a_rbt_insert_adjust(root, node);
545 return A_NULL;
548 a_rbt_node *a_rbt_search(a_rbt const *root, void const *ctx, int (*cmp)(void const *, void const *))
550 for (a_rbt_node *cur = root->node; cur;)
552 int const res = cmp(ctx, cur);
553 if (res < 0)
555 cur = cur->left;
557 else if (res > 0)
559 cur = cur->right;
561 else { return cur; }
563 return A_NULL;
566 a_rbt_node *a_rbt_head(a_rbt const *root)
568 a_rbt_node *node = root->node;
569 if (node)
571 while (node->left) { node = node->left; }
573 return node;
576 a_rbt_node *a_rbt_tail(a_rbt const *root)
578 a_rbt_node *node = root->node;
579 if (node)
581 while (node->right) { node = node->right; }
583 return node;
586 a_rbt_node *a_rbt_next(a_rbt_node *node)
592 / \ / \
593 A C E G
595 if (node->right) /* D -> F -> E */
597 node = node->right;
598 while (node->left) { node = node->left; }
600 else /* C -> B -> D */
602 a_rbt_node *leaf;
603 do {
604 leaf = node;
605 node = a_rbt_parent(node);
606 } while (node && node->left != leaf);
608 return node;
611 a_rbt_node *a_rbt_prev(a_rbt_node *node)
613 if (node->left)
615 node = node->left;
616 while (node->right) { node = node->right; }
618 else
620 a_rbt_node *leaf;
621 do {
622 leaf = node;
623 node = a_rbt_parent(node);
624 } while (node && node->right != leaf);
626 return node;
629 a_rbt_node *a_rbt_pre_next(a_rbt_node *node)
635 / \ / \
636 A C E G
638 a_rbt_node *leaf;
639 if (node->left) { return node->left; }
640 if (node->right) { return node->right; }
641 for (leaf = node, node = a_rbt_parent(node); node;
642 leaf = node, node = a_rbt_parent(node))
644 if (node->right && node->right != leaf)
646 node = node->right;
647 break; /* A -> B -> C */
648 } /* C -> B -> D -> F */
650 return node;
653 a_rbt_node *a_rbt_pre_prev(a_rbt_node *node)
655 a_rbt_node *leaf;
656 if (node->right) { return node->right; }
657 if (node->left) { return node->left; }
658 for (leaf = node, node = a_rbt_parent(node); node;
659 leaf = node, node = a_rbt_parent(node))
661 if (node->left && node->left != leaf)
663 node = node->left;
664 break;
667 return node;
670 #define A_RBT_POST(head, tail) \
671 for (;;) \
673 if (node->head) \
675 node = node->head; \
677 else if (node->tail) \
679 node = node->tail; \
681 else { break; } \
683 (void)0
685 a_rbt_node *a_rbt_post_head(a_rbt const *root)
687 a_rbt_node *node = root->node;
688 if (node) { A_RBT_POST(left, right); }
689 return node;
692 a_rbt_node *a_rbt_post_tail(a_rbt const *root)
694 a_rbt_node *node = root->node;
695 if (node) { A_RBT_POST(right, left); }
696 return node;
699 a_rbt_node *a_rbt_post_next(a_rbt_node *node)
705 / \ / \
706 A C E G
708 a_rbt_node *leaf = node;
709 node = a_rbt_parent(node);
710 if (node && node->right && node->right != leaf)
712 node = node->right; /* B -> D -> F -> E */
713 A_RBT_POST(left, right); /* A -> B -> C */
714 } /* C -> B */
715 return node;
718 a_rbt_node *a_rbt_post_prev(a_rbt_node *node)
720 a_rbt_node *leaf = node;
721 node = a_rbt_parent(node);
722 if (node && node->left && node->left != leaf)
724 node = node->left;
725 A_RBT_POST(right, left);
727 return node;
730 a_rbt_node *a_rbt_tear(a_rbt *root, a_rbt_node **next)
732 a_rbt_node *node = *next;
733 if (!node)
735 node = root->node;
736 if (!node) { return node; }
738 A_RBT_POST(left, right);
739 *next = a_rbt_parent(node);
740 a_rbt_new_child(root, *next, node, A_NULL);
741 return node;