despawn_state.rs

use bevy::prelude::*;
use std::collections::HashSet;

/// Example demonstrating despawning all entities that were spawned during a specific state
fn main() {
    let mut app = App::new();

    app
        .add_state::<AppState>()
        .insert_resource(EntityRecord(HashSet::new()))
        .add_systems(Update, spawn_entities)
        .add_systems(OnEnter(AppState::Main), (print_entities, record_entities))
        .add_systems(OnExit(AppState::Main), (print_entities, despawn_state).chain())
        .add_systems(OnEnter(AppState::Other), print_entities)
    ;

    // flip states a few times to demonstrate how entities behave when spawned from each
    for i in 1..=6 {
        let new_state = if i % 2 == 0 { AppState::Main } else { AppState::Other };
        println!("new state: {new_state:?}");
        app.insert_resource(NextState(Some(new_state)));
        for _ in 0..i {
            app.update();
        }
    }
}

/// Entities spawned outside of our target state will be populated here when we enter the target state
#[derive(Debug, Clone, Default, Resource)]
struct EntityRecord(pub HashSet<Entity>);

/// A random marker component for demonstration
#[derive(Debug, Clone, Default, Component)]
struct Foo;

#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash, States)]
enum AppState {
    #[default]
    Main,
    Other,
}

/// Prints the current entities in the world
fn print_entities(foos: Query<Entity>) {
    print!("foos: [");
    for (i, foo) in foos.iter().enumerate() {
        if i > 0 {
            print!(", ");
        }
        print!("{foo:?}");
    }
    println!("]");
}

fn spawn_entities(mut commands: Commands) {
    println!("spawn stuff");
    commands.spawn(Foo);
    commands.spawn_empty();
}

/// Record **all** entities in the world at the time of state switch
fn record_entities(mut recorded: ResMut<EntityRecord>, world_entities: Query<Entity>) {
    for entity in world_entities.iter() {
        recorded.0.insert(entity);
    }
}

/// Finds all entities in the world that don't exist in our record
/// These are the entities that were spawned during the current state and we wish to despawn
fn despawn_state(mut commands: Commands, mut recorded: ResMut<EntityRecord>, world_entities: Query<Entity>) {
    let mut entities = HashSet::<Entity>::new();
    for entity in world_entities.iter() {
        entities.insert(entity);
    }
    let prev = std::mem::take(&mut recorded.0);
    let diff: HashSet<_> = entities.difference(&prev).collect();
    println!("despawning {} entities created this state", diff.len());
    for entity in diff {
        commands.entity(*entity).despawn();
    }
}