Add the ability to override default command fns

src/core/command_markers.rs

@@ -36,6 +36,7 @@ pub trait AppMarkerExt {
     then these functions will be called for this component during replication
     instead of default [`write`](super::replication_fns::command_fns::write) and
     [`remove`](super::replication_fns::command_fns::remove).
+    See also [`Self::set_write_fn`].
 
     # Safety
 
@@ -122,6 +123,24 @@ pub trait AppMarkerExt {
         write: WriteFn,
         remove: RemoveFn,
     ) -> &mut Self;
+
+    /// Sets default functions for a component when there are no markers.
+    ///
+    /// If there are no markers are present on an entity, then these functions will
+    /// be called for this component during replication instead of default
+    /// [`write`](super::replication_fns::command_fns::write) and
+    /// [`remove`](super::replication_fns::command_fns::remove).
+    /// See also [`Self::set_marker_fns`].
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that passed `write` can be safely called with a
+    /// [`SerdeFns`](super::replication_fns::serde_fns::SerdeFns) created for `C`.
+    unsafe fn set_command_fns<C: Component>(
+        &mut self,
+        write: WriteFn,
+        remove: RemoveFn,
+    ) -> &mut Self;
 }
 
 impl AppMarkerExt for App {
@@ -159,6 +178,20 @@ impl AppMarkerExt for App {
 
         self
     }
+
+    unsafe fn set_command_fns<C: Component>(
+        &mut self,
+        write: WriteFn,
+        remove: RemoveFn,
+    ) -> &mut Self {
+        self.world
+            .resource_scope(|world, mut replication_fns: Mut<ReplicationFns>| unsafe {
+                // SAFETY: The caller ensured that `write` can be safely called with a `SerdeFns` created for `C`.
+                replication_fns.set_command_fns::<C>(world, write, remove);
+            });
+
+        self
+    }
 }
 
 /// Registered markers that override functions if present for

src/core/replication_fns.rs

@@ -77,6 +77,26 @@ impl ReplicationFns {
         command_fns.set_marker_fns(marker_id, write, remove);
     }
 
+    /// Sets default functions for a component when there are no markers.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that passed `write` can be safely called with all
+    /// [`SerdeFns`] registered for `C` with other methods on this struct.
+    pub(super) unsafe fn set_command_fns<C: Component>(
+        &mut self,
+        world: &mut World,
+        write: WriteFn,
+        remove: RemoveFn,
+    ) {
+        let (index, _) = self.init_command_fns::<C>(world);
+        let (command_fns, _) = &mut self.commands[index];
+
+        // SAFETY: `command_fns` was created for `C` and the caller ensured
+        // that `write` can be safely called with a `SerdeFns` created for `C`.
+        command_fns.set_fns(write, remove);
+    }
+
     /// Same as [`Self::register_serde_fns`], but uses default functions for a component.
     ///
     /// If your component contains any [`Entity`] inside, use [`Self::register_mapped_serde_fns`].

src/core/replication_fns/command_fns.rs

@@ -71,6 +71,17 @@ impl CommandFns {
         *fns = Some((write, remove));
     }
 
+    /// Sets default functions when there are no markers.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that passed `write` can be safely called with all
+    /// [`SerdeFns`] created for the same type as this instance.
+    pub(super) unsafe fn set_fns(&mut self, write: WriteFn, remove: RemoveFn) {
+        self.write = write;
+        self.remove = remove;
+    }
+
     /// Calls [`read`] on the type for which this instance was created.
     ///
     /// It's a non-overridable function that is used to restore the erased type from [`Ptr`].

tests/changes.rs

@@ -108,6 +108,69 @@ fn package_size_component() {
     assert_eq!(component.0, BIG_DATA);
 }
 
+#[test]
+fn command_fns() {
+    let mut server_app = App::new();
+    let mut client_app = App::new();
+    for app in [&mut server_app, &mut client_app] {
+        app.add_plugins((
+            MinimalPlugins,
+            RepliconPlugins.set(ServerPlugin {
+                tick_policy: TickPolicy::EveryFrame,
+                ..Default::default()
+            }),
+        ))
+        .replicate::<OriginalComponent>();
+
+        // SAFETY: `replace` can be safely called with a `SerdeFns` created for `OriginalComponent`.
+        unsafe {
+            app.set_command_fns::<OriginalComponent>(
+                replace,
+                command_fns::remove::<ReplacedComponent>,
+            );
+        }
+    }
+
+    server_app.connect_client(&mut client_app);
+
+    let server_entity = server_app
+        .world
+        .spawn((Replication, OriginalComponent(false)))
+        .id();
+
+    let client_entity = client_app
+        .world
+        .spawn((Replication, ReplacedComponent(false)))
+        .id();
+
+    client_app
+        .world
+        .resource_mut::<ServerEntityMap>()
+        .insert(server_entity, client_entity);
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+    server_app.exchange_with_client(&mut client_app);
+
+    // Change value.
+    let mut component = server_app
+        .world
+        .get_mut::<OriginalComponent>(server_entity)
+        .unwrap();
+    component.0 = true;
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+
+    let client_entity = client_app.world.entity(client_entity);
+    assert!(!client_entity.contains::<OriginalComponent>());
+
+    let component = client_entity.get::<ReplacedComponent>().unwrap();
+    assert!(component.0);
+}
+
 #[test]
 fn marker() {
     let mut server_app = App::new();

tests/insertion.rs

@@ -185,6 +185,53 @@ fn mapped_new_entity() {
     assert_eq!(client_app.world.entities().len(), 2);
 }
 
+#[test]
+fn command_fns() {
+    let mut server_app = App::new();
+    let mut client_app = App::new();
+    for app in [&mut server_app, &mut client_app] {
+        app.add_plugins((
+            MinimalPlugins,
+            RepliconPlugins.set(ServerPlugin {
+                tick_policy: TickPolicy::EveryFrame,
+                ..Default::default()
+            }),
+        ))
+        .replicate::<OriginalComponent>();
+
+        // SAFETY: `replace` can be safely called with a `SerdeFns` created for `OriginalComponent`.
+        unsafe {
+            app.set_command_fns::<OriginalComponent>(
+                replace,
+                command_fns::remove::<ReplacedComponent>,
+            );
+        }
+    }
+
+    server_app.connect_client(&mut client_app);
+
+    let server_entity = server_app.world.spawn(Replication).id();
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+    server_app.exchange_with_client(&mut client_app);
+
+    server_app
+        .world
+        .entity_mut(server_entity)
+        .insert(OriginalComponent);
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+
+    client_app
+        .world
+        .query_filtered::<(), (With<ReplacedComponent>, Without<OriginalComponent>)>()
+        .single(&client_app.world);
+}
+
 #[test]
 fn marker() {
     let mut server_app = App::new();

tests/removal.rs

@@ -48,6 +48,60 @@ fn single() {
     assert!(!client_entity.contains::<DummyComponent>());
 }
 
+#[test]
+fn command_fns() {
+    let mut server_app = App::new();
+    let mut client_app = App::new();
+    for app in [&mut server_app, &mut client_app] {
+        app.add_plugins((
+            MinimalPlugins,
+            RepliconPlugins.set(ServerPlugin {
+                tick_policy: TickPolicy::EveryFrame,
+                ..Default::default()
+            }),
+        ))
+        .replicate::<DummyComponent>();
+
+        // SAFETY: `replace` can be safely called with a `SerdeFns` created for `DummyComponent`.
+        unsafe {
+            app.set_command_fns::<DummyComponent>(
+                command_fns::write::<DummyComponent>,
+                command_fns::remove::<RemovingComponent>,
+            );
+        }
+    }
+
+    server_app.connect_client(&mut client_app);
+
+    let server_entity = server_app.world.spawn((Replication, DummyComponent)).id();
+    let client_entity = client_app
+        .world
+        .spawn((Replication, RemovingComponent))
+        .id();
+
+    client_app
+        .world
+        .resource_mut::<ServerEntityMap>()
+        .insert(server_entity, client_entity);
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+    server_app.exchange_with_client(&mut client_app);
+
+    server_app
+        .world
+        .entity_mut(server_entity)
+        .remove::<DummyComponent>();
+
+    server_app.update();
+    server_app.exchange_with_client(&mut client_app);
+    client_app.update();
+
+    let client_entity = client_app.world.entity(client_entity);
+    assert!(!client_entity.contains::<RemovingComponent>());
+}
+
 #[test]
 fn marker() {
     let mut server_app = App::new();