Cleanup: style, use braces for blenlib
[blender.git] / source / blender / blenlib / intern / DLRB_tree.c
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
25 #include "MEM_guardedalloc.h"
26
27 #include "BLI_listbase.h"
28 #include "BLI_dlrbTree.h"
29
30 /* *********************************************** */
31 /* Tree API */
32
33 /* Create a new tree, and initialize as necessary */
34 DLRBT_Tree *BLI_dlrbTree_new(void)
35 {
36         /* just allocate for now */
37         return MEM_callocN(sizeof(DLRBT_Tree), "DLRBT_Tree");
38 }
39
40 /* Just zero out the pointers used */
41 void BLI_dlrbTree_init(DLRBT_Tree *tree)
42 {
43         if (tree == NULL) {
44                 return;
45         }
46
47         tree->first = tree->last = tree->root = NULL;
48 }
49
50 /* Helper for traversing tree and freeing sub-nodes */
51 static void recursive_tree_free_nodes(DLRBT_Node *node)
52 {
53         /* sanity check */
54         if (node == NULL) {
55                 return;
56         }
57
58         /* free child nodes + subtrees */
59         recursive_tree_free_nodes(node->left);
60         recursive_tree_free_nodes(node->right);
61
62         /* free self */
63         MEM_freeN(node);
64 }
65
66 /* Free the given tree's data but not the tree itself */
67 void BLI_dlrbTree_free(DLRBT_Tree *tree)
68 {
69         if (tree == NULL) {
70                 return;
71         }
72
73         /* if the list-base stuff is set, just use that (and assume its set),
74          * otherwise, we'll need to traverse the tree...
75          */
76         if (tree->first) {
77                 /* free list */
78                 BLI_freelistN((ListBase *)tree);
79         }
80         else {
81                 /* traverse tree, freeing sub-nodes */
82                 recursive_tree_free_nodes(tree->root);
83         }
84
85         /* clear pointers */
86         tree->first = tree->last = tree->root = NULL;
87 }
88
89 /* ------- */
90
91 /* Helper function - used for traversing down the tree from the root to add nodes in order */
92 static void linkedlist_sync_add_node(DLRBT_Tree *tree, DLRBT_Node *node)
93 {
94         /* sanity checks */
95         if ((tree == NULL) || (node == NULL)) {
96                 return;
97         }
98
99         /* add left-node (and its subtree) */
100         linkedlist_sync_add_node(tree, node->left);
101
102         /* now add self
103          * - must remove detach from other links first
104          *   (for now, only clear own pointers)
105          */
106         node->prev = node->next = NULL;
107         BLI_addtail((ListBase *)tree, (Link *)node);
108
109         /* finally, add right node (and its subtree) */
110         linkedlist_sync_add_node(tree, node->right);
111 }
112
113 /* Make sure the tree's Double-Linked list representation is valid */
114 void BLI_dlrbTree_linkedlist_sync(DLRBT_Tree *tree)
115 {
116         /* sanity checks */
117         if (tree == NULL) {
118                 return;
119         }
120
121         /* clear list-base pointers so that the new list can be added properly */
122         tree->first = tree->last = NULL;
123
124         /* start adding items from the root */
125         linkedlist_sync_add_node(tree, tree->root);
126 }
127
128 /* *********************************************** */
129 /* Tree Search Utilities */
130
131 /* Find the node which matches or is the closest to the requested node */
132 DLRBT_Node *BLI_dlrbTree_search(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
133 {
134         DLRBT_Node *node = (tree) ? tree->root : NULL;
135         short found = 0;
136
137         /* check that there is a comparator to use */
138         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
139         if (cmp_cb == NULL) {
140                 return NULL;
141         }
142
143         /* iteratively perform this search */
144         while (node && found == 0) {
145                 /* check if traverse further or not
146                  * NOTE: it is assumed that the values will be unit values only
147                  */
148                 switch (cmp_cb(node, search_data)) {
149                         case -1:    /* data less than node */
150                                 if (node->left) {
151                                         node = node->left;
152                                 }
153                                 else {
154                                         found = 1;
155                                 }
156                                 break;
157
158                         case 1:  /* data greater than node */
159                                 if (node->right) {
160                                         node = node->right;
161                                 }
162                                 else {
163                                         found = 1;
164                                 }
165                                 break;
166
167                         default:  /* data equals node */
168                                 found = 1;
169                                 break;
170                 }
171         }
172
173         /* return the nearest matching node */
174         return node;
175 }
176
177 /* Find the node which exactly matches the required data */
178 DLRBT_Node *BLI_dlrbTree_search_exact(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
179 {
180         DLRBT_Node *node = (tree) ? tree->root : NULL;
181         short found = 0;
182
183         /* check that there is a comparator to use */
184         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
185         if (cmp_cb == NULL) {
186                 return NULL;
187         }
188
189         /* iteratively perform this search */
190         while (node && found == 0) {
191                 /* check if traverse further or not
192                  * NOTE: it is assumed that the values will be unit values only
193                  */
194                 switch (cmp_cb(node, search_data)) {
195                         case -1:  /* data less than node */
196                                 if (node->left) {
197                                         node = node->left;
198                                 }
199                                 else {
200                                         found = -1;
201                                 }
202                                 break;
203
204                         case 1:  /* data greater than node */
205                                 if (node->right) {
206                                         node = node->right;
207                                 }
208                                 else {
209                                         found = -1;
210                                 }
211                                 break;
212
213                         default:  /* data equals node */
214                                 found = 1;
215                                 break;
216                 }
217         }
218
219         /* return the exactly matching node */
220         return (found == 1) ? (node) : (NULL);
221 }
222
223 /* Find the node which occurs immediately before the best matching node */
224 DLRBT_Node *BLI_dlrbTree_search_prev(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
225 {
226         DLRBT_Node *node;
227
228         /* check that there is a comparator to use */
229         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
230         if (cmp_cb == NULL) {
231                 return NULL;
232         }
233
234         /* get the node which best matches this description */
235         node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
236
237         if (node) {
238                 /* if the item we're searching for is greater than the node found, we've found the match */
239                 if (cmp_cb(node, search_data) > 0) {
240                         return node;
241                 }
242
243                 /* return the previous node otherwise */
244                 /* NOTE: what happens if there is no previous node? */
245                 return node->prev;
246         }
247
248         /* nothing matching was found */
249         return NULL;
250 }
251
252 /* Find the node which occurs immediately after the best matching node */
253 DLRBT_Node *BLI_dlrbTree_search_next(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
254 {
255         DLRBT_Node *node;
256
257         /* check that there is a comparator to use */
258         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
259         if (cmp_cb == NULL) {
260                 return NULL;
261         }
262
263         /* get the node which best matches this description */
264         node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
265
266         if (node) {
267                 /* if the item we're searching for is less than the node found, we've found the match */
268                 if (cmp_cb(node, search_data) < 0) {
269                         return node;
270                 }
271
272                 /* return the previous node otherwise */
273                 /* NOTE: what happens if there is no previous node? */
274                 return node->next;
275         }
276
277         /* nothing matching was found */
278         return NULL;
279 }
280
281
282 /* Check whether there is a node matching the requested node */
283 short BLI_dlrbTree_contains(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
284 {
285         /* check if an exact search throws up anything... */
286         return (BLI_dlrbTree_search_exact(tree, cmp_cb, search_data) != NULL);
287 }
288
289 /* *********************************************** */
290 /* Tree Relationships Utilities */
291
292 /* get the 'grandparent' - the parent of the parent - of the given node */
293 static DLRBT_Node *get_grandparent(DLRBT_Node *node)
294 {
295         if (node && node->parent) {
296                 return node->parent->parent;
297         }
298         else {
299                 return NULL;
300         }
301 }
302
303 /* get the sibling node (e.g. if node is left child of parent, return right child of parent) */
304 static DLRBT_Node *get_sibling(DLRBT_Node *node)
305 {
306         if (node && node->parent) {
307                 if (node == node->parent->left) {
308                         return node->parent->right;
309                 }
310                 else {
311                         return node->parent->left;
312                 }
313         }
314
315         /* sibling not found */
316         return NULL;
317 }
318
319 /* get the 'uncle' - the sibling of the parent - of the given 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 */
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 */
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) */
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 */
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  */
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  */
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 */
545 DLRBT_Node *BLI_dlrbTree_add(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb,
546                              DLRBT_NAlloc_FP new_cb, DLRBT_NUpdate_FP update_cb, void *data)
547 {
548         DLRBT_Node *parNode, *node = NULL;
549         short new_node = 0;
550
551         /* sanity checks */
552         if (tree == NULL) {
553                 return NULL;
554         }
555
556         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
557         if (cmp_cb == NULL) {
558                 return NULL;
559         }
560         /* TODO: if no allocator is supplied, try using the one supplied with the tree... */
561         if (new_cb == NULL) {
562                 return NULL;
563         }
564         /* TODO: if no updater is supplied, try using the one supplied with the tree... */
565
566         /* try to find the nearest node to this one */
567         parNode = BLI_dlrbTree_search(tree, cmp_cb, data);
568
569         /* add new node to the BST in the 'standard way' as appropriate
570          * NOTE: we do not support duplicates in our tree...
571          */
572         if (parNode) {
573                 /* check how this new node compares with the existing ones
574                  * NOTE: it is assumed that the values will be unit values only
575                  */
576                 switch (cmp_cb(parNode, data)) {
577                         case -1:  /* add new node as left child */
578                         {
579                                 node = new_cb(data);
580                                 new_node = 1;
581
582                                 parNode->left = node;
583                                 node->parent = parNode;
584                                 break;
585                         }
586                         case 1:  /* add new node as right child */
587                         {
588                                 node = new_cb(data);
589                                 new_node = 1;
590
591                                 parNode->right = node;
592                                 node->parent = parNode;
593                                 break;
594                         }
595                         default:  /* update the duplicate node as appropriate */
596                         {
597                                 if (update_cb) {
598                                         update_cb(parNode, data);
599                                 }
600                                 break;
601                         }
602                 }
603         }
604         else {
605                 /* no nodes in the tree yet... add a new node as the root */
606                 node = new_cb(data);
607                 new_node = 1;
608
609                 tree->root = node;
610         }
611
612         /* if a new node was added, it should be tagged as red, and then balanced as appropriate */
613         if (new_node) {
614                 /* tag this new node as being 'red' */
615                 node->tree_col = DLRBT_RED;
616
617                 /* perform BST balancing steps:
618                  * start from case 1, an trek through the tail-recursive insertion checks
619                  */
620                 insert_check_1(tree, node);
621         }
622
623         /* return the node added */
624         return node;
625 }
626
627 /* *********************************************** */
628 /* Remove */
629
630 /* TODO: this hasn't been coded yet, since this functionality was not needed by the author */
631
632 /* *********************************************** */