bevy_map_scatter 0.5.0

//! Bevy plugin for map_scatter providing assets, resources, message types, and systems.
#![forbid(unsafe_code)]

use std::sync::Arc;

#[cfg(feature = "ron")]
pub use assets::ScatterPlanAssetSaver;
pub use assets::{
    ParentDef, SamplingDef, ScatterKindDef, ScatterLayerDef, ScatterPlanAsset,
    ScatterPlanAssetLoader, SelectionStrategyDef,
};
use bevy::prelude::*;
use bevy::tasks::{block_on, AsyncComputeTaskPool, Task};
pub use events::{ChannelSink, ScatterBus, ScatterBusConfig, ScatterEventFilter, ScatterMessage};
use map_scatter::prelude::*;
use rand::rngs::StdRng;
use rand::SeedableRng;
pub use streaming::{
    MapScatterStreamingPlugin, ScatterStreamChunk, ScatterStreamChunks, ScatterStreamPlaced,
    ScatterStreamPlacement, ScatterStreamSettings,
};
pub use textures::{ImageTexture, ImageTextureError};

mod assets;
mod events;
mod streaming;
mod textures;

/// Convenient re-exports for common types. Import with `use bevy_map_scatter::prelude::*;`.
pub mod prelude {
    pub use map_scatter::prelude::*;

    #[cfg(feature = "ron")]
    pub use crate::assets::ScatterPlanAssetSaver;
    pub use crate::assets::{
        ParentDef, SamplingDef, ScatterKindDef, ScatterLayerDef, ScatterPlanAsset,
        ScatterPlanAssetLoader, SelectionStrategyDef,
    };
    pub use crate::events::{
        ChannelSink, ScatterBus, ScatterBusConfig, ScatterEventFilter, ScatterMessage,
    };
    pub use crate::streaming::{
        MapScatterStreamingPlugin, ScatterStreamChunk, ScatterStreamChunks, ScatterStreamPlaced,
        ScatterStreamPlacement, ScatterStreamSettings,
    };
    pub use crate::textures::{ImageTexture, ImageTextureError};
    pub use crate::{MapScatterPlugin, ScatterFinished, ScatterRequest, ScatterTextureRegistry};
}

/// Bevy plugin providing scatter assets, resources, messages, systems, and observers.
///
/// Requests are sent with [`Commands::trigger`] using [`ScatterRequest`]. Completed jobs trigger
/// [`ScatterFinished`], while runner diagnostics are forwarded through [`Messages<ScatterMessage>`].
///
/// Scatter jobs are owned by their request entity. Despawning that entity before the job finishes
/// drops the underlying [`Task`], which cancels the work rather than detaching it.
pub struct MapScatterPlugin;

/// Shared texture registry (read-only) used by all runs.
///
/// In Bevy 0.19, resources are stored as singleton components internally. This type is used only
/// as a resource by the plugin; plugin queries remain narrowed to their own components.
/// Register your textures at startup or via custom systems.
#[derive(Resource, Clone)]
pub struct ScatterTextureRegistry(
    /// Shared texture registry used by scatter runs.
    pub Arc<TextureRegistry>,
);

impl Default for ScatterTextureRegistry {
    fn default() -> Self {
        Self(Arc::new(TextureRegistry::new()))
    }
}

/// Shared field program cache used by async jobs.
///
/// Kept private so it cannot be confused with a user-facing component now that Bevy resources are
/// also ECS components internally.
#[derive(Resource, Clone)]
struct ScatterCache(pub Arc<FieldProgramCache>);

impl Default for ScatterCache {
    fn default() -> Self {
        Self(Arc::new(FieldProgramCache::new()))
    }
}

/// A request to run a scatter plan (by asset handle) with a configuration and RNG seed.
///
/// The request entity owns the spawned scatter job. If that entity is despawned before completion,
/// the job is cancelled by dropping its [`Task`]. Cancelled jobs do not trigger
/// [`ScatterFinished`].
#[non_exhaustive]
#[derive(EntityEvent)]
pub struct ScatterRequest {
    /// Entity used to track the request.
    pub entity: Entity,
    /// Handle to the scatter plan asset.
    pub plan: Handle<ScatterPlanAsset>,
    /// Run configuration for this request.
    pub config: RunConfig,
    /// RNG seed for deterministic scattering.
    pub seed: u64,
}

impl ScatterRequest {
    pub fn new(
        entity: Entity,
        plan: Handle<ScatterPlanAsset>,
        config: RunConfig,
        seed: u64,
    ) -> Self {
        Self {
            entity,
            plan,
            config,
            seed,
        }
    }
}

/// Component holding an async scatter job task.
///
/// This is added to entities with a [`ScatterRequest`] when a job is spawned. The component keeps
/// the task alive while the request entity exists. Removing the component or despawning the entity
/// drops the task and cancels unfinished work.
#[derive(Component)]
struct ScatterJob {
    pub task: Option<Task<RunResult>>,
}

/// [`EntityEvent`] triggered when a scatter run has finished.
///
/// This is emitted only for jobs that complete and are collected by `MapScatterPlugin`. Jobs
/// cancelled by removing the internal job component or despawning the request entity do not emit a
/// cancellation event.
#[non_exhaustive]
#[derive(EntityEvent, Clone)]
pub struct ScatterFinished {
    /// Entity associated with the original request.
    pub entity: Entity,
    /// Result produced by the scatter run.
    pub result: RunResult,
}

impl Plugin for MapScatterPlugin {
    fn build(&self, app: &mut App) {
        app.add_message::<ScatterMessage>()
            .register_type::<ScatterPlanAsset>()
            .register_type::<ScatterLayerDef>()
            .register_type::<ScatterKindDef>()
            .register_type::<SelectionStrategyDef>()
            .register_type::<ParentDef>()
            .register_type::<SamplingDef>()
            .register_type::<Handle<ScatterPlanAsset>>()
            .init_asset::<ScatterPlanAsset>()
            .init_asset_loader::<ScatterPlanAssetLoader>()
            .init_resource::<ScatterBusConfig>()
            .init_resource::<ScatterBus>()
            .init_resource::<ScatterTextureRegistry>()
            .init_resource::<ScatterCache>()
            .add_systems(Update, poll_scatter_jobs)
            .add_systems(Update, drain_scatter_messages)
            .add_observer(spawn_scatter_job);
    }
}

fn spawn_scatter_job(
    request: On<ScatterRequest>,
    mut commands: Commands,
    bus: Res<ScatterBus>,
    cache: Res<ScatterCache>,
    textures: Res<ScatterTextureRegistry>,
    assets: Res<Assets<ScatterPlanAsset>>,
) {
    let pool = AsyncComputeTaskPool::get();
    let tx = bus.sender().clone();
    let filter = bus.filter().clone();
    let entity = request.entity;

    if let Err(err) = request.config.validate() {
        warn!(
            "ScatterRequest config invalid for {:?}: {}",
            request.entity, err
        );
        return;
    }

    let Some(plan) = assets.get(&request.plan) else {
        error!("ScatterPlanAsset not loaded yet: {:?}", request.plan);
        return;
    };

    // Prepare data for the task
    let plan = plan.into();
    let config = request.config.clone();
    let seed = request.seed;
    let textures = textures.0.clone();
    let cache = cache.0.clone();
    let tx = tx.clone();
    let filter = filter.clone();

    // Spawn async job returning the RunResult
    let task = pool.spawn(async move {
        let mut rng = StdRng::seed_from_u64(seed);

        // Stream events through channel sink
        let mut sink = ChannelSink {
            request: entity,
            tx,
            filter,
        };

        let Ok(mut runner) = ScatterRunner::try_new(config.clone(), &textures, cache.as_ref())
        else {
            warn!("Scatter runner failed to initialize for {:?}", entity);
            return RunResult::new();
        };
        runner.run_with_events(&plan, &mut rng, &mut sink)
    });

    // Attach job component to the entity
    commands
        .entity(request.entity)
        .insert(ScatterJob { task: Some(task) });
}

fn poll_scatter_jobs(mut commands: Commands, mut job_query: Query<(Entity, &mut ScatterJob)>) {
    for (entity, mut job) in job_query.iter_mut() {
        let Some(task) = job.task.as_ref() else {
            continue;
        };
        if !task.is_finished() {
            continue;
        }

        let task = job.task.take().expect("finished scatter task should exist");
        let result = block_on(task);

        // Remove job component when done.
        commands.entity(entity).remove::<ScatterJob>();

        // Trigger finished `EntityEvent`.
        commands.trigger(ScatterFinished { entity, result });
    }
}

fn drain_scatter_messages(bus: Res<ScatterBus>, mut messages: ResMut<Messages<ScatterMessage>>) {
    while let Ok(message) = bus.receiver().try_recv() {
        messages.write(message);
    }
}

#[cfg(test)]
mod tests {
    use std::any::TypeId;
    use std::thread;
    use std::time::{Duration, Instant};

    use bevy::asset::{AssetPlugin, ReflectAsset, ReflectHandle};
    use bevy::ecs::reflect::ReflectComponent;
    use bevy::tasks::{TaskPool, TaskPoolBuilder};

    use super::*;

    #[derive(Resource, Default)]
    struct ObservedFinished {
        count: usize,
        entity: Option<Entity>,
        result: Option<RunResult>,
    }

    #[test]
    fn plugins_register_bevy_facing_reflected_types() {
        let mut app = App::new();
        app.add_plugins(AssetPlugin::default())
            .add_plugins((MapScatterPlugin, MapScatterStreamingPlugin));

        let registry = app.world().resource::<AppTypeRegistry>().read();

        assert_registered::<ScatterPlanAsset>(&registry);
        assert_registered::<ScatterLayerDef>(&registry);
        assert_registered::<ScatterKindDef>(&registry);
        assert_registered::<SelectionStrategyDef>(&registry);
        assert_registered::<ParentDef>(&registry);
        assert_registered::<SamplingDef>(&registry);
        assert_registered::<Handle<ScatterPlanAsset>>(&registry);
        assert_registered::<ScatterStreamSettings>(&registry);
        assert_registered::<ScatterStreamChunks>(&registry);
        assert_registered::<ScatterStreamChunk>(&registry);
        assert_registered::<ScatterStreamPlacement>(&registry);

        let asset_registration = registry
            .get(TypeId::of::<ScatterPlanAsset>())
            .expect("ScatterPlanAsset should be registered");
        assert!(asset_registration.data::<ReflectAsset>().is_some());

        let handle_registration = registry
            .get(TypeId::of::<Handle<ScatterPlanAsset>>())
            .expect("Handle<ScatterPlanAsset> should be registered");
        assert!(handle_registration.data::<ReflectHandle>().is_some());

        for type_id in [
            TypeId::of::<ScatterStreamSettings>(),
            TypeId::of::<ScatterStreamChunks>(),
            TypeId::of::<ScatterStreamChunk>(),
            TypeId::of::<ScatterStreamPlacement>(),
        ] {
            let registration = registry.get(type_id).expect("type should be registered");
            assert!(registration.data::<ReflectComponent>().is_some());
        }
    }

    #[test]
    fn unfinished_scatter_job_remains_after_polling() {
        let pool = test_pool();
        let mut app = app_with_polling();
        let task = pool.spawn(async { std::future::pending::<RunResult>().await });
        let entity = app.world_mut().spawn(ScatterJob { task: Some(task) }).id();

        app.update();

        let job = app
            .world()
            .entity(entity)
            .get::<ScatterJob>()
            .expect("unfinished job should remain on the request entity");
        assert!(job.task.is_some());
        assert_eq!(app.world().resource::<ObservedFinished>().count, 0);
    }

    #[test]
    fn finished_scatter_job_triggers_finished_event_and_removes_component() {
        let pool = test_pool();
        let mut app = app_with_polling();
        let task = finished_task(&pool);
        let entity = app.world_mut().spawn(ScatterJob { task: Some(task) }).id();

        app.update();

        assert!(app.world().entity(entity).get::<ScatterJob>().is_none());
        let observed = app.world().resource::<ObservedFinished>();
        assert_eq!(observed.count, 1);
        assert_eq!(observed.entity, Some(entity));
        assert!(observed.result.is_some());
    }

    #[test]
    fn despawned_scatter_job_does_not_panic_or_emit_finished() {
        let pool = test_pool();
        let mut app = app_with_polling();
        let task = pool.spawn(async { std::future::pending::<RunResult>().await });
        let entity = app.world_mut().spawn(ScatterJob { task: Some(task) }).id();
        app.world_mut().entity_mut(entity).despawn();

        app.update();

        assert!(app.world().get_entity(entity).is_err());
        assert_eq!(app.world().resource::<ObservedFinished>().count, 0);
    }

    fn app_with_polling() -> App {
        let mut app = App::new();
        app.init_resource::<ObservedFinished>()
            .add_systems(Update, poll_scatter_jobs)
            .add_observer(record_finished);
        app
    }

    fn record_finished(finished: On<ScatterFinished>, mut observed: ResMut<ObservedFinished>) {
        observed.count += 1;
        observed.entity = Some(finished.entity);
        observed.result = Some(finished.result.clone());
    }

    fn test_pool() -> TaskPool {
        TaskPoolBuilder::new().num_threads(1).build()
    }

    fn finished_task(pool: &TaskPool) -> Task<RunResult> {
        let task = pool.spawn(async { RunResult::new() });
        let timeout = Instant::now() + Duration::from_secs(5);
        while !task.is_finished() {
            assert!(
                Instant::now() < timeout,
                "test scatter task did not finish in time"
            );
            thread::yield_now();
        }
        task
    }

    fn assert_registered<T: 'static>(registry: &bevy::reflect::TypeRegistry) {
        assert!(
            registry.get(TypeId::of::<T>()).is_some(),
            "{} should be registered",
            std::any::type_name::<T>()
        );
    }
}