From 67f3c14b7a13420b418c8ada9aea266d923b6896 Mon Sep 17 00:00:00 2001
From: Oli Scherer <github35764891676564198441@oli-obk.de>
Date: Thu, 7 Mar 2024 11:55:05 +0100
Subject: [PATCH] Add a convenience helper for reading GridCell and Transform
 from a query at the same time (#12)

Co-authored-by: Oliver Scherer <github@oli-obk.de>
Co-authored-by: Aevyrie <aevyrie@gmail.com>
---
 examples/demo.rs   |  12 +++--
 src/camera.rs      |  37 ++++++++-------
 src/grid_cell.rs   |   7 +++
 src/lib.rs         |  26 +++++------
 src/world_query.rs | 109 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 158 insertions(+), 33 deletions(-)
 create mode 100644 src/world_query.rs

diff --git a/examples/demo.rs b/examples/demo.rs
index f7b82e0..7396c19 100644
--- a/examples/demo.rs
+++ b/examples/demo.rs
@@ -5,6 +5,7 @@ use bevy::{
 };
 use big_space::{
     camera::{CameraController, CameraInput},
+    world_query::GridTransformReadOnly,
     FloatingOrigin, GridCell,
 };
 
@@ -153,14 +154,17 @@ fn ui_text_system(
     mut debug_text: Query<&mut Text, (With<BigSpaceDebugText>, Without<FunFactText>)>,
     mut fun_text: Query<&mut Text, (With<FunFactText>, Without<BigSpaceDebugText>)>,
     time: Res<Time>,
-    origin: Query<(&GridCell<i128>, &Transform), With<FloatingOrigin>>,
+    origin: Query<GridTransformReadOnly<i128>, With<FloatingOrigin>>,
     camera: Query<&CameraController>,
     objects: Query<&Transform, With<Handle<Mesh>>>,
 ) {
-    let (cell, transform) = origin.single();
-    let translation = transform.translation;
+    let origin = origin.single();
+    let translation = origin.transform.translation;
 
-    let grid_text = format!("GridCell: {}x, {}y, {}z", cell.x, cell.y, cell.z);
+    let grid_text = format!(
+        "GridCell: {}x, {}y, {}z",
+        origin.cell.x, origin.cell.y, origin.cell.z
+    );
 
     let translation_text = format!(
         "Transform: {:>8.2}x, {:>8.2}y, {:>8.2}z",
diff --git a/src/camera.rs b/src/camera.rs
index 50b63fd..349f43f 100644
--- a/src/camera.rs
+++ b/src/camera.rs
@@ -10,7 +10,11 @@ use bevy::{
     transform::TransformSystem,
 };
 
-use crate::{precision::GridPrecision, FloatingOriginSettings, GridCell};
+use crate::{
+    precision::GridPrecision,
+    world_query::{GridTransform, GridTransformReadOnly},
+    FloatingOriginSettings,
+};
 
 /// Adds the `big_space` camera controller
 #[derive(Default)]
@@ -173,20 +177,21 @@ pub fn default_camera_inputs(
 /// Find the object nearest the camera
 pub fn nearest_objects<T: GridPrecision>(
     settings: Res<FloatingOriginSettings>,
-    objects: Query<(Entity, &GridCell<T>, &Transform, &Aabb)>,
-    mut camera: Query<(&mut CameraController, &GridCell<T>, &Transform)>,
+    objects: Query<(Entity, GridTransformReadOnly<T>, &Aabb)>,
+    mut camera: Query<(&mut CameraController, GridTransformReadOnly<T>)>,
 ) {
-    let (mut camera, cam_cell, cam_transform) = camera.single_mut();
+    let (mut camera, cam_pos) = camera.single_mut();
     let nearest_object = objects
         .iter()
-        .map(|(entity, cell, obj_transform, aabb)| {
-            let pos = settings.grid_position_double(&(cell - cam_cell), obj_transform)
-                - cam_transform.translation.as_dvec3();
-            let dist = pos.length()
-                - (aabb.half_extents.as_dvec3() * obj_transform.scale.as_dvec3())
+        .map(|(entity, obj_pos, aabb)| {
+            let center_distance = settings
+                .grid_position_double(&(*obj_pos.cell - *cam_pos.cell), obj_pos.transform)
+                - cam_pos.transform.translation.as_dvec3();
+            let nearest_distance = center_distance.length()
+                - (aabb.half_extents.as_dvec3() * obj_pos.transform.scale.as_dvec3())
                     .abs()
                     .max_element();
-            (entity, dist)
+            (entity, nearest_distance)
         })
         .filter(|v| v.1.is_finite())
         .reduce(|nearest, this| if this.1 < nearest.1 { this } else { nearest });
@@ -198,9 +203,9 @@ pub fn camera_controller<P: GridPrecision>(
     time: Res<Time>,
     settings: Res<FloatingOriginSettings>,
     mut input: ResMut<CameraInput>,
-    mut camera: Query<(&mut Transform, &mut CameraController, &mut GridCell<P>)>,
+    mut camera: Query<(GridTransform<P>, &mut CameraController)>,
 ) {
-    for (mut cam_transform, mut controller, mut cell) in camera.iter_mut() {
+    for (mut position, mut controller) in camera.iter_mut() {
         let speed = match (controller.nearest_object, controller.slow_near_objects) {
             (Some(nearest), true) => nearest.1.abs(),
             _ => controller.speed,
@@ -215,16 +220,16 @@ pub fn camera_controller<P: GridPrecision>(
         let (vel_t_current, vel_r_current) = (controller.vel_translation, controller.vel_rotation);
         let (vel_t_target, vel_r_target) = input.target_velocity(speed, time.delta_seconds_f64());
 
-        let cam_rot = cam_transform.rotation.as_dquat();
+        let cam_rot = position.transform.rotation.as_dquat();
         let vel_t_next = cam_rot * vel_t_target; // Orients the translation to match the camera
         let vel_t_next = vel_t_current.lerp(vel_t_next, lerp_translation);
         // Convert the high precision translation to a grid cell and low precision translation
         let (cell_offset, new_translation) = settings.translation_to_grid(vel_t_next);
-        *cell += cell_offset;
-        cam_transform.translation += new_translation;
+        *position.cell += cell_offset;
+        position.transform.translation += new_translation;
 
         let new_rotation = vel_r_current.slerp(vel_r_target, lerp_rotation);
-        cam_transform.rotation *= new_rotation.as_quat();
+        position.transform.rotation *= new_rotation.as_quat();
 
         // Store the new velocity to be used in the next frame
         controller.vel_translation = vel_t_next;
diff --git a/src/grid_cell.rs b/src/grid_cell.rs
index 32d777c..9f53b03 100644
--- a/src/grid_cell.rs
+++ b/src/grid_cell.rs
@@ -105,3 +105,10 @@ impl<P: GridPrecision> std::ops::AddAssign for GridCell<P> {
         *self = self.add(rhs);
     }
 }
+
+impl<P: GridPrecision> std::ops::SubAssign for GridCell<P> {
+    fn sub_assign(&mut self, rhs: Self) {
+        use std::ops::Sub;
+        *self = self.sub(rhs);
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index d1540fb..993c4ea 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -89,10 +89,12 @@
 use bevy::{math::DVec3, prelude::*, transform::TransformSystem};
 use propagation::propagate_transforms;
 use std::marker::PhantomData;
+use world_query::{GridTransformReadOnly, GridTransformReadOnlyItem};
 
 pub mod grid_cell;
 pub mod precision;
 pub mod propagation;
+pub mod world_query;
 
 pub use grid_cell::GridCell;
 
@@ -314,41 +316,39 @@ pub fn update_global_from_grid<P: GridPrecision>(
     origin: Query<(Ref<GridCell<P>>, Ref<FloatingOrigin>)>,
     mut entities: ParamSet<(
         Query<
-            (&Transform, &mut GlobalTransform, &GridCell<P>),
+            (GridTransformReadOnly<P>, &mut GlobalTransform),
             Or<(Changed<GridCell<P>>, Changed<Transform>)>,
         >,
-        Query<(&Transform, &mut GlobalTransform, &GridCell<P>)>,
+        Query<(GridTransformReadOnly<P>, &mut GlobalTransform)>,
     )>,
 ) {
     let (origin_cell, floating_origin) = origin.single();
 
     if origin_cell.is_changed() || floating_origin.is_changed() {
         let mut all_entities = entities.p1();
-        all_entities
-            .par_iter_mut()
-            .for_each(|(local, global, entity_cell)| {
-                update_global_from_cell_local(&settings, entity_cell, &origin_cell, local, global);
-            });
+        all_entities.par_iter_mut().for_each(|(local, global)| {
+            update_global_from_cell_local(&settings, &origin_cell, local, global);
+        });
     } else {
         let mut moved_cell_entities = entities.p0();
         moved_cell_entities
             .par_iter_mut()
-            .for_each(|(local, global, entity_cell)| {
-                update_global_from_cell_local(&settings, entity_cell, &origin_cell, local, global);
+            .for_each(|(local, global)| {
+                update_global_from_cell_local(&settings, &origin_cell, local, global);
             });
     }
 }
 
 fn update_global_from_cell_local<P: GridPrecision>(
     settings: &FloatingOriginSettings,
-    entity_cell: &GridCell<P>,
     origin_cell: &GridCell<P>,
-    local: &Transform,
+    local: GridTransformReadOnlyItem<P>,
     mut global: Mut<GlobalTransform>,
 ) {
-    let grid_cell_delta = entity_cell - origin_cell;
+    let grid_cell_delta = *local.cell - *origin_cell;
     *global = local
-        .with_translation(settings.grid_position(&grid_cell_delta, local))
+        .transform
+        .with_translation(settings.grid_position(&grid_cell_delta, local.transform))
         .into();
 }
 
diff --git a/src/world_query.rs b/src/world_query.rs
new file mode 100644
index 0000000..15c0268
--- /dev/null
+++ b/src/world_query.rs
@@ -0,0 +1,109 @@
+//! A helper query argument that ensures you don't forget to handle
+//! the [`GridCell`] when you work with a [`Transform`].
+
+use bevy::ecs::query::QueryData;
+use bevy::math::DVec3;
+use bevy::prelude::*;
+
+use crate::precision::GridPrecision;
+use crate::{FloatingOriginSettings, GridCell};
+
+#[derive(QueryData)]
+#[query_data(mutable)]
+/// A convenience query argument that groups a [`Transform`] with its [`GridCell`].
+/// If you only want to read from the position, use [`GridTransformReadOnly`] instead,
+/// as this will allow the bevy ECS to run multiple queries using [`GridTransformReadOnly`]
+/// at the same time (just like multiple queries with `&Transform` are fine).
+pub struct GridTransform<P: GridPrecision> {
+    /// Grid local transform
+    pub transform: &'static mut Transform,
+    /// The grid to which `transform` is relative to.
+    pub cell: &'static mut GridCell<P>,
+}
+
+impl<'w, P: GridPrecision> GridTransformItem<'w, P> {
+    /// Compute the global position with double precision.
+    pub fn position_double(&self, settings: &FloatingOriginSettings) -> DVec3 {
+        settings.grid_position_double(&self.cell, &self.transform)
+    }
+
+    /// Compute the global position.
+    pub fn position(&self, settings: &FloatingOriginSettings) -> Vec3 {
+        settings.grid_position(&self.cell, &self.transform)
+    }
+
+    /// Get a copy of the fields to work with.
+    pub fn to_owned(&self) -> GridTransformOwned<P> {
+        GridTransformOwned {
+            transform: *self.transform,
+            cell: *self.cell,
+        }
+    }
+}
+
+impl<'w, P: GridPrecision> GridTransformReadOnlyItem<'w, P> {
+    /// Compute the global position with double precision.
+    pub fn position_double(&self, settings: &FloatingOriginSettings) -> DVec3 {
+        settings.grid_position_double(self.cell, self.transform)
+    }
+
+    /// Compute the global position.
+    pub fn position(&self, settings: &FloatingOriginSettings) -> Vec3 {
+        settings.grid_position(self.cell, self.transform)
+    }
+
+    /// Get a copy of the fields to work with.
+    pub fn to_owned(&self) -> GridTransformOwned<P> {
+        GridTransformOwned {
+            transform: *self.transform,
+            cell: *self.cell,
+        }
+    }
+}
+
+/// A convenience wrapper that allows working with grid and transform easily
+#[derive(Copy, Clone)]
+pub struct GridTransformOwned<P: GridPrecision> {
+    /// Grid local transform
+    pub transform: Transform,
+    /// The grid to which `transform` is relative to.
+    pub cell: GridCell<P>,
+}
+
+impl<P: GridPrecision> std::ops::Sub for GridTransformOwned<P> {
+    type Output = Self;
+
+    /// Compute a new transform that maps from `source` to `self`.
+    fn sub(mut self, source: Self) -> Self {
+        self.cell -= source.cell;
+        self.transform.translation -= source.transform.translation;
+        self.transform.scale /= source.transform.scale;
+        self.transform.rotation *= source.transform.rotation.inverse();
+        self
+    }
+}
+
+impl<P: GridPrecision> std::ops::Add for GridTransformOwned<P> {
+    type Output = Self;
+
+    /// Compute a new transform that shifts, scales and rotates `self` by `diff`.
+    fn add(mut self, diff: Self) -> Self {
+        self.cell += diff.cell;
+        self.transform.translation += diff.transform.translation;
+        self.transform.scale *= diff.transform.scale;
+        self.transform.rotation *= diff.transform.rotation;
+        self
+    }
+}
+
+impl<P: GridPrecision> GridTransformOwned<P> {
+    /// Compute the global position with double precision.
+    pub fn position_double(&self, space: &FloatingOriginSettings) -> DVec3 {
+        space.grid_position_double(&self.cell, &self.transform)
+    }
+
+    /// Compute the global position.
+    pub fn position(&self, space: &FloatingOriginSettings) -> Vec3 {
+        space.grid_position(&self.cell, &self.transform)
+    }
+}