cleanup: style
[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];
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])
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         node->index = index;
102         node->d = 0;
103
104 #ifdef DEBUG
105         tree->is_balanced = false;
106 #endif
107 }
108
109 static KDTreeNode *kdtree_balance(KDTreeNode *nodes, unsigned int totnode, unsigned int axis)
110 {
111         KDTreeNode *node;
112         float co;
113         unsigned int left, right, median, i, j;
114
115         if (totnode <= 0)
116                 return NULL;
117         else if (totnode == 1)
118                 return nodes;
119         
120         /* quicksort style sorting around median */
121         left = 0;
122         right = totnode - 1;
123         median = totnode / 2;
124
125         while (right > left) {
126                 co = nodes[right].co[axis];
127                 i = left - 1;
128                 j = right;
129
130                 while (1) {
131                         while (nodes[++i].co[axis] < co) ;
132                         while (nodes[--j].co[axis] > co && j > left) ;
133
134                         if (i >= j)
135                                 break;
136
137                         SWAP(KDTreeNode, nodes[i], nodes[j]);
138                 }
139
140                 SWAP(KDTreeNode, nodes[i], nodes[right]);
141                 if (i >= median)
142                         right = i - 1;
143                 if (i <= median)
144                         left = i + 1;
145         }
146
147         /* set node and sort subnodes */
148         node = &nodes[median];
149         node->d = axis;
150         node->left = kdtree_balance(nodes, median, (axis + 1) % 3);
151         node->right = kdtree_balance(nodes + median + 1, (totnode - (median + 1)), (axis + 1) % 3);
152
153         return node;
154 }
155
156 void BLI_kdtree_balance(KDTree *tree)
157 {
158         tree->root = kdtree_balance(tree->nodes, tree->totnode, 0);
159
160 #ifdef DEBUG
161         tree->is_balanced = true;
162 #endif
163 }
164
165 static float squared_distance(const float v2[3], const float v1[3], const float n2[3])
166 {
167         float d[3], dist;
168
169         d[0] = v2[0] - v1[0];
170         d[1] = v2[1] - v1[1];
171         d[2] = v2[2] - v1[2];
172
173         dist = len_squared_v3(d);
174
175         /* can someone explain why this is done?*/
176         if (n2 && (dot_v3v3(d, n2) < 0.0f)) {
177                 dist *= 10.0f;
178         }
179
180         return dist;
181 }
182
183 static KDTreeNode **realloc_nodes(KDTreeNode **stack, unsigned int *totstack, const bool is_alloc)
184 {
185         KDTreeNode **stack_new = MEM_mallocN((*totstack + KD_NEAR_ALLOC_INC) * sizeof(KDTreeNode *), "KDTree.treestack");
186         memcpy(stack_new, stack, *totstack * sizeof(KDTreeNode *));
187         // memset(stack_new + *totstack, 0, sizeof(KDTreeNode *) * KD_NEAR_ALLOC_INC);
188         if (is_alloc)
189                 MEM_freeN(stack);
190         *totstack += KD_NEAR_ALLOC_INC;
191         return stack_new;
192 }
193
194 /**
195  * Find nearest returns index, and -1 if no node is found.
196  */
197 int BLI_kdtree_find_nearest(
198         KDTree *tree, const float co[3],
199         KDTreeNearest *r_nearest)
200 {
201         KDTreeNode *root, *node, *min_node;
202         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
203         float min_dist, cur_dist;
204         unsigned int totstack, cur = 0;
205
206 #ifdef DEBUG
207         BLI_assert(tree->is_balanced == true);
208 #endif
209
210         if (UNLIKELY(!tree->root))
211                 return -1;
212
213         stack = defaultstack;
214         totstack = KD_STACK_INIT;
215
216         root = tree->root;
217         min_node = root;
218         min_dist = len_squared_v3v3(root->co, co);
219
220         if (co[root->d] < root->co[root->d]) {
221                 if (root->right)
222                         stack[cur++] = root->right;
223                 if (root->left)
224                         stack[cur++] = root->left;
225         }
226         else {
227                 if (root->left)
228                         stack[cur++] = root->left;
229                 if (root->right)
230                         stack[cur++] = root->right;
231         }
232         
233         while (cur--) {
234                 node = stack[cur];
235
236                 cur_dist = node->co[node->d] - co[node->d];
237
238                 if (cur_dist < 0.0f) {
239                         cur_dist = -cur_dist * cur_dist;
240
241                         if (-cur_dist < min_dist) {
242                                 cur_dist = len_squared_v3v3(node->co, co);
243                                 if (cur_dist < min_dist) {
244                                         min_dist = cur_dist;
245                                         min_node = node;
246                                 }
247                                 if (node->left)
248                                         stack[cur++] = node->left;
249                         }
250                         if (node->right)
251                                 stack[cur++] = node->right;
252                 }
253                 else {
254                         cur_dist = cur_dist * cur_dist;
255
256                         if (cur_dist < min_dist) {
257                                 cur_dist = len_squared_v3v3(node->co, co);
258                                 if (cur_dist < min_dist) {
259                                         min_dist = cur_dist;
260                                         min_node = node;
261                                 }
262                                 if (node->right)
263                                         stack[cur++] = node->right;
264                         }
265                         if (node->left)
266                                 stack[cur++] = node->left;
267                 }
268                 if (UNLIKELY(cur + 3 > totstack)) {
269                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
270                 }
271         }
272
273         if (r_nearest) {
274                 r_nearest->index = min_node->index;
275                 r_nearest->dist = sqrtf(min_dist);
276                 copy_v3_v3(r_nearest->co, min_node->co);
277         }
278
279         if (stack != defaultstack)
280                 MEM_freeN(stack);
281
282         return min_node->index;
283 }
284
285 static void add_nearest(KDTreeNearest *ptn, unsigned int *found, unsigned int n, int index,
286                         float dist, const float *co)
287 {
288         unsigned int i;
289
290         if (*found < n) (*found)++;
291
292         for (i = *found - 1; i > 0; i--) {
293                 if (dist >= ptn[i - 1].dist)
294                         break;
295                 else
296                         ptn[i] = ptn[i - 1];
297         }
298
299         ptn[i].index = index;
300         ptn[i].dist = dist;
301         copy_v3_v3(ptn[i].co, co);
302 }
303
304 /**
305  * Find n nearest returns number of points found, with results in nearest.
306  * Normal is optional, but if given will limit results to points in normal direction from co.
307  *
308  * \param r_nearest  An array of nearest, sized at least \a n.
309  */
310 int BLI_kdtree_find_nearest_n__normal(
311         KDTree *tree, const float co[3], const float nor[3],
312         KDTreeNearest r_nearest[],
313         unsigned int n)
314 {
315         KDTreeNode *root, *node = NULL;
316         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
317         float cur_dist;
318         unsigned int totstack, cur = 0;
319         unsigned int i, found = 0;
320
321 #ifdef DEBUG
322         BLI_assert(tree->is_balanced == true);
323 #endif
324
325         if (UNLIKELY(!tree->root || n == 0))
326                 return 0;
327
328         stack = defaultstack;
329         totstack = KD_STACK_INIT;
330
331         root = tree->root;
332
333         cur_dist = squared_distance(root->co, co, nor);
334         add_nearest(r_nearest, &found, n, root->index, cur_dist, root->co);
335         
336         if (co[root->d] < root->co[root->d]) {
337                 if (root->right)
338                         stack[cur++] = root->right;
339                 if (root->left)
340                         stack[cur++] = root->left;
341         }
342         else {
343                 if (root->left)
344                         stack[cur++] = root->left;
345                 if (root->right)
346                         stack[cur++] = root->right;
347         }
348
349         while (cur--) {
350                 node = stack[cur];
351
352                 cur_dist = node->co[node->d] - co[node->d];
353
354                 if (cur_dist < 0.0f) {
355                         cur_dist = -cur_dist * cur_dist;
356
357                         if (found < n || -cur_dist < r_nearest[found - 1].dist) {
358                                 cur_dist = squared_distance(node->co, co, nor);
359
360                                 if (found < n || cur_dist < r_nearest[found - 1].dist)
361                                         add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
362
363                                 if (node->left)
364                                         stack[cur++] = node->left;
365                         }
366                         if (node->right)
367                                 stack[cur++] = node->right;
368                 }
369                 else {
370                         cur_dist = cur_dist * cur_dist;
371
372                         if (found < n || cur_dist < r_nearest[found - 1].dist) {
373                                 cur_dist = squared_distance(node->co, co, nor);
374                                 if (found < n || cur_dist < r_nearest[found - 1].dist)
375                                         add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
376
377                                 if (node->right)
378                                         stack[cur++] = node->right;
379                         }
380                         if (node->left)
381                                 stack[cur++] = node->left;
382                 }
383                 if (UNLIKELY(cur + 3 > totstack)) {
384                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
385                 }
386         }
387
388         for (i = 0; i < found; i++)
389                 r_nearest[i].dist = sqrtf(r_nearest[i].dist);
390
391         if (stack != defaultstack)
392                 MEM_freeN(stack);
393
394         return (int)found;
395 }
396
397 static int range_compare(const void *a, const void *b)
398 {
399         const KDTreeNearest *kda = a;
400         const KDTreeNearest *kdb = b;
401
402         if (kda->dist < kdb->dist)
403                 return -1;
404         else if (kda->dist > kdb->dist)
405                 return 1;
406         else
407                 return 0;
408 }
409 static void add_in_range(
410         KDTreeNearest **r_foundstack,
411         unsigned int   *r_foundstack_tot_alloc,
412         unsigned int      found,
413         const int index, const float dist, const float *co)
414 {
415         KDTreeNearest *to;
416
417         if (UNLIKELY(found >= *r_foundstack_tot_alloc)) {
418                 *r_foundstack = MEM_reallocN_id(
419                         *r_foundstack,
420                         (*r_foundstack_tot_alloc += KD_FOUND_ALLOC_INC) * sizeof(KDTreeNode),
421                         __func__);
422         }
423
424         to = (*r_foundstack) + found;
425
426         to->index = index;
427         to->dist = sqrtf(dist);
428         copy_v3_v3(to->co, co);
429 }
430
431 /**
432  * Range search returns number of points found, with results in nearest
433  * Normal is optional, but if given will limit results to points in normal direction from co.
434  * Remember to free nearest after use!
435  */
436 int BLI_kdtree_range_search__normal(
437         KDTree *tree, const float co[3], const float nor[3],
438         KDTreeNearest **r_nearest, float range)
439 {
440         KDTreeNode *root, *node = NULL;
441         KDTreeNode **stack, *defaultstack[KD_STACK_INIT];
442         KDTreeNearest *foundstack = NULL;
443         float range2 = range * range, dist2;
444         unsigned int totstack, cur = 0, found = 0, totfoundstack = 0;
445
446 #ifdef DEBUG
447         BLI_assert(tree->is_balanced == true);
448 #endif
449
450         if (UNLIKELY(!tree->root))
451                 return 0;
452
453         stack = defaultstack;
454         totstack = KD_STACK_INIT;
455
456         root = tree->root;
457
458         if (co[root->d] + range < root->co[root->d]) {
459                 if (root->left)
460                         stack[cur++] = root->left;
461         }
462         else if (co[root->d] - range > root->co[root->d]) {
463                 if (root->right)
464                         stack[cur++] = root->right;
465         }
466         else {
467                 dist2 = squared_distance(root->co, co, nor);
468                 if (dist2 <= range2)
469                         add_in_range(&foundstack, &totfoundstack, found++, root->index, dist2, root->co);
470
471                 if (root->left)
472                         stack[cur++] = root->left;
473                 if (root->right)
474                         stack[cur++] = root->right;
475         }
476
477         while (cur--) {
478                 node = stack[cur];
479
480                 if (co[node->d] + range < node->co[node->d]) {
481                         if (node->left)
482                                 stack[cur++] = node->left;
483                 }
484                 else if (co[node->d] - range > node->co[node->d]) {
485                         if (node->right)
486                                 stack[cur++] = node->right;
487                 }
488                 else {
489                         dist2 = squared_distance(node->co, co, nor);
490                         if (dist2 <= range2)
491                                 add_in_range(&foundstack, &totfoundstack, found++, node->index, dist2, node->co);
492
493                         if (node->left)
494                                 stack[cur++] = node->left;
495                         if (node->right)
496                                 stack[cur++] = node->right;
497                 }
498
499                 if (UNLIKELY(cur + 3 > totstack)) {
500                         stack = realloc_nodes(stack, &totstack, defaultstack != stack);
501                 }
502         }
503
504         if (stack != defaultstack)
505                 MEM_freeN(stack);
506
507         if (found)
508                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
509
510         *r_nearest = foundstack;
511
512         return (int)found;
513 }