Add a RollbackEventHook Trait

src/ggrs_stage.rs

@@ -1,5 +1,5 @@
-use crate::{world_snapshot::WorldSnapshot, SessionType};
-use bevy::{prelude::*, reflect::TypeRegistry};
+use crate::{world_snapshot::WorldSnapshot, RollbackEventHook, SessionType};
+use bevy::{ecs::entity::EntityMap, prelude::*, reflect::TypeRegistry};
 use ggrs::{
     Config, GGRSError, GGRSRequest, GameStateCell, InputStatus, P2PSession, PlayerHandle,
     SessionState, SpectatorSession, SyncTestSession,
@@ -29,6 +29,7 @@ where
     accumulator: Duration,
     /// boolean to see if we should run slow to let remote clients catch up
     run_slow: bool,
+    hooks: Vec<Box<dyn RollbackEventHook>>,
 }
 
 impl<T: Config + Send + Sync> Stage for GGRSStage<T> {
@@ -81,6 +82,7 @@ impl<T: Config> GGRSStage<T> {
             last_update: Instant::now(),
             accumulator: Duration::ZERO,
             run_slow: false,
+            hooks: Vec::new(),
         }
     }
 
@@ -191,9 +193,33 @@ impl<T: Config> GGRSStage<T> {
     pub(crate) fn handle_requests(&mut self, requests: Vec<GGRSRequest<T>>, world: &mut World) {
         for request in requests {
             match request {
-                GGRSRequest::SaveGameState { cell, frame } => self.save_world(cell, frame, world),
-                GGRSRequest::LoadGameState { frame, .. } => self.load_world(frame, world),
-                GGRSRequest::AdvanceFrame { inputs } => self.advance_frame(inputs, world),
+                GGRSRequest::SaveGameState { cell, frame } => {
+                    for hook in &mut self.hooks {
+                        hook.pre_save(frame, self.snapshots.len(), world);
+                    }
+                    self.save_world(cell, frame, world);
+                    for hook in &mut self.hooks {
+                        hook.post_save(frame, self.snapshots.len(), world);
+                    }
+                }
+                GGRSRequest::LoadGameState { frame, .. } => {
+                    for hook in &mut self.hooks {
+                        hook.pre_load(frame, self.snapshots.len(), world);
+                    }
+                    let entity_map = self.load_world(frame, world);
+                    for hook in &mut self.hooks {
+                        hook.post_load(frame, self.snapshots.len(), &entity_map, world);
+                    }
+                }
+                GGRSRequest::AdvanceFrame { inputs } => {
+                    for hook in &mut self.hooks {
+                        hook.pre_advance(world);
+                    }
+                    self.advance_frame(inputs, world);
+                    for hook in &mut self.hooks {
+                        hook.post_advance(world);
+                    }
+                }
             }
         }
     }
@@ -217,15 +243,15 @@ impl<T: Config> GGRSStage<T> {
         self.snapshots[pos] = snapshot;
     }
 
-    pub(crate) fn load_world(&mut self, frame: i32, world: &mut World) {
+    pub(crate) fn load_world(&mut self, frame: i32, world: &mut World) -> EntityMap {
         self.frame = frame;
 
         // we get the correct snapshot
         let pos = frame as usize % self.snapshots.len();
         let snapshot_to_load = &self.snapshots[pos];
 
         // load the entities
-        snapshot_to_load.write_to_world(world, &self.type_registry);
+        snapshot_to_load.write_to_world(world, &self.type_registry)
     }
 
     pub(crate) fn advance_frame(
@@ -247,6 +273,10 @@ impl<T: Config> GGRSStage<T> {
         self.schedule = schedule;
     }
 
+    pub(crate) fn set_hooks(&mut self, hooks: Vec<Box<dyn RollbackEventHook>>) {
+        self.hooks = hooks;
+    }
+
     pub(crate) fn set_type_registry(&mut self, type_registry: TypeRegistry) {
         self.type_registry = type_registry;
     }

src/lib.rs

@@ -2,10 +2,11 @@
 #![forbid(unsafe_code)] // let us try
 
 use bevy::{
+    ecs::entity::EntityMap,
     prelude::*,
     reflect::{FromType, GetTypeRegistration, TypeRegistry, TypeRegistryInternal},
 };
-use ggrs::{Config, PlayerHandle};
+use ggrs::{Config, Frame, PlayerHandle};
 use ggrs_stage::GGRSStage;
 use parking_lot::RwLock;
 use reflect_resource::ReflectResource;
@@ -73,12 +74,47 @@ impl RollbackIdProvider {
     }
 }
 
+/// An object that may hook into rollback events.
+pub trait RollbackEventHook: Sync + Send + 'static {
+    /// Run before a snapshot is saved
+    fn pre_save(&mut self, frame: Frame, max_snapshots: usize, world: &mut World) {
+        let _ = (frame, max_snapshots, world);
+    }
+    /// Run after a snapshot is saved
+    fn post_save(&mut self, frame: Frame, max_snapshots: usize, world: &mut World) {
+        let _ = (frame, max_snapshots, world);
+    }
+    /// Run before a snapshot is loaded ( restored )
+    fn pre_load(&mut self, frame: Frame, max_snapshots: usize, world: &mut World) {
+        let _ = (frame, max_snapshots, world);
+    }
+    /// Run after a snapshot is loaded ( restored )
+    fn post_load(
+        &mut self,
+        frame: Frame,
+        max_snapshots: usize,
+        entity_map: &EntityMap,
+        world: &mut World,
+    ) {
+        let _ = (frame, max_snapshots, entity_map, world);
+    }
+    /// Run before the game simulation is advanced one frame
+    fn pre_advance(&mut self, world: &mut World) {
+        let _ = world;
+    }
+    /// Run after the game simulation is advanced one frame
+    fn post_advance(&mut self, world: &mut World) {
+        let _ = world;
+    }
+}
+
 /// A builder to configure GGRS for a bevy app.
 pub struct GGRSPlugin<T: Config + Send + Sync> {
     input_system: Option<Box<dyn System<In = PlayerHandle, Out = T::Input>>>,
     fps: usize,
     type_registry: TypeRegistry,
     schedule: Schedule,
+    hooks: Vec<Box<dyn RollbackEventHook>>,
 }
 
 impl<T: Config + Send + Sync> Default for GGRSPlugin<T> {
@@ -103,6 +139,7 @@ impl<T: Config + Send + Sync> Default for GGRSPlugin<T> {
                 })),
             },
             schedule: Default::default(),
+            hooks: Default::default(),
         }
     }
 }
@@ -150,6 +187,13 @@ impl<T: Config + Send + Sync> GGRSPlugin<T> {
         self
     }
 
+    /// Add a [`RollbackEventHook`] object will have the opportunity to modify the world or take
+    /// other actions during snapshot saves, loads, and frame advances.q
+    pub fn add_rollback_hook<H: RollbackEventHook>(mut self, hook: H) -> Self {
+        self.hooks.push(Box::new(hook));
+        self
+    }
+
     /// Consumes the builder and makes changes on the bevy app according to the settings.
     pub fn build(self, app: &mut App) {
         let mut input_system = self
@@ -161,6 +205,7 @@ impl<T: Config + Send + Sync> GGRSPlugin<T> {
         stage.set_update_frequency(self.fps);
         stage.set_schedule(self.schedule);
         stage.set_type_registry(self.type_registry);
+        stage.set_hooks(self.hooks);
         app.add_stage_before(CoreStage::Update, GGRS_UPDATE, stage);
         // other resources
         app.insert_resource(RollbackIdProvider::default());

src/world_snapshot.rs

@@ -128,7 +128,11 @@ impl WorldSnapshot {
         snapshot
     }
 
-    pub(crate) fn write_to_world(&self, world: &mut World, type_registry: &TypeRegistry) {
+    pub(crate) fn write_to_world(
+        &self,
+        world: &mut World,
+        type_registry: &TypeRegistry,
+    ) -> EntityMap {
         let type_registry = type_registry.read();
         let mut rid_map = rollback_id_map(world);
 
@@ -246,5 +250,7 @@ impl WorldSnapshot {
                     .ok();
             }
         }
+
+        entity_map
     }
 }