Cleanup: comments (long lines) in blenlib
[blender.git] / source / blender / blenlib / intern / kdtree_impl.h
1 /*
2  * This program is free software; you can redistribute it and/or
3  * modify it under the terms of the GNU General Public License
4  * as published by the Free Software Foundation; either version 2
5  * of the License, or (at your option) any later version.
6  *
7  * This program is distributed in the hope that it will be useful,
8  * but WITHOUT ANY WARRANTY; without even the implied warranty of
9  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
10  * GNU General Public License for more details.
11  *
12  * You should have received a copy of the GNU General Public License
13  * along with this program; if not, write to the Free Software Foundation,
14  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
15  */
16
17 /** \file
18  * \ingroup bli
19  */
20
21 #include "MEM_guardedalloc.h"
22
23 #include "BLI_math.h"
24 #include "BLI_kdtree_impl.h"
25 #include "BLI_utildefines.h"
26 #include "BLI_strict_flags.h"
27
28 #define _CONCAT_AUX(MACRO_ARG1, MACRO_ARG2) MACRO_ARG1##MACRO_ARG2
29 #define _CONCAT(MACRO_ARG1, MACRO_ARG2) _CONCAT_AUX(MACRO_ARG1, MACRO_ARG2)
30 #define BLI_kdtree_nd_(id) _CONCAT(KDTREE_PREFIX_ID, _##id)
31
32 typedef struct KDTreeNode_head {
33   uint left, right;
34   float co[KD_DIMS];
35   int index;
36 } KDTreeNode_head;
37
38 typedef struct KDTreeNode {
39   uint left, right;
40   float co[KD_DIMS];
41   int index;
42   uint d; /* range is only (0..KD_DIMS - 1) */
43 } KDTreeNode;
44
45 struct KDTree {
46   KDTreeNode *nodes;
47   uint nodes_len;
48   uint root;
49 #ifdef DEBUG
50   bool is_balanced;        /* ensure we call balance first */
51   uint nodes_len_capacity; /* max size of the tree */
52 #endif
53 };
54
55 #define KD_STACK_INIT 100     /* initial size for array (on the stack) */
56 #define KD_NEAR_ALLOC_INC 100 /* alloc increment for collecting nearest */
57 #define KD_FOUND_ALLOC_INC 50 /* alloc increment for collecting nearest */
58
59 #define KD_NODE_UNSET ((uint)-1)
60
61 /** When set we know all values are unbalanced,
62  * otherwise clear them when re-balancing: see T62210. */
63 #define KD_NODE_ROOT_IS_INIT ((uint)-2)
64
65 /* -------------------------------------------------------------------- */
66 /** \name Local Math API
67  * \{ */
68
69 static void copy_vn_vn(float v0[KD_DIMS], const float v1[KD_DIMS])
70 {
71   for (uint j = 0; j < KD_DIMS; j++) {
72     v0[j] = v1[j];
73   }
74 }
75
76 static float len_squared_vnvn(const float v0[KD_DIMS], const float v1[KD_DIMS])
77 {
78   float d = 0.0f;
79   for (uint j = 0; j < KD_DIMS; j++) {
80     d += SQUARE(v0[j] - v1[j]);
81   }
82   return d;
83 }
84
85 static float len_squared_vnvn_cb(const float co_kdtree[KD_DIMS],
86                                  const float co_search[KD_DIMS],
87                                  const void *UNUSED(user_data))
88 {
89   return len_squared_vnvn(co_kdtree, co_search);
90 }
91
92 /** \} */
93
94 /**
95  * Creates or free a kdtree
96  */
97 KDTree *BLI_kdtree_nd_(new)(uint nodes_len_capacity)
98 {
99   KDTree *tree;
100
101   tree = MEM_mallocN(sizeof(KDTree), "KDTree");
102   tree->nodes = MEM_mallocN(sizeof(KDTreeNode) * nodes_len_capacity, "KDTreeNode");
103   tree->nodes_len = 0;
104   tree->root = KD_NODE_ROOT_IS_INIT;
105
106 #ifdef DEBUG
107   tree->is_balanced = false;
108   tree->nodes_len_capacity = nodes_len_capacity;
109 #endif
110
111   return tree;
112 }
113
114 void BLI_kdtree_nd_(free)(KDTree *tree)
115 {
116   if (tree) {
117     MEM_freeN(tree->nodes);
118     MEM_freeN(tree);
119   }
120 }
121
122 /**
123  * Construction: first insert points, then call balance. Normal is optional.
124  */
125 void BLI_kdtree_nd_(insert)(KDTree *tree, int index, const float co[KD_DIMS])
126 {
127   KDTreeNode *node = &tree->nodes[tree->nodes_len++];
128
129 #ifdef DEBUG
130   BLI_assert(tree->nodes_len <= tree->nodes_len_capacity);
131 #endif
132
133   /* note, array isn't calloc'd,
134    * need to initialize all struct members */
135
136   node->left = node->right = KD_NODE_UNSET;
137   copy_vn_vn(node->co, co);
138   node->index = index;
139   node->d = 0;
140
141 #ifdef DEBUG
142   tree->is_balanced = false;
143 #endif
144 }
145
146 static uint kdtree_balance(KDTreeNode *nodes, uint nodes_len, uint axis, const uint ofs)
147 {
148   KDTreeNode *node;
149   float co;
150   uint left, right, median, i, j;
151
152   if (nodes_len <= 0) {
153     return KD_NODE_UNSET;
154   }
155   else if (nodes_len == 1) {
156     return 0 + ofs;
157   }
158
159   /* quicksort style sorting around median */
160   left = 0;
161   right = nodes_len - 1;
162   median = nodes_len / 2;
163
164   while (right > left) {
165     co = nodes[right].co[axis];
166     i = left - 1;
167     j = right;
168
169     while (1) {
170       while (nodes[++i].co[axis] < co) { /* pass */
171       }
172       while (nodes[--j].co[axis] > co && j > left) { /* pass */
173       }
174
175       if (i >= j) {
176         break;
177       }
178
179       SWAP(KDTreeNode_head, *(KDTreeNode_head *)&nodes[i], *(KDTreeNode_head *)&nodes[j]);
180     }
181
182     SWAP(KDTreeNode_head, *(KDTreeNode_head *)&nodes[i], *(KDTreeNode_head *)&nodes[right]);
183     if (i >= median) {
184       right = i - 1;
185     }
186     if (i <= median) {
187       left = i + 1;
188     }
189   }
190
191   /* set node and sort subnodes */
192   node = &nodes[median];
193   node->d = axis;
194   axis = (axis + 1) % KD_DIMS;
195   node->left = kdtree_balance(nodes, median, axis, ofs);
196   node->right = kdtree_balance(
197       nodes + median + 1, (nodes_len - (median + 1)), axis, (median + 1) + ofs);
198
199   return median + ofs;
200 }
201
202 void BLI_kdtree_nd_(balance)(KDTree *tree)
203 {
204   if (tree->root != KD_NODE_ROOT_IS_INIT) {
205     for (uint i = 0; i < tree->nodes_len; i++) {
206       tree->nodes[i].left = KD_NODE_UNSET;
207       tree->nodes[i].right = KD_NODE_UNSET;
208     }
209   }
210
211   tree->root = kdtree_balance(tree->nodes, tree->nodes_len, 0, 0);
212
213 #ifdef DEBUG
214   tree->is_balanced = true;
215 #endif
216 }
217
218 static uint *realloc_nodes(uint *stack, uint *stack_len_capacity, const bool is_alloc)
219 {
220   uint *stack_new = MEM_mallocN((*stack_len_capacity + KD_NEAR_ALLOC_INC) * sizeof(uint),
221                                 "KDTree.treestack");
222   memcpy(stack_new, stack, *stack_len_capacity * sizeof(uint));
223   // memset(stack_new + *stack_len_capacity, 0, sizeof(uint) * KD_NEAR_ALLOC_INC);
224   if (is_alloc) {
225     MEM_freeN(stack);
226   }
227   *stack_len_capacity += KD_NEAR_ALLOC_INC;
228   return stack_new;
229 }
230
231 /**
232  * Find nearest returns index, and -1 if no node is found.
233  */
234 int BLI_kdtree_nd_(find_nearest)(const KDTree *tree,
235                                  const float co[KD_DIMS],
236                                  KDTreeNearest *r_nearest)
237 {
238   const KDTreeNode *nodes = tree->nodes;
239   const KDTreeNode *root, *min_node;
240   uint *stack, stack_default[KD_STACK_INIT];
241   float min_dist, cur_dist;
242   uint stack_len_capacity, cur = 0;
243
244 #ifdef DEBUG
245   BLI_assert(tree->is_balanced == true);
246 #endif
247
248   if (UNLIKELY(tree->root == KD_NODE_UNSET)) {
249     return -1;
250   }
251
252   stack = stack_default;
253   stack_len_capacity = KD_STACK_INIT;
254
255   root = &nodes[tree->root];
256   min_node = root;
257   min_dist = len_squared_vnvn(root->co, co);
258
259   if (co[root->d] < root->co[root->d]) {
260     if (root->right != KD_NODE_UNSET) {
261       stack[cur++] = root->right;
262     }
263     if (root->left != KD_NODE_UNSET) {
264       stack[cur++] = root->left;
265     }
266   }
267   else {
268     if (root->left != KD_NODE_UNSET) {
269       stack[cur++] = root->left;
270     }
271     if (root->right != KD_NODE_UNSET) {
272       stack[cur++] = root->right;
273     }
274   }
275
276   while (cur--) {
277     const KDTreeNode *node = &nodes[stack[cur]];
278
279     cur_dist = node->co[node->d] - co[node->d];
280
281     if (cur_dist < 0.0f) {
282       cur_dist = -cur_dist * cur_dist;
283
284       if (-cur_dist < min_dist) {
285         cur_dist = len_squared_vnvn(node->co, co);
286         if (cur_dist < min_dist) {
287           min_dist = cur_dist;
288           min_node = node;
289         }
290         if (node->left != KD_NODE_UNSET) {
291           stack[cur++] = node->left;
292         }
293       }
294       if (node->right != KD_NODE_UNSET) {
295         stack[cur++] = node->right;
296       }
297     }
298     else {
299       cur_dist = cur_dist * cur_dist;
300
301       if (cur_dist < min_dist) {
302         cur_dist = len_squared_vnvn(node->co, co);
303         if (cur_dist < min_dist) {
304           min_dist = cur_dist;
305           min_node = node;
306         }
307         if (node->right != KD_NODE_UNSET) {
308           stack[cur++] = node->right;
309         }
310       }
311       if (node->left != KD_NODE_UNSET) {
312         stack[cur++] = node->left;
313       }
314     }
315     if (UNLIKELY(cur + KD_DIMS > stack_len_capacity)) {
316       stack = realloc_nodes(stack, &stack_len_capacity, stack_default != stack);
317     }
318   }
319
320   if (r_nearest) {
321     r_nearest->index = min_node->index;
322     r_nearest->dist = sqrtf(min_dist);
323     copy_vn_vn(r_nearest->co, min_node->co);
324   }
325
326   if (stack != stack_default) {
327     MEM_freeN(stack);
328   }
329
330   return min_node->index;
331 }
332
333 /**
334  * A version of #BLI_kdtree_3d_find_nearest which runs a callback
335  * to filter out values.
336  *
337  * \param filter_cb: Filter find results,
338  * Return codes: (1: accept, 0: skip, -1: immediate exit).
339  */
340 int BLI_kdtree_nd_(find_nearest_cb)(
341     const KDTree *tree,
342     const float co[KD_DIMS],
343     int (*filter_cb)(void *user_data, int index, const float co[KD_DIMS], float dist_sq),
344     void *user_data,
345     KDTreeNearest *r_nearest)
346 {
347   const KDTreeNode *nodes = tree->nodes;
348   const KDTreeNode *min_node = NULL;
349
350   uint *stack, stack_default[KD_STACK_INIT];
351   float min_dist = FLT_MAX, cur_dist;
352   uint stack_len_capacity, cur = 0;
353
354 #ifdef DEBUG
355   BLI_assert(tree->is_balanced == true);
356 #endif
357
358   if (UNLIKELY(tree->root == KD_NODE_UNSET)) {
359     return -1;
360   }
361
362   stack = stack_default;
363   stack_len_capacity = ARRAY_SIZE(stack_default);
364
365 #define NODE_TEST_NEAREST(node) \
366   { \
367     const float dist_sq = len_squared_vnvn((node)->co, co); \
368     if (dist_sq < min_dist) { \
369       const int result = filter_cb(user_data, (node)->index, (node)->co, dist_sq); \
370       if (result == 1) { \
371         min_dist = dist_sq; \
372         min_node = node; \
373       } \
374       else if (result == 0) { \
375         /* pass */ \
376       } \
377       else { \
378         BLI_assert(result == -1); \
379         goto finally; \
380       } \
381     } \
382   } \
383   ((void)0)
384
385   stack[cur++] = tree->root;
386
387   while (cur--) {
388     const KDTreeNode *node = &nodes[stack[cur]];
389
390     cur_dist = node->co[node->d] - co[node->d];
391
392     if (cur_dist < 0.0f) {
393       cur_dist = -cur_dist * cur_dist;
394
395       if (-cur_dist < min_dist) {
396         NODE_TEST_NEAREST(node);
397
398         if (node->left != KD_NODE_UNSET) {
399           stack[cur++] = node->left;
400         }
401       }
402       if (node->right != KD_NODE_UNSET) {
403         stack[cur++] = node->right;
404       }
405     }
406     else {
407       cur_dist = cur_dist * cur_dist;
408
409       if (cur_dist < min_dist) {
410         NODE_TEST_NEAREST(node);
411
412         if (node->right != KD_NODE_UNSET) {
413           stack[cur++] = node->right;
414         }
415       }
416       if (node->left != KD_NODE_UNSET) {
417         stack[cur++] = node->left;
418       }
419     }
420     if (UNLIKELY(cur + KD_DIMS > stack_len_capacity)) {
421       stack = realloc_nodes(stack, &stack_len_capacity, stack_default != stack);
422     }
423   }
424
425 #undef NODE_TEST_NEAREST
426
427 finally:
428   if (stack != stack_default) {
429     MEM_freeN(stack);
430   }
431
432   if (min_node) {
433     if (r_nearest) {
434       r_nearest->index = min_node->index;
435       r_nearest->dist = sqrtf(min_dist);
436       copy_vn_vn(r_nearest->co, min_node->co);
437     }
438
439     return min_node->index;
440   }
441   else {
442     return -1;
443   }
444 }
445
446 static void nearest_ordered_insert(KDTreeNearest *nearest,
447                                    uint *nearest_len,
448                                    const uint nearest_len_capacity,
449                                    const int index,
450                                    const float dist,
451                                    const float co[KD_DIMS])
452 {
453   uint i;
454
455   if (*nearest_len < nearest_len_capacity) {
456     (*nearest_len)++;
457   }
458
459   for (i = *nearest_len - 1; i > 0; i--) {
460     if (dist >= nearest[i - 1].dist) {
461       break;
462     }
463     else {
464       nearest[i] = nearest[i - 1];
465     }
466   }
467
468   nearest[i].index = index;
469   nearest[i].dist = dist;
470   copy_vn_vn(nearest[i].co, co);
471 }
472
473 /**
474  * Find \a nearest_len_capacity nearest returns number of points found, with results in nearest.
475  *
476  * \param r_nearest: An array of nearest, sized at least \a nearest_len_capacity.
477  */
478 int BLI_kdtree_nd_(find_nearest_n_with_len_squared_cb)(
479     const KDTree *tree,
480     const float co[KD_DIMS],
481     KDTreeNearest r_nearest[],
482     const uint nearest_len_capacity,
483     float (*len_sq_fn)(const float co_search[KD_DIMS],
484                        const float co_test[KD_DIMS],
485                        const void *user_data),
486     const void *user_data)
487 {
488   const KDTreeNode *nodes = tree->nodes;
489   const KDTreeNode *root;
490   uint *stack, stack_default[KD_STACK_INIT];
491   float cur_dist;
492   uint stack_len_capacity, cur = 0;
493   uint i, nearest_len = 0;
494
495 #ifdef DEBUG
496   BLI_assert(tree->is_balanced == true);
497 #endif
498
499   if (UNLIKELY((tree->root == KD_NODE_UNSET) || nearest_len_capacity == 0)) {
500     return 0;
501   }
502
503   if (len_sq_fn == NULL) {
504     len_sq_fn = len_squared_vnvn_cb;
505     BLI_assert(user_data == NULL);
506   }
507
508   stack = stack_default;
509   stack_len_capacity = ARRAY_SIZE(stack_default);
510
511   root = &nodes[tree->root];
512
513   cur_dist = len_sq_fn(co, root->co, user_data);
514   nearest_ordered_insert(
515       r_nearest, &nearest_len, nearest_len_capacity, root->index, cur_dist, root->co);
516
517   if (co[root->d] < root->co[root->d]) {
518     if (root->right != KD_NODE_UNSET) {
519       stack[cur++] = root->right;
520     }
521     if (root->left != KD_NODE_UNSET) {
522       stack[cur++] = root->left;
523     }
524   }
525   else {
526     if (root->left != KD_NODE_UNSET) {
527       stack[cur++] = root->left;
528     }
529     if (root->right != KD_NODE_UNSET) {
530       stack[cur++] = root->right;
531     }
532   }
533
534   while (cur--) {
535     const KDTreeNode *node = &nodes[stack[cur]];
536
537     cur_dist = node->co[node->d] - co[node->d];
538
539     if (cur_dist < 0.0f) {
540       cur_dist = -cur_dist * cur_dist;
541
542       if (nearest_len < nearest_len_capacity || -cur_dist < r_nearest[nearest_len - 1].dist) {
543         cur_dist = len_sq_fn(co, node->co, user_data);
544
545         if (nearest_len < nearest_len_capacity || cur_dist < r_nearest[nearest_len - 1].dist) {
546           nearest_ordered_insert(
547               r_nearest, &nearest_len, nearest_len_capacity, node->index, cur_dist, node->co);
548         }
549
550         if (node->left != KD_NODE_UNSET) {
551           stack[cur++] = node->left;
552         }
553       }
554       if (node->right != KD_NODE_UNSET) {
555         stack[cur++] = node->right;
556       }
557     }
558     else {
559       cur_dist = cur_dist * cur_dist;
560
561       if (nearest_len < nearest_len_capacity || cur_dist < r_nearest[nearest_len - 1].dist) {
562         cur_dist = len_sq_fn(co, node->co, user_data);
563         if (nearest_len < nearest_len_capacity || cur_dist < r_nearest[nearest_len - 1].dist) {
564           nearest_ordered_insert(
565               r_nearest, &nearest_len, nearest_len_capacity, node->index, cur_dist, node->co);
566         }
567
568         if (node->right != KD_NODE_UNSET) {
569           stack[cur++] = node->right;
570         }
571       }
572       if (node->left != KD_NODE_UNSET) {
573         stack[cur++] = node->left;
574       }
575     }
576     if (UNLIKELY(cur + KD_DIMS > stack_len_capacity)) {
577       stack = realloc_nodes(stack, &stack_len_capacity, stack_default != stack);
578     }
579   }
580
581   for (i = 0; i < nearest_len; i++) {
582     r_nearest[i].dist = sqrtf(r_nearest[i].dist);
583   }
584
585   if (stack != stack_default) {
586     MEM_freeN(stack);
587   }
588
589   return (int)nearest_len;
590 }
591
592 int BLI_kdtree_nd_(find_nearest_n)(const KDTree *tree,
593                                    const float co[KD_DIMS],
594                                    KDTreeNearest r_nearest[],
595                                    const uint nearest_len_capacity)
596 {
597   return BLI_kdtree_nd_(find_nearest_n_with_len_squared_cb)(
598       tree, co, r_nearest, nearest_len_capacity, NULL, NULL);
599 }
600
601 static int nearest_cmp_dist(const void *a, const void *b)
602 {
603   const KDTreeNearest *kda = a;
604   const KDTreeNearest *kdb = b;
605
606   if (kda->dist < kdb->dist) {
607     return -1;
608   }
609   else if (kda->dist > kdb->dist) {
610     return 1;
611   }
612   else {
613     return 0;
614   }
615 }
616 static void nearest_add_in_range(KDTreeNearest **r_nearest,
617                                  uint nearest_index,
618                                  uint *nearest_len_capacity,
619                                  const int index,
620                                  const float dist,
621                                  const float co[KD_DIMS])
622 {
623   KDTreeNearest *to;
624
625   if (UNLIKELY(nearest_index >= *nearest_len_capacity)) {
626     *r_nearest = MEM_reallocN_id(
627         *r_nearest, (*nearest_len_capacity += KD_FOUND_ALLOC_INC) * sizeof(KDTreeNode), __func__);
628   }
629
630   to = (*r_nearest) + nearest_index;
631
632   to->index = index;
633   to->dist = sqrtf(dist);
634   copy_vn_vn(to->co, co);
635 }
636
637 /**
638  * Range search returns number of points nearest_len, with results in nearest
639  *
640  * \param r_nearest: Allocated array of nearest nearest_len (caller is responsible for freeing).
641  */
642 int BLI_kdtree_nd_(range_search_with_len_squared_cb)(
643     const KDTree *tree,
644     const float co[KD_DIMS],
645     KDTreeNearest **r_nearest,
646     const float range,
647     float (*len_sq_fn)(const float co_search[KD_DIMS],
648                        const float co_test[KD_DIMS],
649                        const void *user_data),
650     const void *user_data)
651 {
652   const KDTreeNode *nodes = tree->nodes;
653   uint *stack, stack_default[KD_STACK_INIT];
654   KDTreeNearest *nearest = NULL;
655   const float range_sq = range * range;
656   float dist_sq;
657   uint stack_len_capacity, cur = 0;
658   uint nearest_len = 0, nearest_len_capacity = 0;
659
660 #ifdef DEBUG
661   BLI_assert(tree->is_balanced == true);
662 #endif
663
664   if (UNLIKELY(tree->root == KD_NODE_UNSET)) {
665     return 0;
666   }
667
668   if (len_sq_fn == NULL) {
669     len_sq_fn = len_squared_vnvn_cb;
670     BLI_assert(user_data == NULL);
671   }
672
673   stack = stack_default;
674   stack_len_capacity = ARRAY_SIZE(stack_default);
675
676   stack[cur++] = tree->root;
677
678   while (cur--) {
679     const KDTreeNode *node = &nodes[stack[cur]];
680
681     if (co[node->d] + range < node->co[node->d]) {
682       if (node->left != KD_NODE_UNSET) {
683         stack[cur++] = node->left;
684       }
685     }
686     else if (co[node->d] - range > node->co[node->d]) {
687       if (node->right != KD_NODE_UNSET) {
688         stack[cur++] = node->right;
689       }
690     }
691     else {
692       dist_sq = len_sq_fn(co, node->co, user_data);
693       if (dist_sq <= range_sq) {
694         nearest_add_in_range(
695             &nearest, nearest_len++, &nearest_len_capacity, node->index, dist_sq, node->co);
696       }
697
698       if (node->left != KD_NODE_UNSET) {
699         stack[cur++] = node->left;
700       }
701       if (node->right != KD_NODE_UNSET) {
702         stack[cur++] = node->right;
703       }
704     }
705
706     if (UNLIKELY(cur + KD_DIMS > stack_len_capacity)) {
707       stack = realloc_nodes(stack, &stack_len_capacity, stack_default != stack);
708     }
709   }
710
711   if (stack != stack_default) {
712     MEM_freeN(stack);
713   }
714
715   if (nearest_len) {
716     qsort(nearest, nearest_len, sizeof(KDTreeNearest), nearest_cmp_dist);
717   }
718
719   *r_nearest = nearest;
720
721   return (int)nearest_len;
722 }
723
724 int BLI_kdtree_nd_(range_search)(const KDTree *tree,
725                                  const float co[KD_DIMS],
726                                  KDTreeNearest **r_nearest,
727                                  const float range)
728 {
729   return BLI_kdtree_nd_(range_search_with_len_squared_cb)(tree, co, r_nearest, range, NULL, NULL);
730 }
731
732 /**
733  * A version of #BLI_kdtree_3d_range_search which runs a callback
734  * instead of allocating an array.
735  *
736  * \param search_cb: Called for every node found in \a range,
737  * false return value performs an early exit.
738  *
739  * \note the order of calls isn't sorted based on distance.
740  */
741 void BLI_kdtree_nd_(range_search_cb)(
742     const KDTree *tree,
743     const float co[KD_DIMS],
744     float range,
745     bool (*search_cb)(void *user_data, int index, const float co[KD_DIMS], float dist_sq),
746     void *user_data)
747 {
748   const KDTreeNode *nodes = tree->nodes;
749
750   uint *stack, stack_default[KD_STACK_INIT];
751   float range_sq = range * range, dist_sq;
752   uint stack_len_capacity, cur = 0;
753
754 #ifdef DEBUG
755   BLI_assert(tree->is_balanced == true);
756 #endif
757
758   if (UNLIKELY(tree->root == KD_NODE_UNSET)) {
759     return;
760   }
761
762   stack = stack_default;
763   stack_len_capacity = ARRAY_SIZE(stack_default);
764
765   stack[cur++] = tree->root;
766
767   while (cur--) {
768     const KDTreeNode *node = &nodes[stack[cur]];
769
770     if (co[node->d] + range < node->co[node->d]) {
771       if (node->left != KD_NODE_UNSET) {
772         stack[cur++] = node->left;
773       }
774     }
775     else if (co[node->d] - range > node->co[node->d]) {
776       if (node->right != KD_NODE_UNSET) {
777         stack[cur++] = node->right;
778       }
779     }
780     else {
781       dist_sq = len_squared_vnvn(node->co, co);
782       if (dist_sq <= range_sq) {
783         if (search_cb(user_data, node->index, node->co, dist_sq) == false) {
784           goto finally;
785         }
786       }
787
788       if (node->left != KD_NODE_UNSET) {
789         stack[cur++] = node->left;
790       }
791       if (node->right != KD_NODE_UNSET) {
792         stack[cur++] = node->right;
793       }
794     }
795
796     if (UNLIKELY(cur + KD_DIMS > stack_len_capacity)) {
797       stack = realloc_nodes(stack, &stack_len_capacity, stack_default != stack);
798     }
799   }
800
801 finally:
802   if (stack != stack_default) {
803     MEM_freeN(stack);
804   }
805 }
806
807 /**
808  * Use when we want to loop over nodes ordered by index.
809  * Requires indices to be aligned with nodes.
810  */
811 static uint *kdtree_order(const KDTree *tree)
812 {
813   const KDTreeNode *nodes = tree->nodes;
814   uint *order = MEM_mallocN(sizeof(uint) * tree->nodes_len, __func__);
815   for (uint i = 0; i < tree->nodes_len; i++) {
816     order[nodes[i].index] = i;
817   }
818   return order;
819 }
820
821 /* -------------------------------------------------------------------- */
822 /** \name BLI_kdtree_3d_calc_duplicates_fast
823  * \{ */
824
825 struct DeDuplicateParams {
826   /* Static */
827   const KDTreeNode *nodes;
828   float range;
829   float range_sq;
830   int *duplicates;
831   int *duplicates_found;
832
833   /* Per Search */
834   float search_co[KD_DIMS];
835   int search;
836 };
837
838 static void deduplicate_recursive(const struct DeDuplicateParams *p, uint i)
839 {
840   const KDTreeNode *node = &p->nodes[i];
841   if (p->search_co[node->d] + p->range <= node->co[node->d]) {
842     if (node->left != KD_NODE_UNSET) {
843       deduplicate_recursive(p, node->left);
844     }
845   }
846   else if (p->search_co[node->d] - p->range >= node->co[node->d]) {
847     if (node->right != KD_NODE_UNSET) {
848       deduplicate_recursive(p, node->right);
849     }
850   }
851   else {
852     if ((p->search != node->index) && (p->duplicates[node->index] == -1)) {
853       if (len_squared_vnvn(node->co, p->search_co) <= p->range_sq) {
854         p->duplicates[node->index] = (int)p->search;
855         *p->duplicates_found += 1;
856       }
857     }
858     if (node->left != KD_NODE_UNSET) {
859       deduplicate_recursive(p, node->left);
860     }
861     if (node->right != KD_NODE_UNSET) {
862       deduplicate_recursive(p, node->right);
863     }
864   }
865 }
866
867 /**
868  * Find duplicate points in \a range.
869  * Favors speed over quality since it doesn't find the best target vertex for merging.
870  * Nodes are looped over, duplicates are added when found.
871  * Nevertheless results are predictable.
872  *
873  * \param range: Coordinates in this range are candidates to be merged.
874  * \param use_index_order: Loop over the coordinates ordered by #KDTreeNode.index
875  * At the expense of some performance, this ensures the layout of the tree doesn't influence
876  * the iteration order.
877  * \param duplicates: An array of int's the length of #KDTree.nodes_len
878  * Values initialized to -1 are candidates to me merged.
879  * Setting the index to it's own position in the array prevents it from being touched,
880  * although it can still be used as a target.
881  * \returns The number of merges found (includes any merges already in the \a duplicates array).
882  *
883  * \note Merging is always a single step (target indices wont be marked for merging).
884  */
885 int BLI_kdtree_nd_(calc_duplicates_fast)(const KDTree *tree,
886                                          const float range,
887                                          bool use_index_order,
888                                          int *duplicates)
889 {
890   int found = 0;
891   struct DeDuplicateParams p = {
892       .nodes = tree->nodes,
893       .range = range,
894       .range_sq = SQUARE(range),
895       .duplicates = duplicates,
896       .duplicates_found = &found,
897   };
898
899   if (use_index_order) {
900     uint *order = kdtree_order(tree);
901     for (uint i = 0; i < tree->nodes_len; i++) {
902       const uint node_index = order[i];
903       const int index = (int)i;
904       if (ELEM(duplicates[index], -1, index)) {
905         p.search = index;
906         copy_vn_vn(p.search_co, tree->nodes[node_index].co);
907         int found_prev = found;
908         deduplicate_recursive(&p, tree->root);
909         if (found != found_prev) {
910           /* Prevent chains of doubles. */
911           duplicates[index] = index;
912         }
913       }
914     }
915     MEM_freeN(order);
916   }
917   else {
918     for (uint i = 0; i < tree->nodes_len; i++) {
919       const uint node_index = i;
920       const int index = p.nodes[node_index].index;
921       if (ELEM(duplicates[index], -1, index)) {
922         p.search = index;
923         copy_vn_vn(p.search_co, tree->nodes[node_index].co);
924         int found_prev = found;
925         deduplicate_recursive(&p, tree->root);
926         if (found != found_prev) {
927           /* Prevent chains of doubles. */
928           duplicates[index] = index;
929         }
930       }
931     }
932   }
933   return found;
934 }
935
936 /** \} */
937
938 /* -------------------------------------------------------------------- */
939 /** \name BLI_kdtree_3d_deduplicate
940  * \{ */
941
942 static int kdtree_node_cmp_deduplicate(const void *n0_p, const void *n1_p)
943 {
944   const KDTreeNode *n0 = n0_p;
945   const KDTreeNode *n1 = n1_p;
946   for (uint j = 0; j < KD_DIMS; j++) {
947     if (n0->co[j] < n1->co[j]) {
948       return -1;
949     }
950     else if (n0->co[j] > n1->co[j]) {
951       return 1;
952     }
953   }
954   /* Sort by pointer so the first added will be used.
955    * assignment below ignores const correctness,
956    * however the values aren't used for sorting and are to be discarded. */
957   if (n0 < n1) {
958     ((KDTreeNode *)n1)->d = KD_DIMS; /* tag invalid */
959     return -1;
960   }
961   else {
962     ((KDTreeNode *)n0)->d = KD_DIMS; /* tag invalid */
963     return 1;
964   }
965 }
966
967 /**
968  * Remove exact duplicates (run before before balancing).
969  *
970  * Keep the first element added when duplicates are found.
971  */
972 int BLI_kdtree_nd_(deduplicate)(KDTree *tree)
973 {
974 #ifdef DEBUG
975   tree->is_balanced = false;
976 #endif
977   qsort(tree->nodes, (size_t)tree->nodes_len, sizeof(*tree->nodes), kdtree_node_cmp_deduplicate);
978   uint j = 0;
979   for (uint i = 0; i < tree->nodes_len; i++) {
980     if (tree->nodes[i].d != KD_DIMS) {
981       if (i != j) {
982         tree->nodes[j] = tree->nodes[i];
983       }
984       j++;
985     }
986   }
987   tree->nodes_len = j;
988   return (int)tree->nodes_len;
989 }
990
991 /** \} */