1 files changed, 264 insertions, 0 deletions
diff --git a/azalea/src/client_impl/entity_query.rs b/azalea/src/client_impl/entity_query.rs
new file mode 100644
index 00000000..85e46525
--- /dev/null
+++ b/azalea/src/client_impl/entity_query.rs
@@ -0,0 +1,264 @@
+use std::{any, sync::Arc};
+
+use azalea_core::position::Vec3;
+use azalea_entity::Position;
+use azalea_world::InstanceName;
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    query::{QueryData, QueryEntityError, QueryFilter, QueryItem, ROQueryItem},
+    world::World,
+};
+use parking_lot::Mutex;
+
+use crate::Client;
+
+impl Client {
+    /// A convenience function for getting components from our client's entity.
+    ///
+    /// To query another entity, you can use [`Self::query_entity`].
+    ///
+    /// # Examples
+    /// ```
+    /// # use azalea_world::InstanceName;
+    /// # fn example(mut client: azalea::Client) {
+    /// let is_logged_in = client.query_self::<Option<&InstanceName>, _>(|ins| ins.is_some());
+    /// # }
+    /// ```
+    ///
+    /// # Panics
+    ///
+    /// This will panic if the component doesn't exist on the client.
+    pub fn query_self<D: QueryData, R>(&self, f: impl FnOnce(QueryItem<D>) -> R) -> R {
+        let mut ecs = self.ecs.lock();
+        let mut qs = ecs.query::<D>();
+        let res = qs.get_mut(&mut ecs, self.entity).unwrap_or_else(|_| {
+            panic!(
+                "Our client is missing a required component {:?}",
+                any::type_name::<D>()
+            )
+        });
+        f(res)
+    }
+
+    /// A convenience function for getting components from any entity.
+    ///
+    /// If you're querying the client, you should use [`Self::query_self`].
+    ///
+    /// # Panics
+    ///
+    /// This will panic if the entity doesn't exist or if the query isn't valid
+    /// for the entity. For a non-panicking version, you may use
+    /// [`Self::try_query_entity`].
+    pub fn query_entity<D: QueryData, R>(
+        &self,
+        entity: Entity,
+        f: impl FnOnce(QueryItem<D>) -> R,
+    ) -> R {
+        self.try_query_entity(entity, f).unwrap_or_else(|_| {
+            panic!(
+                "Entity is missing a required component {:?}",
+                any::type_name::<D>()
+            )
+        })
+    }
+
+    /// A convenience function for getting components from any entity, or None
+    /// if the query fails.
+    ///
+    /// If you're sure that the entity exists and that the query will succeed,
+    /// you can use [`Self::query_entity`].
+    pub fn try_query_entity<D: QueryData, R>(
+        &self,
+        entity: Entity,
+        f: impl FnOnce(QueryItem<D>) -> R,
+    ) -> Result<R, QueryEntityError> {
+        let mut ecs = self.ecs.lock();
+        let mut qs = ecs.query::<D>();
+        qs.get_mut(&mut ecs, entity).map(f)
+    }
+
+    /// Quickly returns a lightweight [`Entity`] for an arbitrary entity that
+    /// matches the given predicate function that is in the same
+    /// [`Instance`] as the client.
+    ///
+    /// You can then use [`Self::entity_component`] to get components from this
+    /// entity.
+    ///
+    /// If you want to find the nearest entity, consider using
+    /// [`Self::nearest_entity_by`] instead. If you want to find all entities
+    /// that match the predicate, use [`Self::nearest_entities_by`].
+    ///
+    /// # Example
+    /// ```
+    /// use azalea::{
+    ///     Client,
+    ///     entity::{Position, metadata::Player},
+    ///     player::GameProfileComponent,
+    /// };
+    /// use bevy_ecs::query::With;
+    ///
+    /// # fn example(mut bot: Client, sender_name: String) {
+    /// let entity = bot.any_entity_by::<&GameProfileComponent, With<Player>>(
+    ///     |profile: &GameProfileComponent| profile.name == sender_name,
+    /// );
+    /// if let Some(entity) = entity {
+    ///     let position = bot.entity_component::<Position>(entity);
+    ///     // ...
+    /// }
+    /// # }
+    /// ```
+    ///
+    /// [`Entity`]: bevy_ecs::entity::Entity
+    /// [`Instance`]: azalea_world::Instance
+    pub fn any_entity_by<Q: QueryData, F: QueryFilter>(
+        &self,
+        predicate: impl EntityPredicate<Q, F>,
+    ) -> Option<Entity> {
+        let instance_name = self.get_component::<InstanceName>()?;
+        predicate.find_any(self.ecs.clone(), &instance_name)
+    }
+
+    /// Return a lightweight [`Entity`] for the nearest entity that matches the
+    /// given predicate function.
+    ///
+    /// You can then use [`Self::entity_component`] to get components from this
+    /// entity.
+    ///
+    /// If you don't need the entity to be the nearest one, it may be more
+    /// efficient to use [`Self::any_entity_by`] instead. You can also use
+    /// [`Self::nearest_entities_by`] to get all nearby entities.
+    ///
+    /// ```
+    /// use azalea_entity::{LocalEntity, Position, metadata::Player};
+    /// use bevy_ecs::query::{With, Without};
+    ///
+    /// # fn example(mut bot: azalea::Client, sender_name: String) {
+    /// // get the position of the nearest player
+    /// if let Some(nearest_player) =
+    ///     bot.nearest_entity_by::<(), (With<Player>, Without<LocalEntity>)>(|_: ()| true)
+    /// {
+    ///     let nearest_player_pos = *bot.entity_component::<Position>(nearest_player);
+    ///     bot.chat(format!("You are at {nearest_player_pos}"));
+    /// }
+    /// # }
+    /// ```
+    ///
+    /// [`Entity`]: bevy_ecs::entity::Entity
+    pub fn nearest_entity_by<Q: QueryData, F: QueryFilter>(
+        &self,
+        predicate: impl EntityPredicate<Q, F>,
+    ) -> Option<Entity> {
+        self.nearest_entities_by(predicate).first().copied()
+    }
+
+    /// Similar to [`Self::nearest_entity_by`] but returns a `Vec<Entity>` of
+    /// all entities in our instance that match the predicate.
+    ///
+    /// The first entity is the nearest one.
+    ///
+    /// ```
+    /// # use azalea_entity::{LocalEntity, Position, metadata::Player};
+    /// # use bevy_ecs::query::{With, Without};
+    /// # fn example(mut bot: azalea::Client, sender_name: String) {
+    /// let nearby_players =
+    ///     bot.nearest_entities_by::<(), (With<Player>, Without<LocalEntity>)>(|_: ()| true);
+    /// # }
+    /// ```
+    pub fn nearest_entities_by<Q: QueryData, F: QueryFilter>(
+        &self,
+        predicate: impl EntityPredicate<Q, F>,
+    ) -> Vec<Entity> {
+        let Some(instance_name) = self.get_component::<InstanceName>() else {
+            return vec![];
+        };
+        let Some(position) = self.get_component::<Position>() else {
+            return vec![];
+        };
+        predicate.find_all_sorted(self.ecs.clone(), &instance_name, (&position).into())
+    }
+
+    /// Get a component from an entity.
+    ///
+    /// Note that this will return an owned type (i.e. not a reference) so it
+    /// may be expensive for larger types.
+    ///
+    /// If you're trying to get a component for this client, use
+    /// [`Self::component`].
+    pub fn entity_component<Q: Component + Clone>(&self, entity: Entity) -> Q {
+        let mut ecs = self.ecs.lock();
+        let mut q = ecs.query::<&Q>();
+        let components = q.get(&ecs, entity).unwrap_or_else(|_| {
+            panic!(
+                "Entity is missing a required component {:?}",
+                any::type_name::<Q>()
+            )
+        });
+        components.clone()
+    }
+
+    /// Get a component from an entity, if it exists.
+    ///
+    /// This is similar to [`Self::entity_component`] but returns an `Option`
+    /// instead of panicking if the component isn't present.
+    pub fn get_entity_component<Q: Component + Clone>(&self, entity: Entity) -> Option<Q> {
+        let mut ecs = self.ecs.lock();
+        let mut q = ecs.query::<&Q>();
+        let components = q.get(&ecs, entity).ok();
+        components.cloned()
+    }
+}
+
+pub trait EntityPredicate<Q: QueryData, Filter: QueryFilter> {
+    fn find_any(&self, ecs_lock: Arc<Mutex<World>>, instance_name: &InstanceName)
+    -> Option<Entity>;
+    fn find_all_sorted(
+        &self,
+        ecs_lock: Arc<Mutex<World>>,
+        instance_name: &InstanceName,
+        nearest_to: Vec3,
+    ) -> Vec<Entity>;
+}
+impl<F, Q: QueryData, Filter: QueryFilter> EntityPredicate<Q, Filter> for F
+where
+    F: Fn(ROQueryItem<Q>) -> bool,
+    for<'w, 's> <<Q as QueryData>::ReadOnly as QueryData>::Item<'w, 's>: Copy,
+{
+    fn find_any(
+        &self,
+        ecs_lock: Arc<Mutex<World>>,
+        instance_name: &InstanceName,
+    ) -> Option<Entity> {
+        let mut ecs = ecs_lock.lock();
+        let mut query = ecs.query_filtered::<(Entity, &InstanceName, Q), Filter>();
+        query
+            .iter(&ecs)
+            .find(|(_, e_instance_name, q)| *e_instance_name == instance_name && (self)(*q))
+            .map(|(e, _, _)| e)
+    }
+
+    fn find_all_sorted(
+        &self,
+        ecs_lock: Arc<Mutex<World>>,
+        instance_name: &InstanceName,
+        nearest_to: Vec3,
+    ) -> Vec<Entity> {
+        let mut ecs = ecs_lock.lock();
+        let mut query = ecs.query_filtered::<(Entity, &InstanceName, &Position, Q), Filter>();
+        let mut entities = query
+            .iter(&ecs)
+            .filter(|(_, e_instance_name, _, q)| *e_instance_name == instance_name && (self)(*q))
+            .map(|(e, _, position, _)| (e, Vec3::from(position)))
+            .collect::<Vec<(Entity, Vec3)>>();
+
+        entities.sort_by_cached_key(|(_, position)| {
+            // to_bits is fine here as long as the number is positive
+            position.distance_squared_to(nearest_to).to_bits()
+        });
+
+        entities
+            .into_iter()
+            .map(|(e, _)| e)
+            .collect::<Vec<Entity>>()
+    }
+}