diff options
-rw-r--r-- | source/blender/blenlib/BLI_kdtree.h | 4 | ||||
-rw-r--r-- | source/blender/blenlib/intern/BLI_kdtree.c | 87 |
2 files changed, 91 insertions, 0 deletions
diff --git a/source/blender/blenlib/BLI_kdtree.h b/source/blender/blenlib/BLI_kdtree.h index a3ecfb898c3..d488dbce1fd 100644 --- a/source/blender/blenlib/BLI_kdtree.h +++ b/source/blender/blenlib/BLI_kdtree.h @@ -58,6 +58,10 @@ int BLI_kdtree_find_nearest( #define BLI_kdtree_range_search(tree, co, r_nearest, range) \ BLI_kdtree_range_search__normal(tree, co, NULL, r_nearest, range) +void BLI_kdtree_range_search_cb( + const KDTree *tree, const float co[3], float range, + bool (*search_cb)(void *user_data, int index, const float co[3], float dist_sq), void *user_data); + /* Normal use is deprecated */ /* remove __normal functions when last users drop */ int BLI_kdtree_find_nearest_n__normal( diff --git a/source/blender/blenlib/intern/BLI_kdtree.c b/source/blender/blenlib/intern/BLI_kdtree.c index bf8c78b5ab6..acb1e4dc136 100644 --- a/source/blender/blenlib/intern/BLI_kdtree.c +++ b/source/blender/blenlib/intern/BLI_kdtree.c @@ -524,3 +524,90 @@ int BLI_kdtree_range_search__normal( return (int)found; } + +/** + * A version of #BLI_kdtree_range_search which runs a callback + * instead of allocating an array. + * + * \param search_cb: Called for every node found in \a range, false return value performs an early exit. + * + * \note the order of calls isn't sorted based on distance. + */ +void BLI_kdtree_range_search_cb( + const KDTree *tree, const float co[3], float range, + bool (*search_cb)(void *user_data, int index, const float co[3], float dist_sq), void *user_data) +{ + const KDTreeNode *nodes = tree->nodes; + const KDTreeNode *root; + unsigned int *stack, defaultstack[KD_STACK_INIT]; + float range_sq = range * range, dist_sq; + unsigned int totstack, cur = 0; + +#ifdef DEBUG + BLI_assert(tree->is_balanced == true); +#endif + + if (UNLIKELY(tree->root == KD_NODE_UNSET)) + return false; + + stack = defaultstack; + totstack = KD_STACK_INIT; + + root = &nodes[tree->root]; + + if (co[root->d] + range < root->co[root->d]) { + if (root->left != KD_NODE_UNSET) + stack[cur++] = root->left; + } + else if (co[root->d] - range > root->co[root->d]) { + if (root->right != KD_NODE_UNSET) + stack[cur++] = root->right; + } + else { + dist_sq = len_squared_v3v3(root->co, co); + if (dist_sq <= range_sq) { + if (search_cb(user_data, root->index, root->co, dist_sq) == false) { + goto finally; + } + } + + if (root->left != KD_NODE_UNSET) + stack[cur++] = root->left; + if (root->right != KD_NODE_UNSET) + stack[cur++] = root->right; + } + + while (cur--) { + const KDTreeNode *node = &nodes[stack[cur]]; + + if (co[node->d] + range < node->co[node->d]) { + if (node->left != KD_NODE_UNSET) + stack[cur++] = node->left; + } + else if (co[node->d] - range > node->co[node->d]) { + if (node->right != KD_NODE_UNSET) + stack[cur++] = node->right; + } + else { + dist_sq = len_squared_v3v3(node->co, co); + if (dist_sq <= range_sq) { + if (search_cb(user_data, node->index, node->co, dist_sq) == false) { + goto finally; + } + } + + if (node->left != KD_NODE_UNSET) + stack[cur++] = node->left; + if (node->right != KD_NODE_UNSET) + stack[cur++] = node->right; + } + + if (UNLIKELY(cur + 3 > totstack)) { + stack = realloc_nodes(stack, &totstack, defaultstack != stack); + } + } + +finally: + if (stack != defaultstack) + MEM_freeN(stack); +} |