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