From da26a69283b58c28cfc960e5e86a9aea6e37bb5b Mon Sep 17 00:00:00 2001 From: Mike Hsu Date: Fri, 3 Oct 2025 12:10:37 -0700 Subject: [PATCH 1/2] replace a system --- crates/bevy_ecs/src/schedule/error.rs | 3 ++ crates/bevy_ecs/src/schedule/node.rs | 9 ++++ crates/bevy_ecs/src/schedule/schedule.rs | 64 +++++++++++++++++++++++- 3 files changed, 74 insertions(+), 2 deletions(-) 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..2c8546422f5d0 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; @@ -1540,7 +1557,10 @@ impl ScheduleGraph { .zip(schedule.system_conditions.drain(..)) { 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 let Some(node_conditions) = self.systems.get_conditions_mut(key) { @@ -1990,6 +2010,8 @@ pub struct ScheduleNotInitialized; #[cfg(test)] mod tests { + use alloc::boxed::Box; + use bevy_ecs_macros::ScheduleLabel; use crate::{ @@ -1999,7 +2021,7 @@ mod tests { tests::ResMut, IntoScheduleConfigs, Schedule, ScheduleBuildSettings, ScheduleCleanupPolicy, SystemSet, }, - system::Commands, + system::{Commands, IntoSystem, ScheduleSystem}, world::World, }; @@ -2975,4 +2997,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); + } } From 7a743f1c7683950382bca36d16a740b11692c5cd Mon Sep 17 00:00:00 2001 From: Mike Hsu Date: Fri, 3 Oct 2025 12:14:48 -0700 Subject: [PATCH 2/2] add some comments --- crates/bevy_ecs/src/schedule/schedule.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/bevy_ecs/src/schedule/schedule.rs b/crates/bevy_ecs/src/schedule/schedule.rs index 2c8546422f5d0..581cb2733f1b8 100644 --- a/crates/bevy_ecs/src/schedule/schedule.rs +++ b/crates/bevy_ecs/src/schedule/schedule.rs @@ -1556,6 +1556,7 @@ 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) { // if node.inner.is_some() then the system was replaced if node.inner.is_none() { @@ -1563,6 +1564,7 @@ impl ScheduleGraph { } } + // 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; } @@ -1573,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; }