moonshine_util/
query.rs

1//! Utilities and decorators for working with queries.
2
3use std::marker::PhantomData;
4
5use bevy_ecs::archetype::Archetype;
6use bevy_ecs::change_detection::Tick;
7use bevy_ecs::component::{ComponentId, Components};
8use bevy_ecs::prelude::*;
9use bevy_ecs::query::{FilteredAccess, QueryData, ReadOnlyQueryData, WorldQuery};
10use bevy_ecs::storage::{Table, TableRow};
11
12/// A trait for types that can be constructed from query data.
13///
14/// See [`Get`] for more information on usage.
15pub trait MapQuery {
16    /// The query type which this type can be constructed from.
17    type Query: ReadOnlyQueryData;
18
19    /// The output type that this query will map to.
20    type Output;
21
22    /// Called at the time of query execution to map the query data into `Self`.
23    fn map(data: <Self::Query as QueryData>::Item<'_, '_>) -> Self::Output;
24}
25
26/// A query decorator which maps some query data into `T` using [`MapQuery`].
27///
28/// This is useful for when you want to compute a processed value from some query data.
29///
30/// # Example
31///
32/// ```rust
33/// use bevy::prelude::*;
34/// use moonshine_util::prelude::*;
35///
36/// struct Height;
37///
38/// impl MapQuery for Height {
39///     type Query = &'static GlobalTransform;
40///
41///     type Output = f32;
42///
43///     fn map(data: &GlobalTransform) -> f32 {
44///         data.translation().y
45///     }
46/// }
47///
48/// fn average_height(query: Query<Get<Height>>) -> f32 {
49///     let mut total_height = 0.0;
50///     let mut count = 0;
51///     for h in query.iter() {
52///         total_height += h;
53///         count += 1;
54///     }
55///
56///     total_height / count as f32
57/// }
58///
59/// # bevy_ecs::system::assert_is_system(average_height);
60/// ```
61pub struct Get<T>(PhantomData<T>);
62
63unsafe impl<T: MapQuery> WorldQuery for Get<T> {
64    type Fetch<'a> = <T::Query as WorldQuery>::Fetch<'a>;
65
66    type State = <T::Query as WorldQuery>::State;
67
68    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
69        T::Query::shrink_fetch(fetch)
70    }
71
72    unsafe fn init_fetch<'w>(
73        world: bevy_ecs::world::unsafe_world_cell::UnsafeWorldCell<'w>,
74        state: &Self::State,
75        last_run: Tick,
76        this_run: Tick,
77    ) -> Self::Fetch<'w> {
78        unsafe { T::Query::init_fetch(world, state, last_run, this_run) }
79    }
80
81    const IS_DENSE: bool = T::Query::IS_DENSE;
82
83    unsafe fn set_archetype<'w>(
84        fetch: &mut Self::Fetch<'w>,
85        state: &Self::State,
86        archetype: &'w Archetype,
87        table: &'w Table,
88    ) {
89        unsafe { T::Query::set_archetype(fetch, state, archetype, table) }
90    }
91
92    unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) {
93        unsafe { T::Query::set_table(fetch, state, table) }
94    }
95
96    fn update_component_access(state: &Self::State, access: &mut FilteredAccess) {
97        T::Query::update_component_access(state, access)
98    }
99
100    fn init_state(world: &mut World) -> Self::State {
101        T::Query::init_state(world)
102    }
103
104    fn get_state(components: &Components) -> Option<Self::State> {
105        T::Query::get_state(components)
106    }
107
108    fn matches_component_set(
109        state: &Self::State,
110        set_contains_id: &impl Fn(ComponentId) -> bool,
111    ) -> bool {
112        T::Query::matches_component_set(state, set_contains_id)
113    }
114}
115
116unsafe impl<T: MapQuery> QueryData for Get<T> {
117    type ReadOnly = Self;
118
119    const IS_READ_ONLY: bool = <T::Query as QueryData>::IS_READ_ONLY;
120
121    const IS_ARCHETYPAL: bool = <T::Query as QueryData>::IS_ARCHETYPAL;
122
123    type Item<'w, 's> = T::Output;
124
125    fn shrink<'wlong: 'wshort, 'wshort, 's>(
126        item: Self::Item<'wlong, 's>,
127    ) -> Self::Item<'wshort, 's> {
128        item
129    }
130
131    unsafe fn fetch<'w, 's>(
132        state: &'s Self::State,
133        fetch: &mut Self::Fetch<'w>,
134        entity: Entity,
135        table_row: TableRow,
136    ) -> Option<Self::Item<'w, 's>> {
137        T::Query::fetch(state, fetch, entity, table_row).map(T::map)
138    }
139
140    fn iter_access(
141        state: &Self::State,
142    ) -> impl Iterator<Item = bevy_ecs::query::EcsAccessType<'_>> {
143        T::Query::iter_access(state)
144    }
145}
146
147unsafe impl<T: MapQuery> ReadOnlyQueryData for Get<T> {}
moonshine_util/query.rs

moonshine_util/
query.rs
