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