Add a method to modify all the active dynamic bodies on the RigidBodySet.

1 files changed, 38 insertions, 4 deletions

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,