Add BLI_bvhtree_walk_dfs utility function

This generic function allows callers to walk the tree using callbacks to define behavior.

2 files changed, 105 insertions, 2 deletions

diff --git a/source/blender/blenlib/BLI_kdopbvh.h b/source/blender/blenlib/BLI_kdopbvh.h
index 7f63801699f..2349be11aa7 100644
--- a/source/blender/blenlib/BLI_kdopbvh.h
+++ b/source/blender/blenlib/BLI_kdopbvh.h
@@ -43,6 +43,16 @@ struct BVHTree;
 typedef struct BVHTree BVHTree;
 #define USE_KDOPBVH_WATERTIGHT
 
+typedef struct BVHTreeAxisRange {
+	union {
+		struct {
+			float min, max;
+		};
+		/* alternate access */
+		float range[2];
+	};
+} BVHTreeAxisRange;
+
 typedef struct BVHTreeOverlap {
 	int indexA;
 	int indexB;
@@ -85,8 +95,8 @@ typedef void (*BVHTree_NearestPointCallback)(void *userdata, int index, const fl
 /* callback must update hit in case it finds a nearest successful hit */
 typedef void (*BVHTree_RayCastCallback)(void *userdata, int index, const BVHTreeRay *ray, BVHTreeRayHit *hit);
 
-/* callback must update nearest to ray in case it finds a nearest result */
-typedef void(*BVHTree_NearestToRayCallback)(void *userdata, int index, const BVHTreeRay *ray, BVHTreeNearest *nearest);
+/* callback must update nearest in case it finds a nearest result */
+typedef void (*BVHTree_NearestToRayCallback)(void *userdata, int index, const BVHTreeRay *ray, BVHTreeNearest *nearest);
 
 /* callback to check if 2 nodes overlap (use thread if intersection results need to be stored) */
 typedef bool (*BVHTree_OverlapCallback)(void *userdata, int index_a, int index_b, int thread);
@@ -94,6 +104,16 @@ typedef bool (*BVHTree_OverlapCallback)(void *userdata, int index_a, int index_b
 /* callback to range search query */
 typedef void (*BVHTree_RangeQuery)(void *userdata, int index, float dist_sq);
 
+
+/* callbacks to BLI_bvhtree_walk_dfs */
+/* return true to traverse into this nodes children, else skip. */
+typedef bool (*BVHTree_WalkParentCallback)(const BVHTreeAxisRange *bounds, void *userdata);
+/* return true to keep walking, else early-exit the search. */
+typedef bool (*BVHTree_WalkLeafCallback)(const BVHTreeAxisRange *bounds, int index, void *userdata);
+/* return true to search (min, max) else (max, min). */
+typedef bool (*BVHTree_WalkOrderCallback)(const BVHTreeAxisRange *bounds, char axis, void *userdata);
+
+
 BVHTree *BLI_bvhtree_new(int maxsize, float epsilon, char tree_type, char axis);
 void BLI_bvhtree_free(BVHTree *tree);
 
@@ -145,6 +165,13 @@ float BLI_bvhtree_bb_raycast(const float bv[6], const float light_start[3], cons
 int BLI_bvhtree_range_query(BVHTree *tree, const float co[3], float radius,
                             BVHTree_RangeQuery callback, void *userdata);
 
+void BLI_bvhtree_walk_dfs(
+        BVHTree *tree,
+        BVHTree_WalkParentCallback walk_parent_cb,
+        BVHTree_WalkLeafCallback walk_leaf_cb,
+        BVHTree_WalkOrderCallback walk_order_cb,
+        void *userdata);
+
 
 /* expose for bvh callbacks to use */
 extern const float bvhtree_kdop_axes[13][3];
diff --git a/source/blender/blenlib/intern/BLI_kdopbvh.c b/source/blender/blenlib/intern/BLI_kdopbvh.c
index e8ff84e8f5f..baef5ab06d2 100644
--- a/source/blender/blenlib/intern/BLI_kdopbvh.c
+++ b/source/blender/blenlib/intern/BLI_kdopbvh.c
@@ -1998,3 +1998,79 @@ int BLI_bvhtree_range_query(BVHTree *tree, const float co[3], float radius, BVHT
 
 	return data.hits;
 }
+
+
+/* -------------------------------------------------------------------- */
+
+/** \name BLI_bvhtree_walk_dfs
+ * \{ */
+
+/**
+ * Runs first among nodes children of the first node before going to the next node in the same layer.
+ *
+ * \return false to break out of the search early.
+ */
+static bool bvhtree_walk_dfs_recursive(
+        BVHTree_WalkParentCallback walk_parent_cb,
+        BVHTree_WalkLeafCallback walk_leaf_cb,
+        BVHTree_WalkOrderCallback walk_order_cb,
+        const BVHNode *node, void *userdata)
+{
+	if (node->totnode == 0) {
+		return walk_leaf_cb((const BVHTreeAxisRange *)node->bv, node->index, userdata);
+	}
+	else {
+		/* First pick the closest node to recurse into */
+		if (walk_order_cb((const BVHTreeAxisRange *)node->bv, node->main_axis, userdata)) {
+			for (int i = 0; i != node->totnode; i++) {
+				if (walk_parent_cb((const BVHTreeAxisRange *)node->bv, userdata)) {
+					if (!bvhtree_walk_dfs_recursive(
+					        walk_parent_cb, walk_leaf_cb, walk_order_cb,
+					        node->children[i], userdata))
+					{
+						return false;
+					}
+				}
+			}
+		}
+		else {
+			for (int i = node->totnode - 1; i >= 0; i--) {
+				if (walk_parent_cb((const BVHTreeAxisRange *)node->bv, userdata)) {
+					if (!bvhtree_walk_dfs_recursive(
+					        walk_parent_cb, walk_leaf_cb, walk_order_cb,
+					        node->children[i], userdata))
+					{
+						return false;
+					}
+				}
+			}
+		}
+	}
+	return true;
+}
+
+/**
+ * This is a generic function to perform a depth first search on the BVHTree
+ * where the search order and nodes traversed depend on callbacks passed in.
+ *
+ * \param tree: Tree to walk.
+ * \param walk_parent_cb: Callback on a parents bound-box to test if it should be traversed.
+ * \param walk_leaf_cb: Callback to test leaf nodes, callback must store its own result,
+ * returning false exits early.
+ * \param walk_order_cb: Callback that indicates which direction to search,
+ * either from the node with the lower or higher k-dop axis value.
+ * \param userdata: Argument passed to all callbacks.
+ */
+void BLI_bvhtree_walk_dfs(
+        BVHTree *tree,
+        BVHTree_WalkParentCallback walk_parent_cb,
+        BVHTree_WalkLeafCallback walk_leaf_cb,
+        BVHTree_WalkOrderCallback walk_order_cb, void *userdata)
+{
+	const BVHNode *root = tree->nodes[tree->totleaf];
+	if (root != NULL) {
+		bvhtree_walk_dfs_recursive(walk_parent_cb, walk_leaf_cb, walk_order_cb, root, userdata);
+	}
+}
+
+/** \} */