source/blender/blenlib/intern/DLRB_tree.c

   1 /**
   2  * $Id$
   3  *
   4  * ***** BEGIN GPL LICENSE BLOCK *****
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License
   8  * as published by the Free Software Foundation; either version 2
   9  * of the License, or (at your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14  * GNU General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, write to the Free Software Foundation,
  18  * Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
  19  *
  20  * The Original Code is Copyright (C) 2009 Blender Foundation, Joshua Leung
  21  * All rights reserved.
  22  *
  23  * Contributor(s): Joshua Leung (original author)
  24  *
  25  * ***** END GPL LICENSE BLOCK *****
  26  */
  27
  28 #include <stdlib.h>
  29 #include <stdio.h>
  30 #include <string.h>
  31
  32 #include "MEM_guardedalloc.h"
  33
  34 #include "BLI_blenlib.h"
  35 #include "BLI_dlrbTree.h"
  36
  37 #include "DNA_listBase.h"
  38
  39 /* *********************************************** */
  40 /* Tree API */
  41
  42 /* Create a new tree, and initialise as necessary */
  43 DLRBT_Tree *BLI_dlrbTree_new (void)
  44 {
  45         /* just allocate for now */
  46         return MEM_callocN(sizeof(DLRBT_Tree), "DLRBT_Tree");
  47 }
  48
  49 /* Just zero out the pointers used */
  50 void BLI_dlrbTree_init (DLRBT_Tree *tree)
  51 {
  52         if (tree == NULL)
  53                 return;
  54
  55         tree->first= tree->last= tree->root= NULL;
  56 }
  57
  58 /* Helper for traversing tree and freeing sub-nodes */
  59 static void recursive_tree_free_nodes (DLRBT_Node *node)
  60 {
  61         /* sanity check */
  62         if (node == NULL)
  63                 return;
  64
  65         /* free child nodes + subtrees */
  66         recursive_tree_free_nodes(node->left);
  67         recursive_tree_free_nodes(node->right);
  68
  69         /* free self */
  70         MEM_freeN(node);
  71 }
  72
  73 /* Free the given tree's data but not the tree itself */
  74 void BLI_dlrbTree_free (DLRBT_Tree *tree)
  75 {
  76         if (tree == NULL)
  77                 return;
  78
  79         /* if the list-base stuff is set, just use that (and assume its set),
  80          * otherwise, we'll need to traverse the tree...
  81          */
  82         if (tree->first) {
  83                 /* free list */
  84                 BLI_freelistN((ListBase *)tree);
  85         }
  86         else {
  87                 /* traverse tree, freeing sub-nodes */
  88                 recursive_tree_free_nodes(tree->root);
  89         }
  90
  91         /* clear pointers */
  92         tree->first= tree->last= tree->root= NULL;
  93 }
  94
  95 /* ------- */
  96
  97 /* Helper function - used for traversing down the tree from the root to add nodes in order */
  98 static void linkedlist_sync_add_node (DLRBT_Tree *tree, DLRBT_Node *node)
  99 {
 100         /* sanity checks */
 101         if ((tree == NULL) || (node == NULL))
 102                 return;
 103
 104         /* add left-node (and its subtree) */
 105         linkedlist_sync_add_node(tree, node->left);
 106
 107         /* now add self
 108          *      - must remove detach from other links first
 109          *        (for now, only clear own pointers)
 110          */
 111         node->prev= node->next= NULL;
 112         BLI_addtail((ListBase *)tree, (Link *)node);
 113
 114         /* finally, add right node (and its subtree) */
 115         linkedlist_sync_add_node(tree, node->right);
 116 }
 117
 118 /* Make sure the tree's Double-Linked list representation is valid */
 119 void BLI_dlrbTree_linkedlist_sync (DLRBT_Tree *tree)
 120 {
 121         /* sanity checks */
 122         if (tree == NULL)
 123                 return;
 124
 125         /* clear list-base pointers so that the new list can be added properly */
 126         tree->first= tree->last= NULL;
 127
 128         /* start adding items from the root */
 129         linkedlist_sync_add_node(tree, tree->root);
 130 }
 131
 132 /* *********************************************** */
 133 /* Tree Relationships Utilities */
 134
 135 /* get the 'grandparent' - the parent of the parent - of the given node */
 136 static DLRBT_Node *get_grandparent (DLRBT_Node *node)
 137 {
 138         if (node && node->parent)
 139                 return node->parent->parent;
 140         else
 141                 return NULL;
 142 }
 143
 144 /* get the 'uncle' - the sibling of the parent - of the given node */
 145 static DLRBT_Node *get_uncle (DLRBT_Node *node)
 146 {
 147         DLRBT_Node *gpn= get_grandparent(node);
 148
 149         /* return the child of the grandparent which isn't the node's parent */
 150         if (gpn) {
 151                 if (gpn->left == node->parent)
 152                         return gpn->right;
 153                 else
 154                         return gpn->left;
 155         }
 156
 157         /* not found */
 158         return NULL;
 159 }
 160
 161 /* *********************************************** */
 162 /* Tree Rotation Utilities */
 163
 164 static void rotate_left (DLRBT_Tree *tree, DLRBT_Node *root)
 165 {
 166         DLRBT_Node **root_slot, *pivot;
 167
 168         /* pivot is simply the root's right child, to become the root's parent */
 169         pivot= root->right;
 170         if (pivot == NULL)
 171                 return;
 172
 173         if (root->parent) {
 174                 if (root == root->parent->left)
 175                         root_slot= &root->parent->left;
 176                 else
 177                         root_slot= &root->parent->right;
 178         }
 179         else
 180                 root_slot= ((DLRBT_Node**)&tree->root);//&((DLRBT_Node*)tree->root);
 181
 182         /* - pivot's left child becomes root's right child
 183          * - root now becomes pivot's left child
 184          */
 185         root->right= pivot->left;
 186         if (pivot->left) pivot->left->parent= root;
 187
 188         pivot->left= root;
 189         pivot->parent= root->parent;
 190         root->parent= pivot;
 191
 192         /* make the pivot the new root */
 193         if (root_slot)
 194                 *root_slot= pivot;
 195 }
 196
 197 static void rotate_right (DLRBT_Tree *tree, DLRBT_Node *root)
 198 {
 199         DLRBT_Node **root_slot, *pivot;
 200
 201         /* pivot is simply the root's left child, to become the root's parent */
 202         pivot= root->left;
 203         if (pivot == NULL)
 204                 return;
 205
 206         if (root->parent) {
 207                 if (root == root->parent->left)
 208                         root_slot= &root->parent->left;
 209                 else
 210                         root_slot= &root->parent->right;
 211         }
 212         else
 213                 root_slot= ((DLRBT_Node**)&tree->root);//&((DLRBT_Node*)tree->root);
 214
 215         /* - pivot's right child becomes root's left child
 216          * - root now becomes pivot's right child
 217          */
 218         root->left= pivot->right;
 219         if (pivot->right) pivot->right->parent= root;
 220
 221         pivot->right= root;
 222         pivot->parent= root->parent;
 223         root->parent= pivot;
 224
 225         /* make the pivot the new root */
 226         if (root_slot)
 227                 *root_slot= pivot;
 228 }
 229
 230 /* *********************************************** */
 231 /* Post-Insertion Balancing  */
 232
 233 /* forward defines for insertion checks */
 234 static void insert_check_1(DLRBT_Tree *tree, DLRBT_Node *node);
 235 static void insert_check_2(DLRBT_Tree *tree, DLRBT_Node *node);
 236 static void insert_check_3(DLRBT_Tree *tree, DLRBT_Node *node);
 237
 238 /* ----- */
 239
 240 /* W. 1) Root must be black (so that the 2nd-generation can have a black parent) */
 241 static void insert_check_1 (DLRBT_Tree *tree, DLRBT_Node *node)
 242 {
 243         if (node) {
 244                 /* if this is the root, just ensure that it is black */
 245                 if (node->parent == NULL)
 246                         node->tree_col= DLRBT_BLACK;
 247                 else
 248                         insert_check_2(tree, node);
 249         }
 250 }
 251
 252 /* W. 2+3) Parent of node must be black, otherwise recolor and flush */
 253 static void insert_check_2 (DLRBT_Tree *tree, DLRBT_Node *node)
 254 {
 255         /* if the parent is not black, we need to change that... */
 256         if (node && node->parent && node->parent->tree_col) {
 257                 DLRBT_Node *unc= get_uncle(node);
 258
 259                 /* if uncle and parent are both red, need to change them to black and make
 260                  * the parent black in order to satisfy the criteria of each node having the
 261                  * same number of black nodes to its leaves
 262                  */
 263                 if (unc && unc->tree_col) {
 264                         DLRBT_Node *gp= get_grandparent(node);
 265
 266                         /* make the n-1 generation nodes black */
 267                         node->parent->tree_col= unc->tree_col= DLRBT_BLACK;
 268
 269                         /* - make the grandparent red, so that we maintain alternating red/black property
 270                          *  (it must exist, so no need to check for NULL here),
 271                          * - as the grandparent may now cause inconsistencies with the rest of the tree,
 272                          *   we must flush up the tree and perform checks/rebalancing/repainting, using the
 273                          *      grandparent as the node of interest
 274                          */
 275                         gp->tree_col= DLRBT_RED;
 276                         insert_check_1(tree, gp);
 277                 }
 278                 else {
 279                         /* we've got an unbalanced branch going down the grandparent to the parent,
 280                          * so need to perform some rotations to re-balance the tree
 281                          */
 282                         insert_check_3(tree, node);
 283                 }
 284         }
 285 }
 286
 287 /* W. 4+5) Perform rotation on sub-tree containing the 'new' node, then do any  */
 288 static void insert_check_3 (DLRBT_Tree *tree, DLRBT_Node *node)
 289 {
 290         DLRBT_Node *gp= get_grandparent(node);
 291
 292         /* check that grandparent and node->parent exist (jut in case... really shouldn't happen on a good tree) */
 293         if (node && node->parent && gp) {
 294                 /* a left rotation will switch the roles of node and its parent, assuming that
 295                  * the parent is the left child of the grandparent... otherwise, rotation direction
 296                  * should be swapped
 297                  */
 298                 if ((node == node->parent->right) && (node->parent == gp->left)) {
 299                         rotate_left(tree, node);
 300                         node= node->left;
 301                 }
 302                 else if ((node == node->parent->left) && (node->parent == gp->right)) {
 303                         rotate_right(tree, node);
 304                         node= node->right;
 305                 }
 306
 307                 /* fix old parent's color-tagging, and perform rotation on the old parent in the
 308                  * opposite direction if needed for the current situation
 309                  * NOTE: in the code above, node pointer is changed to point to the old parent
 310                  */
 311                 if (node) {
 312                         /* get 'new' grandparent (i.e. grandparent for old-parent (node)) */
 313                         gp= get_grandparent(node);
 314
 315                         /* modify the coloring of the grandparent and parent so that they still satisfy the constraints */
 316                         node->parent->tree_col= DLRBT_BLACK;
 317                         gp->tree_col= DLRBT_RED;
 318
 319                         /* if there are several nodes that all form a left chain, do a right rotation to correct this
 320                          * (or a rotation in the opposite direction if they all form a right chain)
 321                          */
 322                         if ((node == node->parent->left) && (node->parent == gp->left))
 323                                 rotate_right(tree, gp);
 324                         else //if ((node == node->parent->right) && (node->parent == gp->right))
 325                                 rotate_left(tree, gp);
 326                 }
 327         }
 328 }
 329
 330 /* ----- */
 331
 332 /* Balance the tree after the given element has been added to it
 333  * (using custom code, in the Binary Tree way).
 334  */
 335 void BLI_dlrbTree_insert(DLRBT_Tree *tree, DLRBT_Node *node)
 336 {
 337         /* sanity checks */
 338         if ((tree == NULL) || (node == NULL))
 339                 return;
 340
 341         /* firstly, the node we just added should be red by default */
 342         node->tree_col= DLRBT_RED;
 343
 344         /* start from case 1, an trek through the tail-recursive insertion checks */
 345         insert_check_1(tree, node);
 346 }
 347
 348 /* *********************************************** */
 349 /* Remove */
 350
 351 // TODO: this hasn't been coded yet, since this functionality was not needed by the author
 352
 353 /* *********************************************** */