feat: rework QueryPipeline update API to take less parameters (#647)

* chore: rework QueryPipeline API to take a generic qbvh updater

This allows to pass less parameters depending on the updating mode.

* chore: rework struct and functions names, and docs

---------

Co-authored-by: Sébastien Crozet <sebcrozet@dimforge.com>

4 files changed, 134 insertions, 91 deletions

diff --git a/src/dynamics/ccd/ccd_solver.rs b/src/dynamics/ccd/ccd_solver.rs
index 79c4495..eaabf77 100644
--- a/src/dynamics/ccd/ccd_solver.rs
+++ b/src/dynamics/ccd/ccd_solver.rs
@@ -3,8 +3,8 @@ use crate::dynamics::{IslandManager, RigidBodyHandle, RigidBodySet};
 use crate::geometry::{ColliderParent, ColliderSet, CollisionEvent, NarrowPhase};
 use crate::math::Real;
 use crate::parry::utils::SortedPair;
-use crate::pipeline::{EventHandler, QueryPipeline, QueryPipelineMode};
-use crate::prelude::{ActiveEvents, CollisionEventFlags};
+use crate::pipeline::{EventHandler, QueryPipeline};
+use crate::prelude::{query_pipeline_generators, ActiveEvents, CollisionEventFlags};
 use parry::query::{DefaultQueryDispatcher, QueryDispatcher};
 use parry::utils::hashmap::HashMap;
 use std::collections::BinaryHeap;
@@ -117,10 +117,12 @@ impl CCDSolver {
         narrow_phase: &NarrowPhase,
     ) -> Option<Real> {
         // Update the query pipeline.
-        self.query_pipeline.update_with_mode(
-            bodies,
-            colliders,
-            QueryPipelineMode::SweepTestWithPredictedPosition { dt },
+        self.query_pipeline.update_with_generator(
+            query_pipeline_generators::SweptAabbWithPredictedPosition {
+                bodies,
+                colliders,
+                dt,
+            },
         );
 
         let mut pairs_seen = HashMap::default();
@@ -238,10 +240,8 @@ impl CCDSolver {
         let mut min_overstep = dt;
 
         // Update the query pipeline.
-        self.query_pipeline.update_with_mode(
-            bodies,
-            colliders,
-            QueryPipelineMode::SweepTestWithNextPosition,
+        self.query_pipeline.update_with_generator(
+            query_pipeline_generators::SweptAabbWithNextPosition { bodies, colliders },
         );
 
         /*
diff --git a/src/pipeline/mod.rs b/src/pipeline/mod.rs
index bac67d2..d32d342 100644
--- a/src/pipeline/mod.rs
+++ b/src/pipeline/mod.rs
@@ -4,7 +4,9 @@ pub use collision_pipeline::CollisionPipeline;
 pub use event_handler::{ActiveEvents, ChannelEventCollector, EventHandler};
 pub use physics_hooks::{ActiveHooks, ContactModificationContext, PairFilterContext, PhysicsHooks};
 pub use physics_pipeline::PhysicsPipeline;
-pub use query_pipeline::{QueryFilter, QueryFilterFlags, QueryPipeline, QueryPipelineMode};
+pub use query_pipeline::{
+    generators as query_pipeline_generators, QueryFilter, QueryFilterFlags, QueryPipeline,
+};
 
 #[cfg(feature = "debug-render")]
 pub use self::debug_render_pipeline::{
diff --git a/src/pipeline/query_pipeline/generators.rs b/src/pipeline/query_pipeline/generators.rs
new file mode 100644
index 0000000..5b86d6a
--- /dev/null
+++ b/src/pipeline/query_pipeline/generators.rs
@@ -0,0 +1,109 @@
+//! Structs implementing [`QbvhDataGenerator<ColliderHandle>`] to be used with [`QueryPipeline::update_with_generator`].
+
+use parry::partitioning::QbvhDataGenerator;
+
+use crate::math::Real;
+use crate::prelude::{Aabb, ColliderHandle, ColliderSet, RigidBodySet};
+
+#[cfg(doc)]
+use crate::{
+    dynamics::{IntegrationParameters, RigidBodyPosition},
+    pipeline::QueryPipeline,
+};
+
+/// Generates collider AABBs based on the union of their current AABB and the AABB predicted
+/// from the velocity and forces of their parent rigid-body.
+///
+/// The main purpose of this struct is to be passed as a parameters to
+/// [`QueryPipeline::update_with_generator`] to update the [`QueryPipeline`].
+///
+/// The predicted position is calculated as
+/// `RigidBody::predict_position_using_velocity_and_forces * Collider::position_wrt_parent`.
+pub struct SweptAabbWithPredictedPosition<'a> {
+    /// The rigid bodies of your simulation.
+    pub bodies: &'a RigidBodySet,
+    /// The colliders of your simulation.
+    pub colliders: &'a ColliderSet,
+    /// The delta time to compute predicted position.
+    ///
+    /// You probably want to set it to [`IntegrationParameter::dt`].
+    pub dt: Real,
+}
+impl<'a> QbvhDataGenerator<ColliderHandle> for SweptAabbWithPredictedPosition<'a> {
+    fn size_hint(&self) -> usize {
+        self.colliders.len()
+    }
+
+    #[inline(always)]
+    fn for_each(&mut self, mut f: impl FnMut(ColliderHandle, Aabb)) {
+        for (h, co) in self.colliders.iter_enabled() {
+            if let Some(co_parent) = co.parent {
+                let rb = &self.bodies[co_parent.handle];
+                let predicted_pos = rb
+                    .pos
+                    .integrate_forces_and_velocities(self.dt, &rb.forces, &rb.vels, &rb.mprops);
+
+                let next_position = predicted_pos * co_parent.pos_wrt_parent;
+                f(h, co.shape.compute_swept_aabb(&co.pos, &next_position))
+            } else {
+                f(h, co.shape.compute_aabb(&co.pos))
+            }
+        }
+    }
+}
+
+/// Generates collider AABBs based on the union of their AABB at their current [`Collider::position`]
+/// and the AABB predicted from their parent’s [`RigidBody::next_position`].
+///
+/// The main purpose of this struct is to be passed as a parameters to
+/// [`QueryPipeline::update_with_generator`] to update the [`QueryPipeline`].
+///
+/// The predicted position is calculated as
+/// `RigidBody::next_position * Collider::position_wrt_parent`.
+pub struct SweptAabbWithNextPosition<'a> {
+    /// The rigid bodies of your simulation.
+    pub bodies: &'a RigidBodySet,
+    /// The colliders of your simulation.
+    pub colliders: &'a ColliderSet,
+}
+
+impl<'a> QbvhDataGenerator<ColliderHandle> for SweptAabbWithNextPosition<'a> {
+    fn size_hint(&self) -> usize {
+        self.colliders.len()
+    }
+
+    #[inline(always)]
+    fn for_each(&mut self, mut f: impl FnMut(ColliderHandle, Aabb)) {
+        for (h, co) in self.colliders.iter_enabled() {
+            if let Some(co_parent) = co.parent {
+                let rb_next_pos = &self.bodies[co_parent.handle].pos.next_position;
+                let next_position = rb_next_pos * co_parent.pos_wrt_parent;
+                f(h, co.shape.compute_swept_aabb(&co.pos, &next_position))
+            } else {
+                f(h, co.shape.compute_aabb(&co.pos))
+            }
+        }
+    }
+}
+
+/// Generates collider AABBs based on the AABB at their current [`Collider::position`].
+///
+/// The main purpose of this struct is to be passed as a parameters to
+/// [`QueryPipeline::update_with_generator`] to update the [`QueryPipeline`].
+pub struct CurrentAabb<'a> {
+    /// The colliders of your simulation.
+    pub colliders: &'a ColliderSet,
+}
+
+impl<'a> QbvhDataGenerator<ColliderHandle> for CurrentAabb<'a> {
+    fn size_hint(&self) -> usize {
+        self.colliders.len()
+    }
+
+    #[inline(always)]
+    fn for_each(&mut self, mut f: impl FnMut(ColliderHandle, Aabb)) {
+        for (h, co) in self.colliders.iter_enabled() {
+            f(h, co.shape.compute_aabb(&co.pos))
+        }
+    }
+}
diff --git a/src/pipeline/query_pipeline.rs b/src/pipeline/query_pipeline/mod.rs
index f202107..829499f 100644
--- a/src/pipeline/query_pipeline.rs
+++ b/src/pipeline/query_pipeline/mod.rs
@@ -1,3 +1,5 @@
+pub mod generators;
+
 use crate::dynamics::RigidBodyHandle;
 use crate::geometry::{
     Aabb, Collider, ColliderHandle, InteractionGroups, PointProjection, Qbvh, Ray, RayIntersection,
@@ -227,22 +229,6 @@ impl<'a> QueryFilter<'a> {
     }
 }
 
-/// Indicates how the colliders position should be taken into account when
-/// updating the query pipeline.
-pub enum QueryPipelineMode {
-    /// The `Collider::position` is taken into account.
-    CurrentPosition,
-    /// The `RigidBody::next_position * Collider::position_wrt_parent` is taken into account for
-    /// the colliders positions.
-    SweepTestWithNextPosition,
-    /// The `RigidBody::predict_position_using_velocity_and_forces * Collider::position_wrt_parent`
-    /// is taken into account for the colliders position.
-    SweepTestWithPredictedPosition {
-        /// The time used to integrate the rigid-body's velocity and acceleration.
-        dt: Real,
-    },
-}
-
 impl<'a> TypedSimdCompositeShape for QueryPipelineAsCompositeShape<'a> {
     type PartShape = dyn Shape;
     type PartNormalConstraints = dyn NormalConstraints;
@@ -357,72 +343,18 @@ impl QueryPipeline {
     }
 
     /// Update the acceleration structure on the query pipeline.
-    pub fn update(&mut self, bodies: &RigidBodySet, colliders: &ColliderSet) {
-        self.update_with_mode(bodies, colliders, QueryPipelineMode::CurrentPosition)
+    ///
+    /// Uses [`generators::CurrentAabb`] to update.
+    pub fn update(&mut self, colliders: &ColliderSet) {
+        self.update_with_generator(generators::CurrentAabb { colliders })
     }
 
-    /// Update the acceleration structure on the query pipeline.
-    pub fn update_with_mode(
-        &mut self,
-        bodies: &RigidBodySet,
-        colliders: &ColliderSet,
-        mode: QueryPipelineMode,
-    ) {
-        struct DataGenerator<'a> {
-            bodies: &'a RigidBodySet,
-            colliders: &'a ColliderSet,
-            mode: QueryPipelineMode,
-        }
-
-        impl<'a> QbvhDataGenerator<ColliderHandle> for DataGenerator<'a> {
-            fn size_hint(&self) -> usize {
-                self.colliders.len()
-            }
-
-            #[inline(always)]
-            fn for_each(&mut self, mut f: impl FnMut(ColliderHandle, Aabb)) {
-                match self.mode {
-                    QueryPipelineMode::CurrentPosition => {
-                        for (h, co) in self.colliders.iter_enabled() {
-                            f(h, co.shape.compute_aabb(&co.pos))
-                        }
-                    }
-                    QueryPipelineMode::SweepTestWithNextPosition => {
-                        for (h, co) in self.colliders.iter_enabled() {
-                            if let Some(co_parent) = co.parent {
-                                let rb_next_pos = &self.bodies[co_parent.handle].pos.next_position;
-                                let next_position = rb_next_pos * co_parent.pos_wrt_parent;
-                                f(h, co.shape.compute_swept_aabb(&co.pos, &next_position))
-                            } else {
-                                f(h, co.shape.compute_aabb(&co.pos))
-                            }
-                        }
-                    }
-                    QueryPipelineMode::SweepTestWithPredictedPosition { dt } => {
-                        for (h, co) in self.colliders.iter_enabled() {
-                            if let Some(co_parent) = co.parent {
-                                let rb = &self.bodies[co_parent.handle];
-                                let predicted_pos = rb.pos.integrate_forces_and_velocities(
-                                    dt, &rb.forces, &rb.vels, &rb.mprops,
-                                );
-
-                                let next_position = predicted_pos * co_parent.pos_wrt_parent;
-                                f(h, co.shape.compute_swept_aabb(&co.pos, &next_position))
-                            } else {
-                                f(h, co.shape.compute_aabb(&co.pos))
-                            }
-                        }
-                    }
-                }
-            }
-        }
-
-        let generator = DataGenerator {
-            bodies,
-            colliders,
-            mode,
-        };
-        self.qbvh.clear_and_rebuild(generator, self.dilation_factor);
+    /// Update the acceleration structure on the query pipeline using a custom collider bounding
+    /// volume generator.
+    ///
+    /// See [`generators`] for available generators.
+    pub fn update_with_generator(&mut self, mode: impl QbvhDataGenerator<ColliderHandle>) {
+        self.qbvh.clear_and_rebuild(mode, self.dilation_factor);
     }
 
     /// Find the closest intersection between a ray and a set of collider.