From 2872643accc3100c9609c933d255b2672d6d594d Mon Sep 17 00:00:00 2001
From: Joe Ardent <code@ardent.nebcorp.com>
Date: Sat, 2 Dec 2023 15:44:25 -0800
Subject: [PATCH] add diagnostics, tighten updates

---
 src/lib.rs  | 37 ++++++++++++++++++-------------------
 src/main.rs | 11 ++++++++---
 2 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 21acbbf..ce122b1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,8 @@
 use argh::FromArgs;
-use bevy::{prelude::*, utils::HashMap, utils::HashSet};
+use bevy::{
+    prelude::*,
+    utils::{HashMap, HashSet},
+};
 use bevy_spatial::{kdtree::KDTree3, SpatialAccess};
 
 pub type NNTree = KDTree3<Toid>;
@@ -28,7 +31,7 @@ pub struct Config {
 #[derive(Clone, Debug, Default, Deref, DerefMut, Resource)]
 pub struct Positions(pub HashMap<Entity, Vec3>);
 
-#[derive(Component, Debug, Clone, Default)]
+#[derive(Component, Debug, Deref, DerefMut, Clone, Default)]
 pub struct Buddies(HashSet<Entity>);
 
 #[derive(Component, Debug, Clone, Deref, DerefMut, Default)]
@@ -84,11 +87,11 @@ pub fn update_vel(
         let speed = toid.speed;
         let mut dir = vel.normalize();
         let pos = xform.translation;
-        let original_dir = dir.clone();
+        let original_dir = dir;
 
         // 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 bp = positions.get(buddy).unwrap();
             let bdir = *bp - pos;
             let dist = bdir.length();
             let rot = Quat::from_rotation_arc(dir, bdir.normalize());
@@ -155,7 +158,7 @@ pub fn update_pos(
         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);
+        *positions.entry(entity).or_insert(Vec3::ZERO) = xform.translation;
     }
 }
 
@@ -168,26 +171,22 @@ pub fn update_buddies(
     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);
+        for buddy in buddies.clone().iter() {
+            let bp = positions.get(buddy).unwrap();
             let bd2 = (*bp - pos).length_squared();
             if bd2 > d2 {
-                buddies.0.remove(buddy);
+                buddies.remove(buddy);
             }
         }
 
-        if buddies.0.len() < toid.buddies {
-            let mut neighbors = index
-                .within_distance(pos, BUDDY_RADIUS)
+        if buddies.len() < toid.buddies {
+            let diff = toid.buddies - buddies.len();
+            for (_, neighbor) in index
+                .k_nearest_neighbour(pos, diff + 1)
                 .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());
-                } else {
-                    break;
-                }
+                .filter(|n| n.1.is_some() && n.1.unwrap() != entity)
+            {
+                buddies.insert(neighbor.unwrap());
             }
         }
     }
diff --git a/src/main.rs b/src/main.rs
index 77a1346..8c54ff4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,18 +1,23 @@
 use std::time::Duration;
 
 use audubon::*;
-use bevy::prelude::*;
+use bevy::{
+    diagnostic::{FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
+    prelude::*,
+};
 use bevy_spatial::{AutomaticUpdate, TransformMode};
 
 fn main() {
     let config: audubon::Config = argh::from_env();
     App::new()
         .add_plugins(DefaultPlugins)
-        .add_plugins(
+        .add_plugins((
             AutomaticUpdate::<Toid>::new()
                 .with_transform(TransformMode::GlobalTransform)
                 .with_frequency(Duration::from_millis(150)),
-        )
+            FrameTimeDiagnosticsPlugin,
+            LogDiagnosticsPlugin::default(),
+        ))
         .insert_resource(config)
         .insert_resource(Positions::default())
         .insert_resource(ClearColor(Color::rgb(0.64, 0.745, 0.937))) // a nice light blue