source/blender/blenlib/intern/BLI_kdtree.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) 2001-2002 by NaN Holding BV.
  19  * All rights reserved.
  20  *
  21  * The Original Code is: none of this file.
  22  *
  23  * Contributor(s): Janne Karhu
  24  *                 Brecht Van Lommel
  25  *
  26  * ***** END GPL LICENSE BLOCK *****
  27  */
  28
  29 /** \file blender/blenlib/intern/BLI_kdtree.c
  30  *  \ingroup bli
  31  */
  32
  33
  34
  35 #include "MEM_guardedalloc.h"
  36
  37 #include "BLI_math.h"
  38 #include "BLI_kdtree.h"
  39
  40 #ifndef SWAP
  41 #  define SWAP(type, a, b) { type sw_ap; sw_ap = (a); (a) = (b); (b) = sw_ap; }
  42 #endif
  43
  44 typedef struct KDTreeNode {
  45         struct KDTreeNode *left, *right;
  46         float co[3], nor[3];
  47         int index;
  48         short d;
  49 } KDTreeNode;
  50
  51 struct KDTree {
  52         KDTreeNode *nodes;
  53         int totnode;
  54         KDTreeNode *root;
  55 };
  56
  57 KDTree *BLI_kdtree_new(int maxsize)
  58 {
  59         KDTree *tree;
  60
  61         tree = MEM_callocN(sizeof(KDTree), "KDTree");
  62         tree->nodes = MEM_callocN(sizeof(KDTreeNode) * maxsize, "KDTreeNode");
  63         tree->totnode = 0;
  64
  65         return tree;
  66 }
  67
  68 void BLI_kdtree_free(KDTree *tree)
  69 {
  70         if (tree) {
  71                 MEM_freeN(tree->nodes);
  72                 MEM_freeN(tree);
  73         }
  74 }
  75
  76 void BLI_kdtree_insert(KDTree *tree, int index, float *co, float *nor)
  77 {
  78         KDTreeNode *node = &tree->nodes[tree->totnode++];
  79
  80         node->index = index;
  81         copy_v3_v3(node->co, co);
  82         if (nor) copy_v3_v3(node->nor, nor);
  83 }
  84
  85 static KDTreeNode *kdtree_balance(KDTreeNode *nodes, int totnode, int axis)
  86 {
  87         KDTreeNode *node;
  88         float co;
  89         int left, right, median, i, j;
  90
  91         if (totnode <= 0)
  92                 return NULL;
  93         else if (totnode == 1)
  94                 return nodes;
  95
  96         /* quicksort style sorting around median */
  97         left = 0;
  98         right = totnode - 1;
  99         median = totnode / 2;
 100
 101         while (right > left) {
 102                 co = nodes[right].co[axis];
 103                 i = left - 1;
 104                 j = right;
 105
 106                 while (1) {
 107                         while (nodes[++i].co[axis] < co) ;
 108                         while (nodes[--j].co[axis] > co && j > left) ;
 109
 110                         if (i >= j) break;
 111                         SWAP(KDTreeNode, nodes[i], nodes[j]);
 112                 }
 113
 114                 SWAP(KDTreeNode, nodes[i], nodes[right]);
 115                 if (i >= median)
 116                         right = i - 1;
 117                 if (i <= median)
 118                         left = i + 1;
 119         }
 120
 121         /* set node and sort subnodes */
 122         node = &nodes[median];
 123         node->d = axis;
 124         node->left = kdtree_balance(nodes, median, (axis + 1) % 3);
 125         node->right = kdtree_balance(nodes + median + 1, (totnode - (median + 1)), (axis + 1) % 3);
 126
 127         return node;
 128 }
 129
 130 void BLI_kdtree_balance(KDTree *tree)
 131 {
 132         tree->root = kdtree_balance(tree->nodes, tree->totnode, 0);
 133 }
 134
 135 static float squared_distance(const float v2[3], const float v1[3], float *UNUSED(n1), float *n2)
 136 {
 137         float d[3], dist;
 138
 139         d[0] = v2[0] - v1[0];
 140         d[1] = v2[1] - v1[1];
 141         d[2] = v2[2] - v1[2];
 142
 143         dist = dot_v3v3(d, d);
 144
 145         //if (n1 && n2 && (dot_v3v3(n1, n2) < 0.0f))
 146
 147         /* can someone explain why this is done?*/
 148         if (n2 && (dot_v3v3(d, n2) < 0.0f)) {
 149                 dist *= 10.0f;
 150         }
 151
 152         return dist;
 153 }
 154
 155 int BLI_kdtree_find_nearest(KDTree *tree, float *co, float *nor, KDTreeNearest *nearest)
 156 {
 157         KDTreeNode *root, *node, *min_node;
 158         KDTreeNode **stack, *defaultstack[100];
 159         float min_dist, cur_dist;
 160         int totstack, cur = 0;
 161
 162         if (!tree->root)
 163                 return -1;
 164
 165         stack = defaultstack;
 166         totstack = 100;
 167
 168         root = tree->root;
 169         min_node = root;
 170         min_dist = squared_distance(root->co, co, root->nor, nor);
 171
 172         if (co[root->d] < root->co[root->d]) {
 173                 if (root->right)
 174                         stack[cur++] = root->right;
 175                 if (root->left)
 176                         stack[cur++] = root->left;
 177         }
 178         else {
 179                 if (root->left)
 180                         stack[cur++] = root->left;
 181                 if (root->right)
 182                         stack[cur++] = root->right;
 183         }
 184
 185         while (cur--) {
 186                 node = stack[cur];
 187
 188                 cur_dist = node->co[node->d] - co[node->d];
 189
 190                 if (cur_dist < 0.0f) {
 191                         cur_dist = -cur_dist * cur_dist;
 192
 193                         if (-cur_dist < min_dist) {
 194                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
 195                                 if (cur_dist < min_dist) {
 196                                         min_dist = cur_dist;
 197                                         min_node = node;
 198                                 }
 199                                 if (node->left)
 200                                         stack[cur++] = node->left;
 201                         }
 202                         if (node->right)
 203                                 stack[cur++] = node->right;
 204                 }
 205                 else {
 206                         cur_dist = cur_dist * cur_dist;
 207
 208                         if (cur_dist < min_dist) {
 209                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
 210                                 if (cur_dist < min_dist) {
 211                                         min_dist = cur_dist;
 212                                         min_node = node;
 213                                 }
 214                                 if (node->right)
 215                                         stack[cur++] = node->right;
 216                         }
 217                         if (node->left)
 218                                 stack[cur++] = node->left;
 219                 }
 220                 if (cur + 3 > totstack) {
 221                         KDTreeNode **temp = MEM_callocN((totstack + 100) * sizeof(KDTreeNode *), "psys_treestack");
 222                         memcpy(temp, stack, totstack * sizeof(KDTreeNode *));
 223                         if (stack != defaultstack)
 224                                 MEM_freeN(stack);
 225                         stack = temp;
 226                         totstack += 100;
 227                 }
 228         }
 229
 230         if (nearest) {
 231                 nearest->index = min_node->index;
 232                 nearest->dist = sqrt(min_dist);
 233                 copy_v3_v3(nearest->co, min_node->co);
 234         }
 235
 236         if (stack != defaultstack)
 237                 MEM_freeN(stack);
 238
 239         return min_node->index;
 240 }
 241
 242 static void add_nearest(KDTreeNearest *ptn, int *found, int n, int index, float dist, float *co)
 243 {
 244         int i;
 245
 246         if (*found < n) (*found)++;
 247
 248         for (i = *found - 1; i > 0; i--) {
 249                 if (dist >= ptn[i - 1].dist)
 250                         break;
 251                 else
 252                         ptn[i] = ptn[i - 1];
 253         }
 254
 255         ptn[i].index = index;
 256         ptn[i].dist = dist;
 257         copy_v3_v3(ptn[i].co, co);
 258 }
 259
 260 /* finds the nearest n entries in tree to specified coordinates */
 261 int BLI_kdtree_find_n_nearest(KDTree *tree, int n, float *co, float *nor, KDTreeNearest *nearest)
 262 {
 263         KDTreeNode *root, *node = NULL;
 264         KDTreeNode **stack, *defaultstack[100];
 265         float cur_dist;
 266         int i, totstack, cur = 0, found = 0;
 267
 268         if (!tree->root)
 269                 return 0;
 270
 271         stack = defaultstack;
 272         totstack = 100;
 273
 274         root = tree->root;
 275
 276         cur_dist = squared_distance(root->co, co, root->nor, nor);
 277         add_nearest(nearest, &found, n, root->index, cur_dist, root->co);
 278
 279         if (co[root->d] < root->co[root->d]) {
 280                 if (root->right)
 281                         stack[cur++] = root->right;
 282                 if (root->left)
 283                         stack[cur++] = root->left;
 284         }
 285         else {
 286                 if (root->left)
 287                         stack[cur++] = root->left;
 288                 if (root->right)
 289                         stack[cur++] = root->right;
 290         }
 291
 292         while (cur--) {
 293                 node = stack[cur];
 294
 295                 cur_dist = node->co[node->d] - co[node->d];
 296
 297                 if (cur_dist < 0.0f) {
 298                         cur_dist = -cur_dist * cur_dist;
 299
 300                         if (found < n || -cur_dist < nearest[found - 1].dist) {
 301                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
 302
 303                                 if (found < n || cur_dist < nearest[found - 1].dist)
 304                                         add_nearest(nearest, &found, n, node->index, cur_dist, node->co);
 305
 306                                 if (node->left)
 307                                         stack[cur++] = node->left;
 308                         }
 309                         if (node->right)
 310                                 stack[cur++] = node->right;
 311                 }
 312                 else {
 313                         cur_dist = cur_dist * cur_dist;
 314
 315                         if (found < n || cur_dist < nearest[found - 1].dist) {
 316                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
 317                                 if (found < n || cur_dist < nearest[found - 1].dist)
 318                                         add_nearest(nearest, &found, n, node->index, cur_dist, node->co);
 319
 320                                 if (node->right)
 321                                         stack[cur++] = node->right;
 322                         }
 323                         if (node->left)
 324                                 stack[cur++] = node->left;
 325                 }
 326                 if (cur + 3 > totstack) {
 327                         KDTreeNode **temp = MEM_callocN((totstack + 100) * sizeof(KDTreeNode *), "psys_treestack");
 328                         memcpy(temp, stack, totstack * sizeof(KDTreeNode *));
 329                         if (stack != defaultstack)
 330                                 MEM_freeN(stack);
 331                         stack = temp;
 332                         totstack += 100;
 333                 }
 334         }
 335
 336         for (i = 0; i < found; i++)
 337                 nearest[i].dist = sqrt(nearest[i].dist);
 338
 339         if (stack != defaultstack)
 340                 MEM_freeN(stack);
 341
 342         return found;
 343 }
 344
 345 static int range_compare(const void *a, const void *b)
 346 {
 347         const KDTreeNearest *kda = a;
 348         const KDTreeNearest *kdb = b;
 349
 350         if (kda->dist < kdb->dist)
 351                 return -1;
 352         else if (kda->dist > kdb->dist)
 353                 return 1;
 354         else
 355                 return 0;
 356 }
 357 static void add_in_range(KDTreeNearest **ptn, int found, int *totfoundstack, int index, float dist, float *co)
 358 {
 359         KDTreeNearest *to;
 360
 361         if (found + 1 > *totfoundstack) {
 362                 KDTreeNearest *temp = MEM_callocN((*totfoundstack + 50) * sizeof(KDTreeNode), "psys_treefoundstack");
 363                 memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
 364                 if (*ptn)
 365                         MEM_freeN(*ptn);
 366                 *ptn = temp;
 367                 *totfoundstack += 50;
 368         }
 369
 370         to = (*ptn) + found;
 371
 372         to->index = index;
 373         to->dist = sqrt(dist);
 374         copy_v3_v3(to->co, co);
 375 }
 376 int BLI_kdtree_range_search(KDTree *tree, float range, float *co, float *nor, KDTreeNearest **nearest)
 377 {
 378         KDTreeNode *root, *node = NULL;
 379         KDTreeNode **stack, *defaultstack[100];
 380         KDTreeNearest *foundstack = NULL;
 381         float range2 = range * range, dist2;
 382         int totstack, cur = 0, found = 0, totfoundstack = 0;
 383
 384         if (!tree || !tree->root)
 385                 return 0;
 386
 387         stack = defaultstack;
 388         totstack = 100;
 389
 390         root = tree->root;
 391
 392         if (co[root->d] + range < root->co[root->d]) {
 393                 if (root->left)
 394                         stack[cur++] = root->left;
 395         }
 396         else if (co[root->d] - range > root->co[root->d]) {
 397                 if (root->right)
 398                         stack[cur++] = root->right;
 399         }
 400         else {
 401                 dist2 = squared_distance(root->co, co, root->nor, nor);
 402                 if (dist2  <= range2)
 403                         add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);
 404
 405                 if (root->left)
 406                         stack[cur++] = root->left;
 407                 if (root->right)
 408                         stack[cur++] = root->right;
 409         }
 410
 411         while (cur--) {
 412                 node = stack[cur];
 413
 414                 if (co[node->d] + range < node->co[node->d]) {
 415                         if (node->left)
 416                                 stack[cur++] = node->left;
 417                 }
 418                 else if (co[node->d] - range > node->co[node->d]) {
 419                         if (node->right)
 420                                 stack[cur++] = node->right;
 421                 }
 422                 else {
 423                         dist2 = squared_distance(node->co, co, node->nor, nor);
 424                         if (dist2 <= range2)
 425                                 add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);
 426
 427                         if (node->left)
 428                                 stack[cur++] = node->left;
 429                         if (node->right)
 430                                 stack[cur++] = node->right;
 431                 }
 432
 433                 if (cur + 3 > totstack) {
 434                         KDTreeNode **temp = MEM_callocN((totstack + 100) * sizeof(KDTreeNode *), "psys_treestack");
 435                         memcpy(temp, stack, totstack * sizeof(KDTreeNode *));
 436                         if (stack != defaultstack)
 437                                 MEM_freeN(stack);
 438                         stack = temp;
 439                         totstack += 100;
 440                 }
 441         }
 442
 443         if (stack != defaultstack)
 444                 MEM_freeN(stack);
 445
 446         if (found)
 447                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
 448
 449         *nearest = foundstack;
 450
 451         return found;
 452 }