WIP Despawn using a disabling component

src/snapshot/despawn.rs

@@ -0,0 +1,182 @@
+//! Rollback entity deferred despawning module.
+//!
+//! This module allows rollback entities to have non-rollback components (such as static collision
+//! properties or mesh handles) by deferring the actual despawn until the requested frame has been
+//! confirmed. Despawned entities are instead marked with the disabling component
+//! [`RollbackDespawned`] so that they behave as though they are despawned by default.
+//!
+//! # Examples
+//! ```rust
+//! # use bevy::prelude::*;
+//! # use bevy_ggrs::{prelude::*, ResourceChecksumPlugin};
+//! #
+//! # const FPS: usize = 60;
+//! #
+//! # type MyInputType = u8;
+//! #
+//! # fn read_local_inputs() {}
+//! #
+//! # fn start(session: Session<GgrsConfig<MyInputType>>) {
+//! # let mut app = App::new();
+//! #[derive(Resource, Clone, Hash)]
+//! struct BossHealth(u32);
+//!
+//! // To include something in the checksum, it should also be rolled back
+//! app.rollback_resource_with_clone::<BossHealth>();
+//!
+//! // This will update the checksum every frame to include BossHealth
+//! app.add_plugins(ResourceChecksumPlugin::<BossHealth>::default());
+//! # }
+//! ```
+//!
+//! Entities which have been marked for despawn are disabled using the [`RollbackDespawned`]
+//! component, so they will appear as though they are despawned to normal system queries.
+
+use crate::snapshot::despawn::private::RollbackDespawnCommandExtensionSeal;
+use crate::{
+    AdvanceWorld, AdvanceWorldSystems, ConfirmedFrameCount, LoadWorld, LoadWorldSystems,
+    RollbackFrameCount,
+};
+use bevy::app::{App, Plugin};
+use bevy::prelude::*;
+use ggrs::Frame;
+use std::cmp::Ordering;
+
+/// Marks an entity as despawned, contains the frame that the entity was despawned on.
+///
+/// When an entity is marked with this component, they MUST not be allowed to affect the simulation
+/// in any way, as they may eventually be despawned on different frames for different peers.
+/// This component is registered as "disabling" (see [ecs module docs](bevy_ecs::entity_disabling))
+/// so that the entity will not show up in queries by default.
+#[derive(Component, Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct RollbackDespawned(Frame);
+
+pub struct RollbackDespawnPlugin;
+
+impl Plugin for RollbackDespawnPlugin {
+    fn build(&self, app: &mut App) {
+        app.register_disabling_component::<RollbackDespawned>();
+
+        app.add_systems(
+            LoadWorld,
+            resurrect_entities.in_set(LoadWorldSystems::EntityResurrect),
+        )
+        .add_systems(
+            AdvanceWorld,
+            despawn_confirmed_entities.in_set(AdvanceWorldSystems::DespawnConfirmed),
+        );
+    }
+}
+
+fn resurrect_entities(
+    world: &mut World,
+    despawn_query: &mut QueryState<(Entity, &RollbackDespawned)>,
+) {
+    let rollback_frame = world.resource::<RollbackFrameCount>();
+
+    despawn_query
+        .iter(world)
+        .filter_map(|(entity, despawned_frame)| {
+            // During a world load, entities marked with the current rollback frame were marked
+            // for despawn during that frame, so we only want to resurrect entities despawned after.
+            Some(entity).filter(|_e| despawned_frame > rollback_frame)
+        })
+        .collect::<Vec<_>>()
+        .into_iter()
+        .for_each(|entity| {
+            world.entity_mut(entity).remove::<RollbackDespawned>();
+        });
+}
+
+fn despawn_confirmed_entities(
+    world: &mut World,
+    despawn_query: &mut QueryState<(Entity, &RollbackDespawned)>,
+    mut local: Local<ConfirmedFrameCount>,
+) {
+    let confirmed_frame = world.resource::<ConfirmedFrameCount>();
+    if *confirmed_frame == *local {
+        return; // No work necessary
+    }
+    *local = *confirmed_frame;
+
+    despawn_query
+        .iter(world)
+        .filter_map(|(entity, despawned_frame)| {
+            // Entities marked as despawned on the confirmed frame or earlier can be immediately
+            // despawned.
+            Some(entity).filter(|_e| despawned_frame <= confirmed_frame)
+        })
+        .collect::<Vec<_>>()
+        .into_iter()
+        .for_each(|entity| {
+            world.despawn(entity);
+        });
+}
+
+mod private {
+    pub trait RollbackDespawnCommandExtensionSeal {}
+}
+pub trait RollbackDespawnCommandExtension: private::RollbackDespawnCommandExtensionSeal {
+    /// Despawns this entity and its children recursively using the [`RollbackDespawned`]
+    /// component, such that they can be resurrected following a rollback.
+    ///
+    /// NOTE: This does not yet support [`RelationshipTarget`] with linked spawn mode.
+    fn despawn_rollback(&mut self);
+
+    /// NOTE: Not implemented yet.
+    fn despawn_children_rollback(&mut self) -> &mut Self;
+
+    /// NOTE: Not implemented yet.
+    fn despawn_related_rollback<S>(&mut self) -> &mut Self;
+}
+
+impl RollbackDespawnCommandExtensionSeal for EntityCommands<'_> {}
+
+impl RollbackDespawnCommandExtension for EntityCommands<'_> {
+    fn despawn_rollback(&mut self) {
+        self.queue_silenced(despawn_rollback);
+    }
+
+    fn despawn_children_rollback(&mut self) -> &mut Self {
+        todo!()
+    }
+
+    fn despawn_related_rollback<S>(&mut self) -> &mut Self {
+        todo!()
+    }
+}
+
+fn despawn_rollback(mut entity: EntityWorldMut) {
+    if let Some(&RollbackFrameCount(frame)) = entity.get_resource::<RollbackFrameCount>() {
+        // If we have RollbackFrameCount we should also have ConfirmedFrameCount
+        let &ConfirmedFrameCount(confirmed) = entity.get_resource::<ConfirmedFrameCount>().unwrap();
+
+        // TODO handle wraparound
+        if confirmed < frame {
+            entity.insert_recursive::<Children>(RollbackDespawned(frame));
+            return;
+        }
+    }
+
+    // If current frame is confirmed or rollback sim is not present, we can simply despawn
+    entity.despawn();
+}
+
+macro_rules! newtype_partial_ord {
+    ($i:ident, $j:ident) => {
+        impl PartialEq<$j> for $i {
+            fn eq(&self, other: &$j) -> bool {
+                self.0 == other.0
+            }
+        }
+
+        impl PartialOrd<$j> for $i {
+            fn partial_cmp(&self, other: &$j) -> Option<Ordering> {
+                Some(self.0.cmp(&other.0))
+            }
+        }
+    };
+}
+
+newtype_partial_ord!(RollbackDespawned, RollbackFrameCount);
+newtype_partial_ord!(RollbackDespawned, ConfirmedFrameCount);

src/snapshot/mod.rs

@@ -8,6 +8,7 @@ mod childof_snapshot;
 mod component_checksum;
 mod component_map;
 mod component_snapshot;
+mod despawn;
 mod entity;
 mod entity_checksum;
 mod resource_checksum;
@@ -24,6 +25,7 @@ pub use childof_snapshot::*;
 pub use component_checksum::*;
 pub use component_map::*;
 pub use component_snapshot::*;
+pub use despawn::*;
 pub use entity::*;
 pub use entity_checksum::*;
 pub use resource_checksum::*;
@@ -36,6 +38,7 @@ pub use set::*;
 pub use strategy::*;
 
 pub mod prelude {
+    pub use super::despawn::{RollbackDespawnCommandExtension, RollbackDespawned};
     pub use super::{Checksum, LoadWorldSystems, SaveWorldSystems};
 }
 
@@ -304,6 +307,7 @@ impl Plugin for SnapshotPlugin {
                 EntitySnapshotPlugin,
                 ResourceSnapshotPlugin::<CloneStrategy<RollbackOrdered>>::default(),
                 ChildOfSnapshotPlugin,
+                RollbackDespawnPlugin,
             ));
     }
 }

src/snapshot/set.rs

@@ -7,6 +7,9 @@ use crate::snapshot::{AdvanceWorld, LoadWorld, SaveWorld};
 /// and [`Resource`] snapshots are loaded and applied to the [`World`].
 #[derive(SystemSet, Hash, Debug, PartialEq, Eq, Clone)]
 pub enum LoadWorldSystems {
+    /// Removes any despawn markers if the loaded frame is before they were marked.
+    /// See [`despawn module docs`](`crate::snapshot::despawn`).
+    EntityResurrect,
     /// Recreate the [`Entity`] graph as it was during the frame to be rolled back to.
     /// When this set is complete, all entities that were alive during the snapshot
     /// frame have been recreated, and any that were not have been removed. If the
@@ -19,7 +22,7 @@ pub enum LoadWorldSystems {
     /// When this set is complete, all [`Components`](`Component`) and [`Resources`](`Resource`)
     /// will be rolled back to their exact state during the snapshot.
     ///
-    /// NOTE: At this point, [`Entity`] relationships may be broken, see [`LoadWorldSet::Mapping`]
+    /// NOTE: At this point, [`Entity`] relationships may be broken, see [`LoadWorldSystems::Mapping`]
     /// for when those relationships are fixed.
     Data,
     /// Flush any deferred operations
@@ -49,6 +52,9 @@ pub enum SaveWorldSystems {
 #[derive(SystemSet, Hash, Debug, PartialEq, Eq, Clone)]
 pub enum AdvanceWorldSystems {
     First,
+    /// Despawn any entities if their marked frame has been confirmed.
+    /// See [`despawn module docs`](`crate::snapshot::despawn`).
+    DespawnConfirmed,
     Main,
     Last,
 }
@@ -62,6 +68,7 @@ impl Plugin for SnapshotSetPlugin {
         app.configure_sets(
             LoadWorld,
             (
+                LoadWorldSystems::EntityResurrect,
                 LoadWorldSystems::Entity,
                 LoadWorldSystems::EntityFlush,
                 LoadWorldSystems::Data,
@@ -78,6 +85,7 @@ impl Plugin for SnapshotSetPlugin {
             AdvanceWorld,
             (
                 AdvanceWorldSystems::First,
+                AdvanceWorldSystems::DespawnConfirmed,
                 AdvanceWorldSystems::Main,
                 AdvanceWorldSystems::Last,
             )