KDTree: add BLI_kdtree_range_search_cb

This performs a range search on the kdtree, running a callback instead of allocating an array.

Allows the caller to perform extra checks in the case of overlap,
avoids redundant array allocations, since caller can handle matches.
This commit is contained in:
Campbell Barton 2015-11-18 09:01:55 +11:00
parent 1dc1e9e4ae
commit d6f9ba76a5
2 changed files with 91 additions and 0 deletions

View File

@ -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(

View File

@ -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);
}