BLI_kdtree: refactor boids specific logic into callback

Logic to for boids to avoid head-on collisions was in BLI_kdtree.

Move this into a callback which is now defined in boids.c
so the kdtree code can be kept generic.

3 files changed, 68 insertions, 44 deletions

diff --git a/source/blender/blenkernel/intern/boids.c b/source/blender/blenkernel/intern/boids.c
index 54a8929a7f7..8b71aa0fb69 100644
--- a/source/blender/blenkernel/intern/boids.c
+++ b/source/blender/blenkernel/intern/boids.c
@@ -45,6 +45,23 @@
 
 #include "RNA_enum_types.h"
 
+static float len_squared_v3v3_with_normal_bias(
+        const float co_search[3], const float co_test[3], const void *user_data)
+{
+	const float *normal = user_data;
+	float d[3], dist;
+
+	sub_v3_v3v3(d, co_test, co_search);
+
+	dist = len_squared_v3(d);
+
+	/* Avoid head-on collisions. */
+	if (dot_v3v3(d, normal) < 0.0f) {
+		dist *= 10.0f;
+	}
+	return dist;
+}
+
 typedef struct BoidValues {
 	float max_speed, max_acc;
 	float max_ave, min_speed;
@@ -257,9 +274,9 @@ static int rule_avoid_collision(BoidRule *rule, BoidBrainData *bbd, BoidValues *
 
 	//check boids in own system
 	if (acbr->options & BRULE_ACOLL_WITH_BOIDS) {
-		neighbors = BLI_kdtree_range_search__normal(
-		        bbd->sim->psys->tree, pa->prev_state.co, pa->prev_state.ave,
-		        &ptn, acbr->look_ahead * len_v3(pa->prev_state.vel));
+		neighbors = BLI_kdtree_range_search_with_len_squared_cb(
+		        bbd->sim->psys->tree, pa->prev_state.co, &ptn, acbr->look_ahead * len_v3(pa->prev_state.vel),
+		        len_squared_v3v3_with_normal_bias, pa->prev_state.ave);
 		if (neighbors > 1) for (n=1; n<neighbors; n++) {
 			copy_v3_v3(co1, pa->prev_state.co);
 			copy_v3_v3(vel1, pa->prev_state.vel);
@@ -306,9 +323,9 @@ static int rule_avoid_collision(BoidRule *rule, BoidBrainData *bbd, BoidValues *
 
 		if (epsys) {
 			BLI_assert(epsys->tree != NULL);
-			neighbors = BLI_kdtree_range_search__normal(
-			        epsys->tree, pa->prev_state.co, pa->prev_state.ave,
-			        &ptn, acbr->look_ahead * len_v3(pa->prev_state.vel));
+			neighbors = BLI_kdtree_range_search_with_len_squared_cb(
+			        epsys->tree, pa->prev_state.co, &ptn, acbr->look_ahead * len_v3(pa->prev_state.vel),
+			        len_squared_v3v3_with_normal_bias, pa->prev_state.ave);
 
 			if (neighbors > 0) for (n=0; n<neighbors; n++) {
 				copy_v3_v3(co1, pa->prev_state.co);
@@ -406,7 +423,9 @@ static int rule_flock(BoidRule *UNUSED(rule), BoidBrainData *bbd, BoidValues *UN
 {
 	KDTreeNearest ptn[11];
 	float vec[3] = {0.0f, 0.0f, 0.0f}, loc[3] = {0.0f, 0.0f, 0.0f};
-	int neighbors = BLI_kdtree_find_nearest_n__normal(bbd->sim->psys->tree, pa->state.co, pa->prev_state.ave, ptn, 11);
+	int neighbors = BLI_kdtree_find_nearest_n_with_len_squared_cb(
+	        bbd->sim->psys->tree, pa->state.co, ptn, ARRAY_SIZE(ptn),
+	        len_squared_v3v3_with_normal_bias, pa->prev_state.ave);
 	int n;
 	int ret = 0;
 
diff --git a/source/blender/blenlib/BLI_kdtree.h b/source/blender/blenlib/BLI_kdtree.h
index 864055127fc..3b21ad76b4d 100644
--- a/source/blender/blenlib/BLI_kdtree.h
+++ b/source/blender/blenlib/BLI_kdtree.h
@@ -45,9 +45,9 @@ int BLI_kdtree_find_nearest(
         KDTreeNearest *r_nearest) ATTR_NONNULL(1, 2);
 
 #define BLI_kdtree_find_nearest_n(tree, co, r_nearest, n) \
-        BLI_kdtree_find_nearest_n__normal(tree, co, NULL, r_nearest, n)
+        BLI_kdtree_find_nearest_n_with_len_squared_cb(tree, co, r_nearest, n, NULL, NULL)
 #define BLI_kdtree_range_search(tree, co, r_nearest, range) \
-        BLI_kdtree_range_search__normal(tree, co, NULL, r_nearest, range)
+        BLI_kdtree_range_search_with_len_squared_cb(tree, co, r_nearest, range, NULL, NULL)
 
 int BLI_kdtree_find_nearest_cb(
         const KDTree *tree, const float co[3],
@@ -61,15 +61,18 @@ int BLI_kdtree_calc_duplicates_fast(
         const KDTree *tree, const float range, bool use_index_order,
         int *doubles);
 
-/* Normal use is deprecated */
-/* remove __normal functions when last users drop */
-int BLI_kdtree_find_nearest_n__normal(
-        const KDTree *tree, const float co[3], const float nor[3],
+/* Versions of find/range search that take a squared distance callback to support bias. */
+int BLI_kdtree_find_nearest_n_with_len_squared_cb(
+        const KDTree *tree, const float co[3],
         KDTreeNearest *r_nearest,
-        unsigned int n) ATTR_NONNULL(1, 2, 4);
-int BLI_kdtree_range_search__normal(
-        const KDTree *tree, const float co[3], const float nor[3],
+        unsigned int n,
+        float (*len_sq_fn)(const float co_search[3], const float co_test[3], const void *user_data),
+        const void *user_data) ATTR_NONNULL(1, 2, 3);
+int BLI_kdtree_range_search_with_len_squared_cb(
+        const KDTree *tree, const float co[3],
         KDTreeNearest **r_nearest,
-        float range) ATTR_NONNULL(1, 2, 4) ATTR_WARN_UNUSED_RESULT;
+        float range,
+        float (*len_sq_fn)(const float co_search[3], const float co_test[3], const void *user_data),
+        const void *user_data) ATTR_NONNULL(1, 2) ATTR_WARN_UNUSED_RESULT;
 
 #endif  /* __BLI_KDTREE_H__ */
diff --git a/source/blender/blenlib/intern/BLI_kdtree.c b/source/blender/blenlib/intern/BLI_kdtree.c
index 689b9d9720a..92761a5ec7a 100644
--- a/source/blender/blenlib/intern/BLI_kdtree.c
+++ b/source/blender/blenlib/intern/BLI_kdtree.c
@@ -178,22 +178,9 @@ void BLI_kdtree_balance(KDTree *tree)
 #endif
 }
 
-static float squared_distance(const float v2[3], const float v1[3], const float n2[3])
+static float len_squared_v3v3_cb(const float co_kdtree[3], const float co_search[3], const void *UNUSED(user_data))
 {
-	float d[3], dist;
-
-	d[0] = v2[0] - v1[0];
-	d[1] = v2[1] - v1[1];
-	d[2] = v2[2] - v1[2];
-
-	dist = len_squared_v3(d);
-
-	/* can someone explain why this is done?*/
-	if (n2 && (dot_v3v3(d, n2) < 0.0f)) {
-		dist *= 10.0f;
-	}
-
-	return dist;
+	return len_squared_v3v3(co_kdtree, co_search);
 }
 
 static uint *realloc_nodes(uint *stack, uint *totstack, const bool is_alloc)
@@ -422,8 +409,9 @@ finally:
 	}
 }
 
-static void add_nearest(KDTreeNearest *ptn, uint *found, uint n, int index,
-                        float dist, const float *co)
+static void add_nearest(
+        KDTreeNearest *ptn, uint *found, const uint n, const int index,
+        const float dist, const float co[3])
 {
 	uint i;
 
@@ -451,10 +439,12 @@ static void add_nearest(KDTreeNearest *ptn, uint *found, uint n, int index,
  *
  * \param r_nearest: An array of nearest, sized at least \a n.
  */
-int BLI_kdtree_find_nearest_n__normal(
-        const KDTree *tree, const float co[3], const float nor[3],
+int BLI_kdtree_find_nearest_n_with_len_squared_cb(
+        const KDTree *tree, const float co[3],
         KDTreeNearest r_nearest[],
-        uint n)
+        uint n,
+        float (*len_sq_fn)(const float co_search[3], const float co_test[3], const void *user_data),
+        const void *user_data)
 {
 	const KDTreeNode *nodes = tree->nodes;
 	const KDTreeNode *root;
@@ -471,12 +461,17 @@ int BLI_kdtree_find_nearest_n__normal(
 		return 0;
 	}
 
+	if (len_sq_fn == NULL) {
+		len_sq_fn = len_squared_v3v3_cb;
+		BLI_assert(user_data == NULL);
+	}
+
 	stack = defaultstack;
 	totstack = KD_STACK_INIT;
 
 	root = &nodes[tree->root];
 
-	cur_dist = squared_distance(root->co, co, nor);
+	cur_dist = len_sq_fn(co, root->co, user_data);
 	add_nearest(r_nearest, &found, n, root->index, cur_dist, root->co);
 
 	if (co[root->d] < root->co[root->d]) {
@@ -505,7 +500,7 @@ int BLI_kdtree_find_nearest_n__normal(
 			cur_dist = -cur_dist * cur_dist;
 
 			if (found < n || -cur_dist < r_nearest[found - 1].dist) {
-				cur_dist = squared_distance(node->co, co, nor);
+				cur_dist = len_sq_fn(co, node->co, user_data);
 
 				if (found < n || cur_dist < r_nearest[found - 1].dist) {
 					add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
@@ -523,7 +518,7 @@ int BLI_kdtree_find_nearest_n__normal(
 			cur_dist = cur_dist * cur_dist;
 
 			if (found < n || cur_dist < r_nearest[found - 1].dist) {
-				cur_dist = squared_distance(node->co, co, nor);
+				cur_dist = len_sq_fn(co, node->co, user_data);
 				if (found < n || cur_dist < r_nearest[found - 1].dist) {
 					add_nearest(r_nearest, &found, n, node->index, cur_dist, node->co);
 				}
@@ -594,9 +589,11 @@ static void add_in_range(
  * Normal is optional, but if given will limit results to points in normal direction from co.
  * Remember to free nearest after use!
  */
-int BLI_kdtree_range_search__normal(
-        const KDTree *tree, const float co[3], const float nor[3],
-        KDTreeNearest **r_nearest, float range)
+int BLI_kdtree_range_search_with_len_squared_cb(
+        const KDTree *tree, const float co[3],
+        KDTreeNearest **r_nearest, float range,
+        float (*len_sq_fn)(const float co_search[3], const float co_test[3], const void *user_data),
+        const void *user_data)
 {
 	const KDTreeNode *nodes = tree->nodes;
 	uint *stack, defaultstack[KD_STACK_INIT];
@@ -612,6 +609,11 @@ int BLI_kdtree_range_search__normal(
 		return 0;
 	}
 
+	if (len_sq_fn == NULL) {
+		len_sq_fn = len_squared_v3v3_cb;
+		BLI_assert(user_data == NULL);
+	}
+
 	stack = defaultstack;
 	totstack = KD_STACK_INIT;
 
@@ -631,7 +633,7 @@ int BLI_kdtree_range_search__normal(
 			}
 		}
 		else {
-			dist_sq = squared_distance(node->co, co, nor);
+			dist_sq = len_sq_fn(co, node->co, user_data);
 			if (dist_sq <= range_sq) {
 				add_in_range(&foundstack, &totfoundstack, found++, node->index, dist_sq, node->co);
 			}