Any good patterns to remove entities "tied" to a certain State?

[thick robin]

I'm currently using unit structs that I add to each of the entities that "belong" to a certain app's state, then I remove these "marked" entities on the OnExit schedule. This works for me, for now, but I'm looking for a more flexible pattern that allows me to reuse my app's state variants instead of having to create new structs for each state.

I'll appreciate your ideas and suggestions!

[carmine mesa]

I don;t know if it is the best pattern, but what works for me is running the following system as (one of the) last during each frame:

pub fn despawner(
    mut command: Commands,
    despawn_query: Query<Entity, With<DespawnNow>>,
    state: Res<State<AppState>>,
    next_state: Res<NextState<AppState>>,
    despawn_state_query: Query<Entity, With<DespawnOnStateChange>>,
) {
    // Collect all entities to despawn in a set to avoid duplicates
    let mut despawn_set: HashSet<Entity> = despawn_query.iter().collect();
    if let Some(next_state) = &next_state.0 {
        if *state.get() != *next_state {
            for entity in &despawn_state_query {
                despawn_set.insert(entity);
            }
        }
    }
    for entity in despawn_set {
        command.entity(entity).despawn_recursive();
    }
}

Any entity that needs to despawn whenever a state changes gets a DespawnOnStateChange component.

Any entity that I want to despawn this frame gets a DespawnNow component. I could despawn entities directly, but found it quite hard to do this in a consistent manner so that no missing entity panics would occur.