From b3405e56721dac2ebb22a7709cb7430fde7b4859 Mon Sep 17 00:00:00 2001
From: Crozet Sébastien <developer@crozet.re>
Date: Tue, 23 Feb 2021 15:47:44 +0100
Subject: Add a method to modify all the active dynamic bodies on the
 RigidBodySet.

---
 src/dynamics/rigid_body_set.rs | 42 ++++++++++++++++++++++++++++++++++++++----
 1 file changed, 38 insertions(+), 4 deletions(-)

(limited to 'src/dynamics')

diff --git a/src/dynamics/rigid_body_set.rs b/src/dynamics/rigid_body_set.rs
index 79d5827..f459d4f 100644
--- a/src/dynamics/rigid_body_set.rs
+++ b/src/dynamics/rigid_body_set.rs
@@ -234,13 +234,27 @@ impl RigidBodySet {
         self.bodies.get(handle.0)
     }
 
+    fn mark_as_modified(
+        handle: RigidBodyHandle,
+        rb: &mut RigidBody,
+        modified_bodies: &mut Vec<RigidBodyHandle>,
+        modified_all_bodies: bool,
+    ) {
+        if !modified_all_bodies && !rb.changes.contains(RigidBodyChanges::MODIFIED) {
+            rb.changes = RigidBodyChanges::MODIFIED;
+            modified_bodies.push(handle);
+        }
+    }
+
     /// Gets a mutable reference to the rigid-body with the given handle.
     pub fn get_mut(&mut self, handle: RigidBodyHandle) -> Option<&mut RigidBody> {
         let result = self.bodies.get_mut(handle.0)?;
-        if !self.modified_all_bodies && !result.changes.contains(RigidBodyChanges::MODIFIED) {
-            result.changes = RigidBodyChanges::MODIFIED;
-            self.modified_bodies.push(handle);
-        }
+        Self::mark_as_modified(
+            handle,
+            result,
+            &mut self.modified_bodies,
+            self.modified_all_bodies,
+        );
         Some(result)
     }
 
@@ -300,6 +314,26 @@ impl RigidBodySet {
             .filter_map(move |h| Some((*h, bodies.get(h.0)?)))
     }
 
+    /// Applies the given function on all the active dynamic rigid-bodies
+    /// contained by this set.
+    #[inline(always)]
+    pub fn foreach_active_dynamic_body_mut(
+        &mut self,
+        mut f: impl FnMut(RigidBodyHandle, &mut RigidBody),
+    ) {
+        for handle in &self.active_dynamic_set {
+            if let Some(rb) = self.bodies.get_mut(handle.0) {
+                Self::mark_as_modified(
+                    *handle,
+                    rb,
+                    &mut self.modified_bodies,
+                    self.modified_all_bodies,
+                );
+                f(*handle, rb)
+            }
+        }
+    }
+
     #[inline(always)]
     pub(crate) fn foreach_active_body_mut_internal(
         &mut self,
-- 
cgit