BKI_kdtree: add a find that takes filter callback
Useful when we need to selectively ignore nodes.
This commit is contained in:
parent
ee719e8816
commit
54b95c30ae
|
@ -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)
|
||||
|
||||
int BLI_kdtree_find_nearest_cb(
|
||||
const KDTree *tree, const float co[3],
|
||||
int (*filter_cb)(void *user_data, int index, const float co[3], float dist_sq), void *user_data,
|
||||
KDTreeNearest *r_nearest);
|
||||
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);
|
||||
|
|
|
@ -291,6 +291,112 @@ int BLI_kdtree_find_nearest(
|
|||
return min_node->index;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* A version of #BLI_kdtree_find_nearest which runs a callback
|
||||
* to filter out values.
|
||||
*
|
||||
* \param filter_cb: Filter find results,
|
||||
* Return codes: (1: accept, 0: skip, -1: immediate exit).
|
||||
*/
|
||||
int BLI_kdtree_find_nearest_cb(
|
||||
const KDTree *tree, const float co[3],
|
||||
int (*filter_cb)(void *user_data, int index, const float co[3], float dist_sq), void *user_data,
|
||||
KDTreeNearest *r_nearest)
|
||||
{
|
||||
const KDTreeNode *nodes = tree->nodes;
|
||||
const KDTreeNode *min_node = NULL;
|
||||
|
||||
unsigned int *stack, defaultstack[KD_STACK_INIT];
|
||||
float min_dist = FLT_MAX, cur_dist;
|
||||
unsigned int totstack, cur = 0;
|
||||
|
||||
#ifdef DEBUG
|
||||
BLI_assert(tree->is_balanced == true);
|
||||
#endif
|
||||
|
||||
if (UNLIKELY(tree->root == KD_NODE_UNSET))
|
||||
return -1;
|
||||
|
||||
stack = defaultstack;
|
||||
totstack = KD_STACK_INIT;
|
||||
|
||||
#define NODE_TEST_NEAREST(node) \
|
||||
{ \
|
||||
const float dist_sq = len_squared_v3v3((node)->co, co); \
|
||||
if (dist_sq < min_dist) { \
|
||||
const int result = filter_cb(user_data, (node)->index, (node)->co, dist_sq); \
|
||||
if (result == 1) { \
|
||||
min_dist = dist_sq; \
|
||||
min_node = node; \
|
||||
} \
|
||||
else if (result == 0) { \
|
||||
/* pass */ \
|
||||
} \
|
||||
else { \
|
||||
BLI_assert(result == -1); \
|
||||
goto finally; \
|
||||
} \
|
||||
} \
|
||||
} ((void)0)
|
||||
|
||||
stack[cur++] = tree->root;
|
||||
|
||||
while (cur--) {
|
||||
const KDTreeNode *node = &nodes[stack[cur]];
|
||||
|
||||
cur_dist = node->co[node->d] - co[node->d];
|
||||
|
||||
if (cur_dist < 0.0f) {
|
||||
cur_dist = -cur_dist * cur_dist;
|
||||
|
||||
if (-cur_dist < min_dist) {
|
||||
NODE_TEST_NEAREST(node);
|
||||
|
||||
if (node->left != KD_NODE_UNSET)
|
||||
stack[cur++] = node->left;
|
||||
}
|
||||
if (node->right != KD_NODE_UNSET)
|
||||
stack[cur++] = node->right;
|
||||
}
|
||||
else {
|
||||
cur_dist = cur_dist * cur_dist;
|
||||
|
||||
if (cur_dist < min_dist) {
|
||||
NODE_TEST_NEAREST(node);
|
||||
|
||||
if (node->right != KD_NODE_UNSET)
|
||||
stack[cur++] = node->right;
|
||||
}
|
||||
if (node->left != KD_NODE_UNSET)
|
||||
stack[cur++] = node->left;
|
||||
}
|
||||
if (UNLIKELY(cur + 3 > totstack)) {
|
||||
stack = realloc_nodes(stack, &totstack, defaultstack != stack);
|
||||
}
|
||||
}
|
||||
|
||||
#undef NODE_TEST_NEAREST
|
||||
|
||||
|
||||
finally:
|
||||
if (stack != defaultstack)
|
||||
MEM_freeN(stack);
|
||||
|
||||
if (min_node) {
|
||||
if (r_nearest) {
|
||||
r_nearest->index = min_node->index;
|
||||
r_nearest->dist = sqrtf(min_dist);
|
||||
copy_v3_v3(r_nearest->co, min_node->co);
|
||||
}
|
||||
|
||||
return min_node->index;
|
||||
}
|
||||
else {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
static void add_nearest(KDTreeNearest *ptn, unsigned int *found, unsigned int n, int index,
|
||||
float dist, const float *co)
|
||||
{
|
||||
|
|
Loading…
Reference in New Issue