azalea_client/
entity_query.rs

1use std::{any, sync::Arc};
2
3use azalea_core::position::Vec3;
4use azalea_entity::Position;
5use azalea_world::InstanceName;
6use bevy_ecs::{
7    component::Component,
8    entity::Entity,
9    query::{QueryData, QueryFilter, ROQueryItem},
10    world::World,
11};
12use parking_lot::Mutex;
13
14use crate::Client;
15
16impl Client {
17    /// A convenience function for getting components of our player's entity.
18    ///
19    /// # Examples
20    /// ```
21    /// # use azalea_world::InstanceName;
22    /// # fn example(mut client: azalea_client::Client) {
23    /// let is_logged_in = client
24    ///     .query::<Option<&InstanceName>>(&mut client.ecs.lock())
25    ///     .is_some();
26    /// # }
27    /// ```
28    pub fn query<'w, D: QueryData>(&self, ecs: &'w mut World) -> D::Item<'w> {
29        ecs.query::<D>()
30            .get_mut(ecs, self.entity)
31            .unwrap_or_else(|_| {
32                panic!(
33                    "Our client is missing a required component {:?}",
34                    any::type_name::<D>()
35                )
36            })
37    }
38
39    /// Quickly returns a lightweight [`Entity`] for an arbitrary entity that
40    /// matches the given predicate function that is in the same
41    /// [`Instance`] as the client.
42    ///
43    /// You can then use [`Self::entity_component`] to get components from this
44    /// entity.
45    ///
46    /// If you want to find the nearest entity, consider using
47    /// [`Self::nearest_entity_by`] instead. If you want to find all entities
48    /// that match the predicate, use [`Self::nearest_entities_by`].
49    ///
50    /// # Example
51    /// ```
52    /// use azalea_client::{Client, player::GameProfileComponent};
53    /// use azalea_entity::{Position, metadata::Player};
54    /// use bevy_ecs::query::With;
55    ///
56    /// # fn example(mut bot: Client, sender_name: String) {
57    /// let entity = bot.any_entity_by::<With<Player>, (&GameProfileComponent,)>(
58    ///     |(profile,): &(&GameProfileComponent,)| profile.name == sender_name,
59    /// );
60    /// if let Some(entity) = entity {
61    ///     let position = bot.entity_component::<Position>(entity);
62    ///     // ...
63    /// }
64    /// # }
65    /// ```
66    ///
67    /// [`Entity`]: bevy_ecs::entity::Entity
68    /// [`Instance`]: azalea_world::Instance
69    pub fn any_entity_by<F: QueryFilter, Q: QueryData>(
70        &self,
71        predicate: impl EntityPredicate<Q, F>,
72    ) -> Option<Entity> {
73        let instance_name = self.get_component::<InstanceName>()?;
74        predicate.find_any(self.ecs.clone(), &instance_name)
75    }
76
77    /// Return a lightweight [`Entity`] for the nearest entity that matches the
78    /// given predicate function.
79    ///
80    /// You can then use [`Self::entity_component`] to get components from this
81    /// entity.
82    ///
83    /// If you don't need the entity to be the nearest one, it may be more
84    /// efficient to use [`Self::any_entity_by`] instead. You can also use
85    /// [`Self::nearest_entities_by`] to get all nearby entities.
86    ///
87    /// ```
88    /// use azalea_entity::{LocalEntity, Position, metadata::Player};
89    /// use bevy_ecs::query::{With, Without};
90    ///
91    /// # fn example(mut bot: azalea_client::Client, sender_name: String) {
92    /// // get the position of the nearest player
93    /// if let Some(nearest_player) =
94    ///     bot.nearest_entity_by::<(With<Player>, Without<LocalEntity>), ()>(|_: &()| true)
95    /// {
96    ///     let nearest_player_pos = *bot.entity_component::<Position>(nearest_player);
97    ///     bot.chat(format!("You are at {nearest_player_pos}"));
98    /// }
99    /// # }
100    /// ```
101    ///
102    /// [`Entity`]: bevy_ecs::entity::Entity
103    pub fn nearest_entity_by<F: QueryFilter, Q: QueryData>(
104        &self,
105        predicate: impl EntityPredicate<Q, F>,
106    ) -> Option<Entity> {
107        self.nearest_entities_by(predicate).first().copied()
108    }
109
110    /// Similar to [`Self::nearest_entity_by`] but returns a `Vec<Entity>` of
111    /// all entities in our instance that match the predicate.
112    ///
113    /// The first entity is the nearest one.
114    ///
115    /// ```
116    /// # use azalea_entity::{LocalEntity, Position, metadata::Player};
117    /// # use bevy_ecs::query::{With, Without};
118    /// # fn example(mut bot: azalea_client::Client, sender_name: String) {
119    /// let nearby_players =
120    ///     bot.nearest_entities_by::<(With<Player>, Without<LocalEntity>), ()>(|_: &()| true);
121    /// # }
122    /// ```
123    pub fn nearest_entities_by<F: QueryFilter, Q: QueryData>(
124        &self,
125        predicate: impl EntityPredicate<Q, F>,
126    ) -> Vec<Entity> {
127        let Some(instance_name) = self.get_component::<InstanceName>() else {
128            return vec![];
129        };
130        let Some(position) = self.get_component::<Position>() else {
131            return vec![];
132        };
133        predicate.find_all_sorted(self.ecs.clone(), &instance_name, (&position).into())
134    }
135
136    /// Get a component from an entity. Note that this will return an owned type
137    /// (i.e. not a reference) so it may be expensive for larger types.
138    ///
139    /// If you're trying to get a component for this client, use
140    /// [`Self::component`].
141    pub fn entity_component<Q: Component + Clone>(&self, entity: Entity) -> Q {
142        let mut ecs = self.ecs.lock();
143        let mut q = ecs.query::<&Q>();
144        let components = q.get(&ecs, entity).unwrap_or_else(|_| {
145            panic!(
146                "Entity is missing a required component {:?}",
147                any::type_name::<Q>()
148            )
149        });
150        components.clone()
151    }
152
153    /// Get a component from an entity, if it exists. This is similar to
154    /// [`Self::entity_component`] but returns an `Option` instead of panicking
155    /// if the component isn't present.
156    pub fn get_entity_component<Q: Component + Clone>(&self, entity: Entity) -> Option<Q> {
157        let mut ecs = self.ecs.lock();
158        let mut q = ecs.query::<&Q>();
159        let components = q.get(&ecs, entity).ok();
160        components.cloned()
161    }
162}
163
164pub trait EntityPredicate<Q: QueryData, Filter: QueryFilter> {
165    fn find_any(&self, ecs_lock: Arc<Mutex<World>>, instance_name: &InstanceName)
166    -> Option<Entity>;
167    fn find_all_sorted(
168        &self,
169        ecs_lock: Arc<Mutex<World>>,
170        instance_name: &InstanceName,
171        nearest_to: Vec3,
172    ) -> Vec<Entity>;
173}
174impl<F, Q: QueryData, Filter: QueryFilter> EntityPredicate<Q, Filter> for F
175where
176    F: Fn(&ROQueryItem<Q>) -> bool,
177{
178    fn find_any(
179        &self,
180        ecs_lock: Arc<Mutex<World>>,
181        instance_name: &InstanceName,
182    ) -> Option<Entity> {
183        let mut ecs = ecs_lock.lock();
184        let mut query = ecs.query_filtered::<(Entity, &InstanceName, Q), Filter>();
185        query
186            .iter(&ecs)
187            .find(|(_, e_instance_name, q)| *e_instance_name == instance_name && (self)(q))
188            .map(|(e, _, _)| e)
189    }
190
191    fn find_all_sorted(
192        &self,
193        ecs_lock: Arc<Mutex<World>>,
194        instance_name: &InstanceName,
195        nearest_to: Vec3,
196    ) -> Vec<Entity> {
197        let mut ecs = ecs_lock.lock();
198        let mut query = ecs.query_filtered::<(Entity, &InstanceName, &Position, Q), Filter>();
199        let mut entities = query
200            .iter(&ecs)
201            .filter(|(_, e_instance_name, _, q)| *e_instance_name == instance_name && (self)(q))
202            .map(|(e, _, position, _)| (e, Vec3::from(position)))
203            .collect::<Vec<(Entity, Vec3)>>();
204
205        entities.sort_by_cached_key(|(_, position)| {
206            // to_bits is fine here as long as the number is positive
207            position.distance_squared_to(nearest_to).to_bits()
208        });
209
210        entities
211            .into_iter()
212            .map(|(e, _)| e)
213            .collect::<Vec<Entity>>()
214    }
215}