azalea_client/
entity_query.rs

1use std::sync::Arc;
2
3use bevy_ecs::{
4    component::Component,
5    entity::Entity,
6    query::QueryData,
7    query::{QueryFilter, ROQueryItem},
8    world::World,
9};
10use parking_lot::Mutex;
11
12use crate::Client;
13
14impl Client {
15    /// A convenience function for getting components of our player's entity.
16    ///
17    /// # Examples
18    /// ```
19    /// # use azalea_world::InstanceName;
20    /// # fn example(mut client: azalea_client::Client) {
21    /// let is_logged_in = client
22    ///     .query::<Option<&InstanceName>>(&mut client.ecs.lock())
23    ///     .is_some();
24    /// # }
25    /// ```
26    pub fn query<'w, D: QueryData>(&self, ecs: &'w mut World) -> D::Item<'w> {
27        ecs.query::<D>()
28            .get_mut(ecs, self.entity)
29            .unwrap_or_else(|_| {
30                panic!(
31                    "Our client is missing a required component {:?}",
32                    std::any::type_name::<D>()
33                )
34            })
35    }
36
37    /// Return a lightweight [`Entity`] for the entity that matches the given
38    /// predicate function.
39    ///
40    /// You can then use [`Self::entity_component`] to get components from this
41    /// entity.
42    ///
43    /// # Example
44    /// Note that this will very likely change in the future.
45    /// ```
46    /// use azalea_client::{Client, GameProfileComponent};
47    /// use bevy_ecs::query::With;
48    /// use azalea_entity::{Position, metadata::Player};
49    ///
50    /// # fn example(mut bot: Client, sender_name: String) {
51    /// let entity = bot.entity_by::<With<Player>, (&GameProfileComponent,)>(
52    ///     |(profile,): &(&GameProfileComponent,)| profile.name == sender_name,
53    /// );
54    /// if let Some(entity) = entity {
55    ///     let position = bot.entity_component::<Position>(entity);
56    ///     // ...
57    /// }
58    /// # }
59    /// ```
60    ///
61    /// [`Entity`]: bevy_ecs::entity::Entity
62    pub fn entity_by<F: QueryFilter, Q: QueryData>(
63        &mut self,
64        predicate: impl EntityPredicate<Q, F>,
65    ) -> Option<Entity> {
66        predicate.find(self.ecs.clone())
67    }
68
69    /// Get a component from an entity. Note that this will return an owned type
70    /// (i.e. not a reference) so it may be expensive for larger types.
71    ///
72    /// If you're trying to get a component for this client, use
73    /// [`Self::component`].
74    pub fn entity_component<Q: Component + Clone>(&mut self, entity: Entity) -> Q {
75        let mut ecs = self.ecs.lock();
76        let mut q = ecs.query::<&Q>();
77        let components = q.get(&ecs, entity).unwrap_or_else(|_| {
78            panic!(
79                "Entity is missing a required component {:?}",
80                std::any::type_name::<Q>()
81            )
82        });
83        components.clone()
84    }
85
86    /// Get a component from an entity, if it exists. This is similar to
87    /// [`Self::entity_component`] but returns an `Option` instead of panicking
88    /// if the component isn't present.
89    pub fn get_entity_component<Q: Component + Clone>(&mut self, entity: Entity) -> Option<Q> {
90        let mut ecs = self.ecs.lock();
91        let mut q = ecs.query::<&Q>();
92        let components = q.get(&ecs, entity).ok();
93        components.cloned()
94    }
95}
96
97pub trait EntityPredicate<Q: QueryData, Filter: QueryFilter> {
98    fn find(&self, ecs_lock: Arc<Mutex<World>>) -> Option<Entity>;
99}
100impl<F, Q, Filter> EntityPredicate<Q, Filter> for F
101where
102    F: Fn(&ROQueryItem<Q>) -> bool,
103    Q: QueryData,
104    Filter: QueryFilter,
105{
106    fn find(&self, ecs_lock: Arc<Mutex<World>>) -> Option<Entity> {
107        let mut ecs = ecs_lock.lock();
108        let mut query = ecs.query_filtered::<(Entity, Q), Filter>();
109        let entity = query.iter(&ecs).find(|(_, q)| (self)(q)).map(|(e, _)| e);
110
111        entity
112    }
113}
114
115// impl<'a, F, Q1, Q2> EntityPredicate<'a, (Q1, Q2)> for F
116// where
117//     F: Fn(&<Q1 as WorldQuery>::Item<'_>, &<Q2 as WorldQuery>::Item<'_>) ->
118// bool,     Q1: QueryFilter,
119//     Q2: QueryFilter,
120// {
121//     fn find(&self, ecs: &mut Ecs) -> Option<Entity> {
122//         // (self)(query)
123//         let mut query = ecs.query_filtered::<(Entity, Q1, Q2), ()>();
124//         let entity = query
125//             .iter(ecs)
126//             .find(|(_, q1, q2)| (self)(q1, q2))
127//             .map(|(e, _, _)| e);
128
129//         entity
130//     }
131// }