Support for Custom Checksum Methods

examples/stress_tests/particles.rs

@@ -3,7 +3,10 @@ use bevy_ggrs::{prelude::*, LocalInputs, LocalPlayers};
 use clap::Parser;
 use ggrs::{DesyncDetection, UdpNonBlockingSocket};
 use rand::{Rng, SeedableRng};
-use std::{hash::Hasher, net::SocketAddr};
+use std::{
+    hash::{BuildHasher, Hash, Hasher},
+    net::SocketAddr,
+};
 
 /// Stress test for bevy_ggrs
 ///
@@ -102,7 +105,10 @@ impl std::hash::Hash for Velocity {
     fn hash<H: Hasher>(&self, state: &mut H) {
         // We should have no NaNs or infinite values in our simulation
         // as they're not deterministic.
-        assert!(self.0.is_finite());
+        assert!(
+            self.0.is_finite(),
+            "Hashing is not stable for NaN f32 values."
+        );
 
         self.0.x.to_bits().hash(state);
         self.0.y.to_bits().hash(state);
@@ -190,10 +196,26 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             .rollback_resource_with_clone::<ParticleRng>();
     }
 
-    app.checksum_component_with_hash::<Velocity>()
-        // todo: ideally we'd also be doing checksums for Transforms, but that's
-        // currently very clunky to do.
-        .insert_resource(args)
+    app.insert_resource(args)
+        // Components can be added to the frame checksum automatically if they implement Hash...
+        .checksum_component_with_hash::<Velocity>()
+        // ...or you can provide a custom hashing process
+        .checksum_component::<Transform>(|transform| {
+            let mut hasher = bevy::utils::FixedState.build_hasher();
+
+            // In this demo we only translate particles, so only that value
+            // needs to be tracked.
+            assert!(
+                transform.translation.is_finite(),
+                "Hashing is not stable for NaN f32 values."
+            );
+
+            transform.translation.x.to_bits().hash(&mut hasher);
+            transform.translation.y.to_bits().hash(&mut hasher);
+            transform.translation.z.to_bits().hash(&mut hasher);
+
+            hasher.finish()
+        })
         .add_plugins(DefaultPlugins.set(WindowPlugin {
             primary_window: Some(Window {
                 resolution: WindowResolution::new(720., 720.),

src/lib.rs

@@ -272,6 +272,16 @@ pub trait GgrsApp {
     fn update_resource_with_map_entities<Type>(&mut self) -> &mut Self
     where
         Type: Resource + MapEntities;
+
+    /// Adds a component type to the checksum generation pipeline.
+    fn checksum_component<Type>(&mut self, hasher: for<'a> fn(&'a Type) -> u64) -> &mut Self
+    where
+        Type: Component;
+
+    /// Adds a resource type to the checksum generation pipeline.
+    fn checksum_resource<Type>(&mut self, hasher: for<'a> fn(&'a Type) -> u64) -> &mut Self
+    where
+        Type: Resource;
 }
 
 impl GgrsApp for App {
@@ -328,7 +338,7 @@ impl GgrsApp for App {
     where
         Type: Component + Hash,
     {
-        self.add_plugins(ComponentChecksumHashPlugin::<Type>::default())
+        self.add_plugins(ComponentChecksumPlugin::<Type>::default())
     }
 
     fn update_component_with_map_entities<Type>(&mut self) -> &mut Self
@@ -342,7 +352,7 @@ impl GgrsApp for App {
     where
         Type: Resource + Hash,
     {
-        self.add_plugins(ResourceChecksumHashPlugin::<Type>::default())
+        self.add_plugins(ResourceChecksumPlugin::<Type>::default())
     }
 
     fn update_resource_with_map_entities<Type>(&mut self) -> &mut Self
@@ -351,4 +361,18 @@ impl GgrsApp for App {
     {
         self.add_plugins(ResourceMapEntitiesPlugin::<Type>::default())
     }
+
+    fn checksum_component<Type>(&mut self, hasher: for<'a> fn(&'a Type) -> u64) -> &mut Self
+    where
+        Type: Component,
+    {
+        self.add_plugins(ComponentChecksumPlugin::<Type>(hasher))
+    }
+
+    fn checksum_resource<Type>(&mut self, hasher: for<'a> fn(&'a Type) -> u64) -> &mut Self
+    where
+        Type: Resource,
+    {
+        self.add_plugins(ResourceChecksumPlugin::<Type>(hasher))
+    }
 }

src/snapshot/component_checksum.rs

@@ -0,0 +1,102 @@
+use std::hash::{BuildHasher, Hash, Hasher};
+
+use bevy::prelude::*;
+
+use crate::{ChecksumFlag, ChecksumPart, Rollback, RollbackOrdered, SaveWorld, SaveWorldSet};
+
+/// A [`Plugin`] which will track the [`Component`] `C` on [`Rollback Entities`](`Rollback`) and ensure a
+/// [`ChecksumPart`] is available and updated. This can be used to generate a [`Checksum`](`crate::Checksum`).
+///
+/// # Examples
+/// ```rust
+/// # use bevy::prelude::*;
+/// # use bevy_ggrs::{prelude::*, ComponentChecksumPlugin};
+/// #
+/// # const FPS: usize = 60;
+/// #
+/// # type MyInputType = u8;
+/// #
+/// # fn read_local_inputs() {}
+/// #
+/// # fn start(session: Session<GgrsConfig<MyInputType>>) {
+/// # let mut app = App::new();
+/// #[derive(Component, Clone, Copy, Hash)]
+/// struct Health(u32);
+///
+/// // To include something in the checksum, it should also be rolled back
+/// app.rollback_component_with_clone::<Health>();
+///
+/// // This will update the checksum every frame to include Health on rollback entities
+/// app.add_plugins(ComponentChecksumPlugin::<Health>::default());
+/// # }
+/// ```
+pub struct ComponentChecksumPlugin<C: Component>(pub for<'a> fn(&'a C) -> u64);
+
+fn default_hasher<C: Component + Hash>(component: &C) -> u64 {
+    let mut hasher = bevy::utils::FixedState.build_hasher();
+    component.hash(&mut hasher);
+    hasher.finish()
+}
+
+impl<C> Default for ComponentChecksumPlugin<C>
+where
+    C: Component + Hash,
+{
+    fn default() -> Self {
+        Self(default_hasher::<C>)
+    }
+}
+
+impl<C> Plugin for ComponentChecksumPlugin<C>
+where
+    C: Component,
+{
+    fn build(&self, app: &mut App) {
+        let custom_hasher = self.0;
+
+        let update = move |mut commands: Commands,
+                           rollback_ordered: Res<RollbackOrdered>,
+                           components: Query<
+            (&Rollback, &C),
+            (With<Rollback>, Without<ChecksumFlag<C>>),
+        >,
+                           mut checksum: Query<
+            &mut ChecksumPart,
+            (Without<Rollback>, With<ChecksumFlag<C>>),
+        >| {
+            let mut hasher = bevy::utils::FixedState.build_hasher();
+
+            let mut result = 0;
+
+            for (&rollback, component) in components.iter() {
+                let mut hasher = hasher.clone();
+
+                // Hashing the rollback index ensures this hash is unique and stable
+                rollback_ordered.order(rollback).hash(&mut hasher);
+                custom_hasher(component).hash(&mut hasher);
+
+                // XOR chosen over addition or multiplication as it is closed on u64 and commutative
+                result ^= hasher.finish();
+            }
+
+            // Hash the XOR'ed result to break commutativity with other types
+            result.hash(&mut hasher);
+
+            let result = ChecksumPart(hasher.finish() as u128);
+
+            trace!(
+                "Component {} has checksum {:X}",
+                bevy::utils::get_short_name(std::any::type_name::<C>()),
+                result.0
+            );
+
+            if let Ok(mut checksum) = checksum.get_single_mut() {
+                *checksum = result;
+            } else {
+                commands.spawn((result, ChecksumFlag::<C>::default()));
+            }
+        };
+
+        app.add_systems(SaveWorld, update.in_set(SaveWorldSet::Checksum));
+    }
+}

src/snapshot/entity_checksum.rs

@@ -24,10 +24,7 @@ impl EntityChecksumPlugin {
 
         let result = ChecksumPart(hasher.finish() as u128);
 
-        trace!(
-            "Rollback Entities have checksum {:X}",
-            result.0
-        );
+        trace!("Rollback Entities have checksum {:X}", result.0);
 
         if let Ok(mut checksum) = checksum.get_single_mut() {
             *checksum = result;

src/snapshot/mod.rs

@@ -3,25 +3,25 @@ use bevy::{prelude::*, utils::HashMap};
 use std::{collections::VecDeque, marker::PhantomData};
 
 mod checksum;
-mod component_checksum_hash;
+mod component_checksum;
 mod component_map;
 mod component_snapshot;
 mod entity;
 mod entity_checksum;
-mod resource_checksum_hash;
+mod resource_checksum;
 mod resource_map;
 mod resource_snapshot;
 mod rollback_entity_map;
 mod set;
 mod strategy;
 
 pub use checksum::*;
-pub use component_checksum_hash::*;
+pub use component_checksum::*;
 pub use component_map::*;
 pub use component_snapshot::*;
 pub use entity::*;
 pub use entity_checksum::*;
-pub use resource_checksum_hash::*;
+pub use resource_checksum::*;
 pub use resource_map::*;
 pub use resource_snapshot::*;
 pub use rollback_entity_map::*;