From 3fe152fcdd8ce647ad6c4432dd6d562786ffa886 Mon Sep 17 00:00:00 2001 From: ada Date: Sat, 7 Feb 2026 21:04:40 -0600 Subject: [PATCH] fix: cleanup scoping Cleanup method was being run too late, resulting in errantly cleaning up entities from the next scene. I introduced the LoadQueued and Cleanup states to fix this issue. --- crates/screens/src/data.rs | 64 +++++++++++++++++--- crates/screens/src/scope.rs | 42 ++++++++++--- crates/screens/src/systems.rs | 26 +++++++- crates/screens/tests/screens/entity_scope.rs | 3 + crates/screens/tests/screens/lifecycle.rs | 42 ++++++++----- 5 files changed, 143 insertions(+), 34 deletions(-) diff --git a/crates/screens/src/data.rs b/crates/screens/src/data.rs index 646c54f..8247112 100644 --- a/crates/screens/src/data.rs +++ b/crates/screens/src/data.rs @@ -215,10 +215,11 @@ mod screens { } } - /// Loads the screen. - /// Has no effect if already in Loading or Ready states. - pub fn load(&mut self, tick: Tick) { - if matches!(self.state, ScreenState::Unloaded | ScreenState::Unloading) { + pub(crate) fn load(&mut self, tick: Tick) { + if matches!( + self.state, + ScreenState::Unloaded | ScreenState::Unloading | ScreenState::LoadQueued + ) { if self.skip_load { self.state = ScreenState::Ready } else { @@ -234,7 +235,7 @@ mod screens { pub fn unload(&mut self, tick: Tick) { if matches!(self.state, ScreenState::Loading | ScreenState::Ready) { if self.skip_unload { - self.state = ScreenState::Unloaded + self.state = ScreenState::Cleanup; } else { self.state = ScreenState::Unloading; } @@ -242,6 +243,16 @@ mod screens { self.changed_at = tick; } } + /// Queues the screen to load, unloading any other screens and waiting + /// until their cleanup is complete. + /// Only works if in the `Unloaded` state. + pub fn queue_load(&mut self, tick: Tick) { + if matches!(self.state, ScreenState::Unloaded) { + self.state = ScreenState::LoadQueued; + self.needs_update = true; + self.changed_at = tick; + } + } /// Finishes loading the screen. /// Has no effect if already in Loading or Ready states. pub fn finish_loading(&mut self, tick: Tick) { @@ -255,6 +266,14 @@ mod screens { /// Has no effect if already in Loading or Ready states. pub fn finish_unloading(&mut self, tick: Tick) { if matches!(self.state, ScreenState::Unloading) { + self.state = ScreenState::Cleanup; + self.needs_update = true; + self.changed_at = tick; + } + } + + pub(crate) fn finish_cleanup(&mut self, tick: Tick) { + if matches!(self.state, ScreenState::Cleanup) { self.state = ScreenState::Unloaded; self.needs_update = true; self.changed_at = tick; @@ -344,6 +363,8 @@ mod schedules { Loading, /// Runs on [Update] when the screen has [ScreenState::Unloading] Unloading, + /// For internal use! Runs after [ScreenState::Unloading]. Used to clean up screen-scoped entities. + Cleanup, /// Can also be specified as [on_screen_load] OnLoad, /// Can also be specified as [on_screen_ready] @@ -354,6 +375,7 @@ mod schedules { OnUnloaded, } + // TODO: This should use ScreenId internally. /// Wrapper around [ScreenSchedule]. Needed to make schedules unique per type. #[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct ScreenScheduleLabel { @@ -380,14 +402,18 @@ pub use schedules::*; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] pub enum ScreenState { #[default] - #[allow(missing_docs)] + /// The screen is currently down. Unloaded, - #[allow(missing_docs)] + /// The screen is waiting for cleanup to finish. + LoadQueued, + /// The screen is running any custom load systems. Loading, - #[allow(missing_docs)] + /// The screen is fully loaded and ready to execute its main systems. Ready, - #[allow(missing_docs)] + /// The screen is running any custom unload systems. Unloading, + /// For internal use. Cleaning up screen-scoped entities. + Cleanup, } mod system_params { @@ -472,6 +498,10 @@ mod system_params { let tick = self.change_tick; self.data.finish_unloading(tick); } + pub(crate) fn finish_cleanup(&mut self) { + let tick = self.change_tick; + self.data.finish_cleanup(tick); + } #[allow(missing_docs)] pub fn data(&self) -> &ScreenInfo { self.data.as_ref() @@ -672,6 +702,22 @@ mod helpers { OnScreenLoad(TypeId::of::()) } + /// Label of a schedule which fires when the screen has begun its cleanup schedule. + /// Try to avoid spawning anything during this schedule. + #[derive(ScheduleLabel, Debug, PartialEq, Eq, Hash, Clone, Copy)] + pub struct OnScreenCleanup(pub TypeId); + /// See [OnScreenCleanup] + pub fn on_screen_cleanup() -> impl ScheduleLabel { + OnScreenCleanup(TypeId::of::()) + } + /// Label of a schedule which fires when the screen is waiting to load. + #[derive(ScheduleLabel, Debug, PartialEq, Eq, Hash, Clone, Copy)] + pub struct OnScreenLoadQueued(pub TypeId); + /// See [OnScreenLoadQueued] + pub fn on_screen_load_queued() -> impl ScheduleLabel { + OnScreenLoadQueued(TypeId::of::()) + } + /// Label of a schedule which fires when the screen has finished loading. #[derive(ScheduleLabel, Debug, PartialEq, Eq, Hash, Clone, Copy)] pub struct OnScreenReady(pub TypeId); diff --git a/crates/screens/src/scope.rs b/crates/screens/src/scope.rs index 03c230e..7709bf3 100644 --- a/crates/screens/src/scope.rs +++ b/crates/screens/src/scope.rs @@ -1,7 +1,10 @@ use std::any::TypeId; pub use crate::prelude::*; -use bevy::{ecs::system::ScheduleSystem, platform::collections::HashMap}; +use bevy::{ + ecs::system::{ScheduleSystem, SystemIdMarker}, + platform::collections::HashMap, +}; use strum::IntoEnumIterator; #[allow(missing_docs)] @@ -140,6 +143,7 @@ where app.add_observer(on_switch_screen::); app.add_observer(on_finish_loading::); app.add_observer(on_finish_unloading::); + app.add_systems(on_screen_cleanup::(), clean_up_scoped_entities::); // scope systems for (kind, schedule) in self.schedules.into_iter() { @@ -159,7 +163,7 @@ where ScreenSchedule::OnUnload => { let label = schedule.label(); app.add_systems(on_screen_unload::(), move |mut commands: Commands| { - commands.run_schedule(label) + commands.run_schedule(label); }); } ScreenSchedule::OnUnloaded => { @@ -178,16 +182,25 @@ where // Lifecycle #[cfg(debug_assertions)] { - app.add_systems(on_screen_load::(), || debug!("Loading {:?}", S::name())); - app.add_systems(on_screen_ready::(), || debug!("Ready {:?}", S::name())); + app.add_systems(on_screen_load_queued::(), || { + debug!("LoadQueued {:?}", S::name()) + }); + app.add_systems(on_screen_load::(), || { + debug!(" Loading {:?}", S::name()) + }); + app.add_systems(on_screen_ready::(), || { + debug!(" Ready {:?}", S::name()) + }); app.add_systems(on_screen_unload::(), || { - debug!("Unloading {:?}", S::name()) + debug!(" Unloading {:?}", S::name()) + }); + app.add_systems(on_screen_cleanup::(), || { + debug!(" Cleanup {:?}", S::name()) }); app.add_systems(on_screen_unloaded::(), || { - debug!("Unloaded {:?}", S::name()) + debug!(" Unloaded {:?}", S::name()) }); } - app.add_systems(on_screen_unloaded::(), clean_up_scoped_entities::); debug!("Built {} (id={:?})", S::name(), id); } } @@ -224,13 +237,22 @@ fn clean_up_scoped_entities( )>, ), >, - top_levels: Query, With)>, Without)>, + top_levels: Query< + Entity, + ( + Or<(With, With, With)>, // there are probably others i'm missing + Without, + ), + >, ) { screen_scoped .iter() .filter(|c| !top_levels.iter().contains(c)) .for_each(|e| { - commands.entity(e).detach_all_children().despawn(); + if let Ok(mut cmds) = commands.get_entity(e) { + cmds.clear(); // removes all relationship components + cmds.despawn(); + } }); - screen_data.unload(); + screen_data.finish_cleanup(); } diff --git a/crates/screens/src/systems.rs b/crates/screens/src/systems.rs index acd57f2..9cde586 100644 --- a/crates/screens/src/systems.rs +++ b/crates/screens/src/systems.rs @@ -22,7 +22,7 @@ fn handle_switch_msg( for (key, value) in registry.iter_mut().enumerate() { if let Some(data) = value { if key == *msg_key.0 { - data.load(tick.this_run()); + data.queue_load(tick.this_run()); } else { data.unload(tick.this_run()); } @@ -60,6 +60,13 @@ pub(crate) fn on_finish_unloading( } fn run_schedules(mut data: ResMut, mut commands: Commands, tick: SystemChangeTick) { + let all_clear = data.iter().filter_map(|info| info.as_ref()).all(|info| { + matches!( + info.state(), + ScreenState::Unloaded | ScreenState::LoadQueued + ) + }); + for info in data.iter_mut().filter_map(|info| info.as_mut()) { match info.state() { ScreenState::Unloaded => { @@ -74,6 +81,16 @@ fn run_schedules(mut data: ResMut, mut commands: Commands, tick: Sys info.changed_at = tick.this_run(); } } + ScreenState::LoadQueued => { + if info.needs_update { + commands.run_schedule(OnScreenLoadQueued(info.type_id())); + info.needs_update = false; + info.changed_at = tick.this_run(); + } + if all_clear { + info.load(tick.this_run()); + } + } ScreenState::Loading => { commands.run_schedule(ScreenScheduleLabel::from_id( ScreenSchedule::Loading, @@ -113,6 +130,13 @@ fn run_schedules(mut data: ResMut, mut commands: Commands, tick: Sys info.changed_at = tick.this_run(); } } + ScreenState::Cleanup => { + if info.needs_update { + commands.run_schedule(OnScreenCleanup(info.type_id())); + info.needs_update = false; + info.changed_at = tick.this_run(); + } + } } } } diff --git a/crates/screens/tests/screens/entity_scope.rs b/crates/screens/tests/screens/entity_scope.rs index 73666c5..d04139b 100644 --- a/crates/screens/tests/screens/entity_scope.rs +++ b/crates/screens/tests/screens/entity_scope.rs @@ -6,6 +6,9 @@ use itertools::Itertools; use crate::prelude::*; +// TODO: Test that screens which spawn entities in `load` are able to access those entities. +// Don't run `load` until all the scoped entities are removed! + #[derive(Component, PartialEq, Debug, Copy, Clone)] #[component(on_insert = Self::on_insert)] enum Target { diff --git a/crates/screens/tests/screens/lifecycle.rs b/crates/screens/tests/screens/lifecycle.rs index 3cc05a6..afefa2f 100644 --- a/crates/screens/tests/screens/lifecycle.rs +++ b/crates/screens/tests/screens/lifecycle.rs @@ -53,7 +53,12 @@ macro_rules! impl_test_fns { } } -gen_test_fns!(on_screen_load, on_screen_ready, on_screen_unloaded); +gen_test_fns!( + on_screen_load, + on_screen_ready, + on_screen_unloaded, + on_screen_cleanup +); macro_rules! progress_by { ($name:ident) => { @@ -115,18 +120,6 @@ impl Screen for LifecycleScreen { commands.write_message(AppExit::error()); } }, - |r: Res, r2: Res, mut commands: Commands| { - let ok = r.ok() && r2.ok(); - if ok { - info!("OK!"); - commands.write_message(AppExit::Success); - } else { - error!("Did not reach all expected points."); - error!(?r); - error!(?r2); - commands.write_message(AppExit::error()); - } - }, ) .chain(), ); @@ -141,6 +134,27 @@ fn lifecycle() { app.register_screen::(); app.init_resource::(); app.init_resource::(); - impl_test_fns!(app, on_screen_load, on_screen_ready, on_screen_unloaded); + impl_test_fns!( + app, + on_screen_load, + on_screen_ready, + on_screen_unloaded, + on_screen_cleanup + ); + app.add_systems( + on_screen_ready::(), + |r: Res, r2: Res, mut commands: Commands| { + let ok = r.ok() && r2.ok(); + if ok { + info!("OK!"); + commands.write_message(AppExit::Success); + } else { + error!("Did not reach all expected points."); + error!(?r); + error!(?r2); + commands.write_message(AppExit::error()); + } + }, + ); assert!(app.run().is_success()); } -- 2.51.2