Merge with trunk r39928
[blender-staging.git] / source / blender / blenlib / intern / BLI_kdtree.c
1 /*
2  * $Id$
3  *
4  * ***** BEGIN GPL LICENSE BLOCK *****
5  *
6  * This program is free software; you can redistribute it and/or
7  * modify it under the terms of the GNU General Public License
8  * as published by the Free Software Foundation; either version 2
9  * of the License, or (at your option) any later version.
10  *
11  * This program is distributed in the hope that it will be useful,
12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14  * GNU General Public License for more details.
15  *
16  * You should have received a copy of the GNU General Public License
17  * along with this program; if not, write to the Free Software Foundation,
18  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
19  *
20  * The Original Code is Copyright (C) 2001-2002 by NaN Holding BV.
21  * All rights reserved.
22  *
23  * The Original Code is: none of this file.
24  *
25  * Contributor(s): Janne Karhu
26  *                 Brecht Van Lommel
27  *
28  * ***** END GPL LICENSE BLOCK *****
29  */
30
31 /** \file blender/blenlib/intern/BLI_kdtree.c
32  *  \ingroup bli
33  */
34
35
36
37 #include "MEM_guardedalloc.h"
38
39 #include "BLI_math.h"
40 #include "BLI_kdtree.h"
41
42 #ifndef SWAP
43 #define SWAP(type, a, b) { type sw_ap; sw_ap=(a); (a)=(b); (b)=sw_ap; }
44 #endif
45
46 struct KDTree {
47         KDTreeNode *nodes;
48         int totnode;
49         KDTreeNode *root;
50 };
51
52 KDTree *BLI_kdtree_new(int maxsize)
53 {
54         KDTree *tree;
55
56         tree= MEM_callocN(sizeof(KDTree), "KDTree");
57         tree->nodes= MEM_callocN(sizeof(KDTreeNode)*maxsize, "KDTreeNode");
58         tree->totnode= 0;
59
60         return tree;
61 }
62
63 void BLI_kdtree_free(KDTree *tree)
64 {
65         if(tree) {
66                 MEM_freeN(tree->nodes);
67                 MEM_freeN(tree);
68         }
69 }
70
71 void BLI_kdtree_insert(KDTree *tree, int index, float *co, float *nor)
72 {
73         KDTreeNode *node= &tree->nodes[tree->totnode++];
74
75         node->index= index;
76         copy_v3_v3(node->co, co);
77         if(nor) copy_v3_v3(node->nor, nor);
78 }
79
80 static KDTreeNode *kdtree_balance(KDTreeNode *nodes, int totnode, int axis)
81 {
82         KDTreeNode *node;
83         float co;
84         int left, right, median, i, j;
85
86         if(totnode <= 0)
87                 return NULL;
88         else if(totnode == 1)
89                 return nodes;
90         
91         /* quicksort style sorting around median */
92         left= 0;
93         right= totnode-1;
94         median= totnode/2;
95
96         while(right > left) {
97                 co= nodes[right].co[axis];
98                 i= left-1;
99                 j= right;
100
101                 while(1) {
102                         while(nodes[++i].co[axis] < co);
103                         while(nodes[--j].co[axis] > co && j>left);
104
105                         if(i >= j) break;
106                         SWAP(KDTreeNode, nodes[i], nodes[j]);
107                 }
108
109                 SWAP(KDTreeNode, nodes[i], nodes[right]);
110                 if(i >= median)
111                         right= i-1;
112                 if(i <= median)
113                         left= i+1;
114         }
115
116         /* set node and sort subnodes */
117         node= &nodes[median];
118         node->d= axis;
119         node->left= kdtree_balance(nodes, median, (axis+1)%3);
120         node->right= kdtree_balance(nodes+median+1, (totnode-(median+1)), (axis+1)%3);
121
122         return node;
123 }
124
125 void BLI_kdtree_balance(KDTree *tree)
126 {
127         tree->root= kdtree_balance(tree->nodes, tree->totnode, 0);
128 }
129
130 static float squared_distance(float *v2, float *v1, float *n1, float *n2)
131 {
132         float d[3], dist;
133         (void)n1; /* unused */
134
135         d[0]= v2[0]-v1[0];
136         d[1]= v2[1]-v1[1];
137         d[2]= v2[2]-v1[2];
138
139         dist= d[0]*d[0] + d[1]*d[1] + d[2]*d[2];
140
141         //if(n1 && n2 && n1[0]*n2[0] + n1[1]*n2[1] + n1[2]*n2[2] < 0.0f)
142         if(n2 && d[0]*n2[0] + d[1]*n2[1] + d[2]*n2[2] < 0.0f)
143                 dist *= 10.0f;
144
145         return dist;
146 }
147
148 int     BLI_kdtree_find_nearest(KDTree *tree, float *co, float *nor, KDTreeNearest *nearest)
149 {
150         KDTreeNode *root, *node, *min_node;
151         KDTreeNode **stack, *defaultstack[100];
152         float min_dist, cur_dist;
153         int totstack, cur=0;
154
155         if(!tree->root)
156                 return -1;
157
158         stack= defaultstack;
159         totstack= 100;
160
161         root= tree->root;
162         min_node= root;
163         min_dist= squared_distance(root->co,co,root->nor,nor);
164
165         if(co[root->d] < root->co[root->d]) {
166                 if(root->right)
167                         stack[cur++]=root->right;
168                 if(root->left)
169                         stack[cur++]=root->left;
170         }
171         else {
172                 if(root->left)
173                         stack[cur++]=root->left;
174                 if(root->right)
175                         stack[cur++]=root->right;
176         }
177         
178         while(cur--){
179                 node=stack[cur];
180
181                 cur_dist = node->co[node->d] - co[node->d];
182
183                 if(cur_dist<0.0f){
184                         cur_dist= -cur_dist*cur_dist;
185
186                         if(-cur_dist<min_dist){
187                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
188                                 if(cur_dist<min_dist){
189                                         min_dist=cur_dist;
190                                         min_node=node;
191                                 }
192                                 if(node->left)
193                                         stack[cur++]=node->left;
194                         }
195                         if(node->right)
196                                 stack[cur++]=node->right;
197                 }
198                 else{
199                         cur_dist= cur_dist*cur_dist;
200
201                         if(cur_dist<min_dist){
202                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
203                                 if(cur_dist<min_dist){
204                                         min_dist=cur_dist;
205                                         min_node=node;
206                                 }
207                                 if(node->right)
208                                         stack[cur++]=node->right;
209                         }
210                         if(node->left)
211                                 stack[cur++]=node->left;
212                 }
213                 if(cur+3 > totstack){
214                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
215                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
216                         if(stack != defaultstack)
217                                 MEM_freeN(stack);
218                         stack=temp;
219                         totstack+=100;
220                 }
221         }
222
223         if(nearest) {
224                 nearest->index= min_node->index;
225                 nearest->dist= sqrt(min_dist);
226                 copy_v3_v3(nearest->co, min_node->co);
227         }
228
229         if(stack != defaultstack)
230                 MEM_freeN(stack);
231
232         return min_node->index;
233 }
234
235 static void add_nearest(KDTreeNearest *ptn, int *found, int n, int index, float dist, float *co)
236 {
237         int i;
238
239         if(*found<n) (*found)++;
240
241         for(i=*found-1; i>0; i--) {
242                 if(dist >= ptn[i-1].dist)
243                         break;
244                 else
245                         ptn[i]= ptn[i-1];
246         }
247
248         ptn[i].index= index;
249         ptn[i].dist= dist;
250         copy_v3_v3(ptn[i].co, co);
251 }
252
253 /* finds the nearest n entries in tree to specified coordinates */
254 int     BLI_kdtree_find_n_nearest(KDTree *tree, int n, float *co, float *nor, KDTreeNearest *nearest)
255 {
256         KDTreeNode *root, *node= NULL;
257         KDTreeNode **stack, *defaultstack[100];
258         float cur_dist;
259         int i, totstack, cur=0, found=0;
260
261         if(!tree->root)
262                 return 0;
263
264         stack= defaultstack;
265         totstack= 100;
266
267         root= tree->root;
268
269         cur_dist= squared_distance(root->co,co,root->nor,nor);
270         add_nearest(nearest,&found,n,root->index,cur_dist,root->co);
271         
272         if(co[root->d] < root->co[root->d]) {
273                 if(root->right)
274                         stack[cur++]=root->right;
275                 if(root->left)
276                         stack[cur++]=root->left;
277         }
278         else {
279                 if(root->left)
280                         stack[cur++]=root->left;
281                 if(root->right)
282                         stack[cur++]=root->right;
283         }
284
285         while(cur--){
286                 node=stack[cur];
287
288                 cur_dist = node->co[node->d] - co[node->d];
289
290                 if(cur_dist<0.0f){
291                         cur_dist= -cur_dist*cur_dist;
292
293                         if(found<n || -cur_dist<nearest[found-1].dist){
294                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
295
296                                 if(found<n || cur_dist<nearest[found-1].dist)
297                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
298
299                                 if(node->left)
300                                         stack[cur++]=node->left;
301                         }
302                         if(node->right)
303                                 stack[cur++]=node->right;
304                 }
305                 else{
306                         cur_dist= cur_dist*cur_dist;
307
308                         if(found<n || cur_dist<nearest[found-1].dist){
309                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
310                                 if(found<n || cur_dist<nearest[found-1].dist)
311                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
312
313                                 if(node->right)
314                                         stack[cur++]=node->right;
315                         }
316                         if(node->left)
317                                 stack[cur++]=node->left;
318                 }
319                 if(cur+3 > totstack){
320                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
321                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
322                         if(stack != defaultstack)
323                                 MEM_freeN(stack);
324                         stack=temp;
325                         totstack+=100;
326                 }
327         }
328
329         for(i=0; i<found; i++)
330                 nearest[i].dist= sqrt(nearest[i].dist);
331
332         if(stack != defaultstack)
333                 MEM_freeN(stack);
334
335         return found;
336 }
337
338 static int range_compare(const void * a, const void * b)
339 {
340         const KDTreeNearest *kda = a;
341         const KDTreeNearest *kdb = b;
342
343         if(kda->dist < kdb->dist)
344                 return -1;
345         else if(kda->dist > kdb->dist)
346                 return 1;
347         else
348                 return 0;
349 }
350 static void add_in_range(KDTreeNearest **ptn, int found, int *totfoundstack, int index, float dist, float *co)
351 {
352         KDTreeNearest *to;
353
354         if(found+1 > *totfoundstack) {
355                 KDTreeNearest *temp=MEM_callocN((*totfoundstack+50)*sizeof(KDTreeNode), "psys_treefoundstack");
356                 memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
357                 if(*ptn)
358                         MEM_freeN(*ptn);
359                 *ptn = temp;
360                 *totfoundstack+=50;
361         }
362
363         to = (*ptn) + found;
364
365         to->index = index;
366         to->dist = sqrt(dist);
367         copy_v3_v3(to->co, co);
368 }
369 int BLI_kdtree_range_search(KDTree *tree, float range, float *co, float *nor, KDTreeNearest **nearest)
370 {
371         KDTreeNode *root, *node= NULL;
372         KDTreeNode **stack, *defaultstack[100];
373         KDTreeNearest *foundstack=NULL;
374         float range2 = range*range, dist2;
375         int totstack, cur=0, found=0, totfoundstack=0;
376
377         if(!tree || !tree->root)
378                 return 0;
379
380         stack= defaultstack;
381         totstack= 100;
382
383         root= tree->root;
384
385         if(co[root->d] + range < root->co[root->d]) {
386                 if(root->left)
387                         stack[cur++]=root->left;
388         }
389         else if(co[root->d] - range > root->co[root->d]) {
390                 if(root->right)
391                         stack[cur++]=root->right;
392         }
393         else {
394                 dist2 = squared_distance(root->co, co, root->nor, nor);
395                 if(dist2  <= range2)
396                         add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);
397
398                 if(root->left)
399                         stack[cur++]=root->left;
400                 if(root->right)
401                         stack[cur++]=root->right;
402         }
403
404         while(cur--) {
405                 node=stack[cur];
406
407                 if(co[node->d] + range < node->co[node->d]) {
408                         if(node->left)
409                                 stack[cur++]=node->left;
410                 }
411                 else if(co[node->d] - range > node->co[node->d]) {
412                         if(node->right)
413                                 stack[cur++]=node->right;
414                 }
415                 else {
416                         dist2 = squared_distance(node->co, co, node->nor, nor);
417                         if(dist2 <= range2)
418                                 add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);
419
420                         if(node->left)
421                                 stack[cur++]=node->left;
422                         if(node->right)
423                                 stack[cur++]=node->right;
424                 }
425
426                 if(cur+3 > totstack){
427                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
428                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
429                         if(stack != defaultstack)
430                                 MEM_freeN(stack);
431                         stack=temp;
432                         totstack+=100;
433                 }
434         }
435
436         if(stack != defaultstack)
437                 MEM_freeN(stack);
438
439         if(found)
440                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
441
442         *nearest = foundstack;
443
444         return found;
445 }
446
447
448 static int add_in_range_thread_safe(KDTreeNearest *ptn, int found, int *totfoundstack, int index, float dist, float *co)
449 {
450         KDTreeNearest *to;
451
452         if(found+1 > *totfoundstack)
453                 return 0;
454
455         to = ptn + found;
456
457         to->index = index;
458         to->dist = sqrt(dist);
459         copy_v3_v3(to->co, co);
460         return 1;
461 }
462 int BLI_kdtree_range_search_thread_safe(KDTree *tree, float range, float *co, float *nor, KDTreeNearest *nearest, int limit)
463 {
464         KDTreeNode *root, *node= NULL;
465         KDTreeNode **stack, *defaultstack[500];
466         KDTreeNearest *foundstack=nearest;
467         float range2 = range*range, dist2;
468         int totstack, cur=0, found=0;
469
470         if(!tree || !tree->root)
471                 return 0;
472
473         stack= defaultstack;
474         totstack= 500;
475
476         root= tree->root;
477
478         if(co[root->d] + range < root->co[root->d]) {
479                 if(root->left)
480                         stack[cur++]=root->left;
481         }
482         else if(co[root->d] - range > root->co[root->d]) {
483                 if(root->right)
484                         stack[cur++]=root->right;
485         }
486         else {
487                 dist2 = squared_distance(root->co, co, root->nor, nor);
488                 if(dist2  <= range2)
489                         add_in_range_thread_safe(foundstack, found++, &limit, root->index, dist2, root->co);
490
491                 if(root->left)
492                         stack[cur++]=root->left;
493                 if(root->right)
494                         stack[cur++]=root->right;
495         }
496
497         while(cur--) {
498                 node=stack[cur];
499
500                 if(co[node->d] + range < node->co[node->d]) {
501                         if(node->left)
502                                 stack[cur++]=node->left;
503                 }
504                 else if(co[node->d] - range > node->co[node->d]) {
505                         if(node->right)
506                                 stack[cur++]=node->right;
507                 }
508                 else {
509                         dist2 = squared_distance(node->co, co, node->nor, nor);
510                         if(dist2 <= range2)
511                                 if (!add_in_range_thread_safe(foundstack, found++, &limit, node->index, dist2, node->co))
512                                         break;
513
514                         if(node->left)
515                                 stack[cur++]=node->left;
516                         if(node->right)
517                                 stack[cur++]=node->right;
518                 }
519
520                 /* abort if running out of stack */
521                 if(cur+3 > totstack)
522                         break;
523         }
524
525         if(stack != defaultstack)
526                 MEM_freeN(stack);
527
528         if(found)
529                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
530
531         return found;
532 }