Fast and correct one-shot systems with a system registry resource

crates/bevy_app/src/system_registry.rs

@@ -1,6 +1,8 @@
 use crate::App;
-use bevy_ecs::prelude::*;
-use bevy_ecs::system::{Callback, SystemRegistryError};
+use bevy_ecs::{
+    prelude::*,
+    system::{SystemId, SystemRegistryError},
+};
 
 impl App {
     /// Register a system with any number of [`SystemLabel`]s.
@@ -27,7 +29,7 @@ impl App {
     ///
     /// Calls [`SystemRegistry::run_callback`](bevy_ecs::system::SystemRegistry::run_callback).
     #[inline]
-    pub fn run_callback(&mut self, callback: Callback) -> Result<(), SystemRegistryError> {
-        self.world.run_callback(callback)
+    pub fn run_system_by_id(&mut self, system_id: SystemId) -> Result<(), SystemRegistryError> {
+        self.world.run_system_by_id(system_id)
     }
 }

crates/bevy_ecs/src/system/commands/mod.rs

@@ -5,8 +5,7 @@ use crate::{
     self as bevy_ecs,
     bundle::Bundle,
     entity::{Entities, Entity},
-    schedule::SystemSet,
-    system::{Callback, IntoSystem, RunCallback, RunSystemCommand},
+    system::{IntoSystem, RunSystem, RunSystemById, SystemId},
     world::{FromWorld, World},
 };
 use bevy_ecs_macros::SystemParam;
@@ -517,14 +516,14 @@ impl<'w, 's> Commands<'w, 's> {
         &mut self,
         system: S,
     ) {
-        self.queue.push(RunSystemCommand::new(system));
+        self.queue.push(RunSystem::new(system));
     }
 
     /// Run the systems corresponding to the label stored in the provided [`Callback`]
     ///
     /// Calls [`SystemRegistry::run_callback`](crate::SystemRegistry::run_callback).
-    pub fn run_callback(&mut self, callback: Callback) {
-        self.queue.push(RunCallback { callback });
+    pub fn run_system_by_id(&mut self, system_id: SystemId) {
+        self.queue.push(RunSystemById::new(system_id));
     }
 
     /// Pushes a generic [`Command`] to the command queue.

crates/bevy_ecs/src/system/system_registry.rs

@@ -1,14 +1,12 @@
-use bevy_utils::tracing::warn;
-use bevy_utils::{Entry, HashMap};
+use std::any::TypeId;
 use std::hash::Hash;
 use std::marker::PhantomData;
 
-use crate::schedule::{BoxedSystemSet, IntoSystemConfig, SystemTypeSet};
 use crate::system::{BoxedSystem, Command, IntoSystem};
 use crate::world::{Mut, World};
 // Needed for derive(Component) macro
-use crate as bevy_ecs;
-use bevy_ecs_macros::{Component, Resource};
+use crate::{self as bevy_ecs, TypeIdMap};
+use bevy_ecs_macros::Resource;
 
 /// Stores initialized [`System`]s, so they can be reused and run in an ad-hoc fashion.
 ///
@@ -79,7 +77,8 @@ use bevy_ecs_macros::{Component, Resource};
 /// ```
 #[derive(Resource, Default)]
 pub struct SystemRegistry {
-    systems: HashMap<BoxedSystemSet, BoxedSystem>,
+    systems: Vec<(bool, BoxedSystem)>,
+    indices: TypeIdMap<usize>,
 }
 
 impl SystemRegistry {
@@ -93,23 +92,20 @@ impl SystemRegistry {
     ///
     /// To provide explicit set(s), use [`register_system_with_sets`](SystemRegistry::register_system_with_sets).
     #[inline]
-    fn register_system<M, S: IntoSystem<(), (), M> + 'static>(
+    pub fn register_system<M, S: IntoSystem<(), (), M> + 'static>(
         &mut self,
-        world: &mut World,
         system: S,
-    ) {
-        let boxed_set: BoxedSystemSet = Box::new(SystemTypeSet::<S>::new());
-        // This avoids nasty surprising behavior in case systems are registered twice
-        if !self.systems.contains_key(&boxed_set) {
-            // Initialize the system's state
-            let mut boxed_system: BoxedSystem = Box::new(IntoSystem::into_system(system));
-            boxed_system.initialize(world);
-            // Insert they system keyed by its set
-            self.systems.insert(boxed_set, boxed_system);
-        } else {
-            let type_name = std::any::type_name::<S>();
-            warn!("A system of type {type_name} was registered more than once!");
-        };
+    ) -> SystemId {
+        let type_id = TypeId::of::<S>();
+
+        let index = *self.indices.entry(type_id).or_insert_with(|| {
+            let index = self.systems.len();
+            self.systems
+                .push((false, Box::new(IntoSystem::into_system(system))));
+            index
+        });
+
+        SystemId::new(index)
     }
 
     /// Runs the supplied system on the [`World`] a single time.
@@ -126,17 +122,9 @@ impl SystemRegistry {
         world: &mut World,
         system: S,
     ) {
-        let boxed_set = Box::new(SystemTypeSet::<S>::new());
-        let matching_system = match self.systems.entry(boxed_set) {
-            Entry::Occupied(o) => o.into_mut(),
-            Entry::Vacant(v) => {
-                let mut boxed_system: BoxedSystem = Box::new(IntoSystem::into_system(system));
-                boxed_system.initialize(world);
-                v.insert(boxed_system)
-            }
-        };
-        matching_system.run((), world);
-        matching_system.apply_buffers(world);
+        let id = self.register_system(system);
+        self.run_system_by_id(world, id)
+            .expect("System was registered before running");
     }
 
     /// Run the systems corresponding to the set stored in the provided [`Callback`]
@@ -146,31 +134,49 @@ impl SystemRegistry {
     ///
     /// Systems will be run sequentially in registration order if more than one registered system matches the provided set.
     #[inline]
-    pub fn run_callback(
+    pub fn run_system_by_id(
         &mut self,
         world: &mut World,
-        callback: Callback,
+        id: SystemId,
     ) -> Result<(), SystemRegistryError> {
-        let boxed_set = callback.set;
-        match self.systems.get_mut(&boxed_set) {
-            Some(matching_system) => {
+        match self.systems.get_mut(id.index()) {
+            Some((initialized, matching_system)) => {
+                if !*initialized {
+                    matching_system.initialize(world);
+                    *initialized = true;
+                }
                 matching_system.run((), world);
                 matching_system.apply_buffers(world);
                 Ok(())
             }
-            None => Err(SystemRegistryError::SetNotFound(boxed_set)),
+            None => Err(SystemRegistryError::SystemIdNotRegistered(id)),
         }
     }
 }
 
+#[derive(Debug, Copy, Clone, Hash, Ord, PartialOrd, Eq, PartialEq)]
+pub struct SystemId(usize);
+
+impl SystemId {
+    #[inline]
+    pub const fn new(index: usize) -> SystemId {
+        SystemId(index)
+    }
+
+    #[inline]
+    pub fn index(self) -> usize {
+        self.0
+    }
+}
+
 impl World {
-    /// Register a system with its [`SystemTypeSet`].
-    ///
-    /// Calls [`SystemRegistry::register_system_with_type_set`].
-    pub fn register_system<M, S: IntoSystem<(), (), M> + 'static>(&mut self, system: S) {
-        self.resource_scope(|world, mut registry: Mut<SystemRegistry>| {
-            registry.register_system(world, system);
-        });
+    #[inline]
+    pub fn register_system<M, S: IntoSystem<(), (), M> + 'static>(
+        &mut self,
+        system: S,
+    ) -> SystemId {
+        self.resource_mut::<SystemRegistry>()
+            .register_system(system)
     }
 
     /// Runs the supplied system on the [`World`] a single time.
@@ -187,26 +193,21 @@ impl World {
     ///
     /// Calls [`SystemRegistry::run_callback`].
     #[inline]
-    pub fn run_callback(&mut self, callback: Callback) -> Result<(), SystemRegistryError> {
+    pub fn run_system_by_id(&mut self, id: SystemId) -> Result<(), SystemRegistryError> {
         self.resource_scope(|world, mut registry: Mut<SystemRegistry>| {
-            registry.run_callback(world, callback)
+            registry.run_system_by_id(world, id)
         })
     }
 }
 
 /// The [`Command`] type for [`SystemRegistry::run_system`]
 #[derive(Debug, Clone)]
-pub struct RunSystemCommand<
-    M: Send + Sync + 'static,
-    S: IntoSystem<(), (), M> + Send + Sync + 'static,
-> {
+pub struct RunSystem<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static> {
     _phantom_marker: PhantomData<M>,
     system: S,
 }
 
-impl<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static>
-    RunSystemCommand<M, S>
-{
+impl<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static> RunSystem<M, S> {
     /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands)
     #[inline]
     #[must_use]
@@ -219,7 +220,7 @@ impl<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static>
 }
 
 impl<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static> Command
-    for RunSystemCommand<M, S>
+    for RunSystem<M, S>
 {
     #[inline]
     fn write(self, world: &mut World) {
@@ -229,70 +230,36 @@ impl<M: Send + Sync + 'static, S: IntoSystem<(), (), M> + Send + Sync + 'static>
 
 /// The [`Command`] type for [`SystemRegistry::run_systems_by_set`]
 #[derive(Debug, Clone)]
-pub struct RunCallback {
-    pub callback: Callback,
+pub struct RunSystemById {
+    pub system_id: SystemId,
 }
 
-impl Command for RunCallback {
+impl RunSystemById {
+    pub fn new(system_id: SystemId) -> Self {
+        Self { system_id }
+    }
+}
+
+impl Command for RunSystemById {
     #[inline]
     fn write(self, world: &mut World) {
         world.resource_scope(|world, mut registry: Mut<SystemRegistry>| {
             registry
-                .run_callback(world, self.callback)
+                .run_system_by_id(world, self.system_id)
                 // Ideally this error should be handled more gracefully,
                 // but that's blocked on a full error handling solution for commands
                 .unwrap();
         });
     }
 }
 
-/// A struct that stores a boxed [`SystemSet`], used to cause a [`SystemRegistry`] to run systems.
-///
-/// This might be stored as a component, used as an event, or arranged in a queue stored in a resource.
-/// Unless you need to inspect the list of events or add additional information,
-/// prefer the simpler `commands.run_system` over storing callbacks as events,
-///
-/// Systems must be registered via the `register_system` methods on [`SystemRegistry`], [`World`] or `App`
-/// before they can be run by their set using a callback.
-#[derive(Debug, Component, Clone, Eq)]
-pub struct Callback {
-    /// The set of the system(s) to be run.
-    ///
-    /// By default, this is set to the [`SystemTypeSet`]
-    /// of the system passed in via [`Callback::new()`].
-    pub set: BoxedSystemSet,
-}
-
-impl Callback {
-    /// Creates a new callback from a function that can be used as a system.
-    ///
-    /// Remember that you must register your systems with the `App` / [`World`] before they can be run as callbacks!
-    pub fn new<S: IntoSystemConfig<M> + 'static, M>(_system: S) -> Self {
-        Callback {
-            set: Box::new(SystemTypeSet::<S>::new()),
-        }
-    }
-}
-
-impl PartialEq for Callback {
-    fn eq(&self, other: &Self) -> bool {
-        self.set.dyn_eq(other.set.as_dyn_eq())
-    }
-}
-
-impl Hash for Callback {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.set.dyn_hash(state);
-    }
-}
-
 /// An operation on a [`SystemRegistry`] failed
 #[derive(Debug)]
 pub enum SystemRegistryError {
     /// A system was run by set, but no system with that set was found.
     ///
     /// Did you forget to register it?
-    SetNotFound(BoxedSystemSet),
+    SystemIdNotRegistered(SystemId),
 }
 
 mod tests {

examples/ecs/one_shot_systems.rs

@@ -6,8 +6,10 @@
 //!
 //! See the [`SystemRegistry`](bevy::ecs::SystemRegistry) docs for more details.
 
-use bevy::ecs::system::Callback;
-use bevy::prelude::*;
+use bevy::{
+    ecs::system::{SystemId, SystemRegistry},
+    prelude::*,
+};
 
 fn main() {
     App::new()
@@ -16,11 +18,6 @@ fn main() {
         // One shot systems are interchangeable with ordinarily scheduled systems.
         // Change detection, Local and NonSend all work as expected.
         .add_system(count_entities.in_base_set(CoreSet::PostUpdate))
-        // Registered systems can be called dynamically by their label.
-        // These must be registered in advance, or commands.run_system will panic
-        // as no matching system was found.
-        .register_system(button_pressed)
-        .register_system(slider_toggled)
         // One-shot systems can be used to build complex abstractions
         // to match the needs of your design.
         // Here, we model a very simple component-linked callback architecture.
@@ -35,19 +32,22 @@ fn count_entities(all_entities: Query<()>) {
     dbg!(all_entities.iter().count());
 }
 
+#[derive(Component)]
+struct Callback(SystemId);
+
 #[derive(Component)]
 struct Triggered;
 
-fn setup(mut commands: Commands) {
+fn setup(mut system_registry: ResMut<SystemRegistry>, mut commands: Commands) {
     commands.spawn((
         // The Callback component is defined in bevy_ecs,
         // but wrapping this (or making your own customized variant) is easy.
         // Just stored a boxed SystemLabel!
-        Callback::new(button_pressed),
+        Callback(system_registry.register_system(button_pressed)),
         Triggered,
     ));
     // This entity does not have a Triggered component, so its callback won't run.
-    commands.spawn(Callback::new(slider_toggled));
+    commands.spawn(Callback(system_registry.register_system(slider_toggled)));
     commands.run_system(count_entities);
 }
 
@@ -68,6 +68,6 @@ fn evaluate_callbacks(query: Query<&Callback, With<Triggered>>, mut commands: Co
         // we have to use the layer of indirection provided by system labels.
         // Note that if we had registered multiple systems with the same label,
         // they would all be evaluated here.
-        commands.run_callback(callback.clone());
+        commands.run_system_by_id(callback.0);
     }
 }