From e7e196d9f949a03ef997f0adc629344c3696b1ff Mon Sep 17 00:00:00 2001
From: Bruce Mitchener <bruce.mitchener@gmail.com>
Date: Mon, 16 Sep 2024 21:45:14 +0700
Subject: [PATCH] Improve capacity handling for `ColliderSet`, `RigidBodySet`.
 (#726)

These allow an application to reduce the cost of reallocation when
they know that a large number of colliders or rigid bodies will
be created.
---
 CHANGELOG.md                   |  4 ++++
 src/dynamics/rigid_body_set.rs |  8 ++++++++
 src/geometry/collider_set.rs   | 11 +++++++++++
 3 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 070b778da..71576becf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,10 @@
 - The region key has been replaced by an i64 in the f64 version of rapier, increasing the range before panics occur.
 - Fix `BroadphaseMultiSap` not being able to serialize correctly with serde_json.
 
+### Added
+
+- `RigidBodySet` and `ColliderSet` have a new constructor `with_capacity`.
+
 ### Modified
 
 - `InteractionGroups` default value for `memberships` is now `GROUP_1` (#706)
diff --git a/src/dynamics/rigid_body_set.rs b/src/dynamics/rigid_body_set.rs
index 2ef91be0d..314530039 100644
--- a/src/dynamics/rigid_body_set.rs
+++ b/src/dynamics/rigid_body_set.rs
@@ -43,6 +43,14 @@ impl RigidBodySet {
         }
     }
 
+    /// Create a new set of rigid bodies, with an initial capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        RigidBodySet {
+            bodies: Arena::with_capacity(capacity),
+            modified_bodies: Vec::with_capacity(capacity),
+        }
+    }
+
     pub(crate) fn take_modified(&mut self) -> Vec<RigidBodyHandle> {
         std::mem::take(&mut self.modified_bodies)
     }
diff --git a/src/geometry/collider_set.rs b/src/geometry/collider_set.rs
index 393447243..7a2b6aaf7 100644
--- a/src/geometry/collider_set.rs
+++ b/src/geometry/collider_set.rs
@@ -23,6 +23,17 @@ impl ColliderSet {
         }
     }
 
+    /// Create a new set of colliders, with an initial capacity
+    /// for the set of colliders as well as the tracking of
+    /// modified colliders.
+    pub fn with_capacity(capacity: usize) -> Self {
+        ColliderSet {
+            colliders: Arena::with_capacity(capacity),
+            modified_colliders: Vec::with_capacity(capacity),
+            removed_colliders: Vec::new(),
+        }
+    }
+
     pub(crate) fn take_modified(&mut self) -> Vec<ColliderHandle> {
         std::mem::take(&mut self.modified_colliders)
     }