1 /*
2 Copyright 2011-2019 David Robillard <http://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/common.h"
18 #include "zix/tree.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, p->balance, q->balance, r->balance);
236
237 rotate(q, r);
238 rotate(p, r);
239
240 q->balance -= 1 + MAX(0, r->balance);
241 p->balance += 1 - MIN(MIN(0, r->balance) - 1, r->balance + q->balance);
242 // r->balance += MAX(0, p->balance) + MIN(0, q->balance);
243
244 // p->balance = (p->left && p->right) ? -MIN(r->balance, 0) : 0;
245 // q->balance = - MAX(r->balance, 0);
246 r->balance = 0;
247
248 *height_change = -1;
249
250 ASSERT_BALANCE(p);
251 ASSERT_BALANCE(q);
252 ASSERT_BALANCE(r);
253 return r;
254 }
255
256 /**
257 * Rotate right about `p->right` then right about `p`.
258 *
259 * p r
260 * / \ / \
261 * A q => p q
262 * / \ / \ / \
263 * r D A B C D
264 * / \
265 * B C
266 *
267 */
268 ZIX_PRIVATE ZixTreeNode*
rotate_right_left(ZixTreeNode * p,int * height_change)269 rotate_right_left(ZixTreeNode* p, int* height_change)
270 {
271 ZixTreeNode* const q = p->right;
272 ZixTreeNode* const r = q->left;
273
274 assert(p->balance == 2);
275 assert(q->balance == -1);
276 assert(r->balance == -1 || r->balance == 0 || r->balance == 1);
277
278 DEBUG_PRINTF("RL %ld P: %2d Q: %2d R: %2d\n",
279 (intptr_t)p->data, p->balance, q->balance, r->balance);
280
281 rotate(q, r);
282 rotate(p, r);
283
284 q->balance += 1 - MIN(0, r->balance);
285 p->balance -= 1 + MAX(MAX(0, r->balance) + 1, r->balance + q->balance);
286 // r->balance += MAX(0, q->balance) + MIN(0, p->balance);
287
288 // p->balance = (p->left && p->right) ? -MAX(r->balance, 0) : 0;
289 // q->balance = - MIN(r->balance, 0);
290 r->balance = 0;
291 // assert(r->balance == 0);
292
293 *height_change = -1;
294
295 ASSERT_BALANCE(p);
296 ASSERT_BALANCE(q);
297 ASSERT_BALANCE(r);
298 return r;
299 }
300
301 ZIX_PRIVATE ZixTreeNode*
zix_tree_rebalance(ZixTree * t,ZixTreeNode * node,int * height_change)302 zix_tree_rebalance(ZixTree* t, ZixTreeNode* node, int* height_change)
303 {
304 #ifdef ZIX_TREE_HYPER_VERIFY
305 const size_t old_height = height(node);
306 #endif
307 DEBUG_PRINTF("REBALANCE %ld (%d)\n", (intptr_t)node->data, node->balance);
308 *height_change = 0;
309 const bool is_root = !node->parent;
310 assert((is_root && t->root == node) || (!is_root && t->root != node));
311 ZixTreeNode* replacement = node;
312 if (node->balance == -2) {
313 assert(node->left);
314 if (node->left->balance == 1) {
315 replacement = rotate_left_right(node, height_change);
316 } else {
317 replacement = rotate_right(node, height_change);
318 }
319 } else if (node->balance == 2) {
320 assert(node->right);
321 if (node->right->balance == -1) {
322 replacement = rotate_right_left(node, height_change);
323 } else {
324 replacement = rotate_left(node, height_change);
325 }
326 }
327 if (is_root) {
328 assert(!replacement->parent);
329 t->root = replacement;
330 }
331 DUMP(t);
332 #ifdef ZIX_TREE_HYPER_VERIFY
333 assert(old_height + *height_change == height(replacement));
334 #endif
335 return replacement;
336 }
337
338 ZIX_API ZixStatus
zix_tree_insert(ZixTree * t,void * e,ZixTreeIter ** ti)339 zix_tree_insert(ZixTree* t, void* e, ZixTreeIter** ti)
340 {
341 DEBUG_PRINTF("**** INSERT %ld\n", (intptr_t)e);
342 int cmp = 0;
343 ZixTreeNode* n = t->root;
344 ZixTreeNode* p = NULL;
345
346 // Find the parent p of e
347 while (n) {
348 p = n;
349 cmp = t->cmp(e, n->data, t->cmp_data);
350 if (cmp < 0) {
351 n = n->left;
352 } else if (cmp > 0) {
353 n = n->right;
354 } else if (t->allow_duplicates) {
355 n = n->right;
356 } else {
357 if (ti) {
358 *ti = n;
359 }
360 DEBUG_PRINTF("%ld EXISTS!\n", (intptr_t)e);
361 return ZIX_STATUS_EXISTS;
362 }
363 }
364
365 // Allocate a new node n
366 if (!(n = (ZixTreeNode*)malloc(sizeof(ZixTreeNode)))) {
367 return ZIX_STATUS_NO_MEM;
368 }
369 memset(n, '\0', sizeof(ZixTreeNode));
370 n->data = e;
371 n->balance = 0;
372 if (ti) {
373 *ti = n;
374 }
375
376 bool p_height_increased = false;
377
378 // Make p the parent of n
379 n->parent = p;
380 if (!p) {
381 t->root = n;
382 } else {
383 if (cmp < 0) {
384 assert(!p->left);
385 assert(p->balance == 0 || p->balance == 1);
386 p->left = n;
387 --p->balance;
388 p_height_increased = !p->right;
389 } else {
390 assert(!p->right);
391 assert(p->balance == 0 || p->balance == -1);
392 p->right = n;
393 ++p->balance;
394 p_height_increased = !p->left;
395 }
396 }
397
398 DUMP(t);
399
400 // Rebalance if necessary (at most 1 rotation)
401 assert(!p || p->balance == -1 || p->balance == 0 || p->balance == 1);
402 if (p && p_height_increased) {
403 int height_change = 0;
404 for (ZixTreeNode* i = p; i && i->parent; i = i->parent) {
405 if (i == i->parent->left) {
406 if (--i->parent->balance == -2) {
407 zix_tree_rebalance(t, i->parent, &height_change);
408 break;
409 }
410 } else {
411 assert(i == i->parent->right);
412 if (++i->parent->balance == 2) {
413 zix_tree_rebalance(t, i->parent, &height_change);
414 break;
415 }
416 }
417
418 if (i->parent->balance == 0) {
419 break;
420 }
421 }
422 }
423
424 DUMP(t);
425
426 ++t->size;
427
428 #ifdef ZIX_TREE_VERIFY
429 if (!verify(t, t->root)) {
430 return ZIX_STATUS_ERROR;
431 }
432 #endif
433
434 return ZIX_STATUS_SUCCESS;
435 }
436
437 ZIX_API ZixStatus
zix_tree_remove(ZixTree * t,ZixTreeIter * ti)438 zix_tree_remove(ZixTree* t, ZixTreeIter* ti)
439 {
440 ZixTreeNode* const n = ti;
441 ZixTreeNode** pp = NULL; // parent pointer
442 ZixTreeNode* to_balance = n->parent; // lowest node to balance
443 int8_t d_balance = 0; // delta(balance) for n->parent
444
445 DEBUG_PRINTF("*** REMOVE %ld\n", (intptr_t)n->data);
446
447 if ((n == t->root) && !n->left && !n->right) {
448 t->root = NULL;
449 if (t->destroy) {
450 t->destroy(n->data);
451 }
452 free(n);
453 --t->size;
454 assert(t->size == 0);
455 return ZIX_STATUS_SUCCESS;
456 }
457
458 // Set pp to the parent pointer to n, if applicable
459 if (n->parent) {
460 assert(n->parent->left == n || n->parent->right == n);
461 if (n->parent->left == n) { // n is left child
462 pp = &n->parent->left;
463 d_balance = 1;
464 } else { // n is right child
465 assert(n->parent->right == n);
466 pp = &n->parent->right;
467 d_balance = -1;
468 }
469 }
470
471 assert(!pp || *pp == n);
472
473 int height_change = 0;
474 if (!n->left && !n->right) {
475 // n is a leaf, just remove it
476 if (pp) {
477 *pp = NULL;
478 to_balance = n->parent;
479 height_change = (!n->parent->left && !n->parent->right) ? -1 : 0;
480 }
481 } else if (!n->left) {
482 // Replace n with right (only) child
483 if (pp) {
484 *pp = n->right;
485 to_balance = n->parent;
486 } else {
487 t->root = n->right;
488 }
489 n->right->parent = n->parent;
490 height_change = -1;
491 } else if (!n->right) {
492 // Replace n with left (only) child
493 if (pp) {
494 *pp = n->left;
495 to_balance = n->parent;
496 } else {
497 t->root = n->left;
498 }
499 n->left->parent = n->parent;
500 height_change = -1;
501 } else {
502 // Replace n with in-order successor (leftmost child of right subtree)
503 ZixTreeNode* replace = n->right;
504 while (replace->left) {
505 assert(replace->left->parent == replace);
506 replace = replace->left;
507 }
508
509 // Remove replace from parent (replace_p)
510 if (replace->parent->left == replace) {
511 height_change = replace->parent->right ? 0 : -1;
512 d_balance = 1;
513 to_balance = replace->parent;
514 replace->parent->left = replace->right;
515 } else {
516 assert(replace->parent == n);
517 height_change = replace->parent->left ? 0 : -1;
518 d_balance = -1;
519 to_balance = replace->parent;
520 replace->parent->right = replace->right;
521 }
522
523 if (to_balance == n) {
524 to_balance = replace;
525 }
526
527 if (replace->right) {
528 replace->right->parent = replace->parent;
529 }
530
531 replace->balance = n->balance;
532
533 // Swap node to delete with replace
534 if (pp) {
535 *pp = replace;
536 } else {
537 assert(t->root == n);
538 t->root = replace;
539 }
540 replace->parent = n->parent;
541 replace->left = n->left;
542 n->left->parent = replace;
543 replace->right = n->right;
544 if (n->right) {
545 n->right->parent = replace;
546 }
547
548 assert(!replace->parent
549 || replace->parent->left == replace
550 || replace->parent->right == replace);
551 }
552
553 // Rebalance starting at to_balance upwards.
554 for (ZixTreeNode* i = to_balance; i; i = i->parent) {
555 i->balance += d_balance;
556 if (d_balance == 0 || i->balance == -1 || i->balance == 1) {
557 break;
558 }
559
560 assert(i != n);
561 i = zix_tree_rebalance(t, i, &height_change);
562 if (i->balance == 0) {
563 height_change = -1;
564 }
565
566 if (i->parent) {
567 if (i == i->parent->left) {
568 d_balance = height_change * -1;
569 } else {
570 assert(i == i->parent->right);
571 d_balance = height_change;
572 }
573 }
574 }
575
576 DUMP(t);
577
578 if (t->destroy) {
579 t->destroy(n->data);
580 }
581 free(n);
582
583 --t->size;
584
585 #ifdef ZIX_TREE_VERIFY
586 if (!verify(t, t->root)) {
587 return ZIX_STATUS_ERROR;
588 }
589 #endif
590
591 return ZIX_STATUS_SUCCESS;
592 }
593
594 ZIX_API ZixStatus
zix_tree_find(const ZixTree * t,const void * e,ZixTreeIter ** ti)595 zix_tree_find(const ZixTree* t, const void* e, ZixTreeIter** ti)
596 {
597 ZixTreeNode* n = t->root;
598 while (n) {
599 const int cmp = t->cmp(e, n->data, t->cmp_data);
600 if (cmp == 0) {
601 break;
602 } else if (cmp < 0) {
603 n = n->left;
604 } else {
605 n = n->right;
606 }
607 }
608
609 *ti = n;
610 return (n) ? ZIX_STATUS_SUCCESS : ZIX_STATUS_NOT_FOUND;
611 }
612
613 ZIX_API void*
zix_tree_get(const ZixTreeIter * ti)614 zix_tree_get(const ZixTreeIter* ti)
615 {
616 return ti ? ti->data : NULL;
617 }
618
619 ZIX_API ZixTreeIter*
zix_tree_begin(ZixTree * t)620 zix_tree_begin(ZixTree* t)
621 {
622 if (!t->root) {
623 return NULL;
624 }
625
626 ZixTreeNode* n = t->root;
627 while (n->left) {
628 n = n->left;
629 }
630 return n;
631 }
632
633 ZIX_API ZixTreeIter*
zix_tree_end(ZixTree * t)634 zix_tree_end(ZixTree* t)
635 {
636 return NULL;
637 }
638
639 ZIX_API ZixTreeIter*
zix_tree_rbegin(ZixTree * t)640 zix_tree_rbegin(ZixTree* t)
641 {
642 if (!t->root) {
643 return NULL;
644 }
645
646 ZixTreeNode* n = t->root;
647 while (n->right) {
648 n = n->right;
649 }
650 return n;
651 }
652
653 ZIX_API ZixTreeIter*
zix_tree_rend(ZixTree * t)654 zix_tree_rend(ZixTree* t)
655 {
656 return NULL;
657 }
658
659 ZIX_API bool
zix_tree_iter_is_end(const ZixTreeIter * i)660 zix_tree_iter_is_end(const ZixTreeIter* i)
661 {
662 return !i;
663 }
664
665 ZIX_API bool
zix_tree_iter_is_rend(const ZixTreeIter * i)666 zix_tree_iter_is_rend(const ZixTreeIter* i)
667 {
668 return !i;
669 }
670
671 ZIX_API ZixTreeIter*
zix_tree_iter_next(ZixTreeIter * i)672 zix_tree_iter_next(ZixTreeIter* i)
673 {
674 if (!i) {
675 return NULL;
676 }
677
678 if (i->right) {
679 i = i->right;
680 while (i->left) {
681 i = i->left;
682 }
683 } else {
684 while (i->parent && i->parent->right == i) { // i is a right child
685 i = i->parent;
686 }
687
688 i = i->parent;
689 }
690
691 return i;
692 }
693
694 ZIX_API ZixTreeIter*
zix_tree_iter_prev(ZixTreeIter * i)695 zix_tree_iter_prev(ZixTreeIter* i)
696 {
697 if (!i) {
698 return NULL;
699 }
700
701 if (i->left) {
702 i = i->left;
703 while (i->right) {
704 i = i->right;
705 }
706 } else {
707 while (i->parent && i->parent->left == i) { // i is a left child
708 i = i->parent;
709 }
710
711 i = i->parent;
712 }
713
714 return i;
715 }
716