diff --git a/crates/bevy_ecs/src/schedule/error.rs b/crates/bevy_ecs/src/schedule/error.rs index caa34a1d40a6a..21181a21e21f2 100644 --- a/crates/bevy_ecs/src/schedule/error.rs +++ b/crates/bevy_ecs/src/schedule/error.rs @@ -272,6 +272,9 @@ pub enum ScheduleError { /// Schedule not found #[error("Schedule not found.")] ScheduleNotFound, + /// System not found + #[error("System not found")] + SystemNotFound, /// Error initializing schedule #[error("{0}")] ScheduleBuildError(ScheduleBuildError), diff --git a/crates/bevy_ecs/src/schedule/node.rs b/crates/bevy_ecs/src/schedule/node.rs index 6f07c321e07a7..d346f4c3f69ab 100644 --- a/crates/bevy_ecs/src/schedule/node.rs +++ b/crates/bevy_ecs/src/schedule/node.rs @@ -550,6 +550,15 @@ impl Systems { key } + pub(crate) fn replace(&mut self, key: SystemKey, system: ScheduleSystem) -> bool { + let Some(slot) = self.nodes.get_mut(key) else { + return false; + }; + *slot = SystemNode::new(system); + self.uninit.push(key); + true + } + /// Remove a system with [`SystemKey`] pub(crate) fn remove(&mut self, key: SystemKey) -> bool { let mut found = false; diff --git a/crates/bevy_ecs/src/schedule/schedule.rs b/crates/bevy_ecs/src/schedule/schedule.rs index 87a91c86ae568..581cb2733f1b8 100644 --- a/crates/bevy_ecs/src/schedule/schedule.rs +++ b/crates/bevy_ecs/src/schedule/schedule.rs @@ -1086,6 +1086,23 @@ impl ScheduleGraph { self.ambiguous_with_all.remove(&NodeId::from(key)); } + /// Replace a system in a set. Use [`system_in_set`] to get + /// the [`SystemKey`]. This leaves any run conditions or ordering dependencies + /// in place. + pub fn replace_system_by_key( + &mut self, + key: SystemKey, + system: ScheduleSystem, + ) -> Result<(), ScheduleError> { + if !self.systems.replace(key, system) { + return Err(ScheduleError::SystemNotFound); + } + + self.changed = true; + + Ok(()) + } + /// Update the internal graphs (hierarchy, dependency, ambiguity) by adding a single [`GraphInfo`] fn update_graphs(&mut self, id: NodeId, graph_info: GraphInfo) { self.changed = true; @@ -1539,10 +1556,15 @@ impl ScheduleGraph { .zip(schedule.systems.drain(..)) .zip(schedule.system_conditions.drain(..)) { + // if the node no longer exists it was removed, so do nothing if let Some(node) = self.systems.node_mut(key) { - node.inner = Some(system); + // if node.inner.is_some() then the system was replaced + if node.inner.is_none() { + node.inner = Some(system); + } } + // if the node no longer exists it was removed, so do nothing if let Some(node_conditions) = self.systems.get_conditions_mut(key) { *node_conditions = conditions; } @@ -1553,6 +1575,7 @@ impl ScheduleGraph { .drain(..) .zip(schedule.set_conditions.drain(..)) { + // if the node no longer exists it was removed, so do nothing if let Some(node_conditions) = self.system_sets.get_conditions_mut(key) { *node_conditions = conditions; } @@ -1990,6 +2013,8 @@ pub struct ScheduleNotInitialized; #[cfg(test)] mod tests { + use alloc::boxed::Box; + use bevy_ecs_macros::ScheduleLabel; use crate::{ @@ -1999,7 +2024,7 @@ mod tests { tests::ResMut, IntoScheduleConfigs, Schedule, ScheduleBuildSettings, ScheduleCleanupPolicy, SystemSet, }, - system::Commands, + system::{Commands, IntoSystem, ScheduleSystem}, world::World, }; @@ -2975,4 +3000,42 @@ mod tests { let conflicts = schedule.graph().conflicting_systems(); assert!(conflicts.is_empty()); } + + #[test] + fn replace_a_system() { + #[derive(Resource)] + struct Counter(usize); + + fn system_1(mut counter: ResMut) { + counter.0 = 1; + } + + fn system_2(mut counter: ResMut) { + counter.0 = 2; + } + + let mut schedule = Schedule::default(); + schedule.add_systems(system_1); + + let mut world = World::new(); + world.insert_resource(Counter(0)); + + let boxed: ScheduleSystem = Box::new(IntoSystem::into_system(system_2)); + + schedule.initialize(&mut world).unwrap(); + let Ok(&[key]) = schedule + .graph_mut() + .systems_in_set(system_1.into_system_set().intern()) + else { + panic!("key not found") + }; + assert!(schedule + .graph_mut() + .replace_system_by_key(key, boxed) + .is_ok()); + + schedule.run(&mut world); + + assert_eq!(world.resource::().0, 2); + } }