BKI_kdtree: add a find that takes filter callback

Useful when we need to selectively ignore nodes.

2 files changed, 110 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_kdtree.h b/source/blender/blenlib/BLI_kdtree.h
index d488dbce1fd..aa54e1c823c 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)
 
+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);
diff --git a/source/blender/blenlib/intern/BLI_kdtree.c b/source/blender/blenlib/intern/BLI_kdtree.c
index ae3a1f6ff3d..a81f9b28b83 100644
--- a/source/blender/blenlib/intern/BLI_kdtree.c
+++ b/source/blender/blenlib/intern/BLI_kdtree.c
@@ -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)
 {