From 24fd036cb6441d53fbc7ccad219d528d26a2ea9f Mon Sep 17 00:00:00 2001
From: Thierry Berger <contact@thierryberger.com>
Date: Wed, 26 Jun 2024 18:13:40 +0200
Subject: [PATCH] multiple worlds working + example

example from https://github.com/dimforge/bevy_rapier/pull/328

Co-authored-by: Anthony Tornetta <25857049+AnthonyTornetta@users.noreply.github.com>
---
 CHANGELOG.md                                  |   9 +-
 bevy_rapier3d/examples/multi_world3.rs        | 117 +++
 src/plugin/context/systemparams/mod.rs        |  12 +
 .../rapier_context_access.rs}                 |  33 +-
 src/plugin/mod.rs                             |   2 +-
 src/plugin/plugin.rs                          |  13 +
 src/plugin/systems/character_controller.rs    |  18 +-
 src/plugin/systems/collider.rs                | 125 ++--
 src/plugin/systems/joint.rs                   |  46 +-
 src/plugin/systems/remove.rs                  |  24 +-
 src/plugin/systems/rigid_body.rs              | 677 ++++++++++--------
 11 files changed, 661 insertions(+), 415 deletions(-)
 create mode 100644 bevy_rapier3d/examples/multi_world3.rs
 create mode 100644 src/plugin/context/systemparams/mod.rs
 rename src/plugin/context/{systemparams.rs => systemparams/rapier_context_access.rs} (66%)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index befabb27..b553961e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,11 +7,10 @@
 - `RapierContext` is now a `Component`
   - Rapier now supports multiple worlds.
   - Migration guide: 
-    - `ResMut<mut RapierContext>` -> `RapierContextAccessMut`
-    - `Res<RapierContext>` -> `Query<&RapierContext>`
-    - `ResMut<RapierConfiguration>` -> `Query<&mut RapierConfiguration>`
-    - `Res<RapierConfiguration>` -> `Query<&RapierConfiguration>`
-    - you can reach out for the underlying configuration via .single
+    - `ResMut<mut RapierContext>` -> `DefaultRapierContextAccessMut`
+    - `Res<RapierContext>` -> `DefaultRapierContextAccess`
+    - `ResMut<RapierConfiguration>` -> `DefaultRapierConfigurationMut`
+    - `Res<RapierConfiguration>` -> `DefaultRapierConfiguration`
 
 ## v0.27.0-rc.1 (18 June 2024)
 
diff --git a/bevy_rapier3d/examples/multi_world3.rs b/bevy_rapier3d/examples/multi_world3.rs
new file mode 100644
index 00000000..04f6b29c
--- /dev/null
+++ b/bevy_rapier3d/examples/multi_world3.rs
@@ -0,0 +1,117 @@
+use bevy::prelude::*;
+use bevy_rapier3d::prelude::*;
+
+const N_WORLDS: usize = 2;
+
+fn main() {
+    App::new()
+        .insert_resource(ClearColor(Color::rgb(
+            0xF9 as f32 / 255.0,
+            0xF9 as f32 / 255.0,
+            0xFF as f32 / 255.0,
+        )))
+        .add_plugins((
+            DefaultPlugins,
+            RapierPhysicsPlugin::<NoUserData>::default().with_default_world(None),
+            RapierDebugRenderPlugin::default(),
+        ))
+        .add_systems(
+            Startup,
+            ((create_worlds, setup_physics).chain(), setup_graphics),
+        )
+        .add_systems(Update, move_platforms)
+        // .add_systems(Update, change_world)
+        // .add_systems(Update, despawn_last)
+        .run();
+}
+
+fn create_worlds(mut commands: Commands) {
+    for i in 0..N_WORLDS {
+        commands.spawn((RapierContext::default(), WorldId(i)));
+    }
+}
+
+fn setup_graphics(mut commands: Commands) {
+    commands.spawn(Camera3dBundle {
+        transform: Transform::from_xyz(0.0, 3.0, -10.0)
+            .looking_at(Vec3::new(0.0, 0.0, 0.0), Vec3::Y),
+        ..Default::default()
+    });
+}
+
+#[derive(Component)]
+pub struct WorldId(pub usize);
+
+#[derive(Component)]
+struct Platform {
+    starting_y: f32,
+}
+
+fn move_platforms(time: Res<Time>, mut query: Query<(&mut Transform, &Platform)>) {
+    for (mut transform, platform) in query.iter_mut() {
+        transform.translation.y = platform.starting_y + -time.elapsed_seconds().sin();
+    }
+}
+
+/// Demonstrates despawning an entity removing it from its world
+// fn despawn_last(query: Query<(&PhysicsWorld, Entity)>, mut commands: Commands) {
+//     for (bw, entity) in query.iter() {
+//         if bw.world_id == N_WORLDS - 1 {
+//             commands.entity(entity).despawn_recursive();
+//         }
+//     }
+// }
+
+/// Demonstrates how easy it is to move one entity to another world.
+// fn change_world(mut query: Query<&mut PhysicsWorld>) {
+//     for mut bw in query.iter_mut() {
+//         if bw.world_id == 1 {
+//             bw.world_id = 0;
+//         }
+//     }
+// }
+
+pub fn setup_physics(
+    context: Query<(Entity, &WorldId), With<RapierContext>>,
+    mut commands: Commands,
+) {
+    for (context_entity, id) in context.iter() {
+        let id = id.0;
+
+        let color = [
+            Hsla::hsl(220.0, 1.0, 0.3),
+            Hsla::hsl(180.0, 1.0, 0.3),
+            Hsla::hsl(260.0, 1.0, 0.7),
+        ][id % 3];
+
+        /*
+         * Ground
+         */
+        let ground_size = 5.1;
+        let ground_height = 0.1;
+
+        let starting_y = (id as f32) * -0.5 - ground_height;
+
+        let mut platforms = commands.spawn((
+            TransformBundle::from(Transform::from_xyz(0.0, starting_y, 0.0)),
+            Collider::cuboid(ground_size, ground_height, ground_size),
+            ColliderDebugColor(color),
+            RapierContextEntityLink(context_entity),
+        ));
+        if id == 1 {
+            platforms.insert(Platform { starting_y });
+        }
+
+        /*
+         * Create the cube
+         */
+
+        commands.spawn((
+            TransformBundle::from(Transform::from_xyz(0.0, 1.0 + id as f32 * 5.0, 0.0)),
+            RigidBody::Dynamic,
+            Collider::cuboid(0.5, 0.5, 0.5),
+            ColliderDebugColor(color),
+            RapierContextEntityLink(context_entity),
+        ));
+    }
+}
diff --git a/src/plugin/context/systemparams/mod.rs b/src/plugin/context/systemparams/mod.rs
new file mode 100644
index 00000000..910ad2cf
--- /dev/null
+++ b/src/plugin/context/systemparams/mod.rs
@@ -0,0 +1,12 @@
+mod rapier_context_access;
+
+use bevy::{ecs::query::QueryData, prelude::Entity};
+pub use rapier_context_access::*;
+
+use super::RapierContextEntityLink;
+
+#[derive(QueryData)]
+pub struct RapierEntity {
+    pub entity: Entity,
+    pub rapier_context_link: &'static RapierContextEntityLink,
+}
diff --git a/src/plugin/context/systemparams.rs b/src/plugin/context/systemparams/rapier_context_access.rs
similarity index 66%
rename from src/plugin/context/systemparams.rs
rename to src/plugin/context/systemparams/rapier_context_access.rs
index f9c0a646..e3e88af1 100644
--- a/src/plugin/context/systemparams.rs
+++ b/src/plugin/context/systemparams/rapier_context_access.rs
@@ -2,8 +2,7 @@ use bevy::ecs::system::SystemParam;
 use bevy::prelude::*;
 use std::ops::{Deref, DerefMut};
 
-use super::{DefaultRapierContext, RapierContext, RapierContextEntityLink};
-
+use super::super::{DefaultRapierContext, RapierContext, RapierContextEntityLink};
 /// Utility [`SystemParam`] to easily access the default world [`RapierContext`] immutably
 #[derive(SystemParam)]
 pub struct DefaultRapierContextAccess<'w, 's> {
@@ -50,25 +49,14 @@ impl<'w, 's> DerefMut for DefaultRapierContextAccessMut<'w, 's> {
 #[derive(SystemParam)]
 pub struct RapierContextAccess<'w, 's> {
     rapier_context: Query<'w, 's, &'static RapierContext>,
-    rapier_context_link: Query<'w, 's, &'static RapierContextEntityLink>,
 }
 
 impl<'w, 's> RapierContextAccess<'w, 's> {
-    pub fn link(&self, entity: Entity) -> &RapierContextEntityLink {
-        self.rapier_context_link
-            .get(entity)
-            .expect("RapierContextAccess.link called on an entity without RapierContextEntityLink.")
-    }
-    pub fn follow_link(&self, link: RapierContextEntityLink) -> &'_ RapierContext {
+    pub fn context(&self, link: RapierContextEntityLink) -> &'_ RapierContext {
         self.rapier_context
             .get(link.0)
             .expect("RapierContextEntityLink.0 refers to an entity without RapierContext.")
     }
-
-    pub fn context(&self, entity: Entity) -> &'_ RapierContext {
-        let context_link = self.link(entity);
-        self.follow_link(*context_link)
-    }
 }
 
 impl<'w, 's> Deref for RapierContextAccess<'w, 's> {
@@ -83,35 +71,24 @@ impl<'w, 's> Deref for RapierContextAccess<'w, 's> {
 #[derive(SystemParam)]
 pub struct RapierContextAccessMut<'w, 's> {
     pub rapier_context: Query<'w, 's, &'static mut RapierContext>,
-    pub rapier_context_link: Query<'w, 's, &'static RapierContextEntityLink>,
 }
 
 impl<'w, 's> RapierContextAccessMut<'w, 's> {
-    pub fn link(&self, entity: Entity) -> &RapierContextEntityLink {
-        self.rapier_context_link.get(entity).expect(
-            "RapierContextAccessMut.link called on an entity without RapierContextEntityLink.",
-        )
-    }
-    pub fn follow_link(&mut self, link: RapierContextEntityLink) -> &'_ mut RapierContext {
+    pub fn context(&mut self, link: RapierContextEntityLink) -> &'_ mut RapierContext {
         self.rapier_context
             .get_mut(link.0)
             .expect("RapierContextEntityLink.0 refers to an entity without RapierContext.")
             .into_inner()
     }
-
-    pub fn context(&mut self, entity: Entity) -> &'_ mut RapierContext {
-        let context_link = self.link(entity);
-        self.follow_link(*context_link)
-    }
 }
 
 pub fn try_retrieve_context<'a>(
     link: Option<&RapierContextEntityLink>,
-    context: &Query<(Entity, &mut RapierContext)>,
+    context: &Query<Entity, With<RapierContext>>,
 ) -> Result<Entity, Entity> {
     link.map_or_else(
         || {
-            let context_entity = context.iter().next().unwrap().0;
+            let context_entity = context.iter().next().unwrap();
             Err(context_entity)
         },
         |link| Ok(link.0),
diff --git a/src/plugin/mod.rs b/src/plugin/mod.rs
index 08df87e9..5a4da77d 100644
--- a/src/plugin/mod.rs
+++ b/src/plugin/mod.rs
@@ -4,7 +4,7 @@ pub use self::context::{
         DefaultRapierContextAccess, DefaultRapierContextAccessMut, RapierContextAccess,
         RapierContextAccessMut,
     },
-    RapierContext,
+    RapierContext, RapierContextEntityLink,
 };
 pub use self::plugin::{NoUserData, PhysicsSet, RapierPhysicsPlugin, RapierTransformPropagateSet};
 
diff --git a/src/plugin/plugin.rs b/src/plugin/plugin.rs
index 9b8f0b29..c078e7ae 100644
--- a/src/plugin/plugin.rs
+++ b/src/plugin/plugin.rs
@@ -44,6 +44,18 @@ where
         });
         self
     }
+    /// Specifies a default configuration for the default `RapierContext`
+    ///
+    /// If [`None`], no world will be initialized, you are responsible of creating and maintaining
+    /// a [`RapierContext`] before creating any rapier entities (rigidbodies, colliders, joints),
+    /// and as long as any [`RapierContextEntityLink`] has a reference to its [`RapierContext`].
+    pub fn with_default_world(
+        mut self,
+        default_world_initialization: Option<RapierContextInitialization>,
+    ) -> Self {
+        self.default_world_setup = default_world_initialization;
+        self
+    }
 
     /// Specifies whether the plugin should setup each of its [`PhysicsStages`]
     /// (`true`), or if the user will set them up later (`false`).
@@ -244,6 +256,7 @@ where
     }
 }
 
+#[derive(Resource)]
 pub struct RapierContextInitialization {
     pub length_unit: f32,
 }
diff --git a/src/plugin/systems/character_controller.rs b/src/plugin/systems/character_controller.rs
index 0fec5b50..25233a9b 100644
--- a/src/plugin/systems/character_controller.rs
+++ b/src/plugin/systems/character_controller.rs
@@ -1,6 +1,7 @@
 use crate::control::CharacterCollision;
 use crate::dynamics::RapierRigidBodyHandle;
 use crate::geometry::RapierColliderHandle;
+use crate::plugin::context::RapierContextEntityLink;
 use crate::plugin::RapierConfiguration;
 use crate::plugin::RapierContextAccessMut;
 use crate::prelude::KinematicCharacterController;
@@ -19,6 +20,7 @@ pub fn update_character_controls(
     mut context_access: RapierContextAccessMut,
     mut character_controllers: Query<(
         Entity,
+        &RapierContextEntityLink,
         &mut KinematicCharacterController,
         Option<&mut KinematicCharacterControllerOutput>,
         Option<&RapierColliderHandle>,
@@ -27,12 +29,18 @@ pub fn update_character_controls(
     )>,
     mut transforms: Query<&mut Transform>,
 ) {
-    for (entity, mut controller, output, collider_handle, body_handle, glob_transform) in
-        character_controllers.iter_mut()
+    for (
+        entity,
+        rapier_context_link,
+        mut controller,
+        output,
+        collider_handle,
+        body_handle,
+        glob_transform,
+    ) in character_controllers.iter_mut()
     {
-        let link = *context_access.link(entity);
-        let context = context_access.follow_link(link);
-        let config = config.get(link.0).unwrap();
+        let context = context_access.context(*rapier_context_link);
+        let config = config.get(rapier_context_link.0).unwrap();
         if let (Some(raw_controller), Some(translation)) =
             (controller.to_raw(), controller.translation)
         {
diff --git a/src/plugin/systems/collider.rs b/src/plugin/systems/collider.rs
index 4ee87d52..f5765867 100644
--- a/src/plugin/systems/collider.rs
+++ b/src/plugin/systems/collider.rs
@@ -1,6 +1,6 @@
 use crate::dynamics::ReadMassProperties;
 use crate::geometry::Collider;
-use crate::plugin::context::systemparams::try_retrieve_context;
+use crate::plugin::context::systemparams::{try_retrieve_context, RapierEntity};
 use crate::plugin::context::RapierContextEntityLink;
 use crate::plugin::{RapierConfiguration, RapierContext, RapierContextAccessMut};
 use crate::prelude::{
@@ -45,7 +45,12 @@ pub type ColliderComponents<'a> = (
 pub fn apply_scale(
     config: Query<&RapierConfiguration>,
     mut changed_collider_scales: Query<
-        (&mut Collider, &GlobalTransform, Option<&ColliderScale>),
+        (
+            &mut Collider,
+            &RapierContextEntityLink,
+            &GlobalTransform,
+            Option<&ColliderScale>,
+        ),
         Or<(
             Changed<Collider>,
             Changed<GlobalTransform>,
@@ -53,9 +58,8 @@ pub fn apply_scale(
         )>,
     >,
 ) {
-    let config = &*config.single();
-
-    for (mut shape, transform, custom_scale) in changed_collider_scales.iter_mut() {
+    for (mut shape, link, transform, custom_scale) in changed_collider_scales.iter_mut() {
+        let config = config.get(link.0).unwrap();
         #[cfg(feature = "dim2")]
         let effective_scale = match custom_scale {
             Some(ColliderScale::Absolute(scale)) => *scale,
@@ -83,62 +87,69 @@ pub fn apply_collider_user_changes(
     config: Query<&RapierConfiguration>,
     (changed_collider_transforms, parent_query, transform_query): (
         Query<
-            (Entity, &RapierColliderHandle, &GlobalTransform),
+            (RapierEntity, &RapierColliderHandle, &GlobalTransform),
             (Without<RapierRigidBodyHandle>, Changed<GlobalTransform>),
         >,
         Query<&Parent>,
         Query<&Transform>,
     ),
 
-    changed_shapes: Query<(Entity, &RapierColliderHandle, &Collider), Changed<Collider>>,
+    changed_shapes: Query<(RapierEntity, &RapierColliderHandle, &Collider), Changed<Collider>>,
     changed_active_events: Query<
-        (Entity, &RapierColliderHandle, &ActiveEvents),
+        (RapierEntity, &RapierColliderHandle, &ActiveEvents),
         Changed<ActiveEvents>,
     >,
     changed_active_hooks: Query<
-        (Entity, &RapierColliderHandle, &ActiveHooks),
+        (RapierEntity, &RapierColliderHandle, &ActiveHooks),
         Changed<ActiveHooks>,
     >,
     changed_active_collision_types: Query<
-        (Entity, &RapierColliderHandle, &ActiveCollisionTypes),
+        (RapierEntity, &RapierColliderHandle, &ActiveCollisionTypes),
         Changed<ActiveCollisionTypes>,
     >,
     (changed_friction, changed_restitution, changed_contact_skin): (
-        Query<(Entity, &RapierColliderHandle, &Friction), Changed<Friction>>,
-        Query<(Entity, &RapierColliderHandle, &Restitution), Changed<Restitution>>,
-        Query<(Entity, &RapierColliderHandle, &ContactSkin), Changed<ContactSkin>>,
+        Query<(RapierEntity, &RapierColliderHandle, &Friction), Changed<Friction>>,
+        Query<(RapierEntity, &RapierColliderHandle, &Restitution), Changed<Restitution>>,
+        Query<(RapierEntity, &RapierColliderHandle, &ContactSkin), Changed<ContactSkin>>,
     ),
     changed_collision_groups: Query<
-        (Entity, &RapierColliderHandle, &CollisionGroups),
+        (RapierEntity, &RapierColliderHandle, &CollisionGroups),
         Changed<CollisionGroups>,
     >,
     changed_solver_groups: Query<
-        (Entity, &RapierColliderHandle, &SolverGroups),
+        (RapierEntity, &RapierColliderHandle, &SolverGroups),
         Changed<SolverGroups>,
     >,
-    changed_sensors: Query<(Entity, &RapierColliderHandle, &Sensor), Changed<Sensor>>,
+    changed_sensors: Query<(RapierEntity, &RapierColliderHandle, &Sensor), Changed<Sensor>>,
     changed_disabled: Query<
-        (Entity, &RapierColliderHandle, &ColliderDisabled),
+        (RapierEntity, &RapierColliderHandle, &ColliderDisabled),
         Changed<ColliderDisabled>,
     >,
     changed_contact_force_threshold: Query<
-        (Entity, &RapierColliderHandle, &ContactForceEventThreshold),
+        (
+            RapierEntity,
+            &RapierColliderHandle,
+            &ContactForceEventThreshold,
+        ),
         Changed<ContactForceEventThreshold>,
     >,
     changed_collider_mass_props: Query<
-        (Entity, &RapierColliderHandle, &ColliderMassProperties),
+        (RapierEntity, &RapierColliderHandle, &ColliderMassProperties),
         Changed<ColliderMassProperties>,
     >,
 
     mut mass_modified: EventWriter<MassModifiedEvent>,
 ) {
-    let config = &*config.single();
-
-    for (entity, handle, transform) in changed_collider_transforms.iter() {
-        let context = context.context(entity);
-        if context.collider_parent(entity).is_some() {
-            let (_, collider_position) =
-                collider_offset(entity, &context, &parent_query, &transform_query);
+    for (rapier_entity, handle, transform) in changed_collider_transforms.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
+        let config = config.get(rapier_entity.rapier_context_link.0).unwrap();
+        if context.collider_parent(rapier_entity.entity).is_some() {
+            let (_, collider_position) = collider_offset(
+                rapier_entity.entity,
+                &context,
+                &parent_query,
+                &transform_query,
+            );
 
             if let Some(co) = context.colliders.get_mut(handle.0) {
                 co.set_position_wrt_parent(utils::transform_to_iso(&collider_position));
@@ -148,8 +159,9 @@ pub fn apply_collider_user_changes(
         }
     }
 
-    for (entity, handle, shape) in changed_shapes.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, shape) in changed_shapes.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
+        let config = config.get(rapier_entity.rapier_context_link.0).unwrap();
         if let Some(co) = context.colliders.get_mut(handle.0) {
             let mut scaled_shape = shape.clone();
             scaled_shape.set_scale(shape.scale, config.scaled_shape_subdivision);
@@ -163,87 +175,87 @@ pub fn apply_collider_user_changes(
         }
     }
 
-    for (entity, handle, active_events) in changed_active_events.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, active_events) in changed_active_events.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_active_events((*active_events).into())
         }
     }
 
-    for (entity, handle, active_hooks) in changed_active_hooks.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, active_hooks) in changed_active_hooks.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_active_hooks((*active_hooks).into())
         }
     }
 
-    for (entity, handle, active_collision_types) in changed_active_collision_types.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, active_collision_types) in changed_active_collision_types.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_active_collision_types((*active_collision_types).into())
         }
     }
 
-    for (entity, handle, friction) in changed_friction.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, friction) in changed_friction.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_friction(friction.coefficient);
             co.set_friction_combine_rule(friction.combine_rule.into());
         }
     }
 
-    for (entity, handle, restitution) in changed_restitution.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, restitution) in changed_restitution.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_restitution(restitution.coefficient);
             co.set_restitution_combine_rule(restitution.combine_rule.into());
         }
     }
 
-    for (entity, handle, contact_skin) in changed_contact_skin.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, contact_skin) in changed_contact_skin.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_contact_skin(contact_skin.0);
         }
     }
 
-    for (entity, handle, collision_groups) in changed_collision_groups.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, collision_groups) in changed_collision_groups.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_collision_groups((*collision_groups).into());
         }
     }
 
-    for (entity, handle, solver_groups) in changed_solver_groups.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, solver_groups) in changed_solver_groups.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_solver_groups((*solver_groups).into());
         }
     }
 
-    for (entity, handle, _) in changed_sensors.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, _) in changed_sensors.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_sensor(true);
         }
     }
 
-    for (entity, handle, _) in changed_disabled.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, _) in changed_disabled.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_enabled(false);
         }
     }
 
-    for (entity, handle, threshold) in changed_contact_force_threshold.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, threshold) in changed_contact_force_threshold.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             co.set_contact_force_event_threshold(threshold.0);
         }
     }
 
-    for (entity, handle, mprops) in changed_collider_mass_props.iter() {
-        let context = context.context(entity);
+    for (rapier_entity, handle, mprops) in changed_collider_mass_props.iter() {
+        let context = context.context(*rapier_entity.rapier_context_link);
         if let Some(co) = context.colliders.get_mut(handle.0) {
             match mprops {
                 ColliderMassProperties::Density(density) => co.set_density(*density),
@@ -328,13 +340,16 @@ pub fn init_colliders(
         global_transform,
     ) in colliders.iter()
     {
-        let context_entity =
-            try_retrieve_context(context_link, &context).unwrap_or_else(|context_entity| {
+        let context_entity = context_link.map_or_else(
+            || {
+                let context_entity = context.iter().next().unwrap().0;
                 commands
                     .entity(entity)
                     .insert(RapierContextEntityLink(context_entity));
                 context_entity
-            });
+            },
+            |link| link.0,
+        );
         let context = &mut *context.get_mut(context_entity).unwrap().1;
         let config = config.get(context_entity).unwrap();
         let mut scaled_shape = shape.clone();
diff --git a/src/plugin/systems/joint.rs b/src/plugin/systems/joint.rs
index d08586ec..c361c1b6 100644
--- a/src/plugin/systems/joint.rs
+++ b/src/plugin/systems/joint.rs
@@ -2,20 +2,38 @@ use crate::dynamics::ImpulseJoint;
 use crate::dynamics::MultibodyJoint;
 use crate::dynamics::RapierImpulseJointHandle;
 use crate::dynamics::RapierMultibodyJointHandle;
+use crate::plugin::context::systemparams::try_retrieve_context;
 use crate::plugin::context::RapierContextEntityLink;
+use crate::plugin::RapierContext;
 use crate::plugin::RapierContextAccessMut;
 use bevy::prelude::*;
 
 /// System responsible for creating new Rapier joints from the related `bevy_rapier` components.
 pub fn init_joints(
     mut commands: Commands,
-    mut context: RapierContextAccessMut,
-    impulse_joints: Query<(Entity, &ImpulseJoint), Without<RapierImpulseJointHandle>>,
-    multibody_joints: Query<(Entity, &MultibodyJoint), Without<RapierMultibodyJointHandle>>,
+    mut context: Query<(Entity, &mut RapierContext)>,
+    impulse_joints: Query<
+        (Entity, Option<&RapierContextEntityLink>, &ImpulseJoint),
+        Without<RapierImpulseJointHandle>,
+    >,
+    multibody_joints: Query<
+        (Entity, Option<&RapierContextEntityLink>, &MultibodyJoint),
+        Without<RapierMultibodyJointHandle>,
+    >,
     parent_query: Query<&Parent>,
 ) {
-    for (entity, joint) in impulse_joints.iter() {
-        let context = context.context(entity);
+    for (entity, entity_context_link, joint) in impulse_joints.iter() {
+        let context_entity = entity_context_link.map_or_else(
+            || {
+                let context_entity = context.iter().next().unwrap().0;
+                commands
+                    .entity(entity)
+                    .insert(RapierContextEntityLink(context_entity));
+                context_entity
+            },
+            |link| link.0,
+        );
+        let context = &mut *context.get_mut(context_entity).unwrap().1;
         let mut target = None;
         let mut body_entity = entity;
         while target.is_none() {
@@ -39,8 +57,18 @@ pub fn init_joints(
         }
     }
 
-    for (entity, joint) in multibody_joints.iter() {
-        let context = context.context(entity);
+    for (entity, entity_context_link, joint) in multibody_joints.iter() {
+        let context_entity = entity_context_link.map_or_else(
+            || {
+                let context_entity = context.iter().next().unwrap().0;
+                commands
+                    .entity(entity)
+                    .insert(RapierContextEntityLink(context_entity));
+                context_entity
+            },
+            |link| link.0,
+        );
+        let context = &mut *context.get_mut(context_entity).unwrap().1;
         let target = context.entity2body.get(&entity);
 
         if let (Some(target), Some(source)) = (target, context.entity2body.get(&joint.parent)) {
@@ -83,14 +111,14 @@ pub fn apply_joint_user_changes(
     // TODO: right now, we only support propagating changes made to the joint data.
     //       Re-parenting the joint isn’t supported yet.
     for (link, handle, changed_joint) in changed_impulse_joints.iter() {
-        let context = context.follow_link(*link);
+        let context = context.context(*link);
         if let Some(joint) = context.impulse_joints.get_mut(handle.0) {
             joint.data = changed_joint.data.into_rapier();
         }
     }
 
     for (link, handle, changed_joint) in changed_multibody_joints.iter() {
-        let context = context.follow_link(*link);
+        let context = context.context(*link);
         // TODO: not sure this will always work properly, e.g., if the number of Dofs is changed.
         if let Some((mb, link_id)) = context.multibody_joints.get_mut(handle.0) {
             if let Some(link) = mb.link_mut(link_id) {
diff --git a/src/plugin/systems/remove.rs b/src/plugin/systems/remove.rs
index 1f6256b2..8450c087 100644
--- a/src/plugin/systems/remove.rs
+++ b/src/plugin/systems/remove.rs
@@ -8,6 +8,7 @@ use crate::geometry::Collider;
 use crate::geometry::ColliderDisabled;
 use crate::geometry::RapierColliderHandle;
 use crate::plugin::context::systemparams::RapierContextAccessMut;
+use crate::plugin::context::RapierContextEntityLink;
 use crate::prelude::MassModifiedEvent;
 use crate::prelude::RigidBodyDisabled;
 use crate::prelude::Sensor;
@@ -23,6 +24,7 @@ pub fn sync_removals(
     mut removed_colliders: RemovedComponents<RapierColliderHandle>,
     mut removed_impulse_joints: RemovedComponents<RapierImpulseJointHandle>,
     mut removed_multibody_joints: RemovedComponents<RapierMultibodyJointHandle>,
+    link: Query<&RapierContextEntityLink>,
     orphan_bodies: Query<Entity, (With<RapierRigidBodyHandle>, Without<RigidBody>)>,
     orphan_colliders: Query<Entity, (With<RapierColliderHandle>, Without<Collider>)>,
     orphan_impulse_joints: Query<Entity, (With<RapierImpulseJointHandle>, Without<ImpulseJoint>)>,
@@ -41,7 +43,7 @@ pub fn sync_removals(
      * Rigid-bodies removal detection.
      */
     for entity in removed_bodies.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2body.remove(&entity) {
             let _ = context.last_body_transform_set.remove(&handle);
             context.bodies.remove(
@@ -56,7 +58,7 @@ pub fn sync_removals(
     }
 
     for entity in orphan_bodies.iter() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2body.remove(&entity) {
             let _ = context.last_body_transform_set.remove(&handle);
             context.bodies.remove(
@@ -75,7 +77,7 @@ pub fn sync_removals(
      * Collider removal detection.
      */
     for entity in removed_colliders.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(parent) = context.collider_parent(entity) {
             mass_modified.send(parent.into());
         }
@@ -89,7 +91,7 @@ pub fn sync_removals(
     }
 
     for entity in orphan_colliders.iter() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(parent) = context.collider_parent(entity) {
             mass_modified.send(parent.into());
         }
@@ -107,14 +109,14 @@ pub fn sync_removals(
      * Impulse joint removal detection.
      */
     for entity in removed_impulse_joints.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2impulse_joint.remove(&entity) {
             context.impulse_joints.remove(handle, true);
         }
     }
 
     for entity in orphan_impulse_joints.iter() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2impulse_joint.remove(&entity) {
             context.impulse_joints.remove(handle, true);
         }
@@ -125,14 +127,14 @@ pub fn sync_removals(
      * Multibody joint removal detection.
      */
     for entity in removed_multibody_joints.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2multibody_joint.remove(&entity) {
             context.multibody_joints.remove(handle, true);
         }
     }
 
     for entity in orphan_multibody_joints.iter() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2multibody_joint.remove(&entity) {
             context.multibody_joints.remove(handle, true);
         }
@@ -145,7 +147,7 @@ pub fn sync_removals(
      * Marker components removal detection.
      */
     for entity in removed_sensors.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2collider.get(&entity) {
             if let Some(co) = context.colliders.get_mut(*handle) {
                 co.set_sensor(false);
@@ -154,7 +156,7 @@ pub fn sync_removals(
     }
 
     for entity in removed_colliders_disabled.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2collider.get(&entity) {
             if let Some(co) = context.colliders.get_mut(*handle) {
                 co.set_enabled(true);
@@ -163,7 +165,7 @@ pub fn sync_removals(
     }
 
     for entity in removed_rigid_body_disabled.read() {
-        let context = context.context(entity);
+        let context = context.context(*link.get(entity).unwrap());
         if let Some(handle) = context.entity2body.get(&entity) {
             if let Some(rb) = context.bodies.get_mut(*handle) {
                 rb.set_enabled(true);
diff --git a/src/plugin/systems/rigid_body.rs b/src/plugin/systems/rigid_body.rs
index 0eafb179..38a589b8 100644
--- a/src/plugin/systems/rigid_body.rs
+++ b/src/plugin/systems/rigid_body.rs
@@ -12,6 +12,7 @@ use std::ops::DerefMut;
 /// Components that will be updated after a physics step.
 pub type RigidBodyWritebackComponents<'a> = (
     &'a RapierRigidBodyHandle,
+    &'a RapierContextEntityLink,
     Option<&'a Parent>,
     Option<&'a mut Transform>,
     Option<&'a mut TransformInterpolation>,
@@ -42,238 +43,309 @@ pub type RigidBodyComponents<'a> = (
 pub fn apply_rigid_body_user_changes(
     mut context: Query<&mut RapierContext>,
     config: Query<&RapierConfiguration>,
-    changed_rb_types: Query<(&RapierRigidBodyHandle, &RigidBody), Changed<RigidBody>>,
+    changed_rb_types: Query<
+        (&RapierRigidBodyHandle, &RapierContextEntityLink, &RigidBody),
+        Changed<RigidBody>,
+    >,
     mut changed_transforms: Query<
         (
             &RapierRigidBodyHandle,
+            &RapierContextEntityLink,
             &GlobalTransform,
             Option<&mut TransformInterpolation>,
         ),
         Changed<GlobalTransform>,
     >,
-    changed_velocities: Query<(&RapierRigidBodyHandle, &Velocity), Changed<Velocity>>,
+    changed_velocities: Query<
+        (&RapierRigidBodyHandle, &RapierContextEntityLink, &Velocity),
+        Changed<Velocity>,
+    >,
     changed_additional_mass_props: Query<
-        (Entity, &RapierRigidBodyHandle, &AdditionalMassProperties),
+        (
+            Entity,
+            &RapierContextEntityLink,
+            &RapierRigidBodyHandle,
+            &AdditionalMassProperties,
+        ),
         Changed<AdditionalMassProperties>,
     >,
-    changed_locked_axes: Query<(&RapierRigidBodyHandle, &LockedAxes), Changed<LockedAxes>>,
-    changed_forces: Query<(&RapierRigidBodyHandle, &ExternalForce), Changed<ExternalForce>>,
+    changed_locked_axes: Query<
+        (
+            &RapierRigidBodyHandle,
+            &RapierContextEntityLink,
+            &LockedAxes,
+        ),
+        Changed<LockedAxes>,
+    >,
+    changed_forces: Query<
+        (
+            &RapierRigidBodyHandle,
+            &RapierContextEntityLink,
+            &ExternalForce,
+        ),
+        Changed<ExternalForce>,
+    >,
     mut changed_impulses: Query<
-        (&RapierRigidBodyHandle, &mut ExternalImpulse),
+        (
+            &RapierRigidBodyHandle,
+            &RapierContextEntityLink,
+            &mut ExternalImpulse,
+        ),
         Changed<ExternalImpulse>,
     >,
-    changed_gravity_scale: Query<(&RapierRigidBodyHandle, &GravityScale), Changed<GravityScale>>,
+    changed_gravity_scale: Query<
+        (
+            &RapierRigidBodyHandle,
+            &RapierContextEntityLink,
+            &GravityScale,
+        ),
+        Changed<GravityScale>,
+    >,
     (changed_ccd, changed_soft_ccd): (
-        Query<(&RapierRigidBodyHandle, &Ccd), Changed<Ccd>>,
-        Query<(&RapierRigidBodyHandle, &SoftCcd), Changed<SoftCcd>>,
+        Query<(&RapierRigidBodyHandle, &RapierContextEntityLink, &Ccd), Changed<Ccd>>,
+        Query<(&RapierRigidBodyHandle, &RapierContextEntityLink, &SoftCcd), Changed<SoftCcd>>,
     ),
-    changed_dominance: Query<(&RapierRigidBodyHandle, &Dominance), Changed<Dominance>>,
-    changed_sleeping: Query<(&RapierRigidBodyHandle, &Sleeping), Changed<Sleeping>>,
-    changed_damping: Query<(&RapierRigidBodyHandle, &Damping), Changed<Damping>>,
+    changed_dominance: Query<
+        (&RapierRigidBodyHandle, &RapierContextEntityLink, &Dominance),
+        Changed<Dominance>,
+    >,
+    changed_sleeping: Query<
+        (&RapierRigidBodyHandle, &RapierContextEntityLink, &Sleeping),
+        Changed<Sleeping>,
+    >,
+    changed_damping: Query<
+        (&RapierRigidBodyHandle, &RapierContextEntityLink, &Damping),
+        Changed<Damping>,
+    >,
     (changed_disabled, changed_additional_solver_iterations): (
-        Query<(&RapierRigidBodyHandle, &RigidBodyDisabled), Changed<RigidBodyDisabled>>,
         Query<
-            (&RapierRigidBodyHandle, &AdditionalSolverIterations),
+            (
+                &RapierRigidBodyHandle,
+                &RapierContextEntityLink,
+                &RigidBodyDisabled,
+            ),
+            Changed<RigidBodyDisabled>,
+        >,
+        Query<
+            (
+                &RapierRigidBodyHandle,
+                &RapierContextEntityLink,
+                &AdditionalSolverIterations,
+            ),
             Changed<AdditionalSolverIterations>,
         >,
     ),
     mut mass_modified: EventWriter<MassModifiedEvent>,
 ) {
-    for mut context in context.iter_mut() {
-        let context = context.deref_mut();
-        let config = config.single();
-
-        // Deal with sleeping first, because other changes may then wake-up the
-        // rigid-body again.
-        for (handle, sleeping) in changed_sleeping.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                let activation = rb.activation_mut();
-                activation.normalized_linear_threshold = sleeping.normalized_linear_threshold;
-                activation.angular_threshold = sleeping.angular_threshold;
-
-                if !sleeping.sleeping && activation.sleeping {
-                    rb.wake_up(true);
-                } else if sleeping.sleeping && !activation.sleeping {
-                    rb.sleep();
-                }
+    // Deal with sleeping first, because other changes may then wake-up the
+    // rigid-body again.
+    for (handle, link, sleeping) in changed_sleeping.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            let activation = rb.activation_mut();
+            activation.normalized_linear_threshold = sleeping.normalized_linear_threshold;
+            activation.angular_threshold = sleeping.angular_threshold;
+
+            if !sleeping.sleeping && activation.sleeping {
+                rb.wake_up(true);
+            } else if sleeping.sleeping && !activation.sleeping {
+                rb.sleep();
             }
         }
+    }
 
-        // NOTE: we must change the rigid-body type before updating the
-        //       transform or velocity. Otherwise, if the rigid-body was fixed
-        //       and changed to anything else, the velocity change wouldn’t have any effect.
-        //       Similarly, if the rigid-body was kinematic position-based before and
-        //       changed to anything else, a transform change would modify the next
-        //       position instead of the current one.
-        for (handle, rb_type) in changed_rb_types.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_body_type((*rb_type).into(), true);
-            }
+    // NOTE: we must change the rigid-body type before updating the
+    //       transform or velocity. Otherwise, if the rigid-body was fixed
+    //       and changed to anything else, the velocity change wouldn’t have any effect.
+    //       Similarly, if the rigid-body was kinematic position-based before and
+    //       changed to anything else, a transform change would modify the next
+    //       position instead of the current one.
+    for (handle, link, rb_type) in changed_rb_types.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_body_type((*rb_type).into(), true);
         }
+    }
 
-        // Manually checks if the transform changed.
-        // This is needed for detecting if the user actually changed the rigid-body
-        // transform, or if it was just the change we made in our `writeback_rigid_bodies`
-        // system.
-        let transform_changed_fn =
-            |handle: &RigidBodyHandle,
-             transform: &GlobalTransform,
-             last_transform_set: &HashMap<RigidBodyHandle, GlobalTransform>| {
-                if config.force_update_from_transform_changes {
-                    true
-                } else if let Some(prev) = last_transform_set.get(handle) {
-                    *prev != *transform
-                } else {
-                    true
-                }
-            };
+    // Manually checks if the transform changed.
+    // This is needed for detecting if the user actually changed the rigid-body
+    // transform, or if it was just the change we made in our `writeback_rigid_bodies`
+    // system.
+    let transform_changed_fn =
+        |handle: &RigidBodyHandle,
+         config: &RapierConfiguration,
+         transform: &GlobalTransform,
+         last_transform_set: &HashMap<RigidBodyHandle, GlobalTransform>| {
+            if config.force_update_from_transform_changes {
+                true
+            } else if let Some(prev) = last_transform_set.get(handle) {
+                *prev != *transform
+            } else {
+                true
+            }
+        };
+
+    for (handle, link, global_transform, mut interpolation) in changed_transforms.iter_mut() {
+        let context = &mut *context.get_mut(link.0).unwrap();
+        let config = config.get(link.0).unwrap();
+        // Use an Option<bool> to avoid running the check twice.
+        let mut transform_changed = None;
+
+        if let Some(interpolation) = interpolation.as_deref_mut() {
+            transform_changed = transform_changed.or_else(|| {
+                Some(transform_changed_fn(
+                    &handle.0,
+                    config,
+                    global_transform,
+                    &context.last_body_transform_set,
+                ))
+            });
 
-        for (handle, global_transform, mut interpolation) in changed_transforms.iter_mut() {
-            // Use an Option<bool> to avoid running the check twice.
-            let mut transform_changed = None;
-
-            if let Some(interpolation) = interpolation.as_deref_mut() {
-                transform_changed = transform_changed.or_else(|| {
-                    Some(transform_changed_fn(
-                        &handle.0,
-                        global_transform,
-                        &context.last_body_transform_set,
-                    ))
-                });
-
-                if transform_changed == Some(true) {
-                    // Reset the interpolation so we don’t overwrite
-                    // the user’s input.
-                    interpolation.start = None;
-                    interpolation.end = None;
-                }
+            if transform_changed == Some(true) {
+                // Reset the interpolation so we don’t overwrite
+                // the user’s input.
+                interpolation.start = None;
+                interpolation.end = None;
             }
-            // FIXME: avoid to run multiple times the mutable deref ?
-            let context = &mut *context;
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                transform_changed = transform_changed.or_else(|| {
-                    Some(transform_changed_fn(
-                        &handle.0,
-                        global_transform,
-                        &context.last_body_transform_set,
-                    ))
-                });
-
-                match rb.body_type() {
-                    RigidBodyType::KinematicPositionBased => {
-                        if transform_changed == Some(true) {
-                            rb.set_next_kinematic_position(utils::transform_to_iso(
-                                &global_transform.compute_transform(),
-                            ));
-                            context
-                                .last_body_transform_set
-                                .insert(handle.0, *global_transform);
-                        }
+        }
+        // FIXME: avoid to run multiple times the mutable deref ?
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            transform_changed = transform_changed.or_else(|| {
+                Some(transform_changed_fn(
+                    &handle.0,
+                    config,
+                    global_transform,
+                    &context.last_body_transform_set,
+                ))
+            });
+
+            match rb.body_type() {
+                RigidBodyType::KinematicPositionBased => {
+                    if transform_changed == Some(true) {
+                        rb.set_next_kinematic_position(utils::transform_to_iso(
+                            &global_transform.compute_transform(),
+                        ));
+                        context
+                            .last_body_transform_set
+                            .insert(handle.0, *global_transform);
                     }
-                    _ => {
-                        if transform_changed == Some(true) {
-                            rb.set_position(
-                                utils::transform_to_iso(&global_transform.compute_transform()),
-                                true,
-                            );
-                            context
-                                .last_body_transform_set
-                                .insert(handle.0, *global_transform);
-                        }
+                }
+                _ => {
+                    if transform_changed == Some(true) {
+                        rb.set_position(
+                            utils::transform_to_iso(&global_transform.compute_transform()),
+                            true,
+                        );
+                        context
+                            .last_body_transform_set
+                            .insert(handle.0, *global_transform);
                     }
                 }
             }
         }
+    }
 
-        for (handle, velocity) in changed_velocities.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_linvel(velocity.linvel.into(), true);
-                #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
-                rb.set_angvel(velocity.angvel.into(), true);
-            }
+    for (handle, link, velocity) in changed_velocities.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_linvel(velocity.linvel.into(), true);
+            #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
+            rb.set_angvel(velocity.angvel.into(), true);
         }
+    }
 
-        for (entity, handle, mprops) in changed_additional_mass_props.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                match mprops {
-                    AdditionalMassProperties::MassProperties(mprops) => {
-                        rb.set_additional_mass_properties(mprops.into_rapier(), true);
-                    }
-                    AdditionalMassProperties::Mass(mass) => {
-                        rb.set_additional_mass(*mass, true);
-                    }
+    for (entity, link, handle, mprops) in changed_additional_mass_props.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            match mprops {
+                AdditionalMassProperties::MassProperties(mprops) => {
+                    rb.set_additional_mass_properties(mprops.into_rapier(), true);
+                }
+                AdditionalMassProperties::Mass(mass) => {
+                    rb.set_additional_mass(*mass, true);
                 }
-
-                mass_modified.send(entity.into());
             }
+
+            mass_modified.send(entity.into());
         }
+    }
 
-        for (handle, additional_solver_iters) in changed_additional_solver_iterations.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_additional_solver_iterations(additional_solver_iters.0);
-            }
+    for (handle, link, additional_solver_iters) in changed_additional_solver_iterations.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_additional_solver_iterations(additional_solver_iters.0);
         }
+    }
 
-        for (handle, locked_axes) in changed_locked_axes.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_locked_axes((*locked_axes).into(), true);
-            }
+    for (handle, link, locked_axes) in changed_locked_axes.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_locked_axes((*locked_axes).into(), true);
         }
+    }
 
-        for (handle, forces) in changed_forces.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.reset_forces(true);
-                rb.reset_torques(true);
-                rb.add_force(forces.force.into(), true);
-                #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
-                rb.add_torque(forces.torque.into(), true);
-            }
+    for (handle, link, forces) in changed_forces.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.reset_forces(true);
+            rb.reset_torques(true);
+            rb.add_force(forces.force.into(), true);
+            #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
+            rb.add_torque(forces.torque.into(), true);
         }
+    }
 
-        for (handle, mut impulses) in changed_impulses.iter_mut() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.apply_impulse(impulses.impulse.into(), true);
-                #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
-                rb.apply_torque_impulse(impulses.torque_impulse.into(), true);
-                impulses.reset();
-            }
+    for (handle, link, mut impulses) in changed_impulses.iter_mut() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.apply_impulse(impulses.impulse.into(), true);
+            #[allow(clippy::useless_conversion)] // Need to convert if dim3 enabled
+            rb.apply_torque_impulse(impulses.torque_impulse.into(), true);
+            impulses.reset();
         }
+    }
 
-        for (handle, gravity_scale) in changed_gravity_scale.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_gravity_scale(gravity_scale.0, true);
-            }
+    for (handle, link, gravity_scale) in changed_gravity_scale.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_gravity_scale(gravity_scale.0, true);
         }
+    }
 
-        for (handle, ccd) in changed_ccd.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.enable_ccd(ccd.enabled);
-            }
+    for (handle, link, ccd) in changed_ccd.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.enable_ccd(ccd.enabled);
         }
+    }
 
-        for (handle, soft_ccd) in changed_soft_ccd.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_soft_ccd_prediction(soft_ccd.prediction);
-            }
+    for (handle, link, soft_ccd) in changed_soft_ccd.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_soft_ccd_prediction(soft_ccd.prediction);
         }
+    }
 
-        for (handle, dominance) in changed_dominance.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_dominance_group(dominance.groups);
-            }
+    for (handle, link, dominance) in changed_dominance.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_dominance_group(dominance.groups);
         }
+    }
 
-        for (handle, damping) in changed_damping.iter() {
-            if let Some(rb) = context.bodies.get_mut(handle.0) {
-                rb.set_linear_damping(damping.linear_damping);
-                rb.set_angular_damping(damping.angular_damping);
-            }
+    for (handle, link, damping) in changed_damping.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(rb) = context.bodies.get_mut(handle.0) {
+            rb.set_linear_damping(damping.linear_damping);
+            rb.set_angular_damping(damping.angular_damping);
         }
+    }
 
-        for (handle, _) in changed_disabled.iter() {
-            if let Some(co) = context.bodies.get_mut(handle.0) {
-                co.set_enabled(false);
-            }
+    for (handle, link, _) in changed_disabled.iter() {
+        let mut context = context.get_mut(link.0).unwrap();
+        if let Some(co) = context.bodies.get_mut(handle.0) {
+            co.set_enabled(false);
         }
     }
 }
@@ -291,143 +363,140 @@ pub fn writeback_rigid_bodies(
         (With<RigidBody>, Without<RigidBodyDisabled>),
     >,
 ) {
-    for mut context in context.iter_mut() {
-        let sim_to_render_time = sim_to_render_time.single();
-        let config = config.single();
-
-        if config.physics_pipeline_active {
-            for (handle, parent, transform, mut interpolation, mut velocity, mut sleeping) in
-                writeback.iter_mut()
-            {
-                let handle = handle.0;
-
-                let context = &mut *context;
-                // TODO: do this the other way round: iterate through Rapier’s RigidBodySet on the active bodies,
-                // and update the components accordingly. That way, we don’t have to iterate through the entities that weren’t changed
-                // by physics (for example because they are sleeping).
-                if let Some(rb) = context.bodies.get(handle) {
-                    let mut interpolated_pos = utils::iso_to_transform(rb.position());
-
-                    if let TimestepMode::Interpolated { dt, .. } = *timestep_mode {
-                        if let Some(interpolation) = interpolation.as_deref_mut() {
-                            if interpolation.end.is_none() {
-                                interpolation.end = Some(*rb.position());
-                            }
-
-                            if let Some(interpolated) =
-                                interpolation.lerp_slerp((dt + sim_to_render_time.diff) / dt)
-                            {
-                                interpolated_pos = utils::iso_to_transform(&interpolated);
-                            }
-                        }
+    for (handle, link, parent, transform, mut interpolation, mut velocity, mut sleeping) in
+        writeback.iter_mut()
+    {
+        let config = config.get(link.0).unwrap();
+        if !config.physics_pipeline_active {
+            continue;
+        }
+        let handle = handle.0;
+
+        let context = &mut *context.get_mut(link.0).unwrap();
+        let sim_to_render_time = sim_to_render_time.get(link.0).unwrap();
+        // TODO: do this the other way round: iterate through Rapier’s RigidBodySet on the active bodies,
+        // and update the components accordingly. That way, we don’t have to iterate through the entities that weren’t changed
+        // by physics (for example because they are sleeping).
+        if let Some(rb) = context.bodies.get(handle) {
+            let mut interpolated_pos = utils::iso_to_transform(rb.position());
+
+            if let TimestepMode::Interpolated { dt, .. } = *timestep_mode {
+                if let Some(interpolation) = interpolation.as_deref_mut() {
+                    if interpolation.end.is_none() {
+                        interpolation.end = Some(*rb.position());
                     }
 
-                    if let Some(mut transform) = transform {
-                        // NOTE: Rapier's `RigidBody` doesn't know its own scale as it is encoded
-                        //       directly within its collider, so we have to retrieve it from
-                        //       the scale of its bevy transform.
-                        interpolated_pos = interpolated_pos.with_scale(transform.scale);
-
-                        // NOTE: we query the parent’s global transform here, which is a bit
-                        //       unfortunate (performance-wise). An alternative would be to
-                        //       deduce the parent’s global transform from the current entity’s
-                        //       global transform. However, this makes it nearly impossible
-                        //       (because of rounding errors) to predict the exact next value this
-                        //       entity’s global transform will get after the next transform
-                        //       propagation, which breaks our transform modification detection
-                        //       that we do to detect if the user’s transform has to be written
-                        //       into the rigid-body.
-                        if let Some(parent_global_transform) =
-                            parent.and_then(|p| global_transforms.get(**p).ok())
-                        {
-                            // We need to compute the new local transform such that:
-                            // curr_parent_global_transform * new_transform = interpolated_pos
-                            // new_transform = curr_parent_global_transform.inverse() * interpolated_pos
-                            let (_, inverse_parent_rotation, inverse_parent_translation) =
-                                parent_global_transform
-                                    .affine()
-                                    .inverse()
-                                    .to_scale_rotation_translation();
-                            let new_rotation = inverse_parent_rotation * interpolated_pos.rotation;
-
-                            #[allow(unused_mut)] // mut is needed in 2D but not in 3D.
-                            let mut new_translation = inverse_parent_rotation
-                                * interpolated_pos.translation
-                                + inverse_parent_translation;
-
-                            // In 2D, preserve the transform `z` component that may have been set by the user
-                            #[cfg(feature = "dim2")]
-                            {
-                                new_translation.z = transform.translation.z;
-                            }
-
-                            if transform.rotation != new_rotation
-                                || transform.translation != new_translation
-                            {
-                                // NOTE: we write the new value only if there was an
-                                //       actual change, in order to not trigger bevy’s
-                                //       change tracking when the values didn’t change.
-                                transform.rotation = new_rotation;
-                                transform.translation = new_translation;
-                            }
-
-                            // NOTE: we need to compute the result of the next transform propagation
-                            //       to make sure that our change detection for transforms is exact
-                            //       despite rounding errors.
-                            let new_global_transform =
-                                parent_global_transform.mul_transform(*transform);
-
-                            context
-                                .last_body_transform_set
-                                .insert(handle, new_global_transform);
-                        } else {
-                            // In 2D, preserve the transform `z` component that may have been set by the user
-                            #[cfg(feature = "dim2")]
-                            {
-                                interpolated_pos.translation.z = transform.translation.z;
-                            }
-
-                            if transform.rotation != interpolated_pos.rotation
-                                || transform.translation != interpolated_pos.translation
-                            {
-                                // NOTE: we write the new value only if there was an
-                                //       actual change, in order to not trigger bevy’s
-                                //       change tracking when the values didn’t change.
-                                transform.rotation = interpolated_pos.rotation;
-                                transform.translation = interpolated_pos.translation;
-                            }
-
-                            context
-                                .last_body_transform_set
-                                .insert(handle, GlobalTransform::from(interpolated_pos));
-                        }
+                    if let Some(interpolated) =
+                        interpolation.lerp_slerp((dt + sim_to_render_time.diff) / dt)
+                    {
+                        interpolated_pos = utils::iso_to_transform(&interpolated);
                     }
+                }
+            }
 
-                    if let Some(velocity) = &mut velocity {
-                        let new_vel = Velocity {
-                            linvel: (*rb.linvel()).into(),
-                            #[cfg(feature = "dim3")]
-                            angvel: (*rb.angvel()).into(),
-                            #[cfg(feature = "dim2")]
-                            angvel: rb.angvel(),
-                        };
+            if let Some(mut transform) = transform {
+                // NOTE: Rapier's `RigidBody` doesn't know its own scale as it is encoded
+                //       directly within its collider, so we have to retrieve it from
+                //       the scale of its bevy transform.
+                interpolated_pos = interpolated_pos.with_scale(transform.scale);
+
+                // NOTE: we query the parent’s global transform here, which is a bit
+                //       unfortunate (performance-wise). An alternative would be to
+                //       deduce the parent’s global transform from the current entity’s
+                //       global transform. However, this makes it nearly impossible
+                //       (because of rounding errors) to predict the exact next value this
+                //       entity’s global transform will get after the next transform
+                //       propagation, which breaks our transform modification detection
+                //       that we do to detect if the user’s transform has to be written
+                //       into the rigid-body.
+                if let Some(parent_global_transform) =
+                    parent.and_then(|p| global_transforms.get(**p).ok())
+                {
+                    // We need to compute the new local transform such that:
+                    // curr_parent_global_transform * new_transform = interpolated_pos
+                    // new_transform = curr_parent_global_transform.inverse() * interpolated_pos
+                    let (_, inverse_parent_rotation, inverse_parent_translation) =
+                        parent_global_transform
+                            .affine()
+                            .inverse()
+                            .to_scale_rotation_translation();
+                    let new_rotation = inverse_parent_rotation * interpolated_pos.rotation;
+
+                    #[allow(unused_mut)] // mut is needed in 2D but not in 3D.
+                    let mut new_translation = inverse_parent_rotation
+                        * interpolated_pos.translation
+                        + inverse_parent_translation;
+
+                    // In 2D, preserve the transform `z` component that may have been set by the user
+                    #[cfg(feature = "dim2")]
+                    {
+                        new_translation.z = transform.translation.z;
+                    }
 
+                    if transform.rotation != new_rotation
+                        || transform.translation != new_translation
+                    {
                         // NOTE: we write the new value only if there was an
                         //       actual change, in order to not trigger bevy’s
                         //       change tracking when the values didn’t change.
-                        if **velocity != new_vel {
-                            **velocity = new_vel;
-                        }
+                        transform.rotation = new_rotation;
+                        transform.translation = new_translation;
                     }
 
-                    if let Some(sleeping) = &mut sleeping {
+                    // NOTE: we need to compute the result of the next transform propagation
+                    //       to make sure that our change detection for transforms is exact
+                    //       despite rounding errors.
+                    let new_global_transform = parent_global_transform.mul_transform(*transform);
+
+                    context
+                        .last_body_transform_set
+                        .insert(handle, new_global_transform);
+                } else {
+                    // In 2D, preserve the transform `z` component that may have been set by the user
+                    #[cfg(feature = "dim2")]
+                    {
+                        interpolated_pos.translation.z = transform.translation.z;
+                    }
+
+                    if transform.rotation != interpolated_pos.rotation
+                        || transform.translation != interpolated_pos.translation
+                    {
                         // NOTE: we write the new value only if there was an
                         //       actual change, in order to not trigger bevy’s
                         //       change tracking when the values didn’t change.
-                        if sleeping.sleeping != rb.is_sleeping() {
-                            sleeping.sleeping = rb.is_sleeping();
-                        }
+                        transform.rotation = interpolated_pos.rotation;
+                        transform.translation = interpolated_pos.translation;
                     }
+
+                    context
+                        .last_body_transform_set
+                        .insert(handle, GlobalTransform::from(interpolated_pos));
+                }
+            }
+
+            if let Some(velocity) = &mut velocity {
+                let new_vel = Velocity {
+                    linvel: (*rb.linvel()).into(),
+                    #[cfg(feature = "dim3")]
+                    angvel: (*rb.angvel()).into(),
+                    #[cfg(feature = "dim2")]
+                    angvel: rb.angvel(),
+                };
+
+                // NOTE: we write the new value only if there was an
+                //       actual change, in order to not trigger bevy’s
+                //       change tracking when the values didn’t change.
+                if **velocity != new_vel {
+                    **velocity = new_vel;
+                }
+            }
+
+            if let Some(sleeping) = &mut sleeping {
+                // NOTE: we write the new value only if there was an
+                //       actual change, in order to not trigger bevy’s
+                //       change tracking when the values didn’t change.
+                if sleeping.sleeping != rb.is_sleeping() {
+                    sleeping.sleeping = rb.is_sleeping();
                 }
             }
         }
@@ -533,13 +602,16 @@ pub fn init_rigid_bodies(
             activation.angular_threshold = sleep.angular_threshold;
         }
 
-        let context_entity =
-            try_retrieve_context(entity_context_link, &context).unwrap_or_else(|context_entity| {
+        let context_entity = entity_context_link.map_or_else(
+            || {
+                let context_entity = context.iter().next().unwrap().0;
                 commands
                     .entity(entity)
                     .insert(RapierContextEntityLink(context_entity));
                 context_entity
-            });
+            },
+            |link| link.0,
+        );
         let mut context = context.get_mut(context_entity).unwrap().1;
         let handle = context.bodies.insert(rb);
         commands
@@ -561,10 +633,13 @@ pub fn apply_initial_rigid_body_impulses(
     mut context: RapierContextAccessMut,
     // We can’t use RapierRigidBodyHandle yet because its creation command hasn’t been
     // executed yet.
-    mut init_impulses: Query<(Entity, &mut ExternalImpulse), Without<RapierRigidBodyHandle>>,
+    mut init_impulses: Query<
+        (Entity, &RapierContextEntityLink, &mut ExternalImpulse),
+        Without<RapierRigidBodyHandle>,
+    >,
 ) {
-    for (entity, mut impulse) in init_impulses.iter_mut() {
-        let context = context.context(entity);
+    for (entity, link, mut impulse) in init_impulses.iter_mut() {
+        let context = context.context(*link);
         let bodies = &mut context.bodies;
         if let Some(rb) = context
             .entity2body