bevy_exclusive_with 0.1.0

#![deny(rust_2018_idioms)]
#![allow(elided_lifetimes_in_paths)]
#![warn(missing_docs)]
#![doc = include_str!("../README.md")]

use bevy::{
    ecs::query::{QueryFilter, WorldQuery},
    prelude::*,
};
use std::marker::PhantomData;

/// A [`QueryFilter`] that ensures that the queried entity has exactly the component `C`
/// from the exclusive set `S`.
///
/// Use [`exclusive_set!`] to define exclusive sets.
#[derive(Debug)]
pub struct ExclusiveWith<S, C>(PhantomData<(S, C)>);

impl<S, C> Default for ExclusiveWith<S, C> {
    fn default() -> Self {
        Self(PhantomData)
    }
}

unsafe impl<S, C> QueryFilter for ExclusiveWith<S, C>
where
    C: ExclusiveSet<S>,
{
    const IS_ARCHETYPAL: bool = <C::Filter as QueryFilter>::IS_ARCHETYPAL;

    unsafe fn filter_fetch(
        state: &Self::State,
        fetch: &mut Self::Fetch<'_>,
        entity: Entity,
        table_row: bevy::ecs::storage::TableRow,
    ) -> bool {
        <C::Filter as QueryFilter>::filter_fetch(state, fetch, entity, table_row)
    }
}

unsafe impl<S, C> WorldQuery for ExclusiveWith<S, C>
where
    C: ExclusiveSet<S>,
{
    type Fetch<'w> = <C::Filter as WorldQuery>::Fetch<'w>;

    type State = <C::Filter as WorldQuery>::State;

    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
        <C::Filter as WorldQuery>::shrink_fetch(fetch)
    }

    unsafe fn init_fetch<'w>(
        world: bevy::ecs::world::unsafe_world_cell::UnsafeWorldCell<'w>,
        state: &Self::State,
        last_run: bevy::ecs::change_detection::Tick,
        this_run: bevy::ecs::change_detection::Tick,
    ) -> Self::Fetch<'w> {
        <C::Filter as WorldQuery>::init_fetch(world, state, last_run, this_run)
    }

    const IS_DENSE: bool = <C::Filter as WorldQuery>::IS_DENSE;

    unsafe fn set_archetype<'w>(
        fetch: &mut Self::Fetch<'w>,
        state: &Self::State,
        archetype: &'w bevy::ecs::archetype::Archetype,
        table: &'w bevy::ecs::storage::Table,
    ) {
        <C::Filter as WorldQuery>::set_archetype(fetch, state, archetype, table)
    }

    unsafe fn set_table<'w>(
        fetch: &mut Self::Fetch<'w>,
        state: &Self::State,
        table: &'w bevy::ecs::storage::Table,
    ) {
        <C::Filter as WorldQuery>::set_table(fetch, state, table)
    }

    fn update_component_access(state: &Self::State, access: &mut bevy::ecs::query::FilteredAccess) {
        <C::Filter as WorldQuery>::update_component_access(state, access)
    }

    fn init_state(world: &mut World) -> Self::State {
        <C::Filter as WorldQuery>::init_state(world)
    }

    fn get_state(components: &bevy::ecs::component::Components) -> Option<Self::State> {
        <C::Filter as WorldQuery>::get_state(components)
    }

    fn matches_component_set(
        state: &Self::State,
        set_contains_id: &impl Fn(bevy::ecs::component::ComponentId) -> bool,
    ) -> bool {
        <C::Filter as WorldQuery>::matches_component_set(state, set_contains_id)
    }
}

#[doc(hidden)]
pub trait ExclusiveSet<S> {
    type Filter: QueryFilter;
}

/// Macro to define an exclusive set of components. Requires at least two components. Components are
/// not required to be defined locally and may originate from external crates.
///
/// First parameter is the type that represents the exclusive set. Subsequent parameters are the
/// component types that belong to the exclusive set.
///
/// # Examples
///
/// ```
/// # use bevy_exclusive_with::exclusive_set;
/// # use bevy::prelude::*;
/// # #[derive(Component)]
/// # struct ArcherTower;
/// # #[derive(Component)]
/// # struct CannonTower;
/// # #[derive(Component)]
/// # struct MagicTower;
/// # struct Towers;
/// exclusive_set!(Towers: ArcherTower, CannonTower, MagicTower);
/// ```
#[macro_export]
macro_rules! exclusive_set {
    ($set:ty: $($comp:ty),+ $(,)?) => {
        $crate::exclusive_set_impl!($set: $($comp),+; $($comp,)+);
    };
}

#[macro_export]
#[doc(hidden)]
macro_rules! exclusive_set_impl {
    ($set:ty: $comp_first:ty,$($comp:ty),+; $rec_first:ty,$($rec:ty,)*) => {
        impl $crate::ExclusiveSet<$set> for $comp_first {
            type Filter = (With<$comp_first>, $(Without<$comp>,)+);
        }
        $crate::exclusive_set_impl!($set: $($comp),+,$comp_first; $($rec,)*);
    };
    ($set:ty: $comp_first:ty,$($comp:ty),+;) => {};
}

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

    #[derive(Debug, Component)]
    struct MyTransformA;

    #[derive(Debug, Component)]
    struct MyTransformB;

    #[derive(Debug)]
    struct Transforms;

    exclusive_set!(Transforms: Transform, MyTransformA, MyTransformB);

    fn update_exclusive_with(
        _transforms: Query<&mut GlobalTransform, ExclusiveWith<Transforms, Transform>>,
        _transforms_a: Query<&mut GlobalTransform, ExclusiveWith<Transforms, MyTransformA>>,
        _transforms_b: Query<&mut GlobalTransform, ExclusiveWith<Transforms, MyTransformB>>,
        mut app_exit: MessageWriter<AppExit>,
    ) {
        app_exit.write(AppExit::Success);
    }

    fn update_with(
        _transforms: Query<&mut GlobalTransform, With<Transform>>,
        _transforms_a: Query<&mut GlobalTransform, With<MyTransformA>>,
        _transforms_b: Query<&mut GlobalTransform, With<MyTransformB>>,
        mut app_exit: MessageWriter<AppExit>,
    ) {
        app_exit.write(AppExit::Success);
    }

    #[test]
    fn exclusive_with() {
        App::new()
            .add_plugins(MinimalPlugins)
            .add_systems(Update, update_exclusive_with)
            .run();
    }

    #[test]
    #[should_panic]
    fn with() {
        App::new()
            .add_plugins(MinimalPlugins)
            .add_systems(Update, update_with)
            .run();
    }
}