Improve error for when you try running a one-shot system with wrong SystemId type

crates/bevy_ecs/src/system/system_registry.rs

@@ -2,7 +2,7 @@
 use crate::reflect::ReflectComponent;
 use crate::{
     change_detection::Mut,
-    entity::Entity,
+    entity::{Entity, EntityHashSet},
     system::{input::SystemInput, BoxedSystem, IntoSystem, SystemParamValidationError},
     world::World,
 };
@@ -11,6 +11,7 @@ use bevy_ecs_macros::{Component, Resource};
 #[cfg(feature = "bevy_reflect")]
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use core::marker::PhantomData;
+use derive_more::derive::{Deref, DerefMut};
 use thiserror::Error;
 
 /// A small wrapper for [`BoxedSystem`] that also keeps track whether or not the system has been initialized.
@@ -36,6 +37,10 @@ impl<I, O> RegisteredSystem<I, O> {
 #[cfg_attr(feature = "bevy_reflect", reflect(Component, Default))]
 pub struct SystemIdMarker;
 
+/// Set of systems currently running
+#[derive(Default, Resource, Deref, DerefMut)]
+struct RunSystemStack(EntityHashSet);
+
 /// A system that has been removed from the registry.
 /// It contains the system and whether or not it has been initialized.
 ///
@@ -332,25 +337,37 @@ impl World {
         I: SystemInput + 'static,
         O: 'static,
     {
+        self.init_resource::<RunSystemStack>();
+
+        if self.resource::<RunSystemStack>().contains(&id.entity) {
+            return Err(RegisteredSystemError::Recursive(id));
+        }
+
         // Lookup
         let mut entity = self
             .get_entity_mut(id.entity)
             .map_err(|_| RegisteredSystemError::SystemIdNotRegistered(id))?;
 
         // Take ownership of system trait object
-        let RegisteredSystem {
+        let Some(RegisteredSystem {
             mut initialized,
             mut system,
-        } = entity
-            .take::<RegisteredSystem<I, O>>()
-            .ok_or(RegisteredSystemError::Recursive(id))?;
+        }) = entity.take::<RegisteredSystem<I, O>>()
+        else {
+            return Err(RegisteredSystemError::MaybeIncorrectType(id));
+        };
 
-        // Run the system
+        // Initialize the system
         if !initialized {
             system.initialize(self);
             initialized = true;
         }
 
+        let Some(mut sys_stack) = self.get_resource_mut::<RunSystemStack>() else {
+            panic!("`RunSystemStack` should never be removed during the execution of a one-shot system.");
+        };
+        sys_stack.insert(id.entity);
+
         let result = system
             .validate_param(self)
             .map_err(|err| RegisteredSystemError::InvalidParams { system: id, err })
@@ -362,6 +379,11 @@ impl World {
                 ret
             });
 
+        let Some(mut sys_stack) = self.get_resource_mut::<RunSystemStack>() else {
+            panic!("`RunSystemStack` should never be removed during the execution of a one-shot system.");
+        };
+        sys_stack.remove(&id.entity);
+
         // Return ownership of system trait object (if entity still exists)
         if let Ok(mut entity) = self.get_entity_mut(id.entity) {
             entity.insert::<RegisteredSystem<I, O>>(RegisteredSystem {
@@ -501,6 +523,10 @@ pub enum RegisteredSystemError<I: SystemInput = (), O = ()> {
         /// The returned parameter validation error.
         err: SystemParamValidationError,
     },
+    /// System could not be run due to parameters that failed validation.
+    /// This should not be considered an error if [`field@SystemParamValidationError::skipped`] is `true`.
+    #[error("System {:?} was not found on entity {:?}. Maybe it has wrong type", core::any::type_name::<SystemId<I, O>>(), core::convert::identity(.0))]
+    MaybeIncorrectType(SystemId<I, O>),
 }
 
 impl<I: SystemInput, O> core::fmt::Debug for RegisteredSystemError<I, O> {
@@ -517,6 +543,9 @@ impl<I: SystemInput, O> core::fmt::Debug for RegisteredSystemError<I, O> {
                 .field("system", system)
                 .field("err", err)
                 .finish(),
+            Self::MaybeIncorrectType(arg0) => {
+                f.debug_tuple("MaybeIncorrectType").field(arg0).finish()
+            }
         }
     }
 }
@@ -527,7 +556,12 @@ mod tests {
 
     use bevy_utils::default;
 
-    use crate::{prelude::*, system::SystemId};
+    use crate::{
+        prelude::*,
+        system::{RegisteredSystemError, SystemId},
+    };
+
+    use super::{RegisteredSystem, RunSystemStack};
 
     #[derive(Resource, Default, PartialEq, Debug)]
     struct Counter(u8);
@@ -907,6 +941,23 @@ mod tests {
         assert_eq!(INVOCATIONS_LEFT.get(), 0);
     }
 
+    #[test]
+    fn world_run_system_recursive() {
+        fn world_recursive(world: &mut World) {
+            match world.run_system_cached(world_recursive) {
+                Ok(_) => panic!("Calling an already running system should fail."),
+                Err(RegisteredSystemError::Recursive(_)) => (),
+                Err(err) => panic!(
+                    "Should've failed with `Recursive`, but failed with `{}`.",
+                    err
+                ),
+            };
+        }
+
+        let mut world = World::new();
+        world.run_system_cached(world_recursive).unwrap();
+    }
+
     #[test]
     fn run_system_exclusive_adapters() {
         let mut world = World::new();
@@ -915,4 +966,29 @@ mod tests {
         world.run_system_cached(system.pipe(system)).unwrap();
         world.run_system_cached(system.map(|()| {})).unwrap();
     }
+
+    #[test]
+    #[should_panic]
+    fn removing_system_stack_with_one_shot() {
+        let mut world = World::new();
+        world
+            .run_system_cached(|world: &mut World| {
+                world.remove_resource::<RunSystemStack>();
+            })
+            .expect_err("Removing RunSystemStack should error.");
+    }
+
+    #[test]
+    #[should_panic]
+    fn removing_system_stack_with_observer() {
+        let mut world = World::new();
+        world.add_observer(
+            |_trigger: Trigger<OnRemove, RegisteredSystem<(), ()>>, world: &mut World| {
+                world.remove_resource::<RunSystemStack>();
+            },
+        );
+        world
+            .run_system_cached(|_world: &mut World| {})
+            .expect_err("Removing RunSystemStack should error.");
+    }
 }