BLI_math_rotation: properly name the quaternion power function.
[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_listbase.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                 /* 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                                 else
153                                         found = 1;
154                                 break;
155
156                         case 1:  /* data greater than node */
157                                 if (node->right)
158                                         node = node->right;
159                                 else
160                                         found = 1;
161                                 break;
162
163                         default:  /* data equals node */
164                                 found = 1;
165                                 break;
166                 }
167         }
168
169         /* return the nearest matching node */
170         return node;
171 }
172
173 /* Find the node which exactly matches the required data */
174 DLRBT_Node *BLI_dlrbTree_search_exact(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
175 {
176         DLRBT_Node *node = (tree) ? tree->root : NULL;
177         short found = 0;
178
179         /* check that there is a comparator to use */
180         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
181         if (cmp_cb == NULL)
182                 return NULL;
183
184         /* iteratively perform this search */
185         while (node && found == 0) {
186                 /* check if traverse further or not
187                  * NOTE: it is assumed that the values will be unit values only
188                  */
189                 switch (cmp_cb(node, search_data)) {
190                         case -1:  /* data less than node */
191                                 if (node->left)
192                                         node = node->left;
193                                 else
194                                         found = -1;
195                                 break;
196
197                         case 1:  /* data greater than node */
198                                 if (node->right)
199                                         node = node->right;
200                                 else
201                                         found = -1;
202                                 break;
203
204                         default:  /* data equals node */
205                                 found = 1;
206                                 break;
207                 }
208         }
209
210         /* return the exactly matching node */
211         return (found == 1) ? (node) : (NULL);
212 }
213
214 /* Find the node which occurs immediately before the best matching node */
215 DLRBT_Node *BLI_dlrbTree_search_prev(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
216 {
217         DLRBT_Node *node;
218
219         /* check that there is a comparator to use */
220         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
221         if (cmp_cb == NULL)
222                 return NULL;
223
224         /* get the node which best matches this description */
225         node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
226
227         if (node) {
228                 /* if the item we're searching for is greater than the node found, we've found the match */
229                 if (cmp_cb(node, search_data) > 0)
230                         return node;
231
232                 /* return the previous node otherwise */
233                 /* NOTE: what happens if there is no previous node? */
234                 return node->prev;
235         }
236
237         /* nothing matching was found */
238         return NULL;
239 }
240
241 /* Find the node which occurs immediately after the best matching node */
242 DLRBT_Node *BLI_dlrbTree_search_next(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
243 {
244         DLRBT_Node *node;
245
246         /* check that there is a comparator to use */
247         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
248         if (cmp_cb == NULL)
249                 return NULL;
250
251         /* get the node which best matches this description */
252         node = BLI_dlrbTree_search(tree, cmp_cb, search_data);
253
254         if (node) {
255                 /* if the item we're searching for is less than the node found, we've found the match */
256                 if (cmp_cb(node, search_data) < 0)
257                         return node;
258
259                 /* return the previous node otherwise */
260                 /* NOTE: what happens if there is no previous node? */
261                 return node->next;
262         }
263
264         /* nothing matching was found */
265         return NULL;
266 }
267
268
269 /* Check whether there is a node matching the requested node */
270 short BLI_dlrbTree_contains(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb, void *search_data)
271 {
272         /* check if an exact search throws up anything... */
273         return (BLI_dlrbTree_search_exact(tree, cmp_cb, search_data) != NULL);
274 }
275
276 /* *********************************************** */
277 /* Tree Relationships Utilities */
278
279 /* get the 'grandparent' - the parent of the parent - of the given node */
280 static DLRBT_Node *get_grandparent(DLRBT_Node *node)
281 {
282         if (node && node->parent)
283                 return node->parent->parent;
284         else
285                 return NULL;
286 }
287
288 /* get the sibling node (e.g. if node is left child of parent, return right child of parent) */
289 static DLRBT_Node *get_sibling(DLRBT_Node *node)
290 {
291         if (node && node->parent) {
292                 if (node == node->parent->left)
293                         return node->parent->right;
294                 else
295                         return node->parent->left;
296         }
297
298         /* sibling not found */
299         return NULL;
300 }
301
302 /* get the 'uncle' - the sibling of the parent - of the given node */
303 static DLRBT_Node *get_uncle(DLRBT_Node *node)
304 {
305         if (node)
306                 /* return the child of the grandparent which isn't the node's parent */
307                 return get_sibling(node->parent);
308
309         /* uncle not found */
310         return NULL;
311 }
312
313 /* *********************************************** */
314 /* Tree Rotation Utilities */
315
316 /* make right child of 'root' the new root */
317 static void rotate_left(DLRBT_Tree *tree, DLRBT_Node *root)
318 {
319         DLRBT_Node **root_slot, *pivot;
320
321         /* pivot is simply the root's right child, to become the root's parent */
322         pivot = root->right;
323         if (pivot == NULL)
324                 return;
325
326         if (root->parent) {
327                 if (root == root->parent->left)
328                         root_slot = &root->parent->left;
329                 else
330                         root_slot = &root->parent->right;
331         }
332         else
333                 root_slot = ((DLRBT_Node **)&tree->root);  /* &((DLRBT_Node *)tree->root); */
334
335         /* - pivot's left child becomes root's right child
336          * - root now becomes pivot's left child
337          */
338         root->right = pivot->left;
339         if (pivot->left) pivot->left->parent = root;
340
341         pivot->left = root;
342         pivot->parent = root->parent;
343         root->parent = pivot;
344
345         /* make the pivot the new root */
346         if (root_slot)
347                 *root_slot = pivot;
348 }
349
350 /* make the left child of the 'root' the new root */
351 static void rotate_right(DLRBT_Tree *tree, DLRBT_Node *root)
352 {
353         DLRBT_Node **root_slot, *pivot;
354
355         /* pivot is simply the root's left child, to become the root's parent */
356         pivot = root->left;
357         if (pivot == NULL)
358                 return;
359
360         if (root->parent) {
361                 if (root == root->parent->left)
362                         root_slot = &root->parent->left;
363                 else
364                         root_slot = &root->parent->right;
365         }
366         else
367                 root_slot = ((DLRBT_Node **)&tree->root);  /* &((DLRBT_Node *)tree->root); */
368
369         /* - pivot's right child becomes root's left child
370          * - root now becomes pivot's right child
371          */
372         root->left = pivot->right;
373         if (pivot->right) pivot->right->parent = root;
374
375         pivot->right = root;
376         pivot->parent = root->parent;
377         root->parent = pivot;
378
379         /* make the pivot the new root */
380         if (root_slot)
381                 *root_slot = pivot;
382 }
383
384 /* *********************************************** */
385 /* Post-Insertion Balancing  */
386
387 /* forward defines for insertion checks */
388 static void insert_check_1(DLRBT_Tree *tree, DLRBT_Node *node);
389 static void insert_check_2(DLRBT_Tree *tree, DLRBT_Node *node);
390 static void insert_check_3(DLRBT_Tree *tree, DLRBT_Node *node);
391
392 /* ----- */
393
394 /* W. 1) Root must be black (so that the 2nd-generation can have a black parent) */
395 static void insert_check_1(DLRBT_Tree *tree, DLRBT_Node *node)
396 {
397         if (node) {
398                 /* if this is the root, just ensure that it is black */
399                 if (node->parent == NULL)
400                         node->tree_col = DLRBT_BLACK;
401                 else
402                         insert_check_2(tree, node);
403         }
404 }
405
406 /* W. 2+3) Parent of node must be black, otherwise recolor and flush */
407 static void insert_check_2(DLRBT_Tree *tree, DLRBT_Node *node)
408 {
409         /* if the parent is not black, we need to change that... */
410         if (node && node->parent && node->parent->tree_col) {
411                 DLRBT_Node *unc = get_uncle(node);
412
413                 /* if uncle and parent are both red, need to change them to black and make
414                  * the parent black in order to satisfy the criteria of each node having the
415                  * same number of black nodes to its leaves
416                  */
417                 if (unc && unc->tree_col) {
418                         DLRBT_Node *gp = get_grandparent(node);
419
420                         /* make the n-1 generation nodes black */
421                         node->parent->tree_col = unc->tree_col = DLRBT_BLACK;
422
423                         /* - make the grandparent red, so that we maintain alternating red/black property
424                          *  (it must exist, so no need to check for NULL here),
425                          * - as the grandparent may now cause inconsistencies with the rest of the tree,
426                          *   we must flush up the tree and perform checks/re-balancing/re-painting, using the
427                          *   grandparent as the node of interest
428                          */
429                         gp->tree_col = DLRBT_RED;
430                         insert_check_1(tree, gp);
431                 }
432                 else {
433                         /* we've got an unbalanced branch going down the grandparent to the parent,
434                          * so need to perform some rotations to re-balance the tree
435                          */
436                         insert_check_3(tree, node);
437                 }
438         }
439 }
440
441 /* W. 4+5) Perform rotation on sub-tree containing the 'new' node, then do any  */
442 static void insert_check_3(DLRBT_Tree *tree, DLRBT_Node *node)
443 {
444         DLRBT_Node *gp = get_grandparent(node);
445
446         /* check that grandparent and node->parent exist (jut in case... really shouldn't happen on a good tree) */
447         if (node && node->parent && gp) {
448                 /* a left rotation will switch the roles of node and its parent, assuming that
449                  * the parent is the left child of the grandparent... otherwise, rotation direction
450                  * should be swapped
451                  */
452                 if ((node == node->parent->right) && (node->parent == gp->left)) {
453                         rotate_left(tree, node);
454                         node = node->left;
455                 }
456                 else if ((node == node->parent->left) && (node->parent == gp->right)) {
457                         rotate_right(tree, node);
458                         node = node->right;
459                 }
460
461                 /* fix old parent's color-tagging, and perform rotation on the old parent in the
462                  * opposite direction if needed for the current situation
463                  * NOTE: in the code above, node pointer is changed to point to the old parent
464                  */
465                 if (node) {
466                         /* get 'new' grandparent (i.e. grandparent for old-parent (node)) */
467                         gp = get_grandparent(node);
468
469                         /* modify the coloring of the grandparent and parent so that they still satisfy the constraints */
470                         node->parent->tree_col = DLRBT_BLACK;
471                         gp->tree_col = DLRBT_RED;
472
473                         /* if there are several nodes that all form a left chain, do a right rotation to correct this
474                          * (or a rotation in the opposite direction if they all form a right chain)
475                          */
476                         if ((node == node->parent->left) && (node->parent == gp->left))
477                                 rotate_right(tree, gp);
478                         else //if ((node == node->parent->right) && (node->parent == gp->right))
479                                 rotate_left(tree, gp);
480                 }
481         }
482 }
483
484 /* ----- */
485
486 /* Balance the tree after the given element has been added to it
487  * (using custom code, in the Binary Tree way).
488  */
489 void BLI_dlrbTree_insert(DLRBT_Tree *tree, DLRBT_Node *node)
490 {
491         /* sanity checks */
492         if ((tree == NULL) || (node == NULL))
493                 return;
494
495         /* firstly, the node we just added should be red by default */
496         node->tree_col = DLRBT_RED;
497
498         /* start from case 1, an trek through the tail-recursive insertion checks */
499         insert_check_1(tree, node);
500 }
501
502 /* ----- */
503
504 /* Add the given data to the tree, and return the node added */
505 /* NOTE: for duplicates, the update_cb is called (if available), and the existing node is returned */
506 DLRBT_Node *BLI_dlrbTree_add(DLRBT_Tree *tree, DLRBT_Comparator_FP cmp_cb,
507                              DLRBT_NAlloc_FP new_cb, DLRBT_NUpdate_FP update_cb, void *data)
508 {
509         DLRBT_Node *parNode, *node = NULL;
510         short new_node = 0;
511
512         /* sanity checks */
513         if (tree == NULL)
514                 return NULL;
515
516         /* TODO: if no comparator is supplied, try using the one supplied with the tree... */
517         if (cmp_cb == NULL)
518                 return NULL;
519         /* TODO: if no allocator is supplied, try using the one supplied with the tree... */
520         if (new_cb == NULL)
521                 return NULL;
522         /* TODO: if no updater is supplied, try using the one supplied with the tree... */
523
524         /* try to find the nearest node to this one */
525         parNode = BLI_dlrbTree_search(tree, cmp_cb, data);
526
527         /* add new node to the BST in the 'standard way' as appropriate
528          * NOTE: we do not support duplicates in our tree...
529          */
530         if (parNode) {
531                 /* check how this new node compares with the existing ones
532                  * NOTE: it is assumed that the values will be unit values only
533                  */
534                 switch (cmp_cb(parNode, data)) {
535                         case -1:  /* add new node as left child */
536                         {
537                                 node = new_cb(data);
538                                 new_node = 1;
539
540                                 parNode->left = node;
541                                 node->parent = parNode;
542                                 break;
543                         }
544                         case 1:  /* add new node as right child */
545                         {
546                                 node = new_cb(data);
547                                 new_node = 1;
548
549                                 parNode->right = node;
550                                 node->parent = parNode;
551                                 break;
552                         }
553                         default:  /* update the duplicate node as appropriate */
554                         {
555                                 if (update_cb)
556                                         update_cb(parNode, data);
557                                 break;
558                         }
559                 }
560         }
561         else {
562                 /* no nodes in the tree yet... add a new node as the root */
563                 node = new_cb(data);
564                 new_node = 1;
565
566                 tree->root = node;
567         }
568
569         /* if a new node was added, it should be tagged as red, and then balanced as appropriate */
570         if (new_node) {
571                 /* tag this new node as being 'red' */
572                 node->tree_col = DLRBT_RED;
573
574                 /* perform BST balancing steps:
575                  * start from case 1, an trek through the tail-recursive insertion checks
576                  */
577                 insert_check_1(tree, node);
578         }
579
580         /* return the node added */
581         return node;
582 }
583
584 /* *********************************************** */
585 /* Remove */
586
587 /* TODO: this hasn't been coded yet, since this functionality was not needed by the author */
588
589 /* *********************************************** */