bevy_pipe_affect 0.3.0

//! [`Effect`]s that queue `Commands`.
use std::marker::PhantomData;

use bevy::prelude::*;

use crate::Effect;

/// [`Effect`] that pushes a generic command to the command queue.
///
/// This effect is mostly intended to "fill the gaps" of `bevy_pipe_affect`, as not all commands or
/// world mutations provided by bevy have an equivalent effect.
/// Using this typically requires writing some impure code anyway, so users should also consider
/// writing their own custom [`Effect`] instead.
/// (Then, if it's general-purpose enough, consider contributing it upstream!)
///
/// Can be constructed with [`command_queue`].
///
/// # Example
/// In this example, a system is written that gives all `Player`s a `Score` of 0 if they do not
/// already have a score.
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// # #[derive(proptest_derive::Arbitrary)]
/// struct Player;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// # #[derive(proptest_derive::Arbitrary)]
/// struct Score(u32);
///
/// #[derive(Clone, Debug, Default, PartialEq, Eq)]
/// struct InitScores(Vec<Entity>);
///
/// impl Command for InitScores {
///     fn apply(self, world: &mut World) {
///         world.insert_batch_if_new(self.0.into_iter().map(|entity| (entity, Score(0))));
///     }
/// }
///
/// /// Pure-ish system using effects.
/// fn init_player_score_pureish(query: Query<Entity, With<Player>>) -> CommandQueue<InitScores> {
///     command_queue(InitScores(query.iter().collect()))
/// }
///
/// /// Equivalent impure system.
/// fn init_player_score_impure(query: Query<Entity, With<Player>>, mut commands: Commands) {
///     commands.queue(InitScores(query.iter().collect()));
/// }
/// #
/// # use proptest::prelude::*;
/// #
/// # fn app_setup(component_table: Vec<(Option<Player>, Option<Score>)>) -> App {
/// #     let mut app = App::new();
/// #
/// #     component_table.into_iter().for_each(|(player, score)| {
/// #         let mut entity = app.world_mut().spawn_empty();
/// #         if let Some(player) = player {
/// #             entity.insert(player);
/// #         }
/// #         if let Some(score) = score {
/// #             entity.insert(score);
/// #         }
/// #     });
/// #
/// #     app
/// # }
/// #
/// # fn test_state(world: &mut World) -> Vec<(Entity, Option<&Player>, Option<&Score>)> {
/// #     let mut query = world.query::<(Entity, Option<&Player>, Option<&Score>)>();
/// #     query.iter(world).collect()
/// # }
/// #
/// # proptest! {
/// #     fn main(component_table: Vec<(Option<Player>, Option<Score>)>) {
/// #          let mut pure_app = app_setup(component_table.clone());
/// #          pure_app.add_systems(Update, init_player_score_pureish.pipe(affect));
/// #
/// #          let mut impure_app = app_setup(component_table.clone());
/// #          impure_app.add_systems(Update, init_player_score_impure);
/// #
/// #          for _ in 0..3 {
/// #              prop_assert_eq!(test_state(pure_app.world_mut()), test_state(impure_app.world_mut()));
/// #              pure_app.update();
/// #              impure_app.update();
/// #          }
/// #     }
/// # }
/// ```
#[doc = include_str!("defer_command_note.md")]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CommandQueue<C>
where
    C: Command,
{
    /// The command to push onto the queue.
    pub command: C,
}

/// Construct a new [`CommandQueue`] [`Effect`].
pub fn command_queue<C>(command: C) -> CommandQueue<C>
where
    C: Command,
{
    CommandQueue { command }
}

impl<C> Effect for CommandQueue<C>
where
    C: Command,
{
    type MutParam = Commands<'static, 'static>;

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        param.queue(self.command)
    }
}

/// [`Effect`] that queues a command for inserting the provided `Resource` in the `World`.
///
/// Can be constucted with [`command_insert_resource`].
///
/// # Example
/// In this example, a sysstem is written that inserts a `Score` of 0.
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Resource)]
/// # #[derive(proptest_derive::Arbitrary)]
/// struct Score(u32);
///
/// /// Pure system using effects.
/// fn init_score_pure() -> CommandInsertResource<Score> {
///     command_insert_resource(Score(0))
/// }
///
/// /// Equivalent impure system.
/// fn init_score_impure(mut commands: Commands) {
///     commands.insert_resource(Score(0))
/// }
/// #
/// # use proptest::prelude::*;
/// #
/// # fn app_setup(score: Option<Score>) -> App {
/// #     let mut app = App::new();
/// #     if let Some(score) = score {
/// #         app.insert_resource(score);
/// #     }
/// #
/// #     app
/// # }
/// #
/// # fn resource_state(world: &World) -> Option<&Score> {
/// #     world.get_resource::<Score>()
/// # }
/// #
/// # proptest! {
/// #     fn main(score: Option<Score>) {
/// #         let mut pure_app = app_setup(score);
/// #         pure_app.add_systems(Update, init_score_pure.pipe(affect));
/// #
/// #         let mut impure_app = app_setup(score);
/// #         impure_app.add_systems(Update, init_score_impure);
/// #
/// #         for _ in 0..3 {
/// #              prop_assert_eq!(resource_state(pure_app.world_mut()), resource_state(impure_app.world_mut()));
/// #              pure_app.update();
/// #              impure_app.update();
/// #         }
/// #     }
/// # }
/// ```
#[doc = include_str!("defer_command_note.md")]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CommandInsertResource<R>
where
    R: Resource,
{
    /// The initial value of the inserted resource.
    pub resource: R,
}

/// Construct a new [`CommandInsertResource`] [`Effect`].
pub fn command_insert_resource<R>(resource: R) -> CommandInsertResource<R>
where
    R: Resource,
{
    CommandInsertResource { resource }
}

impl<R> Effect for CommandInsertResource<R>
where
    R: Resource,
{
    type MutParam = Commands<'static, 'static>;

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        param.insert_resource(self.resource);
    }
}

/// [`Effect`] that queues a command for removing a `Resource` from the `World`.
///
/// Can be constructed with [`command_remove_resource`].
///
/// # Example
/// In this example, a system is written that removes the `GoToLevel` resource (presumably, after
/// some level transition has completed).
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Resource)]
/// # #[derive(proptest_derive::Arbitrary)]
/// struct GoToLevel(usize);
///
/// /// Pure system using effects.
/// fn complete_level_transition_pure() -> CommandRemoveResource<GoToLevel> {
///     command_remove_resource::<GoToLevel>()
/// }
///
/// /// Equivalent impure system.
/// fn complete_level_transition_impure(mut commands: Commands) {
///     commands.remove_resource::<GoToLevel>()
/// }
/// #
/// # use proptest::prelude::*;
/// #
/// # fn app_setup(resource: Option<GoToLevel>) -> App {
/// #     let mut app = App::new();
/// #
/// #     if let Some(resource) = resource {
/// #         app.insert_resource(resource);
/// #     }
/// #
/// #     app
/// # }
/// #
/// # fn test_state(world: &World) -> Option<&GoToLevel> {
/// #     world.get_resource::<GoToLevel>()
/// # }
/// #
/// # proptest! {
/// #     fn main(resource: Option<GoToLevel>) {
/// #         let mut pure_app = app_setup(resource);
/// #         pure_app.add_systems(Update, complete_level_transition_pure.pipe(affect));
/// #
/// #         let mut impure_app = app_setup(resource);
/// #         impure_app.add_systems(Update, complete_level_transition_impure);
/// #
/// #         for _ in 0..3 {
/// #              prop_assert_eq!(test_state(pure_app.world_mut()), test_state(impure_app.world_mut()));
/// #              pure_app.update();
/// #              impure_app.update();
/// #         }
/// #     }
/// # }
/// ```
#[doc = include_str!("defer_command_note.md")]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CommandRemoveResource<R>
where
    R: Resource,
{
    resource: PhantomData<R>,
}

impl<R> CommandRemoveResource<R>
where
    R: Resource,
{
    /// Construct a new [`CommandRemoveResource`]
    pub fn new() -> Self {
        CommandRemoveResource {
            resource: PhantomData,
        }
    }
}

/// Construct a new [`CommandRemoveResource`] [`Effect`].
pub fn command_remove_resource<R>() -> CommandRemoveResource<R>
where
    R: Resource,
{
    CommandRemoveResource::new()
}

impl<R> Effect for CommandRemoveResource<R>
where
    R: Resource,
{
    type MutParam = Commands<'static, 'static>;

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        param.remove_resource::<R>();
    }
}

/// [`Effect`] that queues a command for spawning an entity with the provided `Bundle`.
///
/// See [`CommandSpawnAnd`] if you need to produce an extra effect with the spawned `Entity` id.
///
/// Can be constructed with [`command_spawn`].
///
/// # Example
/// In this example, a system is written that spawns an `Enemy`.
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// struct Enemy;
///
/// /// Pure system using effects.
/// fn spawn_enemy_pure() -> CommandSpawn<Enemy> {
///     command_spawn(Enemy)
/// }
///
/// /// Equivalent impure system.
/// fn spawn_enemy_impure(mut commands: Commands) {
///     commands.spawn(Enemy);
/// }
/// #
/// # fn app_setup() -> App {
/// #     App::new()
/// # }
/// #
/// # fn test_state(world: &mut World) -> Vec<(Entity, Option<&Enemy>)> {
/// #     let mut query = world.query::<(Entity, Option<&Enemy>)>();
/// #     query.iter(world).collect()
/// # }
/// #
/// # fn main() {
/// #     let mut pure_app = app_setup();
/// #     pure_app.add_systems(Update, spawn_enemy_pure.pipe(affect));
/// #
/// #     let mut impure_app = app_setup();
/// #     impure_app.add_systems(Update, spawn_enemy_impure);
/// #
/// #     for _ in 0..32 {
/// #         assert_eq!(
/// #             test_state(pure_app.world_mut()),
/// #             test_state(impure_app.world_mut())
/// #         );
/// #         pure_app.update();
/// #         impure_app.update();
/// #     }
/// # }
/// ```
#[doc = include_str!("defer_command_note.md")]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CommandSpawn<B>
where
    B: Bundle,
{
    /// The bundle to spawn.
    pub bundle: B,
}

/// Construct a new [`CommandSpawn`] [`Effect`], without an extra effect.
pub fn command_spawn<B>(bundle: B) -> CommandSpawn<B>
where
    B: Bundle,
{
    CommandSpawn { bundle }
}

impl<B> Effect for CommandSpawn<B>
where
    B: Bundle,
{
    type MutParam = Commands<'static, 'static>;

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        param.spawn(self.bundle);
    }
}

/// [`Effect`] that queues a command for spawning an entity with the provided `Bundle`, then
/// supplies the entity id to the provided effect-producing function to cause another effect.
///
/// See [`CommandSpawn`] if you do not need to produce an extra effect.
///
/// Can be constructed with [`command_spawn_and`].
///
/// # Example
/// In this example, a system is written that spawns a `Player`, and a `Sword` as a child of the
/// player.
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// struct Player;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// struct Sword;
///
/// /// Pure system using effects.
/// fn spawn_armed_player_pure() -> CommandSpawnAnd<Player, CommandSpawn<(Sword, ChildOf)>> {
///     command_spawn_and(Player, |player_entity| {
///         command_spawn((Sword, ChildOf(player_entity)))
///     })
/// }
///
/// /// Equivalent impure system.
/// fn spawn_armed_player_impure(mut commands: Commands) {
///     commands.spawn(Player).with_children(|parent| {
///         parent.spawn(Sword);
///     });
/// }
/// #
/// # fn app_setup() -> App {
/// #     App::new()
/// # }
/// #
/// # fn test_state(
/// #     world: &mut World,
/// # ) -> Vec<(Entity, Option<&Player>, Option<&Sword>, Option<&ChildOf>)> {
/// #     let mut query =
/// #         world.query::<(Entity, Option<&Player>, Option<&Sword>, Option<&ChildOf>)>();
/// #     query.iter(world).collect()
/// # }
/// #
/// # fn main() {
/// #     let mut pure_app = app_setup();
/// #     pure_app.add_systems(Update, spawn_armed_player_pure.pipe(affect));
/// #
/// #     let mut impure_app = app_setup();
/// #     impure_app.add_systems(Update, spawn_armed_player_impure);
/// #
/// #     for _ in 0..32 {
/// #         assert_eq!(
/// #             test_state(pure_app.world_mut()),
/// #             test_state(impure_app.world_mut())
/// #         );
/// #         pure_app.update();
/// #         impure_app.update();
/// #     }
/// # }
/// ```
///
/// Not shown...
/// - In this example, [`CommandSpawn`] is used as the additional [`Effect`], but any other effect
/// could be produced.
#[doc = include_str!("defer_command_note.md")]
#[derive(derive_more::Debug)]
pub struct CommandSpawnAnd<B, E>
where
    B: Bundle,
    E: Effect,
{
    /// The bundle to spawn.
    pub bundle: B,
    /// The `Entity -> Effect` function that may cause another effect.
    #[debug("Entity -> {}", std::any::type_name::<E>())]
    pub f: Box<dyn FnOnce(Entity) -> E>,
}

/// Construct a new [`CommandSpawnAnd`] [`Effect`], with an extra effect using the `Entity`.
pub fn command_spawn_and<B, F, E>(bundle: B, f: F) -> CommandSpawnAnd<B, E>
where
    B: Bundle,
    F: FnOnce(Entity) -> E + 'static,
    E: Effect,
{
    CommandSpawnAnd {
        bundle,
        f: Box::new(f),
    }
}

impl<B, E> Default for CommandSpawnAnd<B, E>
where
    B: Bundle + Default,
    E: Effect + Default,
{
    fn default() -> Self {
        CommandSpawnAnd {
            bundle: default(),
            f: Box::new(|_| default()),
        }
    }
}

impl<B, E> Effect for CommandSpawnAnd<B, E>
where
    B: Bundle,
    E: Effect,
{
    type MutParam = (Commands<'static, 'static>, E::MutParam);

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        let entity = param.0.spawn(self.bundle).id();

        (self.f)(entity).affect(&mut param.1);
    }
}

/// [`Effect`] that queues a command for triggering the given event.
///
/// Can be constructed with [`command_trigger`].
///
/// # Example
/// In this example, a system is written that triggers a `Winner` event if any entity has a score
/// above 100.
/// ```
/// use bevy::prelude::*;
/// use bevy_pipe_affect::prelude::*;
///
/// #[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Component)]
/// # #[derive(proptest_derive::Arbitrary)]
/// struct Score(u8);
///
/// #[derive(Copy, Clone, Debug, PartialEq, Eq, Event)]
/// struct Winner(Entity);
///
/// /// Pure system using effects.
/// fn declare_winner_pure(query: Query<(Entity, &Score)>) -> Option<CommandTrigger<Winner>> {
///     query
///         .iter()
///         .find(|(_, score)| score.0 >= 100)
///         .map(|(entity, _)| command_trigger(Winner(entity)))
/// }
///
/// /// Equivalent impure system.
/// fn declare_winner_impure(query: Query<(Entity, &Score)>, mut commands: Commands) {
///     if let Some((entity, _)) = query.iter().find(|(_, score)| score.0 >= 100) {
///         commands.trigger(Winner(entity))
///     }
/// }
/// #
/// # use proptest::prelude::*;
/// #
/// # fn reset_winner_score(winner: On<Winner>) -> QueryEntityAffect<ComponentSet<Score>> {
/// #     query_entity_affect(winner.0, component_set(Score(0)))
/// # }
/// #
/// # fn app_setup(component_table: Vec<Option<Score>>) -> App {
/// #     let mut app = App::new();
/// #     app.add_systems(
/// #         Update,
/// #         (|| command_spawn(Observer::new(reset_winner_score.pipe(affect)))).pipe(affect),
/// #     );
/// #     component_table.into_iter().for_each(|score| {
/// #         let mut entity = app.world_mut().spawn_empty();
/// #         if let Some(score) = score {
/// #             entity.insert(score);
/// #         }
/// #     });
/// #
/// #     app
/// # }
/// #
/// # fn test_state(world: &mut World) -> Vec<(Entity, Option<&Score>)> {
/// #     let mut query = world.query::<(Entity, Option<&Score>)>();
/// #     query.iter(world).collect()
/// # }
/// #
/// # proptest! {
/// #     fn main(component_table: Vec<Option<Score>>) {
/// #         let mut pure_app = app_setup(component_table.clone());
/// #         pure_app.add_systems(Update, declare_winner_pure.pipe(affect));
/// #
/// #         let mut impure_app = app_setup(component_table.clone());
/// #         impure_app.add_systems(Update, declare_winner_impure);
/// #
/// #         for _ in 0..3 {
/// #             prop_assert_eq!(test_state(pure_app.world_mut()), test_state(impure_app.world_mut()));
/// #             pure_app.update();
/// #             impure_app.update();
/// #         }
/// #     }
/// # }
/// ```
#[doc = include_str!("defer_command_note.md")]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub struct CommandTrigger<E>
where
    E: Event,
    for<'a> E::Trigger<'a>: Default,
{
    /// The event being triggered.
    pub event: E,
}

/// Construct a new [`CommandTrigger`] [`Effect`].
pub fn command_trigger<E>(event: E) -> CommandTrigger<E>
where
    E: Event,
    for<'a> E::Trigger<'a>: Default,
{
    CommandTrigger { event }
}

impl<E> Effect for CommandTrigger<E>
where
    E: Event,
    for<'a> E::Trigger<'a>: Default,
{
    type MutParam = Commands<'static, 'static>;

    fn affect(self, param: &mut <Self::MutParam as bevy::ecs::system::SystemParam>::Item<'_, '_>) {
        param.trigger(self.event);
    }
}

#[cfg(test)]
mod tests {
    use proptest::prelude::*;

    use super::*;
    use crate::effects::number_data::{NumberComponent, NumberEvent, NumberResource};
    use crate::prelude::affect;

    proptest! {
        #[test]
        fn command_queue_can_spawn_entities_non_exclusively(component in any::<NumberComponent<0>>()) {
            let mut app = App::new();

            let component_count = app.world_mut().query_filtered::<(), With<NumberComponent<0>>>().iter(app.world()).count();

            assert_eq!(component_count, 0);

            let spawn_component_system = move || {
                command_queue(move |world: &mut World| {
                    world.spawn(component.clone());
                })
            };


            assert!(!IntoSystem::into_system(spawn_component_system.pipe(affect)).is_exclusive());

            app.add_systems(
                Update,
                spawn_component_system.pipe(affect),
            );

            app.update();

            let component_count = app.world_mut().query_filtered::<(), With<NumberComponent<0>>>().iter(app.world()).count();

            assert_eq!(component_count, 1);

            app.update();

            let component_count = app.world_mut().query_filtered::<(), With<NumberComponent<0>>>().iter(app.world()).count();

            assert_eq!(component_count, 2);
        }

        #[test]
        fn resource_commands_correctly_insert_and_remove(resource in any::<NumberResource>()) {
            let mut app = App::new();

            assert!(app.world().get_resource::<NumberResource>().is_none());

            #[derive(Debug, Clone, PartialEq, Eq, Hash, SystemSet)]
            struct InsertSystem;

            app.add_systems(
                Update,
                (move || command_insert_resource(resource)).pipe(affect).in_set(InsertSystem),
            );

            app.update();

            assert_eq!(app.world().get_resource::<NumberResource>(), Some(&resource));

            app.add_systems(
                Update,
                (move || command_remove_resource::<NumberResource>()).pipe(affect).after(InsertSystem),
            );

            app.update();

            assert!(app.world().get_resource::<NumberResource>().is_none());
        }

        #[test]
        fn command_trigger_correctly_triggers_observers(event in any::<NumberEvent>()) {
            let mut app = App::new();

            app.add_systems(
                Startup,
                (move || command_spawn(Observer::new((|event: On<NumberEvent>| command_insert_resource(NumberResource(event.0))).pipe(affect)))).pipe(affect),
            );

            app.update();
            assert!(app.world().get_resource::<NumberResource>().is_none());

            app.update();
            assert!(app.world().get_resource::<NumberResource>().is_none());

            app.add_systems(
                Update,
                (move || command_trigger(event)).pipe(affect)
            );

            app.update();

            assert_eq!(app.world().get_resource::<NumberResource>(), Some(&NumberResource(event.0)));
        }
    }

    #[test]
    fn command_spawn_effect_can_create_parent_child_relationship() {
        let mut app = App::new();

        let children_count = app
            .world_mut()
            .query::<&ChildOf>()
            .iter(app.world())
            .count();

        assert_eq!(children_count, 0);

        #[derive(Resource)]
        struct ParentEntity(Entity);

        app.add_systems(
            Update,
            (move || {
                command_spawn_and((), move |parent| {
                    (
                        command_spawn(ChildOf(parent)),
                        command_insert_resource(ParentEntity(parent)),
                    )
                })
            })
            .pipe(affect),
        );

        app.update();

        let children_count = app
            .world_mut()
            .query::<&ChildOf>()
            .iter(app.world())
            .count();

        assert_eq!(children_count, 1);

        let parent_entity = app.world().resource::<ParentEntity>().0;

        app.world_mut()
            .query::<&ChildOf>()
            .iter(app.world())
            .for_each(|child_of| {
                assert_eq!(child_of.0, parent_entity);
            });

        let children_of_parent_count = app
            .world()
            .entity(parent_entity)
            .get::<Children>()
            .iter()
            .count();

        assert_eq!(children_of_parent_count, 1);
    }
}