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}