Map Entity IDs When Writing Snapshot to World

src/lib.rs

@@ -7,9 +7,9 @@ use bevy::{
 };
 use ggrs::{Config, PlayerHandle};
 use ggrs_stage::GGRSStage;
+use parking_lot::RwLock;
 use reflect_resource::ReflectResource;
 use std::sync::Arc;
-use parking_lot::RwLock;
 
 pub use ggrs;
 
@@ -87,7 +87,20 @@ impl<T: Config + Send + Sync> Default for GGRSPlugin<T> {
             input_system: None,
             fps: DEFAULT_FPS,
             type_registry: TypeRegistry {
-                internal: Arc::new(RwLock::new(TypeRegistryInternal::empty())),
+                internal: Arc::new(RwLock::new({
+                    let mut r = TypeRegistryInternal::empty();
+                    // `Parent` and `Children` must be regisrered so that their `ReflectMapEntities`
+                    // data may be used.
+                    //
+                    // While this is a little bit of a weird spot to register these, are the only
+                    // Bevy core types implementing `MapEntities`, so for now it's probably fine to
+                    // just manually register these here.
+                    //
+                    // The user can still register any custom types with `register_rollback_type()`.
+                    r.register::<Parent>();
+                    r.register::<Children>();
+                    r
+                })),
             },
             schedule: Default::default(),
         }

src/world_snapshot.rs

@@ -1,4 +1,5 @@
 use bevy::{
+    ecs::{entity::EntityMap, reflect::ReflectMapEntities},
     prelude::*,
     reflect::{Reflect, TypeRegistry},
     utils::HashMap,
@@ -87,8 +88,7 @@ impl WorldSnapshot {
                         .filter(|&&entity| world.get::<Rollback>(entity).is_some())
                         .enumerate()
                     {
-                        if let Some(component) = reflect_component.reflect(world, *entity)
-                        {
+                        if let Some(component) = reflect_component.reflect(world, *entity) {
                             assert_eq!(*entity, snapshot.entities[entities_offset + i].entity);
                             // add the hash value of that component to the shapshot checksum, if that component supports hashing
                             if let Some(hash) = component.reflect_hash() {
@@ -132,6 +132,9 @@ impl WorldSnapshot {
         let type_registry = type_registry.read();
         let mut rid_map = rollback_id_map(world);
 
+        // Mapping of the old entity ids ( when snapshot was taken ) to new entity ids
+        let mut entity_map = EntityMap::default();
+
         // first, we write all entities
         for rollback_entity in self.entities.iter() {
             // find the corresponding current entity or create new entity, if it doesn't exist
@@ -146,6 +149,9 @@ impl WorldSnapshot {
                         .id()
                 });
 
+            // Add the mapping from the old entity ID to the new entity ID
+            entity_map.insert(rollback_entity.entity, entity);
+
             // for each registered type, check what we need to do
             for registration in type_registry.iter() {
                 let type_id = registration.type_id();
@@ -161,9 +167,7 @@ impl WorldSnapshot {
                         .find(|comp| comp.type_name() == registration.type_name())
                     {
                         // if we have data saved in the snapshot, overwrite the world
-                        Some(component) => {
-                            reflect_component.apply(world, entity, &**component)
-                        }
+                        Some(component) => reflect_component.apply(world, entity, &**component),
                         // if we don't have any data saved, we need to remove that component from the entity
                         None => reflect_component.remove(world, entity),
                     }
@@ -230,5 +234,17 @@ impl WorldSnapshot {
                 }
             }
         }
+
+        // For every type that reflects `MapEntities`, map the entities so that they reference the
+        // new IDs after applying the snapshot.
+        for registration in type_registry.iter() {
+            if let Some(map_entities_reflect) = registration.data::<ReflectMapEntities>() {
+                map_entities_reflect
+                    .map_entities(world, &entity_map)
+                    // This may fail if an entity is not found in the entity map, but that's fine,
+                    // because if it's not found in the map, then the entity may remain un-mapped.
+                    .ok();
+            }
+        }
     }
 }

tests/entity_maping.rs

@@ -0,0 +1,152 @@
+use bevy::prelude::*;
+
+use bevy_ggrs::*;
+use ggrs::*;
+use instant::Duration;
+
+pub struct GGRSConfig;
+impl Config for GGRSConfig {
+    type Input = u8;
+    type State = u8;
+    type Address = usize;
+}
+
+#[derive(Reflect, Component, Default)]
+struct ChildEntity;
+
+#[derive(Reflect, Component, Default)]
+struct ParentEntity;
+
+#[derive(Reflect, Component, Default, Debug)]
+struct FrameCounter(u16);
+
+fn input_system(_: In<PlayerHandle>, mut delete_events: EventReader<DeleteChildEntityEvent>) -> u8 {
+    u8::from(delete_events.iter().count() > 0)
+}
+
+fn setup_system(mut commands: Commands) {
+    commands
+        .spawn()
+        .insert(Rollback::new(0))
+        .insert(ParentEntity)
+        .with_children(|parent| {
+            parent.spawn().insert(Rollback::new(1)).insert(ChildEntity);
+        });
+}
+
+fn delete_child_system(
+    mut commands: Commands,
+    inputs: Res<Vec<(u8, InputStatus)>>,
+    parent: Query<&Children, With<ParentEntity>>,
+    child: Query<Entity, With<ChildEntity>>,
+) {
+    println!("Inputs: {:?}", *inputs);
+
+    println!("Parent's children: {:?}", parent.single());
+
+    if let Ok(child) = child.get_single() {
+        println!("Child exists: {child:?}");
+    }
+
+    if inputs[0].0 == 1 {
+        println!("Despawning child");
+        let child_entity = parent.single()[0];
+        commands.entity(child_entity).despawn();
+    }
+}
+
+fn frame_counter(mut counter: ResMut<FrameCounter>) {
+    println!("==== Frame {} ====", counter.0);
+    counter.0 = counter.0.wrapping_add(1);
+}
+
+struct DeleteChildEntityEvent;
+
+/// This test makes sure that we correctly map entities stored in resource and components during
+/// snapshot and restore.
+#[test]
+fn entity_mapping() {
+    let mut app = App::new();
+
+    app.add_plugins(MinimalPlugins)
+        .add_plugin(TransformPlugin)
+        .add_event::<DeleteChildEntityEvent>()
+        .init_resource::<FrameCounter>()
+        .add_startup_system(setup_system)
+        // Insert the GGRS session
+        .insert_resource(
+            SessionBuilder::<GGRSConfig>::new()
+                .with_num_players(1)
+                .with_check_distance(2)
+                .add_player(PlayerType::Local, 0)
+                .unwrap()
+                .start_synctest_session()
+                .unwrap(),
+        )
+        .insert_resource(SessionType::SyncTestSession);
+
+    GGRSPlugin::<GGRSConfig>::new()
+        .with_update_frequency(60)
+        .with_input_system(input_system)
+        .register_rollback_type::<ChildEntity>()
+        .register_rollback_type::<ParentEntity>()
+        .register_rollback_type::<FrameCounter>()
+        .with_rollback_schedule(
+            Schedule::default().with_stage(
+                "default",
+                SystemStage::single_threaded()
+                    .with_system(delete_child_system)
+                    .with_system(frame_counter.before(delete_child_system)),
+            ),
+        )
+        .build(&mut app);
+
+    // Sleep helper that will make sure at least one frame should be executed by the GGRS fixed
+    // update loop.
+    let sleep = || std::thread::sleep(Duration::from_secs_f32(1.0 / 60.0));
+
+    // Re-usable queries
+    let get_queries = |app: &mut App| {
+        (
+            app.world.query::<(&ChildEntity, &Parent)>(),
+            app.world.query::<(&ParentEntity, &Children)>(),
+        )
+    };
+
+    // Update once, the world should now be setup
+    app.update();
+    let (mut child_query, mut parent_query) = get_queries(&mut app);
+    assert!(
+        child_query.get_single(&app.world).is_ok(),
+        "Child doesn't exist"
+    );
+    assert!(
+        parent_query.get_single(&app.world).is_ok(),
+        "Parent doesn't exist"
+    );
+
+    sleep();
+    app.update();
+
+    // Send the event to delete the child entity
+    app.world
+        .resource_mut::<Events<DeleteChildEntityEvent>>()
+        .send(DeleteChildEntityEvent);
+
+    // Run for a number of times to make sure we get some rollbacks to happen
+    for _ in 0..5 {
+        sleep();
+        app.update();
+    }
+
+    // Make sure the child is delete and the parent still exists
+    let (mut child_query, mut parent_query) = get_queries(&mut app);
+    assert!(
+        child_query.get_single(&app.world).is_err(),
+        "Child exists after deletion"
+    );
+    assert!(
+        parent_query.get_single(&app.world).is_ok(),
+        "Parent doesn't exist"
+    );
+}