Merge with 2.5 -r 21619:21756.
[blender.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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, 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 #include <stdlib.h>
32 #include <string.h>
33 #include <math.h>
34
35 #include "MEM_guardedalloc.h"
36
37 #include "BLI_arithb.h"
38 #include "BLI_kdtree.h"
39
40 #define SWAP(type, a, b) { type sw_ap; sw_ap=(a); (a)=(b); (b)=sw_ap; }
41
42 typedef struct KDTreeNode {
43         struct KDTreeNode *left, *right;
44         float co[3], nor[3];
45         int index;
46         short d;
47 } KDTreeNode;
48
49 struct KDTree {
50         KDTreeNode *nodes;
51         int totnode;
52         KDTreeNode *root;
53 };
54
55 KDTree *BLI_kdtree_new(int maxsize)
56 {
57         KDTree *tree;
58
59         tree= MEM_callocN(sizeof(KDTree), "KDTree");
60         tree->nodes= MEM_callocN(sizeof(KDTreeNode)*maxsize, "KDTreeNode");
61         tree->totnode= 0;
62
63         return tree;
64 }
65
66 void BLI_kdtree_free(KDTree *tree)
67 {
68         if(tree) {
69                 MEM_freeN(tree->nodes);
70                 MEM_freeN(tree);
71         }
72 }
73
74 void BLI_kdtree_insert(KDTree *tree, int index, float *co, float *nor)
75 {
76         KDTreeNode *node= &tree->nodes[tree->totnode++];
77
78         node->index= index;
79         VecCopyf(node->co, co);
80         if(nor) VecCopyf(node->nor, nor);
81 }
82
83 static KDTreeNode *kdtree_balance(KDTreeNode *nodes, int totnode, int axis)
84 {
85         KDTreeNode *node;
86         float co;
87         int left, right, median, i, j;
88
89         if(totnode <= 0)
90                 return NULL;
91         else if(totnode == 1)
92                 return nodes;
93         
94         /* quicksort style sorting around median */
95         left= 0;
96         right= totnode-1;
97         median= totnode/2;
98
99         while(right > left) {
100                 co= nodes[right].co[axis];
101                 i= left-1;
102                 j= right;
103
104                 while(1) {
105                         while(nodes[++i].co[axis] < co);
106                         while(nodes[--j].co[axis] > co && j>left);
107
108                         if(i >= j) break;
109                         SWAP(KDTreeNode, nodes[i], nodes[j]);
110                 }
111
112                 SWAP(KDTreeNode, nodes[i], nodes[right]);
113                 if(i >= median)
114                         right= i-1;
115                 if(i <= median)
116                         left= i+1;
117         }
118
119         /* set node and sort subnodes */
120         node= &nodes[median];
121         node->d= axis;
122         node->left= kdtree_balance(nodes, median, (axis+1)%3);
123         node->right= kdtree_balance(nodes+median+1, (totnode-(median+1)), (axis+1)%3);
124
125         return node;
126 }
127
128 void BLI_kdtree_balance(KDTree *tree)
129 {
130         tree->root= kdtree_balance(tree->nodes, tree->totnode, 0);
131 }
132
133 static float squared_distance(float *v2, float *v1, float *n1, float *n2)
134 {
135         float d[3], dist;
136
137         d[0]= v2[0]-v1[0];
138         d[1]= v2[1]-v1[1];
139         d[2]= v2[2]-v1[2];
140
141         dist= d[0]*d[0] + d[1]*d[1] + d[2]*d[2];
142
143         //if(n1 && n2 && n1[0]*n2[0] + n1[1]*n2[1] + n1[2]*n2[2] < 0.0f)
144         if(n2 && d[0]*n2[0] + d[1]*n2[1] + d[2]*n2[2] < 0.0f)
145                 dist *= 10.0f;
146
147         return dist;
148 }
149
150 int     BLI_kdtree_find_nearest(KDTree *tree, float *co, float *nor, KDTreeNearest *nearest)
151 {
152         KDTreeNode *root, *node, *min_node;
153         KDTreeNode **stack, *defaultstack[100];
154         float min_dist, cur_dist;
155         int totstack, cur=0;
156
157         if(!tree->root)
158                 return -1;
159
160         stack= defaultstack;
161         totstack= 100;
162
163         root= tree->root;
164         min_node= root;
165         min_dist= squared_distance(root->co,co,root->nor,nor);
166
167         if(co[root->d] < root->co[root->d]) {
168                 if(root->right)
169                         stack[cur++]=root->right;
170                 if(root->left)
171                         stack[cur++]=root->left;
172         }
173         else {
174                 if(root->left)
175                         stack[cur++]=root->left;
176                 if(root->right)
177                         stack[cur++]=root->right;
178         }
179         
180         while(cur--){
181                 node=stack[cur];
182
183                 cur_dist = node->co[node->d] - co[node->d];
184
185                 if(cur_dist<0.0){
186                         cur_dist= -cur_dist*cur_dist;
187
188                         if(-cur_dist<min_dist){
189                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
190                                 if(cur_dist<min_dist){
191                                         min_dist=cur_dist;
192                                         min_node=node;
193                                 }
194                                 if(node->left)
195                                         stack[cur++]=node->left;
196                         }
197                         if(node->right)
198                                 stack[cur++]=node->right;
199                 }
200                 else{
201                         cur_dist= cur_dist*cur_dist;
202
203                         if(cur_dist<min_dist){
204                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
205                                 if(cur_dist<min_dist){
206                                         min_dist=cur_dist;
207                                         min_node=node;
208                                 }
209                                 if(node->right)
210                                         stack[cur++]=node->right;
211                         }
212                         if(node->left)
213                                 stack[cur++]=node->left;
214                 }
215                 if(cur+3 > totstack){
216                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
217                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
218                         if(stack != defaultstack)
219                                 MEM_freeN(stack);
220                         stack=temp;
221                         totstack+=100;
222                 }
223         }
224
225         if(nearest) {
226                 nearest->index= min_node->index;
227                 nearest->dist= sqrt(min_dist);
228                 VecCopyf(nearest->co, min_node->co);
229         }
230
231         if(stack != defaultstack)
232                 MEM_freeN(stack);
233
234         return min_node->index;
235 }
236
237 static void add_nearest(KDTreeNearest *ptn, int *found, int n, int index, float dist, float *co)
238 {
239         int i;
240
241         if(*found<n) (*found)++;
242
243         for(i=*found-1; i>0; i--) {
244                 if(dist >= ptn[i-1].dist)
245                         break;
246                 else
247                         ptn[i]= ptn[i-1];
248         }
249
250         ptn[i].index= index;
251         ptn[i].dist= dist;
252         VecCopyf(ptn[i].co, co);
253 }
254
255 /* finds the nearest n entries in tree to specified coordinates */
256 int     BLI_kdtree_find_n_nearest(KDTree *tree, int n, float *co, float *nor, KDTreeNearest *nearest)
257 {
258         KDTreeNode *root, *node=0;
259         KDTreeNode **stack, *defaultstack[100];
260         float cur_dist;
261         int i, totstack, cur=0, found=0;
262
263         if(!tree->root)
264                 return 0;
265
266         stack= defaultstack;
267         totstack= 100;
268
269         root= tree->root;
270
271         cur_dist= squared_distance(root->co,co,root->nor,nor);
272         add_nearest(nearest,&found,n,root->index,cur_dist,root->co);
273         
274         if(co[root->d] < root->co[root->d]) {
275                 if(root->right)
276                         stack[cur++]=root->right;
277                 if(root->left)
278                         stack[cur++]=root->left;
279         }
280         else {
281                 if(root->left)
282                         stack[cur++]=root->left;
283                 if(root->right)
284                         stack[cur++]=root->right;
285         }
286
287         while(cur--){
288                 node=stack[cur];
289
290                 cur_dist = node->co[node->d] - co[node->d];
291
292                 if(cur_dist<0.0){
293                         cur_dist= -cur_dist*cur_dist;
294
295                         if(found<n || -cur_dist<nearest[found-1].dist){
296                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
297
298                                 if(found<n || cur_dist<nearest[found-1].dist)
299                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
300
301                                 if(node->left)
302                                         stack[cur++]=node->left;
303                         }
304                         if(node->right)
305                                 stack[cur++]=node->right;
306                 }
307                 else{
308                         cur_dist= cur_dist*cur_dist;
309
310                         if(found<n || cur_dist<nearest[found-1].dist){
311                                 cur_dist=squared_distance(node->co,co,node->nor,nor);
312                                 if(found<n || cur_dist<nearest[found-1].dist)
313                                         add_nearest(nearest,&found,n,node->index,cur_dist,node->co);
314
315                                 if(node->right)
316                                         stack[cur++]=node->right;
317                         }
318                         if(node->left)
319                                 stack[cur++]=node->left;
320                 }
321                 if(cur+3 > totstack){
322                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
323                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
324                         if(stack != defaultstack)
325                                 MEM_freeN(stack);
326                         stack=temp;
327                         totstack+=100;
328                 }
329         }
330
331         for(i=0; i<found; i++)
332                 nearest[i].dist= sqrt(nearest[i].dist);
333
334         if(stack != defaultstack)
335                 MEM_freeN(stack);
336
337         return found;
338 }
339
340 int range_compare(const void * a, const void * b)
341 {
342         const KDTreeNearest *kda = a;
343         const KDTreeNearest *kdb = b;
344
345         if(kda->dist < kdb->dist)
346                 return -1;
347         else if(kda->dist > kdb->dist)
348                 return 1;
349         else
350                 return 0;
351 }
352 static void add_in_range(KDTreeNearest **ptn, int found, int *totfoundstack, int index, float dist, float *co)
353 {
354         KDTreeNearest *to;
355
356         if(found+1 > *totfoundstack) {
357                 KDTreeNearest *temp=MEM_callocN((*totfoundstack+50)*sizeof(KDTreeNode), "psys_treefoundstack");
358                 memcpy(temp, *ptn, *totfoundstack * sizeof(KDTreeNearest));
359                 if(*ptn)
360                         MEM_freeN(*ptn);
361                 *ptn = temp;
362                 *totfoundstack+=50;
363         }
364
365         to = (*ptn) + found;
366
367         to->index = index;
368         to->dist = sqrt(dist);
369         VecCopyf(to->co, co);
370 }
371 int BLI_kdtree_range_search(KDTree *tree, float range, float *co, float *nor, KDTreeNearest **nearest)
372 {
373         KDTreeNode *root, *node=0;
374         KDTreeNode **stack, *defaultstack[100];
375         KDTreeNearest *foundstack=NULL;
376         float range2 = range*range, dist2;
377         int i, totstack, cur=0, found=0, totfoundstack=0;
378
379         if(!tree || !tree->root)
380                 return 0;
381
382         stack= defaultstack;
383         totstack= 100;
384
385         root= tree->root;
386
387         if(co[root->d] + range < root->co[root->d]) {
388                 if(root->left)
389                         stack[cur++]=root->left;
390         }
391         else if(co[root->d] - range > root->co[root->d]) {
392                 if(root->right)
393                         stack[cur++]=root->right;
394         }
395         else {
396                 dist2 = squared_distance(root->co, co, root->nor, nor);
397                 if(dist2  <= range2)
398                         add_in_range(&foundstack, found++, &totfoundstack, root->index, dist2, root->co);
399
400                 if(root->left)
401                         stack[cur++]=root->left;
402                 if(root->right)
403                         stack[cur++]=root->right;
404         }
405
406         while(cur--) {
407                 node=stack[cur];
408
409                 if(co[node->d] + range < node->co[node->d]) {
410                         if(node->left)
411                                 stack[cur++]=node->left;
412                 }
413                 else if(co[node->d] - range > node->co[node->d]) {
414                         if(node->right)
415                                 stack[cur++]=node->right;
416                 }
417                 else {
418                         dist2 = squared_distance(node->co, co, node->nor, nor);
419                         if(dist2 <= range2)
420                                 add_in_range(&foundstack, found++, &totfoundstack, node->index, dist2, node->co);
421
422                         if(node->left)
423                                 stack[cur++]=node->left;
424                         if(node->right)
425                                 stack[cur++]=node->right;
426                 }
427
428                 if(cur+3 > totstack){
429                         KDTreeNode **temp=MEM_callocN((totstack+100)*sizeof(KDTreeNode*), "psys_treestack");
430                         memcpy(temp,stack,totstack*sizeof(KDTreeNode*));
431                         if(stack != defaultstack)
432                                 MEM_freeN(stack);
433                         stack=temp;
434                         totstack+=100;
435                 }
436         }
437
438         if(stack != defaultstack)
439                 MEM_freeN(stack);
440
441         if(found)
442                 qsort(foundstack, found, sizeof(KDTreeNearest), range_compare);
443
444         *nearest = foundstack;
445
446         return found;
447 }