Add helper functions to map from local world to remote world

lightyear/src/client/connection.rs

@@ -1,5 +1,6 @@
 //! Specify how a Client sends/receives messages with a Server
 use bevy::ecs::component::Tick as BevyTick;
+use bevy::ecs::entity::MapEntities;
 use bevy::prelude::{Mut, Resource, World};
 use bevy::utils::{Duration, HashMap};
 use bytes::Bytes;
@@ -33,6 +34,7 @@ use crate::shared::message::MessageSend;
 use crate::shared::ping::manager::{PingConfig, PingManager};
 use crate::shared::ping::message::{Ping, Pong};
 use crate::shared::replication::delta::DeltaManager;
+use crate::shared::replication::entity_map::EntityMap;
 use crate::shared::replication::network_target::NetworkTarget;
 use crate::shared::replication::receive::ReplicationReceiver;
 use crate::shared::replication::send::ReplicationSender;
@@ -226,6 +228,18 @@ impl ConnectionManager {
         Ok(())
     }
 
+    // TODO: we need `&mut self` because MapEntities requires `&mut EntityMapper` even though it's not needed here
+    /// Convert entities in the message to be compatible with the remote world
+    pub fn map_entities_to_remote<M: Message + MapEntities>(&mut self, message: &mut M) {
+        let mapper = &mut self.replication_receiver.remote_entity_map.local_to_remote;
+        message.map_entities(mapper);
+    }
+
+    /// Map from the local entities to the remote entities
+    pub fn local_to_remote_map(&mut self) -> &mut EntityMap {
+        &mut self.replication_receiver.remote_entity_map.local_to_remote
+    }
+
     /// Send a [`Message`] to the server using a specific [`Channel`]
     pub fn send_message<C: Channel, M: Message>(&mut self, message: &M) -> Result<(), ClientError> {
         self.send_message_to_target::<C, M>(message, NetworkTarget::None)
@@ -574,3 +588,57 @@ impl ReplicationSend for ConnectionManager {
         self.delta_manager.tick_cleanup(tick);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::prelude::{client, server, ClientConnectionManager};
+    use crate::tests::protocol::Message2;
+    use crate::tests::stepper::{BevyStepper, Step};
+
+    /// Check that we can map entities from the local world to the remote world
+    /// using the ConnectionManager
+    #[test]
+    fn test_map_entities_to_remote() {
+        let mut stepper = BevyStepper::default();
+
+        // spawn an entity on server
+        let server_entity = stepper
+            .server_app
+            .world_mut()
+            .spawn(server::Replicate::default())
+            .id();
+        stepper.frame_step();
+        stepper.frame_step();
+
+        // check that the entity was spawned
+        let client_entity = *stepper
+            .client_app
+            .world()
+            .resource::<client::ConnectionManager>()
+            .replication_receiver
+            .remote_entity_map
+            .get_local(server_entity)
+            .expect("entity was not replicated to client");
+
+        // spawn an entity on the client which contains a link to another entity
+        // we need to map that entity to the remote world
+        assert_eq!(
+            *stepper
+                .client_app
+                .world_mut()
+                .resource_mut::<ClientConnectionManager>()
+                .local_to_remote_map()
+                .get(&client_entity)
+                .unwrap(),
+            server_entity
+        );
+
+        let mut message = Message2(client_entity);
+        stepper
+            .client_app
+            .world_mut()
+            .resource_mut::<ClientConnectionManager>()
+            .map_entities_to_remote(&mut message);
+        assert_eq!(message.0, server_entity);
+    }
+}

lightyear/src/server/connection.rs

@@ -1,6 +1,6 @@
 //! Specify how a Server sends/receives messages with a Client
 use bevy::ecs::component::Tick as BevyTick;
-use bevy::ecs::entity::EntityHash;
+use bevy::ecs::entity::{EntityHash, MapEntities};
 use bevy::prelude::{Component, Entity, Mut, Resource, World};
 use bevy::ptr::Ptr;
 use bevy::utils::{Duration, HashMap};
@@ -45,6 +45,7 @@ use crate::shared::ping::manager::{PingConfig, PingManager};
 use crate::shared::ping::message::{Ping, Pong};
 use crate::shared::replication::components::ReplicationGroupId;
 use crate::shared::replication::delta::DeltaManager;
+use crate::shared::replication::entity_map::EntityMap;
 use crate::shared::replication::network_target::NetworkTarget;
 use crate::shared::replication::receive::ReplicationReceiver;
 use crate::shared::replication::send::ReplicationSender;
@@ -121,6 +122,22 @@ impl ConnectionManager {
         self.connections.keys().copied()
     }
 
+    // TODO: we need `&mut self` because MapEntities requires `&mut EntityMapper` even though it's not needed here
+    /// Convert entities in the message to be compatible with the remote world of the provided client
+    pub fn map_entities_to_remote<M: Message + MapEntities>(
+        &mut self,
+        message: &mut M,
+        client_id: ClientId,
+    ) -> Result<(), ServerError> {
+        let mapper = &mut self
+            .connection_mut(client_id)?
+            .replication_receiver
+            .remote_entity_map
+            .local_to_remote;
+        message.map_entities(mapper);
+        Ok(())
+    }
+
     /// Queues up a message to be sent to all clients matching the specific [`NetworkTarget`]
     pub fn send_message_to_target<C: Channel, M: Message>(
         &mut self,
@@ -487,6 +504,11 @@ impl Connection {
         self.is_local_client
     }
 
+    /// Map from the local entities to the remote entities
+    pub fn local_to_remote_map(&mut self) -> &mut EntityMap {
+        &mut self.replication_receiver.remote_entity_map.local_to_remote
+    }
+
     /// Return the latest estimate of rtt
     pub fn rtt(&self) -> Duration {
         self.ping_manager.rtt()

lightyear/src/tests/protocol.rs

@@ -22,7 +22,13 @@ use crate::shared::replication::delta::Diffable;
 pub struct Message1(pub String);
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Clone, Reflect)]
-pub struct Message2(pub u32);
+pub struct Message2(pub Entity);
+
+impl MapEntities for Message2 {
+    fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
+        self.0 = entity_mapper.map_entity(self.0);
+    }
+}
 
 // Components
 #[derive(Component, Serialize, Deserialize, Clone, Debug, PartialEq, Reflect)]
@@ -194,7 +200,8 @@ impl Plugin for ProtocolPlugin {
     fn build(&self, app: &mut App) {
         // messages
         app.register_message::<Message1>(ChannelDirection::Bidirectional);
-        app.register_message::<Message2>(ChannelDirection::Bidirectional);
+        app.register_message::<Message2>(ChannelDirection::Bidirectional)
+            .add_map_entities();
         // inputs
         app.add_plugins(InputPlugin::<MyInput>::default());
         // components