bevy-trait-query 0.19.0

use bevy_ecs::change_detection::{Mut, Ref, Tick};
use bevy_ecs::entity::Entity;
use bevy_ecs::prelude::World;
use bevy_ecs::ptr::UnsafeCellDeref;
use bevy_ecs::query::{IterQueryData, SingleEntityQueryData};
use bevy_ecs::{
    component::{ComponentId, Components},
    query::{QueryData, QueryItem, ReadOnlyQueryData, WorldQuery},
    storage::TableRow,
    world::unsafe_world_cell::UnsafeWorldCell,
};

use crate::TraitImplMeta;
use crate::{
    OneTraitFetch, TraitQuery, TraitQueryState, debug_unreachable, one::FetchStorage, zip_exact,
};

/// [`WorldQuery`] adapter that fetches entities with exactly one component implementing a trait.
///
/// Depending on whether you requested shared or exclusive access to the trait objects, iterating
/// over these queries yields types with different capacities
///
/// - `Query<One<&dyn Trait>>` yields a [`Ref`] object
/// - `Query<One<&mut dyn Trait>>` yields a [`Mut`] object
pub struct One<T>(pub T);

unsafe impl<Trait: ?Sized + TraitQuery> IterQueryData for One<&Trait> {}
unsafe impl<Trait: ?Sized + TraitQuery> SingleEntityQueryData for One<&Trait> {}
unsafe impl<Trait: ?Sized + TraitQuery> QueryData for One<&Trait> {
    type ReadOnly = Self;

    const IS_READ_ONLY: bool = true;
    const IS_ARCHETYPAL: bool = false;

    type Item<'w, 's> = Ref<'w, Trait>;

    #[inline]
    fn shrink<'wlong: 'wshort, 'wshort, 's>(
        item: QueryItem<'wlong, 's, Self>,
    ) -> QueryItem<'wshort, 's, Self> {
        item
    }

    #[inline]
    unsafe fn fetch<'w, 's>(
        _state: &'s Self::State,
        fetch: &mut Self::Fetch<'w>,
        entity: Entity,
        table_row: TableRow,
    ) -> Option<Self::Item<'w, 's>> {
        unsafe {
            let table_row = table_row.index();
            let (dyn_ctor, ptr, added, changed, location) = match fetch.storage {
                // SAFETY: This function must have been called after `set_archetype`,
                // so we know that `self.storage` has been initialized.
                FetchStorage::Uninit => debug_unreachable(),
                FetchStorage::Table {
                    column,
                    added_ticks,
                    changed_ticks,
                    location,
                    meta,
                } => {
                    let ptr = column.byte_add(table_row * meta.size_bytes);
                    (
                        meta.dyn_ctor,
                        ptr,
                        // SAFETY: We have read access to the component, so by extension
                        // we have access to the corresponding `ComponentTicks`.
                        added_ticks.get_unchecked(table_row).deref(),
                        changed_ticks.get_unchecked(table_row).deref(),
                        location,
                    )
                }
                FetchStorage::SparseSet { components, meta } => {
                    let (ptr, ticks) = components
                        .get_with_ticks(entity)
                        .unwrap_or_else(|| debug_unreachable());
                    (
                        meta.dyn_ctor,
                        ptr,
                        // SAFETY: We have read access to the component, so by extension
                        // we have access to the corresponding `ComponentTicks`.
                        ticks.added.deref(),
                        ticks.changed.deref(),
                        ticks.changed_by,
                    )
                }
            };

            Some(Ref::new(
                dyn_ctor.cast(ptr),
                added,
                changed,
                fetch.last_run,
                fetch.this_run,
                location.map(|loc| loc.deref()),
            ))
        }
    }

    fn iter_access(
        _state: &Self::State,
    ) -> impl Iterator<Item = bevy_ecs::query::EcsAccessType<'_>> {
        core::iter::empty()
    }
}

unsafe impl<Trait: ?Sized + TraitQuery> ReadOnlyQueryData for One<&Trait> {}

// SAFETY: We only access the components registered in TraitQueryState.
// This same set of components is used to match archetypes, and used to register world access.
unsafe impl<Trait: ?Sized + TraitQuery> WorldQuery for One<&Trait> {
    type Fetch<'w> = OneTraitFetch<'w, Trait>;
    type State = TraitQueryState<Trait>;

    #[inline]
    unsafe fn init_fetch<'w>(
        world: UnsafeWorldCell<'w>,
        _state: &Self::State,
        _last_run: Tick,
        _this_run: Tick,
    ) -> OneTraitFetch<'w, Trait> {
        unsafe {
            OneTraitFetch {
                storage: FetchStorage::Uninit,
                last_run: Tick::new(0),
                sparse_sets: &world.storages().sparse_sets,
                this_run: Tick::new(0),
            }
        }
    }

    const IS_DENSE: bool = false;
    // const IS_ARCHETYPAL: bool = false;

    #[inline]
    unsafe fn set_archetype<'w>(
        fetch: &mut OneTraitFetch<'w, Trait>,
        state: &Self::State,
        _archetype: &'w bevy_ecs::archetype::Archetype,
        table: &'w bevy_ecs::storage::Table,
    ) {
        unsafe {
            // Search for a registered trait impl that is present in the archetype.
            // We check the table components first since it is faster to retrieve data of this type.
            //
            // without loss of generality we use the zero-th row since we only care about whether the
            // component exists in the table
            let row = TableRow::new(0_u16.into());
            for (&component_id, &meta) in zip_exact(&*state.components, &*state.meta) {
                if let Some(table_storage) = get_table_fetch_data(table, component_id, row, meta) {
                    fetch.storage = table_storage;
                    return;
                }
            }
            for (&component, &meta) in zip_exact(&*state.components, &*state.meta) {
                if let Some(sparse_set) = fetch.sparse_sets.get(component) {
                    fetch.storage = FetchStorage::SparseSet {
                        components: sparse_set,
                        meta,
                    };
                    return;
                }
            }
            // At least one of the components must be present in the table/sparse set.
            debug_unreachable()
        }
    }

    #[inline]
    unsafe fn set_table<'w>(
        fetch: &mut OneTraitFetch<'w, Trait>,
        state: &Self::State,
        table: &'w bevy_ecs::storage::Table,
    ) {
        unsafe {
            // Search for a registered trait impl that is present in the table.
            //
            // without loss of generality we use the zero-th row since we only care about whether the
            // component exists in the table
            let row = TableRow::new(0_u16.into());
            for (&component_id, &meta) in core::iter::zip(&*state.components, &*state.meta) {
                if let Some(table_storage) = get_table_fetch_data(table, component_id, row, meta) {
                    fetch.storage = table_storage;
                    return;
                }
            }
            // At least one of the components must be present in the table.
            debug_unreachable()
        }
    }

    #[inline]
    fn update_component_access(state: &Self::State, access: &mut bevy_ecs::query::FilteredAccess) {
        let mut new_access = access.clone();
        let mut not_first = false;
        for &component in &*state.components {
            assert!(
                !access.access().has_write(component),
                "&{} conflicts with a previous access in this query. Shared access cannot coincide with exclusive access.",
                core::any::type_name::<Trait>(),
            );
            if not_first {
                let mut intermediate = access.clone();
                intermediate.add_read(component);
                new_access.append_or(&intermediate);
                new_access.extend_access(&intermediate);
            } else {
                new_access.and_with(component);
                new_access.access_mut().add_read(component);
                not_first = true;
            }
        }
        *access = new_access;
    }

    #[inline]
    fn init_state(world: &mut World) -> Self::State {
        TraitQueryState::init(world)
    }

    #[inline]
    fn get_state(_: &Components) -> Option<Self::State> {
        // TODO: fix this https://github.com/bevyengine/bevy/issues/13798
        panic!(
            "transmuting and any other operations concerning the state of a query are currently broken and shouldn't be used. See https://github.com/JoJoJet/bevy-trait-query/issues/59"
        );
    }

    #[inline]
    fn matches_component_set(
        state: &Self::State,
        set_contains_id: &impl Fn(ComponentId) -> bool,
    ) -> bool {
        state.matches_component_set_one(set_contains_id)
    }

    #[inline]
    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
        fetch
    }
}

unsafe impl<Trait: ?Sized + TraitQuery> IterQueryData for One<&mut Trait> {}
unsafe impl<Trait: ?Sized + TraitQuery> SingleEntityQueryData for One<&mut Trait> {}
unsafe impl<'a, Trait: ?Sized + TraitQuery> QueryData for One<&'a mut Trait> {
    type ReadOnly = One<&'a Trait>;

    const IS_READ_ONLY: bool = false;
    const IS_ARCHETYPAL: bool = false;

    type Item<'w, 's> = Mut<'w, Trait>;

    #[inline]
    fn shrink<'wlong: 'wshort, 'wshort, 's>(
        item: QueryItem<'wlong, 's, Self>,
    ) -> QueryItem<'wshort, 's, Self> {
        item
    }

    #[inline]
    unsafe fn fetch<'w>(
        _state: &Self::State,
        fetch: &mut Self::Fetch<'w>,
        entity: Entity,
        table_row: TableRow,
    ) -> Option<Mut<'w, Trait>> {
        unsafe {
            let table_row = table_row.index();
            let (dyn_ctor, ptr, added, changed, location) = match fetch.storage {
                // SAFETY: This function must have been called after `set_archetype`,
                // so we know that `self.storage` has been initialized.
                FetchStorage::Uninit => debug_unreachable(),
                FetchStorage::Table {
                    column,
                    added_ticks,
                    changed_ticks,
                    location,
                    meta,
                } => {
                    let ptr = column.byte_add(table_row * meta.size_bytes);
                    (
                        meta.dyn_ctor,
                        // SAFETY: `column` allows for shared mutable access.
                        // So long as the caller does not invoke this function twice with the same archetype_index,
                        // this pointer will never be aliased.
                        ptr.assert_unique(),
                        // SAFETY: We have exclusive access to the component, so by extension
                        // we have exclusive access to the corresponding `ComponentTicks`.
                        added_ticks.get_unchecked(table_row).deref_mut(),
                        changed_ticks.get_unchecked(table_row).deref_mut(),
                        location,
                    )
                }
                FetchStorage::SparseSet { components, meta } => {
                    let (ptr, ticks) = components
                        .get_with_ticks(entity)
                        .unwrap_or_else(|| debug_unreachable());
                    (
                        meta.dyn_ctor,
                        // SAFETY: We have exclusive access to the sparse set `components`.
                        // So long as the caller does not invoke this function twice with the same archetype_index,
                        // this pointer will never be aliased.
                        ptr.assert_unique(),
                        // SAFETY: We have exclusive access to the component, so by extension
                        // we have exclusive access to the corresponding `ComponentTicks`.
                        ticks.added.deref_mut(),
                        ticks.changed.deref_mut(),
                        ticks.changed_by,
                    )
                }
            };

            Some(Mut::new(
                dyn_ctor.cast_mut(ptr),
                added,
                changed,
                fetch.last_run,
                fetch.this_run,
                location.map(|loc| loc.deref_mut()),
            ))
        }
    }

    fn iter_access(
        _state: &Self::State,
    ) -> impl Iterator<Item = bevy_ecs::query::EcsAccessType<'_>> {
        core::iter::empty()
    }
}

// SAFETY: We only access the components registered in TraitQueryState.
// This same set of components is used to match archetypes, and used to register world access.
unsafe impl<Trait: ?Sized + TraitQuery> WorldQuery for One<&mut Trait> {
    type Fetch<'w> = OneTraitFetch<'w, Trait>;
    type State = TraitQueryState<Trait>;

    #[inline]
    unsafe fn init_fetch<'w>(
        world: UnsafeWorldCell<'w>,
        _state: &Self::State,
        last_run: Tick,
        this_run: Tick,
    ) -> OneTraitFetch<'w, Trait> {
        unsafe {
            OneTraitFetch {
                storage: FetchStorage::Uninit,
                sparse_sets: &world.storages().sparse_sets,
                last_run,
                this_run,
            }
        }
    }

    const IS_DENSE: bool = false;

    #[inline]
    unsafe fn set_archetype<'w>(
        fetch: &mut OneTraitFetch<'w, Trait>,
        state: &Self::State,
        _archetype: &'w bevy_ecs::archetype::Archetype,
        table: &'w bevy_ecs::storage::Table,
    ) {
        unsafe {
            // Search for a registered trait impl that is present in the archetype.
            //
            // without loss of generality we use the zero-th row since we only care about whether the
            // component exists in the table
            let row = TableRow::new(0_u16.into());
            for (&component_id, &meta) in zip_exact(&*state.components, &*state.meta) {
                if let Some(table_storage) = get_table_fetch_data(table, component_id, row, meta) {
                    fetch.storage = table_storage;
                    return;
                }
            }
            for (&component, &meta) in zip_exact(&*state.components, &*state.meta) {
                if let Some(sparse_set) = fetch.sparse_sets.get(component) {
                    fetch.storage = FetchStorage::SparseSet {
                        components: sparse_set,
                        meta,
                    };
                    return;
                }
            }
            // At least one of the components must be present in the table/sparse set.
            debug_unreachable()
        }
    }

    #[inline]
    unsafe fn set_table<'w>(
        fetch: &mut OneTraitFetch<'w, Trait>,
        state: &Self::State,
        table: &'w bevy_ecs::storage::Table,
    ) {
        unsafe {
            // Search for a registered trait impl that is present in the table.
            //
            // without loss of generality we use the zero-th row since we only care about whether the
            // component exists in the table
            let row = TableRow::new(0_u16.into());
            for (&component_id, &meta) in core::iter::zip(&*state.components, &*state.meta) {
                if let Some(table_storage) = get_table_fetch_data(table, component_id, row, meta) {
                    fetch.storage = table_storage;
                    return;
                }
            }
            // At least one of the components must be present in the table.
            debug_unreachable()
        }
    }

    #[inline]
    fn update_component_access(state: &Self::State, access: &mut bevy_ecs::query::FilteredAccess) {
        let mut new_access = access.clone();
        let mut not_first = false;
        for &component in &*state.components {
            assert!(
                !access.access().has_write(component),
                "&mut {} conflicts with a previous access in this query. Mutable component access must be unique.",
                core::any::type_name::<Trait>(),
            );
            if not_first {
                let mut intermediate = access.clone();
                intermediate.add_write(component);
                new_access.append_or(&intermediate);
                new_access.extend_access(&intermediate);
            } else {
                new_access.and_with(component);
                new_access.access_mut().add_write(component);
                not_first = true;
            }
        }
        *access = new_access;
    }

    #[inline]
    fn init_state(world: &mut World) -> Self::State {
        TraitQueryState::init(world)
    }

    #[inline]
    fn get_state(_: &Components) -> Option<Self::State> {
        // TODO: fix this https://github.com/bevyengine/bevy/issues/13798
        panic!(
            "transmuting and any other operations concerning the state of a query are currently broken and shouldn't be used. See https://github.com/JoJoJet/bevy-trait-query/issues/59"
        );
    }

    #[inline]
    fn matches_component_set(
        state: &Self::State,
        set_contains_id: &impl Fn(ComponentId) -> bool,
    ) -> bool {
        state.matches_component_set_one(set_contains_id)
    }

    #[inline]
    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
        fetch
    }
}

// gets all the relevant data repeatingly used for the table storage
#[inline]
unsafe fn get_table_fetch_data<Trait: ?Sized + TraitQuery>(
    table: &'_ bevy_ecs::storage::Table,
    component_id: ComponentId,
    row: TableRow,
    meta: TraitImplMeta<Trait>,
) -> Option<FetchStorage<'_, Trait>> {
    unsafe {
        let ptr = table.get_component(component_id, row)?;
        let location = table.get_changed_by(component_id, row).transpose()?;
        let added = table.get_added_ticks_slice_for(component_id)?;
        let changed = table.get_changed_ticks_slice_for(component_id)?;
        Some(FetchStorage::Table {
            column: ptr,
            added_ticks: added.into(),
            changed_ticks: changed.into(),
            location,
            meta,
        })
    }
}