From 0a841ba9c1ea6350e6d640bc91788bdc9aeb530c Mon Sep 17 00:00:00 2001
From: Ben Frankel <ben.frankel7@gmail.com>
Date: Mon, 3 Mar 2025 23:31:10 -0800
Subject: [PATCH] Cache systems by `S` instead of `S::System` (#16694)

# Objective

- Fixes the issue described in this comment:
https://github.com/bevyengine/bevy/issues/16680#issuecomment-2522764239.

## Solution

- Cache one-shot systems by `S: IntoSystem` (which is const-asserted to
be a ZST) rather than `S::System`.

## Testing

Added a new unit test named `cached_system_into_same_system_type` to
`system_registry.rs`.

---

## Migration Guide

The `CachedSystemId` resource has been changed:

```rust
// Before:
let cached_id = CachedSystemId::<S::System>(id);
assert!(id == cached_id.0);

// After:
let cached_id = CachedSystemId::<S>::new(id);
assert!(id == SystemId::from_entity(cached_id.entity));
```
---
 crates/bevy_ecs/src/system/system_registry.rs | 73 ++++++++++++++++---
 1 file changed, 62 insertions(+), 11 deletions(-)

diff --git a/crates/bevy_ecs/src/system/system_registry.rs b/crates/bevy_ecs/src/system/system_registry.rs
index 2d849dda95..cf42d0d875 100644
--- a/crates/bevy_ecs/src/system/system_registry.rs
+++ b/crates/bevy_ecs/src/system/system_registry.rs
@@ -3,7 +3,7 @@ use crate::reflect::ReflectComponent;
 use crate::{
     change_detection::Mut,
     entity::Entity,
-    system::{input::SystemInput, BoxedSystem, IntoSystem, System},
+    system::{input::SystemInput, BoxedSystem, IntoSystem},
     world::World,
 };
 use alloc::boxed::Box;
@@ -126,7 +126,21 @@ impl<I: SystemInput, O> core::fmt::Debug for SystemId<I, O> {
 ///
 /// This resource is inserted by [`World::register_system_cached`].
 #[derive(Resource)]
-pub struct CachedSystemId<S: System>(pub SystemId<S::In, S::Out>);
+pub struct CachedSystemId<S> {
+    /// The cached `SystemId` as an `Entity`.
+    pub entity: Entity,
+    _marker: PhantomData<fn() -> S>,
+}
+
+impl<S> CachedSystemId<S> {
+    /// Creates a new `CachedSystemId` struct given a `SystemId`.
+    pub fn new<I: SystemInput, O>(id: SystemId<I, O>) -> Self {
+        Self {
+            entity: id.entity(),
+            _marker: PhantomData,
+        }
+    }
+}
 
 impl World {
     /// Registers a system and returns a [`SystemId`] so it can later be called by [`World::run_system`].
@@ -154,7 +168,7 @@ impl World {
     /// Similar to [`Self::register_system`], but allows passing in a [`BoxedSystem`].
     ///
     ///  This is useful if the [`IntoSystem`] implementor has already been turned into a
-    /// [`System`] trait object and put in a [`Box`].
+    /// [`System`](crate::system::System) trait object and put in a [`Box`].
     pub fn register_boxed_system<I, O>(&mut self, system: BoxedSystem<I, O>) -> SystemId<I, O>
     where
         I: SystemInput + 'static,
@@ -389,23 +403,23 @@ impl World {
             );
         }
 
-        if !self.contains_resource::<CachedSystemId<S::System>>() {
+        if !self.contains_resource::<CachedSystemId<S>>() {
             let id = self.register_system(system);
-            self.insert_resource(CachedSystemId::<S::System>(id));
+            self.insert_resource(CachedSystemId::<S>::new(id));
             return id;
         }
 
-        self.resource_scope(|world, mut id: Mut<CachedSystemId<S::System>>| {
-            if let Ok(mut entity) = world.get_entity_mut(id.0.entity()) {
+        self.resource_scope(|world, mut id: Mut<CachedSystemId<S>>| {
+            if let Ok(mut entity) = world.get_entity_mut(id.entity) {
                 if !entity.contains::<RegisteredSystem<I, O>>() {
                     entity.insert(RegisteredSystem::new(Box::new(IntoSystem::into_system(
                         system,
                     ))));
                 }
             } else {
-                id.0 = world.register_system(system);
+                id.entity = world.register_system(system).entity();
             }
-            id.0
+            SystemId::from_entity(id.entity)
         })
     }
 
@@ -422,9 +436,9 @@ impl World {
         S: IntoSystem<I, O, M> + 'static,
     {
         let id = self
-            .remove_resource::<CachedSystemId<S::System>>()
+            .remove_resource::<CachedSystemId<S>>()
             .ok_or(RegisteredSystemError::SystemNotCached)?;
-        self.unregister_system(id.0)
+        self.unregister_system(SystemId::<I, O>::from_entity(id.entity))
     }
 
     /// Runs a cached system, registering it if necessary.
@@ -754,6 +768,43 @@ mod tests {
         assert!(matches!(output, Ok(8)));
     }
 
+    #[test]
+    fn cached_system_into_same_system_type() {
+        use crate::result::Result;
+
+        struct Foo;
+        impl IntoSystem<(), Result<()>, ()> for Foo {
+            type System = ApplyDeferred;
+            fn into_system(_: Self) -> Self::System {
+                ApplyDeferred
+            }
+        }
+
+        struct Bar;
+        impl IntoSystem<(), Result<()>, ()> for Bar {
+            type System = ApplyDeferred;
+            fn into_system(_: Self) -> Self::System {
+                ApplyDeferred
+            }
+        }
+
+        let mut world = World::new();
+        let foo1 = world.register_system_cached(Foo);
+        let foo2 = world.register_system_cached(Foo);
+        let bar1 = world.register_system_cached(Bar);
+        let bar2 = world.register_system_cached(Bar);
+
+        // The `S: IntoSystem` types are different, so they should be cached
+        // as separate systems, even though the `<S as IntoSystem>::System`
+        // types / values are the same (`ApplyDeferred`).
+        assert_ne!(foo1, bar1);
+
+        // But if the `S: IntoSystem` types are the same, they'll be cached
+        // as the same system.
+        assert_eq!(foo1, foo2);
+        assert_eq!(bar1, bar2);
+    }
+
     #[test]
     fn system_with_input_ref() {
         fn with_ref(InRef(input): InRef<u8>, mut counter: ResMut<Counter>) {