a265505cc8f4d6c240d08fd9a944f75f534bd75b
[blender.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  * 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; } (void)0
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, const float co[3], const float nor[3])
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], const float *UNUSED(n1), const 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, const float co[3], const float nor[3], 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, const float co[3], const float nor[3], 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 }