Merge with trunk r41625
[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; }
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(float *v2, float *v1, float *n1, float *n2)
136 {
137         float d[3], dist;
138         (void)n1; /* unused */
139
140         d[0]= v2[0]-v1[0];
141         d[1]= v2[1]-v1[1];
142         d[2]= v2[2]-v1[2];
143
144         dist= d[0]*d[0] + d[1]*d[1] + d[2]*d[2];
145
146         //if(n1 && n2 && n1[0]*n2[0] + n1[1]*n2[1] + n1[2]*n2[2] < 0.0f)
147         if(n2 && d[0]*n2[0] + d[1]*n2[1] + d[2]*n2[2] < 0.0f)
148                 dist *= 10.0f;
149
150         return dist;
151 }
152
153 int     BLI_kdtree_find_nearest(KDTree *tree, float *co, float *nor, KDTreeNearest *nearest)
154 {
155         KDTreeNode *root, *node, *min_node;
156         KDTreeNode **stack, *defaultstack[100];
157         float min_dist, cur_dist;
158         int totstack, cur=0;
159
160         if(!tree->root)
161                 return -1;
162
163         stack= defaultstack;
164         totstack= 100;
165
166         root= tree->root;
167         min_node= root;
168         min_dist= squared_distance(root->co,co,root->nor,nor);
169
170         if(co[root->d] < root->co[root->d]) {
171                 if(root->right)
172                         stack[cur++]=root->right;
173                 if(root->left)
174                         stack[cur++]=root->left;
175         }
176         else {
177                 if(root->left)
178                         stack[cur++]=root->left;
179                 if(root->right)
180                         stack[cur++]=root->right;
181         }
182         
183         while(cur--){
184                 node=stack[cur];
185
186                 cur_dist = node->co[node->d] - co[node->d];
187
188                 if(cur_dist<0.0f){
189                         cur_dist= -cur_dist*cur_dist;
190
191                         if(-cur_dist<min_dist){
192                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
193                                 if(cur_dist<min_dist){
194                                         min_dist=cur_dist;
195                                         min_node=node;
196                                 }
197                                 if(node->left)
198                                         stack[cur++]=node->left;
199                         }
200                         if(node->right)
201                                 stack[cur++]=node->right;
202                 }
203                 else{
204                         cur_dist= cur_dist*cur_dist;
205
206                         if(cur_dist<min_dist){
207                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
208                                 if(cur_dist<min_dist){
209                                         min_dist=cur_dist;
210                                         min_node=node;
211                                 }
212                                 if(node->right)
213                                         stack[cur++]=node->right;
214                         }
215                         if(node->left)
216                                 stack[cur++]=node->left;
217                 }
218                 if(cur+3 > totstack){
219                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
220                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
221                         if(stack != defaultstack)
222                                 MEM_freeN(stack);
223                         stack=temp;
224                         totstack+=100;
225                 }
226         }
227
228         if(nearest) {
229                 nearest->index= min_node->index;
230                 nearest->dist= sqrt(min_dist);
231                 copy_v3_v3(nearest->co, min_node->co);
232         }
233
234         if(stack != defaultstack)
235                 MEM_freeN(stack);
236
237         return min_node->index;
238 }
239
240 static void add_nearest(KDTreeNearest *ptn, int *found, int n, int index, float dist, float *co)
241 {
242         int i;
243
244         if(*found<n) (*found)++;
245
246         for(i=*found-1; i>0; i--) {
247                 if(dist >= ptn[i-1].dist)
248                         break;
249                 else
250                         ptn[i]= ptn[i-1];
251         }
252
253         ptn[i].index= index;
254         ptn[i].dist= dist;
255         copy_v3_v3(ptn[i].co, co);
256 }
257
258 /* finds the nearest n entries in tree to specified coordinates */
259 int     BLI_kdtree_find_n_nearest(KDTree *tree, int n, float *co, float *nor, KDTreeNearest *nearest)
260 {
261         KDTreeNode *root, *node= NULL;
262         KDTreeNode **stack, *defaultstack[100];
263         float cur_dist;
264         int i, totstack, cur=0, found=0;
265
266         if(!tree->root)
267                 return 0;
268
269         stack= defaultstack;
270         totstack= 100;
271
272         root= tree->root;
273
274         cur_dist= squared_distance(root->co,co,root->nor,nor);
275         add_nearest(nearest,&found,n,root->index,cur_dist,root->co);
276         
277         if(co[root->d] < root->co[root->d]) {
278                 if(root->right)
279                         stack[cur++]=root->right;
280                 if(root->left)
281                         stack[cur++]=root->left;
282         }
283         else {
284                 if(root->left)
285                         stack[cur++]=root->left;
286                 if(root->right)
287                         stack[cur++]=root->right;
288         }
289
290         while(cur--){
291                 node=stack[cur];
292
293                 cur_dist = node->co[node->d] - co[node->d];
294
295                 if(cur_dist<0.0f){
296                         cur_dist= -cur_dist*cur_dist;
297
298                         if(found<n || -cur_dist<nearest[found-1].dist){
299                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
300
301                                 if(found<n || cur_dist<nearest[found-1].dist)
302                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
303
304                                 if(node->left)
305                                         stack[cur++]=node->left;
306                         }
307                         if(node->right)
308                                 stack[cur++]=node->right;
309                 }
310                 else{
311                         cur_dist= cur_dist*cur_dist;
312
313                         if(found<n || cur_dist<nearest[found-1].dist){
314                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
315                                 if(found<n || cur_dist<nearest[found-1].dist)
316                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
317
318                                 if(node->right)
319                                         stack[cur++]=node->right;
320                         }
321                         if(node->left)
322                                 stack[cur++]=node->left;
323                 }
324                 if(cur+3 > totstack){
325                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
326                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
327                         if(stack != defaultstack)
328                                 MEM_freeN(stack);
329                         stack=temp;
330                         totstack+=100;
331                 }
332         }
333
334         for(i=0; i<found; i++)
335                 nearest[i].dist= sqrt(nearest[i].dist);
336
337         if(stack != defaultstack)
338                 MEM_freeN(stack);
339
340         return found;
341 }
342
343 static int range_compare(const void * a, const void * b)
344 {
345         const KDTreeNearest *kda = a;
346         const KDTreeNearest *kdb = b;
347
348         if(kda->dist < kdb->dist)
349                 return -1;
350         else if(kda->dist > kdb->dist)
351                 return 1;
352         else
353                 return 0;
354 }
355 static void add_in_range(KDTreeNearest **ptn, int found, int *totfoundstack, int index, float dist, float *co)
356 {
357         KDTreeNearest *to;
358
359         if(found+1 > *totfoundstack) {
360                 KDTreeNearest *temp=MEM_callocN((*totfoundstack+50)*sizeof(KDTreeNode), "psys_treefoundstack");
361                 memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
362                 if(*ptn)
363                         MEM_freeN(*ptn);
364                 *ptn = temp;
365                 *totfoundstack+=50;
366         }
367
368         to = (*ptn) + found;
369
370         to->index = index;
371         to->dist = sqrt(dist);
372         copy_v3_v3(to->co, co);
373 }
374 int BLI_kdtree_range_search(KDTree *tree, float range, float *co, float *nor, KDTreeNearest **nearest)
375 {
376         KDTreeNode *root, *node= NULL;
377         KDTreeNode **stack, *defaultstack[100];
378         KDTreeNearest *foundstack=NULL;
379         float range2 = range*range, dist2;
380         int totstack, cur=0, found=0, totfoundstack=0;
381
382         if(!tree || !tree->root)
383                 return 0;
384
385         stack= defaultstack;
386         totstack= 100;
387
388         root= tree->root;
389
390         if(co[root->d] + range < root->co[root->d]) {
391                 if(root->left)
392                         stack[cur++]=root->left;
393         }
394         else if(co[root->d] - range > root->co[root->d]) {
395                 if(root->right)
396                         stack[cur++]=root->right;
397         }
398         else {
399                 dist2 = squared_distance(root->co, co, root->nor, nor);
400                 if(dist2  <= range2)
401                         add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);
402
403                 if(root->left)
404                         stack[cur++]=root->left;
405                 if(root->right)
406                         stack[cur++]=root->right;
407         }
408
409         while(cur--) {
410                 node=stack[cur];
411
412                 if(co[node->d] + range < node->co[node->d]) {
413                         if(node->left)
414                                 stack[cur++]=node->left;
415                 }
416                 else if(co[node->d] - range > node->co[node->d]) {
417                         if(node->right)
418                                 stack[cur++]=node->right;
419                 }
420                 else {
421                         dist2 = squared_distance(node->co, co, node->nor, nor);
422                         if(dist2 <= range2)
423                                 add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);
424
425                         if(node->left)
426                                 stack[cur++]=node->left;
427                         if(node->right)
428                                 stack[cur++]=node->right;
429                 }
430
431                 if(cur+3 > totstack){
432                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
433                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
434                         if(stack != defaultstack)
435                                 MEM_freeN(stack);
436                         stack=temp;
437                         totstack+=100;
438                 }
439         }
440
441         if(stack != defaultstack)
442                 MEM_freeN(stack);
443
444         if(found)
445                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
446
447         *nearest = foundstack;
448
449         return found;
450 }