1 /*
2 * This program is free software; you can redistribute it and/or
3 * modify it under the terms of the GNU General Public License
4 * as published by the Free Software Foundation; either version 2
5 * of the License, or (at your option) any later version.
6 *
7 * This program is distributed in the hope that it will be useful,
8 * but WITHOUT ANY WARRANTY; without even the implied warranty of
9 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
10 * GNU General Public License for more details.
11 *
12 * You should have received a copy of the GNU General Public License
13 * along with this program; if not, write to the Free Software Foundation,
14 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
15 *
16 * The Original Code is Copyright (C) 2009 Blender Foundation, Joshua Leung
17 * All rights reserved.
18 */
19
20 /** \file
21 * \ingroup bli
22 */
23
24 #include "MEM_guardedalloc.h"
25
26 #include "BLI_dlrbTree.h"
27 #include "BLI_listbase.h"
28
29 /* *********************************************** */
30 /* Tree API */
31
32 /* Create a new tree, and initialize as necessary */
BLI_dlrbTree_new(void)33 DLRBT_Tree *BLI_dlrbTree_new(void)
34 {
35 /* just allocate for now */
36 return MEM_callocN(sizeof(DLRBT_Tree), "DLRBT_Tree");
37 }
38
39 /* Just zero out the pointers used */
BLI_dlrbTree_init(DLRBT_Tree * tree)40 void BLI_dlrbTree_init(DLRBT_Tree *tree)
41 {
42 if (tree == NULL) {
43 return;
44 }
45
46 tree->first = tree->last = tree->root = NULL;
47 }
48
49 /* Helper for traversing tree and freeing sub-nodes */
recursive_tree_free_nodes(DLRBT_Node * node)50 static void recursive_tree_free_nodes(DLRBT_Node *node)
51 {
52 /* sanity check */
53 if (node == NULL) {
54 return;
55 }
56
57 /* free child nodes + subtrees */
58 recursive_tree_free_nodes(node->left);
59 recursive_tree_free_nodes(node->right);
60
61 /* free self */
62 MEM_freeN(node);
63 }
64
65 /* Free the given tree's data but not the tree itself */
BLI_dlrbTree_free(DLRBT_Tree * tree)66 void BLI_dlrbTree_free(DLRBT_Tree *tree)
67 {
68 if (tree == NULL) {
69 return;
70 }
71
72 /* if the list-base stuff is set, just use that (and assume its set),
73 * otherwise, we'll need to traverse the tree...
74 */
75 if (tree->first) {
76 /* free list */
77 BLI_freelistN((ListBase *)tree);
78 }
79 else {
80 /* traverse tree, freeing sub-nodes */
81 recursive_tree_free_nodes(tree->root);
82 }
83
84 /* clear pointers */
85 tree->first = tree->last = tree->root = NULL;
86 }
87
88 /* ------- */
89
90 /* Helper function - used for traversing down the tree from the root to add nodes in order */
linkedlist_sync_add_node(DLRBT_Tree * tree,DLRBT_Node * node)91 static void linkedlist_sync_add_node(DLRBT_Tree *tree, DLRBT_Node *node)
92 {
93 /* sanity checks */
94 if ((tree == NULL) || (node == NULL)) {
95 return;
96 }
97
98 /* add left-node (and its subtree) */
99 linkedlist_sync_add_node(tree, node->left);
100
101 /* now add self
102 * - must remove detach from other links first
103 * (for now, only clear own pointers)
104 */
105 node->prev = node->next = NULL;
106 BLI_addtail((ListBase *)tree, (Link *)node);
107
108 /* finally, add right node (and its subtree) */
109 linkedlist_sync_add_node(tree, node->right);
110 }
111
112 /* Make sure the tree's Double-Linked list representation is valid */
BLI_dlrbTree_linkedlist_sync(DLRBT_Tree * tree)113 void BLI_dlrbTree_linkedlist_sync(DLRBT_Tree *tree)
114 {
115 /* sanity checks */
116 if (tree == NULL) {
117 return;
118 }
119
120 /* clear list-base pointers so that the new list can be added properly */
121 tree->first = tree->last = NULL;
122
123 /* start adding items from the root */
124 linkedlist_sync_add_node(tree, tree->root);
125 }
126
127 /* *********************************************** */
128 /* Tree Search Utilities */
129
130 /* Find the node which matches or is the closest to the requested node */
BLI_dlrbTree_search(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,void * search_data)131 DLRBT_Node *BLI_dlrbTree_search(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
132 {
133 DLRBT_Node *node = (tree) ? tree->root : NULL;
134 short found = 0;
135
136 /* check that there is a comparator to use */
137 /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
138 if (cmp_cb == NULL) {
139 return NULL;
140 }
141
142 /* iteratively perform this search */
143 while (node && found == 0) {
144 /* check if traverse further or not
145 * NOTE: it is assumed that the values will be unit values only
146 */
147 switch (cmp_cb(node, search_data)) {
148 case -1: /* data less than node */
149 if (node->left) {
150 node = node->left;
151 }
152 else {
153 found = 1;
154 }
155 break;
156
157 case 1: /* data greater than node */
158 if (node->right) {
159 node = node->right;
160 }
161 else {
162 found = 1;
163 }
164 break;
165
166 default: /* data equals node */
167 found = 1;
168 break;
169 }
170 }
171
172 /* return the nearest matching node */
173 return node;
174 }
175
176 /* Find the node which exactly matches the required data */
BLI_dlrbTree_search_exact(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,void * search_data)177 DLRBT_Node *BLI_dlrbTree_search_exact(DLRBT_Tree *tree,
178 DLRBT_Comparator_FP cmp_cb,
179 void *search_data)
180 {
181 DLRBT_Node *node = (tree) ? tree->root : NULL;
182 short found = 0;
183
184 /* check that there is a comparator to use */
185 /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
186 if (cmp_cb == NULL) {
187 return NULL;
188 }
189
190 /* iteratively perform this search */
191 while (node && found == 0) {
192 /* check if traverse further or not
193 * NOTE: it is assumed that the values will be unit values only
194 */
195 switch (cmp_cb(node, search_data)) {
196 case -1: /* data less than node */
197 if (node->left) {
198 node = node->left;
199 }
200 else {
201 found = -1;
202 }
203 break;
204
205 case 1: /* data greater than node */
206 if (node->right) {
207 node = node->right;
208 }
209 else {
210 found = -1;
211 }
212 break;
213
214 default: /* data equals node */
215 found = 1;
216 break;
217 }
218 }
219
220 /* return the exactly matching node */
221 return (found == 1) ? (node) : (NULL);
222 }
223
224 /* Find the node which occurs immediately before the best matching node */
BLI_dlrbTree_search_prev(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,void * search_data)225 DLRBT_Node *BLI_dlrbTree_search_prev(DLRBT_Tree *tree,
226 DLRBT_Comparator_FP cmp_cb,
227 void *search_data)
228 {
229 DLRBT_Node *node;
230
231 /* check that there is a comparator to use */
232 /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
233 if (cmp_cb == NULL) {
234 return NULL;
235 }
236
237 /* get the node which best matches this description */
238 node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
239
240 if (node) {
241 /* if the item we're searching for is greater than the node found, we've found the match */
242 if (cmp_cb(node, search_data) > 0) {
243 return node;
244 }
245
246 /* return the previous node otherwise */
247 /* NOTE: what happens if there is no previous node? */
248 return node->prev;
249 }
250
251 /* nothing matching was found */
252 return NULL;
253 }
254
255 /* Find the node which occurs immediately after the best matching node */
BLI_dlrbTree_search_next(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,void * search_data)256 DLRBT_Node *BLI_dlrbTree_search_next(DLRBT_Tree *tree,
257 DLRBT_Comparator_FP cmp_cb,
258 void *search_data)
259 {
260 DLRBT_Node *node;
261
262 /* check that there is a comparator to use */
263 /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
264 if (cmp_cb == NULL) {
265 return NULL;
266 }
267
268 /* get the node which best matches this description */
269 node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
270
271 if (node) {
272 /* if the item we're searching for is less than the node found, we've found the match */
273 if (cmp_cb(node, search_data) < 0) {
274 return node;
275 }
276
277 /* return the previous node otherwise */
278 /* NOTE: what happens if there is no previous node? */
279 return node->next;
280 }
281
282 /* nothing matching was found */
283 return NULL;
284 }
285
286 /* Check whether there is a node matching the requested node */
BLI_dlrbTree_contains(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,void * search_data)287 short BLI_dlrbTree_contains(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
288 {
289 /* check if an exact search throws up anything... */
290 return (BLI_dlrbTree_search_exact(tree, cmp_cb, search_data) != NULL);
291 }
292
293 /* *********************************************** */
294 /* Tree Relationships Utilities */
295
296 /* get the 'grandparent' - the parent of the parent - of the given node */
get_grandparent(DLRBT_Node * node)297 static DLRBT_Node *get_grandparent(DLRBT_Node *node)
298 {
299 if (node && node->parent) {
300 return node->parent->parent;
301 }
302 return NULL;
303 }
304
305 /* get the sibling node (e.g. if node is left child of parent, return right child of parent) */
get_sibling(DLRBT_Node * node)306 static DLRBT_Node *get_sibling(DLRBT_Node *node)
307 {
308 if (node && node->parent) {
309 if (node == node->parent->left) {
310 return node->parent->right;
311 }
312 return node->parent->left;
313 }
314
315 /* sibling not found */
316 return NULL;
317 }
318
319 /* get the 'uncle' - the sibling of the parent - of the given node */
get_uncle(DLRBT_Node * node)320 static DLRBT_Node *get_uncle(DLRBT_Node *node)
321 {
322 if (node) {
323 /* return the child of the grandparent which isn't the node's parent */
324 return get_sibling(node->parent);
325 }
326
327 /* uncle not found */
328 return NULL;
329 }
330
331 /* *********************************************** */
332 /* Tree Rotation Utilities */
333
334 /* make right child of 'root' the new root */
rotate_left(DLRBT_Tree * tree,DLRBT_Node * root)335 static void rotate_left(DLRBT_Tree *tree, DLRBT_Node *root)
336 {
337 DLRBT_Node **root_slot, *pivot;
338
339 /* pivot is simply the root's right child, to become the root's parent */
340 pivot = root->right;
341 if (pivot == NULL) {
342 return;
343 }
344
345 if (root->parent) {
346 if (root == root->parent->left) {
347 root_slot = &root->parent->left;
348 }
349 else {
350 root_slot = &root->parent->right;
351 }
352 }
353 else {
354 root_slot = ((DLRBT_Node **)&tree->root); /* &((DLRBT_Node *)tree->root); */
355 }
356
357 /* - pivot's left child becomes root's right child
358 * - root now becomes pivot's left child
359 */
360 root->right = pivot->left;
361 if (pivot->left) {
362 pivot->left->parent = root;
363 }
364
365 pivot->left = root;
366 pivot->parent = root->parent;
367 root->parent = pivot;
368
369 /* make the pivot the new root */
370 if (root_slot) {
371 *root_slot = pivot;
372 }
373 }
374
375 /* make the left child of the 'root' the new root */
rotate_right(DLRBT_Tree * tree,DLRBT_Node * root)376 static void rotate_right(DLRBT_Tree *tree, DLRBT_Node *root)
377 {
378 DLRBT_Node **root_slot, *pivot;
379
380 /* pivot is simply the root's left child, to become the root's parent */
381 pivot = root->left;
382 if (pivot == NULL) {
383 return;
384 }
385
386 if (root->parent) {
387 if (root == root->parent->left) {
388 root_slot = &root->parent->left;
389 }
390 else {
391 root_slot = &root->parent->right;
392 }
393 }
394 else {
395 root_slot = ((DLRBT_Node **)&tree->root); /* &((DLRBT_Node *)tree->root); */
396 }
397
398 /* - pivot's right child becomes root's left child
399 * - root now becomes pivot's right child
400 */
401 root->left = pivot->right;
402 if (pivot->right) {
403 pivot->right->parent = root;
404 }
405
406 pivot->right = root;
407 pivot->parent = root->parent;
408 root->parent = pivot;
409
410 /* make the pivot the new root */
411 if (root_slot) {
412 *root_slot = pivot;
413 }
414 }
415
416 /* *********************************************** */
417 /* Post-Insertion Balancing */
418
419 /* forward defines for insertion checks */
420 static void insert_check_1(DLRBT_Tree *tree, DLRBT_Node *node);
421 static void insert_check_2(DLRBT_Tree *tree, DLRBT_Node *node);
422 static void insert_check_3(DLRBT_Tree *tree, DLRBT_Node *node);
423
424 /* ----- */
425
426 /* W. 1) Root must be black (so that the 2nd-generation can have a black parent) */
insert_check_1(DLRBT_Tree * tree,DLRBT_Node * node)427 static void insert_check_1(DLRBT_Tree *tree, DLRBT_Node *node)
428 {
429 if (node) {
430 /* if this is the root, just ensure that it is black */
431 if (node->parent == NULL) {
432 node->tree_col = DLRBT_BLACK;
433 }
434 else {
435 insert_check_2(tree, node);
436 }
437 }
438 }
439
440 /* W. 2+3) Parent of node must be black, otherwise recolor and flush */
insert_check_2(DLRBT_Tree * tree,DLRBT_Node * node)441 static void insert_check_2(DLRBT_Tree *tree, DLRBT_Node *node)
442 {
443 /* if the parent is not black, we need to change that... */
444 if (node && node->parent && node->parent->tree_col) {
445 DLRBT_Node *unc = get_uncle(node);
446
447 /* if uncle and parent are both red, need to change them to black and make
448 * the parent black in order to satisfy the criteria of each node having the
449 * same number of black nodes to its leaves
450 */
451 if (unc && unc->tree_col) {
452 DLRBT_Node *gp = get_grandparent(node);
453
454 /* make the n-1 generation nodes black */
455 node->parent->tree_col = unc->tree_col = DLRBT_BLACK;
456
457 /* - make the grandparent red, so that we maintain alternating red/black property
458 * (it must exist, so no need to check for NULL here),
459 * - as the grandparent may now cause inconsistencies with the rest of the tree,
460 * we must flush up the tree and perform checks/re-balancing/re-painting, using the
461 * grandparent as the node of interest
462 */
463 gp->tree_col = DLRBT_RED;
464 insert_check_1(tree, gp);
465 }
466 else {
467 /* we've got an unbalanced branch going down the grandparent to the parent,
468 * so need to perform some rotations to re-balance the tree
469 */
470 insert_check_3(tree, node);
471 }
472 }
473 }
474
475 /* W. 4+5) Perform rotation on sub-tree containing the 'new' node, then do any */
insert_check_3(DLRBT_Tree * tree,DLRBT_Node * node)476 static void insert_check_3(DLRBT_Tree *tree, DLRBT_Node *node)
477 {
478 DLRBT_Node *gp = get_grandparent(node);
479
480 /* check that grandparent and node->parent exist
481 * (jut in case... really shouldn't happen on a good tree) */
482 if (node && node->parent && gp) {
483 /* a left rotation will switch the roles of node and its parent, assuming that
484 * the parent is the left child of the grandparent... otherwise, rotation direction
485 * should be swapped
486 */
487 if ((node == node->parent->right) && (node->parent == gp->left)) {
488 rotate_left(tree, node);
489 node = node->left;
490 }
491 else if ((node == node->parent->left) && (node->parent == gp->right)) {
492 rotate_right(tree, node);
493 node = node->right;
494 }
495
496 /* fix old parent's color-tagging, and perform rotation on the old parent in the
497 * opposite direction if needed for the current situation
498 * NOTE: in the code above, node pointer is changed to point to the old parent
499 */
500 if (node) {
501 /* get 'new' grandparent (i.e. grandparent for old-parent (node)) */
502 gp = get_grandparent(node);
503
504 /* modify the coloring of the grandparent and parent
505 * so that they still satisfy the constraints */
506 node->parent->tree_col = DLRBT_BLACK;
507 gp->tree_col = DLRBT_RED;
508
509 /* if there are several nodes that all form a left chain, do a right rotation to correct
510 * this (or a rotation in the opposite direction if they all form a right chain) */
511 if ((node == node->parent->left) && (node->parent == gp->left)) {
512 rotate_right(tree, gp);
513 }
514 else { // if ((node == node->parent->right) && (node->parent == gp->right))
515 rotate_left(tree, gp);
516 }
517 }
518 }
519 }
520
521 /* ----- */
522
523 /* Balance the tree after the given element has been added to it
524 * (using custom code, in the Binary Tree way).
525 */
BLI_dlrbTree_insert(DLRBT_Tree * tree,DLRBT_Node * node)526 void BLI_dlrbTree_insert(DLRBT_Tree *tree, DLRBT_Node *node)
527 {
528 /* sanity checks */
529 if ((tree == NULL) || (node == NULL)) {
530 return;
531 }
532
533 /* firstly, the node we just added should be red by default */
534 node->tree_col = DLRBT_RED;
535
536 /* start from case 1, an trek through the tail-recursive insertion checks */
537 insert_check_1(tree, node);
538 }
539
540 /* ----- */
541
542 /* Add the given data to the tree, and return the node added */
543 /* NOTE: for duplicates, the update_cb is called (if available),
544 * and the existing node is returned */
BLI_dlrbTree_add(DLRBT_Tree * tree,DLRBT_Comparator_FP cmp_cb,DLRBT_NAlloc_FP new_cb,DLRBT_NUpdate_FP update_cb,void * data)545 DLRBT_Node *BLI_dlrbTree_add(DLRBT_Tree *tree,
546 DLRBT_Comparator_FP cmp_cb,
547 DLRBT_NAlloc_FP new_cb,
548 DLRBT_NUpdate_FP update_cb,
549 void *data)
550 {
551 DLRBT_Node *parNode, *node = NULL;
552 short new_node = 0;
553
554 /* sanity checks */
555 if (tree == NULL) {
556 return NULL;
557 }
558
559 /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
560 if (cmp_cb == NULL) {
561 return NULL;
562 }
563 /* TODO: if no allocator is supplied, try using the one supplied with the tree... */
564 if (new_cb == NULL) {
565 return NULL;
566 }
567 /* TODO: if no updater is supplied, try using the one supplied with the tree... */
568
569 /* try to find the nearest node to this one */
570 parNode = BLI_dlrbTree_search(tree, cmp_cb, data);
571
572 /* add new node to the BST in the 'standard way' as appropriate
573 * NOTE: we do not support duplicates in our tree...
574 */
575 if (parNode) {
576 /* check how this new node compares with the existing ones
577 * NOTE: it is assumed that the values will be unit values only
578 */
579 switch (cmp_cb(parNode, data)) {
580 case -1: /* add new node as left child */
581 {
582 node = new_cb(data);
583 new_node = 1;
584
585 parNode->left = node;
586 node->parent = parNode;
587 break;
588 }
589 case 1: /* add new node as right child */
590 {
591 node = new_cb(data);
592 new_node = 1;
593
594 parNode->right = node;
595 node->parent = parNode;
596 break;
597 }
598 default: /* update the duplicate node as appropriate */
599 {
600 if (update_cb) {
601 update_cb(parNode, data);
602 }
603 break;
604 }
605 }
606 }
607 else {
608 /* no nodes in the tree yet... add a new node as the root */
609 node = new_cb(data);
610 new_node = 1;
611
612 tree->root = node;
613 }
614
615 /* if a new node was added, it should be tagged as red, and then balanced as appropriate */
616 if (new_node) {
617 /* tag this new node as being 'red' */
618 node->tree_col = DLRBT_RED;
619
620 /* perform BST balancing steps:
621 * start from case 1, an trek through the tail-recursive insertion checks
622 */
623 insert_check_1(tree, node);
624 }
625
626 /* return the node added */
627 return node;
628 }
629
630 /* *********************************************** */
631 /* Remove */
632
633 /* TODO: this hasn't been coded yet, since this functionality was not needed by the author */
634
635 /* *********************************************** */
636