c8ee0c11fc705ba5ad79d667d8e597aa3e29ef47
[blender-staging.git] / 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  * Contributor(s): Janne Karhu
19  *                 Brecht Van Lommel
20  *
21  * ***** END GPL LICENSE BLOCK *****
22  */
23
24 /** \file blender/blenlib/intern/BLI_kdtree.c
25  *  \ingroup bli
26  */
27
28 #include "MEM_guardedalloc.h"
29
30 #include "BLI_math.h"
31 #include "BLI_kdtree.h"
32 #include "BLI_utildefines.h"
33 #include "BLI_strict_flags.h"
34
35
36 typedef struct KDTreeNode {
37         struct KDTreeNode *left, *right;
38         float co[3], nor[3];
39         int index;
40         unsigned int d;  /* range is only (0-2) */
41 } KDTreeNode;
42
43 struct KDTree {
44         KDTreeNode *nodes;
45         unsigned int totnode;
46         KDTreeNode *root;
47 #ifdef DEBUG
48         bool is_balanced;  /* ensure we call balance first */
49         unsigned int maxsize;   /* max size of the tree */
50 #endif
51 };
52
53 #define KD_STACK_INIT 100      /* initial size for array (on the stack) */
54 #define KD_NEAR_ALLOC_INC 100  /* alloc increment for collecting nearest */
55 #define KD_FOUND_ALLOC_INC 50  /* alloc increment for collecting nearest */
56
57 /**
58  * Creates or free a kdtree
59  */
60 KDTree *BLI_kdtree_new(unsigned int maxsize)
61 {
62         KDTree *tree;
63
64         tree = MEM_mallocN(sizeof(KDTree), "KDTree");
65         tree->nodes = MEM_mallocN(sizeof(KDTreeNode) * maxsize, "KDTreeNode");
66         tree->totnode = 0;
67         tree->root = NULL;
68
69 #ifdef DEBUG
70         tree->is_balanced = false;
71         tree->maxsize = maxsize;
72 #endif
73
74         return tree;
75 }
76
77 void BLI_kdtree_free(KDTree *tree)
78 {
79         if (tree) {
80                 MEM_freeN(tree->nodes);
81                 MEM_freeN(tree);
82         }
83 }
84
85 /**
86  * Construction: first insert points, then call balance. Normal is optional.
87  */
88 void BLI_kdtree_insert(KDTree *tree, int index, const float co[3], const float nor[3])
89 {
90         KDTreeNode *node = &tree->nodes[tree->totnode++];
91
92 #ifdef DEBUG
93         BLI_assert(tree->totnode <= tree->maxsize);
94 #endif
95
96         /* note, array isn't calloc'd,
97          * need to initialize all struct members */
98
99         node->left = node->right = NULL;
100         copy_v3_v3(node->co, co);
101         if (nor)
102                 copy_v3_v3(node->nor, nor);
103         else
104                 zero_v3(node->nor);
105
106         node->index = index;
107         node->d = 0;
108
109 #ifdef DEBUG
110         tree->is_balanced = false;
111 #endif
112 }
113
114 static KDTreeNode *kdtree_balance(KDTreeNode *nodes, unsigned int totnode, unsigned int axis)
115 {
116         KDTreeNode *node;
117         float co;
118         unsigned int left, right, median, i, j;
119
120         if (totnode <= 0)
121                 return NULL;
122         else if (totnode == 1)
123                 return nodes;
124         
125         /* quicksort style sorting around median */
126         left = 0;
127         right = totnode - 1;
128         median = totnode / 2;
129
130         while (right > left) {
131                 co = nodes[right].co[axis];
132                 i = left - 1;
133                 j = right;
134
135                 while (1) {
136                         while (nodes[++i].co[axis] < co) ;
137                         while (nodes[--j].co[axis] > co && j > left) ;
138
139                         if (i >= j)
140                                 break;
141
142                         SWAP(KDTreeNode, nodes[i], nodes[j]);
143                 }
144
145                 SWAP(KDTreeNode, nodes[i], nodes[right]);
146                 if (i >= median)
147                         right = i - 1;
148                 if (i <= median)
149                         left = i + 1;
150         }
151
152         /* set node and sort subnodes */
153         node = &nodes[median];
154         node->d = axis;
155         node->left = kdtree_balance(nodes, median, (axis + 1) % 3);
156         node->right = kdtree_balance(nodes + median + 1, (totnode - (median + 1)), (axis + 1) % 3);
157
158         return node;
159 }
160
161 void BLI_kdtree_balance(KDTree *tree)
162 {
163         tree->root = kdtree_balance(tree->nodes, tree->totnode, 0);
164
165 #ifdef DEBUG
166         tree->is_balanced = true;
167 #endif
168 }
169
170 static float squared_distance(const float v2[3], const float v1[3], const float UNUSED(n1[3]), const float n2[3])
171 {
172         float d[3], dist;
173
174         d[0] = v2[0] - v1[0];
175         d[1] = v2[1] - v1[1];
176         d[2] = v2[2] - v1[2];
177
178         dist = dot_v3v3(d, d);
179
180         //if (n1 && n2 && (dot_v3v3(n1, n2) < 0.0f))
181
182         /* can someone explain why this is done?*/
183         if (n2 && (dot_v3v3(d, n2) < 0.0f)) {
184                 dist *= 10.0f;
185         }
186
187         return dist;
188 }
189
190 static KDTreeNode **realloc_nodes(KDTreeNode **stack, unsigned int *totstack, const bool is_alloc)
191 {
192         KDTreeNode **stack_new = MEM_mallocN((*totstack + KD_NEAR_ALLOC_INC) * sizeof(KDTreeNode *), "KDTree.treestack");
193         memcpy(stack_new, stack, *totstack * sizeof(KDTreeNode *));
194         // memset(stack_new + *totstack, 0, sizeof(KDTreeNode *) * KD_NEAR_ALLOC_INC);
195         if (is_alloc)
196                 MEM_freeN(stack);
197         *totstack += KD_NEAR_ALLOC_INC;
198         return stack_new;
199 }
200
201 /**
202  * Find nearest returns index, and -1 if no node is found.
203  */
204 int BLI_kdtree_find_nearest(KDTree *tree, const float co[3], const float nor[3],
205                             KDTreeNearest *r_nearest)
206 {
207         KDTreeNode *root, *node, *min_node;
208         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
209         float min_dist, cur_dist;
210         unsigned int totstack, cur = 0;
211
212 #ifdef DEBUG
213         BLI_assert(tree->is_balanced == true);
214 #endif
215
216         if (!tree->root)
217                 return -1;
218
219         stack = defaultstack;
220         totstack = KD_STACK_INIT;
221
222         root = tree->root;
223         min_node = root;
224         min_dist = squared_distance(root->co, co, root->nor, nor);
225
226         if (co[root->d] < root->co[root->d]) {
227                 if (root->right)
228                         stack[cur++] = root->right;
229                 if (root->left)
230                         stack[cur++] = root->left;
231         }
232         else {
233                 if (root->left)
234                         stack[cur++] = root->left;
235                 if (root->right)
236                         stack[cur++] = root->right;
237         }
238         
239         while (cur--) {
240                 node = stack[cur];
241
242                 cur_dist = node->co[node->d] - co[node->d];
243
244                 if (cur_dist < 0.0f) {
245                         cur_dist = -cur_dist * cur_dist;
246
247                         if (-cur_dist < min_dist) {
248                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
249                                 if (cur_dist < min_dist) {
250                                         min_dist = cur_dist;
251                                         min_node = node;
252                                 }
253                                 if (node->left)
254                                         stack[cur++] = node->left;
255                         }
256                         if (node->right)
257                                 stack[cur++] = node->right;
258                 }
259                 else {
260                         cur_dist = cur_dist * cur_dist;
261
262                         if (cur_dist < min_dist) {
263                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
264                                 if (cur_dist < min_dist) {
265                                         min_dist = cur_dist;
266                                         min_node = node;
267                                 }
268                                 if (node->right)
269                                         stack[cur++] = node->right;
270                         }
271                         if (node->left)
272                                 stack[cur++] = node->left;
273                 }
274                 if (UNLIKELY(cur + 3 > totstack)) {
275                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
276                 }
277         }
278
279         if (r_nearest) {
280                 r_nearest->index = min_node->index;
281                 r_nearest->dist = sqrtf(min_dist);
282                 copy_v3_v3(r_nearest->co, min_node->co);
283         }
284
285         if (stack != defaultstack)
286                 MEM_freeN(stack);
287
288         return min_node->index;
289 }
290
291 static void add_nearest(KDTreeNearest *ptn, unsigned int *found, unsigned int n, int index,
292                         float dist, const float *co)
293 {
294         unsigned int i;
295
296         if (*found < n) (*found)++;
297
298         for (i = *found - 1; i > 0; i--) {
299                 if (dist >= ptn[i - 1].dist)
300                         break;
301                 else
302                         ptn[i] = ptn[i - 1];
303         }
304
305         ptn[i].index = index;
306         ptn[i].dist = dist;
307         copy_v3_v3(ptn[i].co, co);
308 }
309
310 /**
311  * Find n nearest returns number of points found, with results in nearest.
312  * Normal is optional, but if given will limit results to points in normal direction from co.
313  *
314  * \param r_nearest  An array of nearest, sized at least \a n.
315  */
316 int BLI_kdtree_find_nearest_n(KDTree *tree, const float co[3], const float nor[3],
317                               KDTreeNearest r_nearest[],
318                               unsigned int n)
319 {
320         KDTreeNode *root, *node = NULL;
321         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
322         float cur_dist;
323         unsigned int totstack, cur = 0;
324         unsigned int i, found = 0;
325
326 #ifdef DEBUG
327         BLI_assert(tree->is_balanced == true);
328 #endif
329
330         if (!tree->root || n == 0)
331                 return 0;
332
333         stack = defaultstack;
334         totstack = KD_STACK_INIT;
335
336         root = tree->root;
337
338         cur_dist = squared_distance(root->co, co, root->nor, nor);
339         add_nearest(r_nearest, &found, n, root->index, cur_dist, root->co);
340         
341         if (co[root->d] < root->co[root->d]) {
342                 if (root->right)
343                         stack[cur++] = root->right;
344                 if (root->left)
345                         stack[cur++] = root->left;
346         }
347         else {
348                 if (root->left)
349                         stack[cur++] = root->left;
350                 if (root->right)
351                         stack[cur++] = root->right;
352         }
353
354         while (cur--) {
355                 node = stack[cur];
356
357                 cur_dist = node->co[node->d] - co[node->d];
358
359                 if (cur_dist < 0.0f) {
360                         cur_dist = -cur_dist * cur_dist;
361
362                         if (found < n || -cur_dist < r_nearest[found - 1].dist) {
363                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
364
365                                 if (found < n || cur_dist < r_nearest[found - 1].dist)
366                                         add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
367
368                                 if (node->left)
369                                         stack[cur++] = node->left;
370                         }
371                         if (node->right)
372                                 stack[cur++] = node->right;
373                 }
374                 else {
375                         cur_dist = cur_dist * cur_dist;
376
377                         if (found < n || cur_dist < r_nearest[found - 1].dist) {
378                                 cur_dist = squared_distance(node->co, co, node->nor, nor);
379                                 if (found < n || cur_dist < r_nearest[found - 1].dist)
380                                         add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
381
382                                 if (node->right)
383                                         stack[cur++] = node->right;
384                         }
385                         if (node->left)
386                                 stack[cur++] = node->left;
387                 }
388                 if (UNLIKELY(cur + 3 > totstack)) {
389                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
390                 }
391         }
392
393         for (i = 0; i < found; i++)
394                 r_nearest[i].dist = sqrtf(r_nearest[i].dist);
395
396         if (stack != defaultstack)
397                 MEM_freeN(stack);
398
399         return (int)found;
400 }
401
402 static int range_compare(const void *a, const void *b)
403 {
404         const KDTreeNearest *kda = a;
405         const KDTreeNearest *kdb = b;
406
407         if (kda->dist < kdb->dist)
408                 return -1;
409         else if (kda->dist > kdb->dist)
410                 return 1;
411         else
412                 return 0;
413 }
414 static void add_in_range(KDTreeNearest **ptn, unsigned int found, unsigned int *totfoundstack, int index, float dist, float *co)
415 {
416         KDTreeNearest *to;
417
418         if (found >= *totfoundstack) {
419                 KDTreeNearest *temp = MEM_mallocN((*totfoundstack + KD_FOUND_ALLOC_INC) * sizeof(KDTreeNode), "KDTree.treefoundstack");
420                 memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
421                 if (*ptn)
422                         MEM_freeN(*ptn);
423                 *ptn = temp;
424                 *totfoundstack += KD_FOUND_ALLOC_INC;
425         }
426
427         to = (*ptn) + found;
428
429         to->index = index;
430         to->dist = sqrtf(dist);
431         copy_v3_v3(to->co, co);
432 }
433
434 /**
435  * Range search returns number of points found, with results in nearest
436  * Normal is optional, but if given will limit results to points in normal direction from co.
437  * Remember to free nearest after use!
438  */
439 int BLI_kdtree_range_search(KDTree *tree, const float co[3], const float nor[3],
440                             KDTreeNearest **r_nearest, float range)
441 {
442         KDTreeNode *root, *node = NULL;
443         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
444         KDTreeNearest *foundstack = NULL;
445         float range2 = range * range, dist2;
446         unsigned int totstack, cur = 0, found = 0, totfoundstack = 0;
447
448 #ifdef DEBUG
449         BLI_assert(tree->is_balanced == true);
450 #endif
451
452         if (!tree->root)
453                 return 0;
454
455         stack = defaultstack;
456         totstack = KD_STACK_INIT;
457
458         root = tree->root;
459
460         if (co[root->d] + range < root->co[root->d]) {
461                 if (root->left)
462                         stack[cur++] = root->left;
463         }
464         else if (co[root->d] - range > root->co[root->d]) {
465                 if (root->right)
466                         stack[cur++] = root->right;
467         }
468         else {
469                 dist2 = squared_distance(root->co, co, root->nor, nor);
470                 if (dist2 <= range2)
471                         add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);
472
473                 if (root->left)
474                         stack[cur++] = root->left;
475                 if (root->right)
476                         stack[cur++] = root->right;
477         }
478
479         while (cur--) {
480                 node = stack[cur];
481
482                 if (co[node->d] + range < node->co[node->d]) {
483                         if (node->left)
484                                 stack[cur++] = node->left;
485                 }
486                 else if (co[node->d] - range > node->co[node->d]) {
487                         if (node->right)
488                                 stack[cur++] = node->right;
489                 }
490                 else {
491                         dist2 = squared_distance(node->co, co, node->nor, nor);
492                         if (dist2 <= range2)
493                                 add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);
494
495                         if (node->left)
496                                 stack[cur++] = node->left;
497                         if (node->right)
498                                 stack[cur++] = node->right;
499                 }
500
501                 if (UNLIKELY(cur + 3 > totstack)) {
502                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
503                 }
504         }
505
506         if (stack != defaultstack)
507                 MEM_freeN(stack);
508
509         if (found)
510                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
511
512         *r_nearest = foundstack;
513
514         return (int)found;
515 }