From c6abd2fd8eee13a3239595768fee4ed2ea2e518a Mon Sep 17 00:00:00 2001
From: Joe Ardent <code@ardent.nebcorp.com>
Date: Sun, 26 Nov 2023 13:46:13 -0800
Subject: [PATCH] use bevy_spatial for nearest neighbor searches

---
 Cargo.lock  | 193 +++++++++++++++++++++++++++++-----------------------
 Cargo.toml  |   2 +-
 src/lib.rs  | 132 +++++++++++++++++++++--------------
 src/main.rs |   7 +-
 4 files changed, 191 insertions(+), 143 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 04a764c..4369040 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -322,15 +322,6 @@ version = "4.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b4eb2cdb97421e01129ccb49169d8279ed21e829929144f4a22a6e54ac549ca1"
 
-[[package]]
-name = "atomic-polyfill"
-version = "0.1.11"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e3ff7eb3f316534d83a8a2c3d1674ace8a5a71198eba31e2e2b597833f699b28"
-dependencies = [
- "critical-section",
-]
-
 [[package]]
 name = "atomic-waker"
 version = "1.1.2"
@@ -343,8 +334,8 @@ version = "0.1.0"
 dependencies = [
  "argh",
  "bevy",
+ "bevy_spatial",
  "rand",
- "rstar",
 ]
 
 [[package]]
@@ -929,6 +920,18 @@ dependencies = [
  "uuid",
 ]
 
+[[package]]
+name = "bevy_spatial"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "56e846c2c81c0f82e098b5b3b0b25e74b6f110e87930008ed7359a20c47d832a"
+dependencies = [
+ "bevy",
+ "kd-tree",
+ "num-traits",
+ "typenum",
+]
+
 [[package]]
 name = "bevy_sprite"
 version = "0.12.0"
@@ -1481,12 +1484,6 @@ dependencies = [
  "cfg-if",
 ]
 
-[[package]]
-name = "critical-section"
-version = "1.1.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7059fff8937831a9ae6f0fe4d658ffabf58f2ca96aa9dec1c889f936f705f216"
-
 [[package]]
 name = "crossbeam-channel"
 version = "0.5.8"
@@ -1497,6 +1494,30 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ce6fd6f855243022dcecf8702fef0c297d4338e226845fe067f6341ad9fa0cef"
+dependencies = [
+ "cfg-if",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae211234986c545741a7dc064309f67ee1e5ad243d0e48335adc0484d960bcc7"
+dependencies = [
+ "autocfg",
+ "cfg-if",
+ "crossbeam-utils",
+ "memoffset",
+ "scopeguard",
+]
+
 [[package]]
 name = "crossbeam-utils"
 version = "0.8.16"
@@ -1541,6 +1562,12 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9ea835d29036a4087793836fa931b08837ad5e957da9e23886b29586fb9b6650"
 
+[[package]]
+name = "either"
+version = "1.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07"
+
 [[package]]
 name = "encase"
 version = "0.6.1"
@@ -1948,15 +1975,6 @@ dependencies = [
  "svg_fmt",
 ]
 
-[[package]]
-name = "hash32"
-version = "0.2.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67"
-dependencies = [
- "byteorder",
-]
-
 [[package]]
 name = "hashbrown"
 version = "0.12.3"
@@ -1989,19 +2007,6 @@ dependencies = [
  "winapi",
 ]
 
-[[package]]
-name = "heapless"
-version = "0.7.16"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "db04bc24a18b9ea980628ecf00e6c0264f3c1426dac36c00cb49b6fbad8b0743"
-dependencies = [
- "atomic-polyfill",
- "hash32",
- "rustc_version",
- "spin",
- "stable_deref_trait",
-]
-
 [[package]]
 name = "hexasphere"
 version = "9.1.0"
@@ -2158,6 +2163,20 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "kd-tree"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f54287107c30b23cf7a32e394c68b9472220a090c897d916dfdf467caef88864"
+dependencies = [
+ "num-traits",
+ "ordered-float",
+ "paste",
+ "pdqselect",
+ "rayon",
+ "typenum",
+]
+
 [[package]]
 name = "khronos-egl"
 version = "4.1.0"
@@ -2227,12 +2246,6 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
-[[package]]
-name = "libm"
-version = "0.2.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
-
 [[package]]
 name = "libredox"
 version = "0.0.2"
@@ -2303,6 +2316,15 @@ version = "2.6.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f665ee40bc4a3c5590afb1e9677db74a508659dfd71e126420da8274909a0167"
 
+[[package]]
+name = "memoffset"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a634b1c61a95585bd15607c6ab0c4e5b226e695ff2800ba0cdccddf208c406c"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "metal"
 version = "0.26.0"
@@ -2512,7 +2534,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
 dependencies = [
  "autocfg",
- "libm",
 ]
 
 [[package]]
@@ -2658,6 +2679,15 @@ dependencies = [
  "libredox",
 ]
 
+[[package]]
+name = "ordered-float"
+version = "3.9.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1e1c390732d15f1d48471625cd92d154e66db2c56645e29a9cd26f4699f72dc"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "overload"
 version = "0.1.1"
@@ -2708,6 +2738,12 @@ version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "de3145af08024dea9fa9914f381a17b8fc6034dfb00f3a84013f7ff43f29ed4c"
 
+[[package]]
+name = "pdqselect"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7778906d9321dd56cde1d1ffa69a73e59dcf5fda6d366f62727adf2bd4193aee"
+
 [[package]]
 name = "peeking_take_while"
 version = "0.1.2"
@@ -2863,6 +2899,26 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f2ff9a1f06a88b01621b7ae906ef0211290d1c8a168a15542486a8f61c0833b9"
 
+[[package]]
+name = "rayon"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "rectangle-pack"
 version = "0.4.2"
@@ -2965,17 +3021,6 @@ dependencies = [
  "serde_derive",
 ]
 
-[[package]]
-name = "rstar"
-version = "0.11.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "73111312eb7a2287d229f06c00ff35b51ddee180f017ab6dec1f69d62ac098d6"
-dependencies = [
- "heapless",
- "num-traits",
- "smallvec",
-]
-
 [[package]]
 name = "rustc-demangle"
 version = "0.1.23"
@@ -2988,15 +3033,6 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
-[[package]]
-name = "rustc_version"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366"
-dependencies = [
- "semver",
-]
-
 [[package]]
 name = "ruzstd"
 version = "0.4.0"
@@ -3029,12 +3065,6 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
 
-[[package]]
-name = "semver"
-version = "1.0.20"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "836fa6a3e1e547f9a2c4040802ec865b5d85f4014efe00555d7090a3dcaa1090"
-
 [[package]]
 name = "serde"
 version = "1.0.193"
@@ -3123,15 +3153,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "spin"
-version = "0.9.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
-dependencies = [
- "lock_api",
-]
-
 [[package]]
 name = "spirv"
 version = "0.2.0+1.5.4"
@@ -3142,12 +3163,6 @@ dependencies = [
  "num-traits",
 ]
 
-[[package]]
-name = "stable_deref_trait"
-version = "1.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
-
 [[package]]
 name = "static_assertions"
 version = "1.1.0"
@@ -3409,6 +3424,12 @@ dependencies = [
  "static_assertions",
 ]
 
+[[package]]
+name = "typenum"
+version = "1.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825"
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.12"
diff --git a/Cargo.toml b/Cargo.toml
index 16aff24..3546182 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,5 +8,5 @@ edition = "2021"
 [dependencies]
 argh = "0.1.12"
 bevy = "0.12.0"
+bevy_spatial = "0.7.0"
 rand = "0.8.5"
-rstar = "0.11.0"
diff --git a/src/lib.rs b/src/lib.rs
index 2195abd..b3758eb 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,11 @@
 use argh::FromArgs;
-use bevy::prelude::*;
+use bevy::{prelude::*, utils::HashMap, utils::HashSet};
+use bevy_spatial::{kdtree::KDTree3, SpatialAccess};
+
+pub type NNTree = KDTree3<Toid>;
 
 // toid
-const SPEED: f32 = 10.0;
+const SPEED: f32 = 1.0;
 const SPEED_DIFF_RANGE: f32 = 0.08; // +/- 8%
 
 // how far from origin before really wanting to come back
@@ -11,10 +14,10 @@ const RADIUS: f32 = 800.0;
 // how close to try to stay to your buddies
 const BUDDY_RADIUS: f32 = 15.0;
 
-const MIN_ALTITUDE: f32 = 1.5;
+// how much room to try to maintain between toids
+const PERSONAL_SPACE: f32 = 1.5;
 
-pub type Point = [f32; 3];
-pub type Toint = rstar::primitives::GeomWithData<Point, Entity>;
+const MIN_ALTITUDE: f32 = 1.5;
 
 #[derive(Debug, FromArgs, Resource)]
 /// Toid Watching
@@ -24,11 +27,11 @@ pub struct Config {
     pub toids: usize,
 }
 
-#[derive(Resource, Deref, DerefMut, Default)]
-pub struct Index(pub rstar::RTree<Toint>);
+#[derive(Clone, Debug, Default, Deref, DerefMut, Resource)]
+pub struct Positions(pub HashMap<Entity, Vec3>);
 
-#[derive(Component, Debug, Clone, Deref, DerefMut, Default)]
-pub struct Buddies(Vec<Toint>);
+#[derive(Component, Debug, Clone, Default)]
+pub struct Buddies(HashSet<Entity>);
 
 #[derive(Component, Debug, Clone, Deref, DerefMut, Default)]
 pub struct Velocity(Vec3);
@@ -48,8 +51,11 @@ pub fn turkey_time(
     let speed = SPEED + (SPEED * speed_diff);
     let vel = unit_vec(r) * speed;
     let buddies = r.gen_range(6..=8);
+    let x = r.gen_range(-5.0..=5.0);
+    let z = r.gen_range(-5.0..=5.0);
+    let y = r.gen_range(0.1..=2.5);
     let spatial_bundle = SpatialBundle {
-        transform: Transform::from_xyz(0.0, MIN_ALTITUDE + 0.1, 0.0),
+        transform: Transform::from_xyz(x, MIN_ALTITUDE + y, z),
         ..Default::default()
     };
     commands
@@ -68,40 +74,11 @@ pub fn turkey_time(
         .id()
 }
 
-pub fn add_gizmos(mut gizmos: Gizmos, toids: Query<(&Transform, Entity), With<Toid>>) {
-    let gizmos = &mut gizmos;
-    for (pos, _entity) in toids.iter() {
-        let nudge = pos.up() * 0.15;
-        let rpos = pos.translation + nudge;
-        gizmos.ray(rpos, pos.forward(), Color::RED);
-    }
-}
-
-pub fn update_pos(
-    mut toids: Query<(&mut Transform, &Velocity, Entity)>,
-    time: Res<Time>,
-    mut index: ResMut<Index>,
-) {
-    let mut positions = Vec::with_capacity(toids.iter().len());
-
-    let dt = time.delta_seconds();
-    for (mut xform, vel, entity) in toids.iter_mut() {
-        let look_at = xform.translation + vel.0;
-        xform.translation += vel.0 * dt;
-        xform.look_at(look_at, Vec3::Y);
-        let toint = Toint::new(xform.translation.to_array(), entity);
-        positions.push(toint);
-    }
-    let tree = rstar::RTree::bulk_load(positions);
-    **index = tree;
-}
-
 pub fn update_vel(
-    mut toids: Query<(&Transform, &mut Velocity, &Toid), With<Toid>>,
-    time: Res<Time>,
+    mut toids: Query<(&Transform, &mut Velocity, &Toid, &Buddies)>,
+    positions: Res<Positions>,
 ) {
-    let dt = time.delta_seconds();
-    for (xform, mut vel, toid) in toids.iter_mut() {
+    for (xform, mut vel, toid, buddies) in toids.iter_mut() {
         let speed = toid.speed;
         let mut dir = vel.normalize();
         let pos = xform.translation;
@@ -111,29 +88,80 @@ pub fn update_vel(
             let dh = MIN_ALTITUDE - pos.y;
             let pct = (dh / MIN_ALTITUDE).min(1.0).powi(2);
             let rot = Quat::from_rotation_arc(dir, Vec3::Y);
-            let (axis, ang) = rot.to_axis_angle();
-            let rot = Quat::from_axis_angle(axis, ang * pct).normalize();
-            dir = rot.mul_vec3(dir);
+            let rot = Quat::IDENTITY.slerp(rot, pct);
+            dir = rot.mul_vec3(dir).normalize();
         }
 
         // nudge toward origin if too far
         {
             let dist = pos.length();
             let toward_origin = -pos.normalize();
-            let forward = dir.normalize();
-            let rot = Quat::from_rotation_arc(forward, toward_origin);
+            let rot = Quat::from_rotation_arc(dir, toward_origin);
             let pct = (dist / RADIUS).min(1.0);
-            let (axis, ang) = rot.to_axis_angle();
-            let rot = Quat::from_axis_angle(axis, ang * pct).normalize();
-            dir = rot.mul_vec3(dir);
+            let rot = Quat::IDENTITY.slerp(rot, pct);
+            dir = rot.mul_vec3(dir).normalize();
+        }
+
+        // find buddies and orient; point more towards further-away buddies
+        for buddy in buddies.0.iter() {
+            let bp = positions.get(buddy).unwrap_or(&pos);
+            let bdir = *bp - pos;
+            let dist = bdir.length();
+            let rot = Quat::from_rotation_arc(dir, bdir.normalize());
+            let s = (dist / (BUDDY_RADIUS * 1.2)).min(1.0);
+            let rot = Quat::IDENTITY.slerp(rot, s);
+            dir = rot.mul_vec3(dir).normalize();
         }
-        // find buddies and orient: point more towards further-away buddies
 
         **vel = dir * speed;
     }
 }
 
-pub fn update_buddies(mut toids: Query<(&Transform, Entity, &mut Buddies)>, index: Res<Index>) {}
+pub fn update_pos(
+    mut toids: Query<(&mut Transform, &Velocity, Entity), With<Toid>>,
+    mut positions: ResMut<Positions>,
+    time: Res<Time>,
+) {
+    let dt = time.delta_seconds();
+    for (mut xform, vel, entity) in toids.iter_mut() {
+        let look_at = xform.translation + vel.0;
+        xform.translation += vel.0 * dt;
+        xform.look_at(look_at, Vec3::Y);
+        let _ = positions.insert(entity, xform.translation);
+    }
+}
+
+pub fn update_buddies(
+    mut toids: Query<(&Transform, Entity, &Toid, &mut Buddies)>,
+    index: Res<NNTree>,
+    positions: Res<Positions>,
+) {
+    let d2 = (BUDDY_RADIUS * 1.5).powi(2);
+    for (xform, entity, toid, mut buddies) in toids.iter_mut() {
+        let pos = xform.translation;
+
+        for buddy in buddies.0.clone().iter() {
+            let bp = positions.get(buddy).unwrap_or(&pos);
+            let bd2 = (*bp - pos).length_squared();
+            if bd2 > d2 {
+                buddies.0.remove(buddy);
+            }
+        }
+
+        if buddies.0.len() < toid.buddies {
+            let mut neighbors = index
+                .within_distance(pos, BUDDY_RADIUS)
+                .into_iter()
+                .filter(|n| n.1.is_some() && n.1.unwrap() != entity);
+            let diff = toid.buddies - buddies.0.len();
+            for _ in 0..diff {
+                if let Some(neighbor) = neighbors.next() {
+                    buddies.0.insert(neighbor.1.unwrap());
+                }
+            }
+        }
+    }
+}
 
 //-************************************************************************
 // util
diff --git a/src/main.rs b/src/main.rs
index 0e72a20..2ed7d5e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,5 @@
-//use std::f32::consts::PI;
-
 use bevy::prelude::*;
+use bevy_spatial::{AutomaticUpdate, TransformMode};
 
 use audubon::*;
 
@@ -8,14 +7,14 @@ fn main() {
     let config: audubon::Config = argh::from_env();
     App::new()
         .add_plugins(DefaultPlugins)
-        .insert_resource(Index(rstar::RTree::new()))
+        .add_plugins(AutomaticUpdate::<Toid>::new().with_transform(TransformMode::GlobalTransform))
         .insert_resource(config)
+        .insert_resource(Positions::default())
         .add_systems(Startup, setup)
         .add_systems(
             Update,
             (
                 update_vel,
-                //add_gizmos,
                 update_pos,
                 update_buddies,
                 rotate_camera,