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