aboutsummaryrefslogtreecommitdiff
path: root/azalea-client/src/entity_query.rs
blob: 5b19bb248f540d39993253c0ae093671b6d95b03 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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, QueryFilter, ROQueryItem},
    world::World,
};
use parking_lot::Mutex;

use crate::Client;

impl Client {
    /// A convenience function for getting components of our player's entity.
    ///
    /// # Examples
    /// ```
    /// # use azalea_world::InstanceName;
    /// # fn example(mut client: azalea_client::Client) {
    /// let is_logged_in = client
    ///     .query::<Option<&InstanceName>>(&mut client.ecs.lock())
    ///     .is_some();
    /// # }
    /// ```
    pub fn query<'w, D: QueryData>(&self, ecs: &'w mut World) -> D::Item<'w> {
        ecs.query::<D>()
            .get_mut(ecs, self.entity)
            .unwrap_or_else(|_| {
                panic!(
                    "Our client is missing a required component {:?}",
                    any::type_name::<D>()
                )
            })
    }

    /// Return 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.
    ///
    /// Also see [`Self::entities_by`] which will return all entities that match
    /// the predicate and sorts them by distance (unlike `entity_by`).
    ///
    /// # Example
    /// ```
    /// use azalea_client::{Client, player::GameProfileComponent};
    /// use azalea_entity::{Position, metadata::Player};
    /// use bevy_ecs::query::With;
    ///
    /// # fn example(mut bot: Client, sender_name: String) {
    /// let entity = bot.entity_by::<With<Player>, (&GameProfileComponent,)>(
    ///     |(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 entity_by<F: QueryFilter, Q: QueryData>(
        &self,
        predicate: impl EntityPredicate<Q, F>,
    ) -> Option<Entity> {
        let instance_name = self.get_component::<InstanceName>()?;
        predicate.find_any(self.ecs.clone(), &instance_name)
    }

    /// Similar to [`Self::entity_by`] but returns a `Vec<Entity>` of all
    /// entities in our instance that match the predicate.
    ///
    /// Unlike `entity_by`, the result is sorted by distance to our client's
    /// position, so the closest entity is first.
    pub fn entities_by<F: QueryFilter, Q: QueryData>(
        &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,
{
    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>>()
    }
}