1 /*
2 Copyright 2011-2019 David Robillard <d@drobilla.net>
3
4 Permission to use, copy, modify, and/or distribute this software for any
5 purpose with or without fee is hereby granted, provided that the above
6 copyright notice and this permission notice appear in all copies.
7
8 THIS SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9 WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10 MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11 ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12 WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13 ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14 OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15 */
16
17 #include "zix/tree.h"
18 #include "zix/common.h"
19
20 #include <assert.h>
21 #include <stdint.h>
22 #include <stdlib.h>
23 #include <string.h>
24
25 typedef struct ZixTreeNodeImpl ZixTreeNode;
26
27 struct ZixTreeImpl {
28 ZixTreeNode* root;
29 ZixDestroyFunc destroy;
30 ZixComparator cmp;
31 void* cmp_data;
32 size_t size;
33 bool allow_duplicates;
34 };
35
36 struct ZixTreeNodeImpl {
37 void* data;
38 struct ZixTreeNodeImpl* left;
39 struct ZixTreeNodeImpl* right;
40 struct ZixTreeNodeImpl* parent;
41 int_fast8_t balance;
42 };
43
44 #define MIN(a, b) (((a) < (b)) ? (a) : (b))
45 #define MAX(a, b) (((a) > (b)) ? (a) : (b))
46
47 // Uncomment these for debugging features
48 // #define ZIX_TREE_DUMP 1
49 // #define ZIX_TREE_VERIFY 1
50 // #define ZIX_TREE_HYPER_VERIFY 1
51
52 #if defined(ZIX_TREE_VERIFY) || defined(ZIX_TREE_HYPER_VERIFY)
53 # include "tree_debug.h"
54 # define ASSERT_BALANCE(n) assert(verify_balance(n))
55 #else
56 # define ASSERT_BALANCE(n)
57 #endif
58
59 #ifdef ZIX_TREE_DUMP
60 # include "tree_debug.h"
61 # define DUMP(t) zix_tree_print(t->root, 0)
62 # define DEBUG_PRINTF(fmt, ...) printf(fmt, __VA_ARGS__)
63 #else
64 # define DUMP(t)
65 # define DEBUG_PRINTF(fmt, ...)
66 #endif
67
68 ZIX_API ZixTree*
zix_tree_new(bool allow_duplicates,ZixComparator cmp,void * cmp_data,ZixDestroyFunc destroy)69 zix_tree_new(bool allow_duplicates,
70 ZixComparator cmp,
71 void* cmp_data,
72 ZixDestroyFunc destroy)
73 {
74 ZixTree* t = (ZixTree*)malloc(sizeof(ZixTree));
75 t->root = NULL;
76 t->destroy = destroy;
77 t->cmp = cmp;
78 t->cmp_data = cmp_data;
79 t->size = 0;
80 t->allow_duplicates = allow_duplicates;
81 return t;
82 }
83
84 ZIX_PRIVATE void
zix_tree_free_rec(ZixTree * t,ZixTreeNode * n)85 zix_tree_free_rec(ZixTree* t, ZixTreeNode* n)
86 {
87 if (n) {
88 zix_tree_free_rec(t, n->left);
89 zix_tree_free_rec(t, n->right);
90 if (t->destroy) {
91 t->destroy(n->data);
92 }
93 free(n);
94 }
95 }
96
97 ZIX_API void
zix_tree_free(ZixTree * t)98 zix_tree_free(ZixTree* t)
99 {
100 if (t) {
101 zix_tree_free_rec(t, t->root);
102 free(t);
103 }
104 }
105
106 ZIX_API size_t
zix_tree_size(const ZixTree * t)107 zix_tree_size(const ZixTree* t)
108 {
109 return t->size;
110 }
111
112 ZIX_PRIVATE void
rotate(ZixTreeNode * p,ZixTreeNode * q)113 rotate(ZixTreeNode* p, ZixTreeNode* q)
114 {
115 assert(q->parent == p);
116 assert(p->left == q || p->right == q);
117
118 q->parent = p->parent;
119 if (q->parent) {
120 if (q->parent->left == p) {
121 q->parent->left = q;
122 } else {
123 q->parent->right = q;
124 }
125 }
126
127 if (p->right == q) {
128 // Rotate left
129 p->right = q->left;
130 q->left = p;
131 if (p->right) {
132 p->right->parent = p;
133 }
134 } else {
135 // Rotate right
136 assert(p->left == q);
137 p->left = q->right;
138 q->right = p;
139 if (p->left) {
140 p->left->parent = p;
141 }
142 }
143
144 p->parent = q;
145 }
146
147 /**
148 * Rotate left about `p`.
149 *
150 * p q
151 * / \ / \
152 * A q => p C
153 * / \ / \
154 * B C A B
155 */
156 ZIX_PRIVATE ZixTreeNode*
rotate_left(ZixTreeNode * p,int * height_change)157 rotate_left(ZixTreeNode* p, int* height_change)
158 {
159 ZixTreeNode* const q = p->right;
160 *height_change = (q->balance == 0) ? 0 : -1;
161
162 DEBUG_PRINTF("LL %ld\n", (intptr_t)p->data);
163
164 assert(p->balance == 2);
165 assert(q->balance == 0 || q->balance == 1);
166
167 rotate(p, q);
168
169 // p->balance -= 1 + MAX(0, q->balance);
170 // q->balance -= 1 - MIN(0, p->balance);
171 --q->balance;
172 p->balance = -(q->balance);
173
174 ASSERT_BALANCE(p);
175 ASSERT_BALANCE(q);
176 return q;
177 }
178
179 /**
180 * Rotate right about `p`.
181 *
182 * p q
183 * / \ / \
184 * q C => A p
185 * / \ / \
186 * A B B C
187 *
188 */
189 ZIX_PRIVATE ZixTreeNode*
rotate_right(ZixTreeNode * p,int * height_change)190 rotate_right(ZixTreeNode* p, int* height_change)
191 {
192 ZixTreeNode* const q = p->left;
193 *height_change = (q->balance == 0) ? 0 : -1;
194
195 DEBUG_PRINTF("RR %ld\n", (intptr_t)p->data);
196
197 assert(p->balance == -2);
198 assert(q->balance == 0 || q->balance == -1);
199
200 rotate(p, q);
201
202 // p->balance += 1 - MIN(0, q->balance);
203 // q->balance += 1 + MAX(0, p->balance);
204 ++q->balance;
205 p->balance = -(q->balance);
206
207 ASSERT_BALANCE(p);
208 ASSERT_BALANCE(q);
209 return q;
210 }
211
212 /**
213 * Rotate left about `p->left` then right about `p`.
214 *
215 * p r
216 * / \ / \
217 * q D => q p
218 * / \ / \ / \
219 * A r A B C D
220 * / \
221 * B C
222 *
223 */
224 ZIX_PRIVATE ZixTreeNode*
rotate_left_right(ZixTreeNode * p,int * height_change)225 rotate_left_right(ZixTreeNode* p, int* height_change)
226 {
227 ZixTreeNode* const q = p->left;
228 ZixTreeNode* const r = q->right;
229
230 assert(p->balance == -2);
231 assert(q->balance == 1);
232 assert(r->balance == -1 || r->balance == 0 || r->balance == 1);
233
234 DEBUG_PRINTF("LR %ld P: %2d Q: %2d R: %2d\n",
235 (intptr_t)p->data,
236 p->balance,
237 q->balance,
238 r->balance);
239
240 rotate(q, r);
241 rotate(p, r);
242
243 q->balance -= 1 + MAX(0, r->balance);
244 p->balance += 1 - MIN(MIN(0, r->balance) - 1, r->balance + q->balance);
245 // r->balance += MAX(0, p->balance) + MIN(0, q->balance);
246
247 // p->balance = (p->left && p->right) ? -MIN(r->balance, 0) : 0;
248 // q->balance = - MAX(r->balance, 0);
249 r->balance = 0;
250
251 *height_change = -1;
252
253 ASSERT_BALANCE(p);
254 ASSERT_BALANCE(q);
255 ASSERT_BALANCE(r);
256 return r;
257 }
258
259 /**
260 * Rotate right about `p->right` then right about `p`.
261 *
262 * p r
263 * / \ / \
264 * A q => p q
265 * / \ / \ / \
266 * r D A B C D
267 * / \
268 * B C
269 *
270 */
271 ZIX_PRIVATE ZixTreeNode*
rotate_right_left(ZixTreeNode * p,int * height_change)272 rotate_right_left(ZixTreeNode* p, int* height_change)
273 {
274 ZixTreeNode* const q = p->right;
275 ZixTreeNode* const r = q->left;
276
277 assert(p->balance == 2);
278 assert(q->balance == -1);
279 assert(r->balance == -1 || r->balance == 0 || r->balance == 1);
280
281 DEBUG_PRINTF("RL %ld P: %2d Q: %2d R: %2d\n",
282 (intptr_t)p->data,
283 p->balance,
284 q->balance,
285 r->balance);
286
287 rotate(q, r);
288 rotate(p, r);
289
290 q->balance += 1 - MIN(0, r->balance);
291 p->balance -= 1 + MAX(MAX(0, r->balance) + 1, r->balance + q->balance);
292 // r->balance += MAX(0, q->balance) + MIN(0, p->balance);
293
294 // p->balance = (p->left && p->right) ? -MAX(r->balance, 0) : 0;
295 // q->balance = - MIN(r->balance, 0);
296 r->balance = 0;
297 // assert(r->balance == 0);
298
299 *height_change = -1;
300
301 ASSERT_BALANCE(p);
302 ASSERT_BALANCE(q);
303 ASSERT_BALANCE(r);
304 return r;
305 }
306
307 ZIX_PRIVATE ZixTreeNode*
zix_tree_rebalance(ZixTree * t,ZixTreeNode * node,int * height_change)308 zix_tree_rebalance(ZixTree* t, ZixTreeNode* node, int* height_change)
309 {
310 #ifdef ZIX_TREE_HYPER_VERIFY
311 const size_t old_height = height(node);
312 #endif
313 DEBUG_PRINTF("REBALANCE %ld (%d)\n", (intptr_t)node->data, node->balance);
314 *height_change = 0;
315 const bool is_root = !node->parent;
316 assert((is_root && t->root == node) || (!is_root && t->root != node));
317 ZixTreeNode* replacement = node;
318 if (node->balance == -2) {
319 assert(node->left);
320 if (node->left->balance == 1) {
321 replacement = rotate_left_right(node, height_change);
322 } else {
323 replacement = rotate_right(node, height_change);
324 }
325 } else if (node->balance == 2) {
326 assert(node->right);
327 if (node->right->balance == -1) {
328 replacement = rotate_right_left(node, height_change);
329 } else {
330 replacement = rotate_left(node, height_change);
331 }
332 }
333 if (is_root) {
334 assert(!replacement->parent);
335 t->root = replacement;
336 }
337 DUMP(t);
338 #ifdef ZIX_TREE_HYPER_VERIFY
339 assert(old_height + *height_change == height(replacement));
340 #endif
341 return replacement;
342 }
343
344 ZIX_API ZixStatus
zix_tree_insert(ZixTree * t,void * e,ZixTreeIter ** ti)345 zix_tree_insert(ZixTree* t, void* e, ZixTreeIter** ti)
346 {
347 DEBUG_PRINTF("**** INSERT %ld\n", (intptr_t)e);
348 int cmp = 0;
349 ZixTreeNode* n = t->root;
350 ZixTreeNode* p = NULL;
351
352 // Find the parent p of e
353 while (n) {
354 p = n;
355 cmp = t->cmp(e, n->data, t->cmp_data);
356 if (cmp < 0) {
357 n = n->left;
358 } else if (cmp > 0) {
359 n = n->right;
360 } else if (t->allow_duplicates) {
361 n = n->right;
362 } else {
363 if (ti) {
364 *ti = n;
365 }
366 DEBUG_PRINTF("%ld EXISTS!\n", (intptr_t)e);
367 return ZIX_STATUS_EXISTS;
368 }
369 }
370
371 // Allocate a new node n
372 if (!(n = (ZixTreeNode*)malloc(sizeof(ZixTreeNode)))) {
373 return ZIX_STATUS_NO_MEM;
374 }
375 memset(n, '\0', sizeof(ZixTreeNode));
376 n->data = e;
377 n->balance = 0;
378 if (ti) {
379 *ti = n;
380 }
381
382 bool p_height_increased = false;
383
384 // Make p the parent of n
385 n->parent = p;
386 if (!p) {
387 t->root = n;
388 } else {
389 if (cmp < 0) {
390 assert(!p->left);
391 assert(p->balance == 0 || p->balance == 1);
392 p->left = n;
393 --p->balance;
394 p_height_increased = !p->right;
395 } else {
396 assert(!p->right);
397 assert(p->balance == 0 || p->balance == -1);
398 p->right = n;
399 ++p->balance;
400 p_height_increased = !p->left;
401 }
402 }
403
404 DUMP(t);
405
406 // Rebalance if necessary (at most 1 rotation)
407 assert(!p || p->balance == -1 || p->balance == 0 || p->balance == 1);
408 if (p && p_height_increased) {
409 int height_change = 0;
410 for (ZixTreeNode* i = p; i && i->parent; i = i->parent) {
411 if (i == i->parent->left) {
412 if (--i->parent->balance == -2) {
413 zix_tree_rebalance(t, i->parent, &height_change);
414 break;
415 }
416 } else {
417 assert(i == i->parent->right);
418 if (++i->parent->balance == 2) {
419 zix_tree_rebalance(t, i->parent, &height_change);
420 break;
421 }
422 }
423
424 if (i->parent->balance == 0) {
425 break;
426 }
427 }
428 }
429
430 DUMP(t);
431
432 ++t->size;
433
434 #ifdef ZIX_TREE_VERIFY
435 if (!verify(t, t->root)) {
436 return ZIX_STATUS_ERROR;
437 }
438 #endif
439
440 return ZIX_STATUS_SUCCESS;
441 }
442
443 ZIX_API ZixStatus
zix_tree_remove(ZixTree * t,ZixTreeIter * ti)444 zix_tree_remove(ZixTree* t, ZixTreeIter* ti)
445 {
446 ZixTreeNode* const n = ti;
447 ZixTreeNode** pp = NULL; // parent pointer
448 ZixTreeNode* to_balance = n->parent; // lowest node to balance
449 int8_t d_balance = 0; // delta(balance) for n->parent
450
451 DEBUG_PRINTF("*** REMOVE %ld\n", (intptr_t)n->data);
452
453 if ((n == t->root) && !n->left && !n->right) {
454 t->root = NULL;
455 if (t->destroy) {
456 t->destroy(n->data);
457 }
458 free(n);
459 --t->size;
460 assert(t->size == 0);
461 return ZIX_STATUS_SUCCESS;
462 }
463
464 // Set pp to the parent pointer to n, if applicable
465 if (n->parent) {
466 assert(n->parent->left == n || n->parent->right == n);
467 if (n->parent->left == n) { // n is left child
468 pp = &n->parent->left;
469 d_balance = 1;
470 } else { // n is right child
471 assert(n->parent->right == n);
472 pp = &n->parent->right;
473 d_balance = -1;
474 }
475 }
476
477 assert(!pp || *pp == n);
478
479 int height_change = 0;
480 if (!n->left && !n->right) {
481 // n is a leaf, just remove it
482 if (pp) {
483 *pp = NULL;
484 to_balance = n->parent;
485 height_change = (!n->parent->left && !n->parent->right) ? -1 : 0;
486 }
487 } else if (!n->left) {
488 // Replace n with right (only) child
489 if (pp) {
490 *pp = n->right;
491 to_balance = n->parent;
492 } else {
493 t->root = n->right;
494 }
495 n->right->parent = n->parent;
496 height_change = -1;
497 } else if (!n->right) {
498 // Replace n with left (only) child
499 if (pp) {
500 *pp = n->left;
501 to_balance = n->parent;
502 } else {
503 t->root = n->left;
504 }
505 n->left->parent = n->parent;
506 height_change = -1;
507 } else {
508 // Replace n with in-order successor (leftmost child of right subtree)
509 ZixTreeNode* replace = n->right;
510 while (replace->left) {
511 assert(replace->left->parent == replace);
512 replace = replace->left;
513 }
514
515 // Remove replace from parent (replace_p)
516 if (replace->parent->left == replace) {
517 height_change = replace->parent->right ? 0 : -1;
518 d_balance = 1;
519 to_balance = replace->parent;
520 replace->parent->left = replace->right;
521 } else {
522 assert(replace->parent == n);
523 height_change = replace->parent->left ? 0 : -1;
524 d_balance = -1;
525 to_balance = replace->parent;
526 replace->parent->right = replace->right;
527 }
528
529 if (to_balance == n) {
530 to_balance = replace;
531 }
532
533 if (replace->right) {
534 replace->right->parent = replace->parent;
535 }
536
537 replace->balance = n->balance;
538
539 // Swap node to delete with replace
540 if (pp) {
541 *pp = replace;
542 } else {
543 assert(t->root == n);
544 t->root = replace;
545 }
546 replace->parent = n->parent;
547 replace->left = n->left;
548 n->left->parent = replace;
549 replace->right = n->right;
550 if (n->right) {
551 n->right->parent = replace;
552 }
553
554 assert(!replace->parent || replace->parent->left == replace ||
555 replace->parent->right == replace);
556 }
557
558 // Rebalance starting at to_balance upwards.
559 for (ZixTreeNode* i = to_balance; i; i = i->parent) {
560 i->balance += d_balance;
561 if (d_balance == 0 || i->balance == -1 || i->balance == 1) {
562 break;
563 }
564
565 assert(i != n);
566 i = zix_tree_rebalance(t, i, &height_change);
567 if (i->balance == 0) {
568 height_change = -1;
569 }
570
571 if (i->parent) {
572 if (i == i->parent->left) {
573 d_balance = height_change * -1;
574 } else {
575 assert(i == i->parent->right);
576 d_balance = height_change;
577 }
578 }
579 }
580
581 DUMP(t);
582
583 if (t->destroy) {
584 t->destroy(n->data);
585 }
586 free(n);
587
588 --t->size;
589
590 #ifdef ZIX_TREE_VERIFY
591 if (!verify(t, t->root)) {
592 return ZIX_STATUS_ERROR;
593 }
594 #endif
595
596 return ZIX_STATUS_SUCCESS;
597 }
598
599 ZIX_API ZixStatus
zix_tree_find(const ZixTree * t,const void * e,ZixTreeIter ** ti)600 zix_tree_find(const ZixTree* t, const void* e, ZixTreeIter** ti)
601 {
602 ZixTreeNode* n = t->root;
603 while (n) {
604 const int cmp = t->cmp(e, n->data, t->cmp_data);
605 if (cmp == 0) {
606 break;
607 }
608
609 if (cmp < 0) {
610 n = n->left;
611 } else {
612 n = n->right;
613 }
614 }
615
616 *ti = n;
617 return (n) ? ZIX_STATUS_SUCCESS : ZIX_STATUS_NOT_FOUND;
618 }
619
620 ZIX_API void*
zix_tree_get(const ZixTreeIter * ti)621 zix_tree_get(const ZixTreeIter* ti)
622 {
623 return ti ? ti->data : NULL;
624 }
625
626 ZIX_API ZixTreeIter*
zix_tree_begin(ZixTree * t)627 zix_tree_begin(ZixTree* t)
628 {
629 if (!t->root) {
630 return NULL;
631 }
632
633 ZixTreeNode* n = t->root;
634 while (n->left) {
635 n = n->left;
636 }
637 return n;
638 }
639
640 ZIX_API ZixTreeIter*
zix_tree_end(ZixTree * t)641 zix_tree_end(ZixTree* t)
642 {
643 return NULL;
644 }
645
646 ZIX_API ZixTreeIter*
zix_tree_rbegin(ZixTree * t)647 zix_tree_rbegin(ZixTree* t)
648 {
649 if (!t->root) {
650 return NULL;
651 }
652
653 ZixTreeNode* n = t->root;
654 while (n->right) {
655 n = n->right;
656 }
657 return n;
658 }
659
660 ZIX_API ZixTreeIter*
zix_tree_rend(ZixTree * t)661 zix_tree_rend(ZixTree* t)
662 {
663 return NULL;
664 }
665
666 ZIX_API bool
zix_tree_iter_is_end(const ZixTreeIter * i)667 zix_tree_iter_is_end(const ZixTreeIter* i)
668 {
669 return !i;
670 }
671
672 ZIX_API bool
zix_tree_iter_is_rend(const ZixTreeIter * i)673 zix_tree_iter_is_rend(const ZixTreeIter* i)
674 {
675 return !i;
676 }
677
678 ZIX_API ZixTreeIter*
zix_tree_iter_next(ZixTreeIter * i)679 zix_tree_iter_next(ZixTreeIter* i)
680 {
681 if (!i) {
682 return NULL;
683 }
684
685 if (i->right) {
686 i = i->right;
687 while (i->left) {
688 i = i->left;
689 }
690 } else {
691 while (i->parent && i->parent->right == i) { // i is a right child
692 i = i->parent;
693 }
694
695 i = i->parent;
696 }
697
698 return i;
699 }
700
701 ZIX_API ZixTreeIter*
zix_tree_iter_prev(ZixTreeIter * i)702 zix_tree_iter_prev(ZixTreeIter* i)
703 {
704 if (!i) {
705 return NULL;
706 }
707
708 if (i->left) {
709 i = i->left;
710 while (i->right) {
711 i = i->right;
712 }
713 } else {
714 while (i->parent && i->parent->left == i) { // i is a left child
715 i = i->parent;
716 }
717
718 i = i->parent;
719 }
720
721 return i;
722 }
723