BLI_kdopbvh: add an option to use a priority queue in find_nearest.

Simple find_nearest relies on a heuristic for efficient culling of
the BVH tree, which involves a fast callback that always updates the
result, and the caller reusing the result of the previous find_nearest
to prime the process for the next vertex.

If the callback is slow and/or applies significant restrictions on
what kind of nodes can qualify for the result, the heuristic can't
work. Thus for such tasks it is necessary to order and prune nodes
before the callback at BVH tree level using a priority queue.

Since, according to code history, for simple find_nearest the
heuristic approach is faster, this mode has to be an option.

source/blender/blenlib/BLI_kdopbvh.h

@@ -84,6 +84,10 @@ typedef struct BVHTreeRayHit {
 	float dist;         /* distance to the hit point */
 } BVHTreeRayHit;
 
+enum {
+	/* Use a priority queue to process nodes in the optimal order (for slow callbacks) */
+	BVH_NEAREST_OPTIMAL_ORDER   = (1 << 0),
+};
 enum {
 	/* calculate IsectRayPrecalc data */
 	BVH_RAYCAST_WATERTIGHT		= (1 << 0),
@@ -144,6 +148,10 @@ float BLI_bvhtree_get_epsilon(const BVHTree *tree);
 
 /* find nearest node to the given coordinates
  * (if nearest is given it will only search nodes where square distance is smaller than nearest->dist) */
+int BLI_bvhtree_find_nearest_ex(
+        BVHTree *tree, const float co[3], BVHTreeNearest *nearest,
+        BVHTree_NearestPointCallback callback, void *userdata,
+        int flag);
 int BLI_bvhtree_find_nearest(
         BVHTree *tree, const float co[3], BVHTreeNearest *nearest,
         BVHTree_NearestPointCallback callback, void *userdata);

source/blender/blenlib/intern/BLI_kdopbvh.c

@@ -56,6 +56,7 @@
 #include "BLI_kdopbvh.h"
 #include "BLI_math.h"
 #include "BLI_task.h"
+#include "BLI_heap_simple.h"
 
 #include "BLI_strict_flags.h"
 
@@ -1277,7 +1278,7 @@ static float calc_nearest_point_squared(const float proj[3], BVHNode *node, floa
 	return len_squared_v3v3(proj, nearest);
 }
 
-/* TODO: use a priority queue to reduce the number of nodes looked on */
+/* Depth first search method */
 static void dfs_find_nearest_dfs(BVHNearestData *data, BVHNode *node)
 {
 	if (node->totnode == 0) {
@@ -1321,9 +1322,53 @@ static void dfs_find_nearest_begin(BVHNearestData *data, BVHNode *node)
 	dfs_find_nearest_dfs(data, node);
 }
 
-int BLI_bvhtree_find_nearest(
+/* Priority queue method */
+static void heap_find_nearest_inner(BVHNearestData *data, HeapSimple *heap, BVHNode *node)
+{
+	if (node->totnode == 0) {
+		if (data->callback)
+			data->callback(data->userdata, node->index, data->co, &data->nearest);
+		else {
+			data->nearest.index = node->index;
+			data->nearest.dist_sq = calc_nearest_point_squared(data->proj, node, data->nearest.co);
+		}
+	}
+	else {
+		float nearest[3];
+
+		for (int i = 0; i != node->totnode; i++) {
+			float dist_sq = calc_nearest_point_squared(data->proj, node->children[i], nearest);
+
+			if (dist_sq < data->nearest.dist_sq) {
+				BLI_heapsimple_insert(heap, dist_sq, node->children[i]);
+			}
+		}
+	}
+}
+
+static void heap_find_nearest_begin(BVHNearestData *data, BVHNode *root)
+{
+	float nearest[3];
+	float dist_sq = calc_nearest_point_squared(data->proj, root, nearest);
+
+	if (dist_sq < data->nearest.dist_sq) {
+		HeapSimple *heap = BLI_heapsimple_new_ex(32);
+
+		heap_find_nearest_inner(data, heap, root);
+
+		while (!BLI_heapsimple_is_empty(heap) && BLI_heapsimple_top_value(heap) < data->nearest.dist_sq) {
+			BVHNode *node = BLI_heapsimple_pop_min(heap);
+			heap_find_nearest_inner(data, heap, node);
+		}
+
+		BLI_heapsimple_free(heap, NULL);
+	}
+}
+
+int BLI_bvhtree_find_nearest_ex(
         BVHTree *tree, const float co[3], BVHTreeNearest *nearest,
-        BVHTree_NearestPointCallback callback, void *userdata)
+        BVHTree_NearestPointCallback callback, void *userdata,
+        int flag)
 {
 	axis_t axis_iter;
 
@@ -1350,8 +1395,14 @@ int BLI_bvhtree_find_nearest(
 	}
 
 	/* dfs search */
-	if (root)
-		dfs_find_nearest_begin(&data, root);
+	if (root) {
+		if (flag & BVH_NEAREST_OPTIMAL_ORDER) {
+			heap_find_nearest_begin(&data, root);
+		}
+		else {
+			dfs_find_nearest_begin(&data, root);
+		}
+	}
 
 	/* copy back results */
 	if (nearest) {
@@ -1361,6 +1412,13 @@ int BLI_bvhtree_find_nearest(
 	return data.nearest.index;
 }
 
+int BLI_bvhtree_find_nearest(
+        BVHTree *tree, const float co[3], BVHTreeNearest *nearest,
+        BVHTree_NearestPointCallback callback, void *userdata)
+{
+	return BLI_bvhtree_find_nearest_ex(tree, co, nearest, callback, userdata, 0);
+}
+
 /** \} */
 
 

tests/gtests/blenlib/BLI_kdopbvh_test.cc

@@ -52,11 +52,23 @@ TEST(kdopbvh, Single)
 	BLI_bvhtree_free(tree);
 }
 
+void optimal_check_callback(void *userdata, int index, const float co[3], BVHTreeNearest *nearest)
+{
+	float (*points)[3] = (float (*)[3])userdata;
+
+	/* BVH_NEAREST_OPTIMAL_ORDER should hit the right node on the first try */
+	EXPECT_EQ(nearest->index, -1);
+	EXPECT_EQ_ARRAY(co, points[index], 3);
+
+	nearest->index = index;
+	nearest->dist_sq = len_squared_v3v3(co, points[index]);
+}
+
 /**
  * Note that a small epsilon is added to the BVH nodes bounds, even if we pass in zero.
  * Use rounding to ensure very close nodes don't cause the wrong node to be found as nearest.
  */
-static void find_nearest_points_test(int points_len, float scale, int round, int random_seed)
+static void find_nearest_points_test(int points_len, float scale, int round, int random_seed, bool optimal = false)
 {
 	struct RNG *rng = BLI_rng_new(random_seed);
 	BVHTree *tree = BLI_bvhtree_new(points_len, 0.0, 8, 8);
@@ -69,9 +81,13 @@ static void find_nearest_points_test(int points_len, float scale, int round, int
 		BLI_bvhtree_insert(tree, i, points[i], 1);
 	}
 	BLI_bvhtree_balance(tree);
+
 	/* first find each point */
+	BVHTree_NearestPointCallback callback = optimal ? optimal_check_callback : NULL;
+	int flags = optimal ? BVH_NEAREST_OPTIMAL_ORDER : 0;
+
 	for (int i = 0; i < points_len; i++) {
-		const int j = BLI_bvhtree_find_nearest(tree, points[i], NULL, NULL, NULL);
+		const int j = BLI_bvhtree_find_nearest_ex(tree, points[i], NULL, callback, points, flags);
 		if (j != i) {
 #if 0
 			const float dist = len_v3v3(points[i], points[j]);
@@ -95,3 +111,7 @@ static void find_nearest_points_test(int points_len, float scale, int round, int
 TEST(kdopbvh, FindNearest_1)		{ find_nearest_points_test(1, 1.0, 1000, 1234); }
 TEST(kdopbvh, FindNearest_2)		{ find_nearest_points_test(2, 1.0, 1000, 123); }
 TEST(kdopbvh, FindNearest_500)		{ find_nearest_points_test(500, 1.0, 1000, 12); }
+
+TEST(kdopbvh, OptimalFindNearest_1)		{ find_nearest_points_test(1, 1.0, 1000, 1234, true); }
+TEST(kdopbvh, OptimalFindNearest_2)		{ find_nearest_points_test(2, 1.0, 1000, 123, true); }
+TEST(kdopbvh, OptimalFindNearest_500)		{ find_nearest_points_test(500, 1.0, 1000, 12, true); }