azalea_entity/plugin/
indexing.rs

1//! Stuff related to entity indexes and keeping track of entities in the world.
2
3use std::{
4    collections::{HashMap, HashSet},
5    fmt::{self, Debug},
6};
7
8use azalea_core::position::ChunkPos;
9use azalea_world::{Instance, InstanceContainer, InstanceName, MinecraftEntityId};
10use bevy_ecs::prelude::*;
11use derive_more::{Deref, DerefMut};
12use nohash_hasher::IntMap;
13use tracing::{debug, trace, warn};
14use uuid::Uuid;
15
16use super::LoadedBy;
17use crate::{EntityUuid, LocalEntity, Position};
18
19#[derive(Resource, Default)]
20pub struct EntityUuidIndex {
21    /// An index of entities by their UUIDs
22    entity_by_uuid: HashMap<Uuid, Entity>,
23}
24impl EntityUuidIndex {
25    pub fn new() -> Self {
26        Self {
27            entity_by_uuid: HashMap::default(),
28        }
29    }
30
31    pub fn get(&self, uuid: &Uuid) -> Option<Entity> {
32        self.entity_by_uuid.get(uuid).copied()
33    }
34
35    pub fn contains_key(&self, uuid: &Uuid) -> bool {
36        self.entity_by_uuid.contains_key(uuid)
37    }
38
39    pub fn insert(&mut self, uuid: Uuid, entity: Entity) {
40        self.entity_by_uuid.insert(uuid, entity);
41    }
42
43    pub fn remove(&mut self, uuid: &Uuid) -> Option<Entity> {
44        self.entity_by_uuid.remove(uuid)
45    }
46}
47
48/// An index of Minecraft entity IDs to Azalea ECS entities. This is a
49/// `Component` so local players can keep track of entity IDs independently from
50/// the instance.
51///
52/// If you need a per-instance instead of per-client version of this, you can
53/// use [`Instance::entity_by_id`].
54#[derive(Component, Default)]
55pub struct EntityIdIndex {
56    /// An index of entities by their MinecraftEntityId
57    entity_by_id: IntMap<MinecraftEntityId, Entity>,
58    id_by_entity: HashMap<Entity, MinecraftEntityId>,
59}
60
61impl EntityIdIndex {
62    pub fn get_by_minecraft_entity(&self, id: MinecraftEntityId) -> Option<Entity> {
63        self.entity_by_id.get(&id).copied()
64    }
65    pub fn get_by_ecs_entity(&self, entity: Entity) -> Option<MinecraftEntityId> {
66        self.id_by_entity.get(&entity).copied()
67    }
68
69    pub fn contains_minecraft_entity(&self, id: MinecraftEntityId) -> bool {
70        self.entity_by_id.contains_key(&id)
71    }
72    pub fn contains_ecs_entity(&self, id: Entity) -> bool {
73        self.id_by_entity.contains_key(&id)
74    }
75
76    pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
77        self.entity_by_id.insert(id, entity);
78        self.id_by_entity.insert(entity, id);
79        trace!("Inserted {id} -> {entity:?} into a client's EntityIdIndex");
80    }
81
82    pub fn remove_by_minecraft_entity(&mut self, id: MinecraftEntityId) -> Option<Entity> {
83        if let Some(entity) = self.entity_by_id.remove(&id) {
84            trace!(
85                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove)"
86            );
87            self.id_by_entity.remove(&entity);
88            Some(entity)
89        } else {
90            trace!(
91                "Failed to remove {id} from a client's EntityIdIndex (using EntityIdIndex::remove)"
92            );
93            None
94        }
95    }
96
97    pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
98        if let Some(id) = self.id_by_entity.remove(&entity) {
99            trace!(
100                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity)."
101            );
102            self.entity_by_id.remove(&id);
103            Some(id)
104        } else {
105            // this is expected to happen when despawning entities if it was already
106            // despawned for another reason (like because the client received a
107            // remove_entities packet, or if we're in a shared instance where entity ids are
108            // different for each client)
109            trace!(
110                "Failed to remove {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity). This may be expected behavior."
111            );
112            None
113        }
114    }
115}
116
117impl Debug for EntityUuidIndex {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        f.debug_struct("EntityUuidIndex").finish()
120    }
121}
122
123/// The chunk position that an entity is currently in.
124#[derive(Component, Debug, Deref, DerefMut)]
125pub struct EntityChunkPos(pub ChunkPos);
126
127/// Update the chunk position indexes in [`Instance::entities_by_chunk`].
128///
129/// [`Instance::entities_by_chunk`]: azalea_world::Instance::entities_by_chunk
130pub fn update_entity_chunk_positions(
131    mut query: Query<(Entity, &Position, &InstanceName, &mut EntityChunkPos), Changed<Position>>,
132    instance_container: Res<InstanceContainer>,
133) {
134    for (entity, pos, instance_name, mut entity_chunk_pos) in query.iter_mut() {
135        let old_chunk = **entity_chunk_pos;
136        let new_chunk = ChunkPos::from(*pos);
137        if old_chunk != new_chunk {
138            **entity_chunk_pos = new_chunk;
139
140            if old_chunk != new_chunk {
141                let Some(instance_lock) = instance_container.get(instance_name) else {
142                    continue;
143                };
144                let mut instance = instance_lock.write();
145
146                // move the entity from the old chunk to the new one
147                if let Some(entities) = instance.entities_by_chunk.get_mut(&old_chunk) {
148                    entities.remove(&entity);
149                }
150                instance
151                    .entities_by_chunk
152                    .entry(new_chunk)
153                    .or_default()
154                    .insert(entity);
155                trace!("Entity {entity:?} moved from {old_chunk:?} to {new_chunk:?}");
156            }
157        }
158    }
159}
160
161/// Insert new entities into [`Instance::entities_by_chunk`].
162pub fn insert_entity_chunk_position(
163    query: Query<(Entity, &Position, &InstanceName), Added<EntityChunkPos>>,
164    instance_container: Res<InstanceContainer>,
165) {
166    for (entity, pos, world_name) in query.iter() {
167        let Some(instance_lock) = instance_container.get(world_name) else {
168            // entity must've been despawned already
169            continue;
170        };
171        let mut instance = instance_lock.write();
172
173        let chunk = ChunkPos::from(*pos);
174        instance
175            .entities_by_chunk
176            .entry(chunk)
177            .or_default()
178            .insert(entity);
179    }
180}
181
182/// Despawn entities that aren't being loaded by anything.
183#[allow(clippy::type_complexity)]
184pub fn remove_despawned_entities_from_indexes(
185    mut commands: Commands,
186    mut entity_uuid_index: ResMut<EntityUuidIndex>,
187    instance_container: Res<InstanceContainer>,
188    query: Query<
189        (
190            Entity,
191            &EntityUuid,
192            &MinecraftEntityId,
193            &Position,
194            &InstanceName,
195            &LoadedBy,
196        ),
197        (Changed<LoadedBy>, Without<LocalEntity>),
198    >,
199    mut entity_id_index_query: Query<&mut EntityIdIndex>,
200) {
201    for (entity, uuid, minecraft_id, position, instance_name, loaded_by) in &query {
202        let Some(instance_lock) = instance_container.get(instance_name) else {
203            // the instance isn't even loaded by us, so we can safely delete the entity
204            debug!(
205                "Despawned entity {entity:?} because it's in an instance that isn't loaded anymore"
206            );
207            if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
208                warn!(
209                    "Tried to remove entity {entity:?} from the uuid index but it was not there."
210                );
211            }
212            // and now remove the entity from the ecs
213            commands.entity(entity).despawn();
214
215            continue;
216        };
217
218        let mut instance = instance_lock.write();
219
220        // if the entity is being loaded by any of our clients, don't despawn it
221        if !loaded_by.is_empty() {
222            continue;
223        }
224
225        // remove the entity from the chunk index
226        let chunk = ChunkPos::from(position);
227        match instance.entities_by_chunk.get_mut(&chunk) {
228            Some(entities_in_chunk) => {
229                if entities_in_chunk.remove(&entity) {
230                    // remove the chunk if there's no entities in it anymore
231                    if entities_in_chunk.is_empty() {
232                        instance.entities_by_chunk.remove(&chunk);
233                    }
234                } else {
235                    // search all the other chunks for it :(
236                    let mut found_in_other_chunks = HashSet::new();
237                    for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
238                        if entities_in_other_chunk.remove(&entity) {
239                            found_in_other_chunks.insert(other_chunk);
240                        }
241                    }
242                    if found_in_other_chunks.is_empty() {
243                        warn!(
244                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there or in any other chunks."
245                        );
246                    } else {
247                        warn!(
248                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there. Found in and removed from other chunk(s): {found_in_other_chunks:?}"
249                        );
250                    }
251                }
252            }
253            _ => {
254                let mut found_in_other_chunks = HashSet::new();
255                for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
256                    if entities_in_other_chunk.remove(&entity) {
257                        found_in_other_chunks.insert(other_chunk);
258                    }
259                }
260                if found_in_other_chunks.is_empty() {
261                    warn!(
262                        "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found and the entity wasn't in any other chunks."
263                    );
264                } else {
265                    warn!(
266                        "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found. Entity found in and removed from other chunk(s): {found_in_other_chunks:?}"
267                    );
268                }
269            }
270        }
271        // remove it from the uuid index
272        if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
273            warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
274        }
275        if instance.entity_by_id.remove(minecraft_id).is_none() {
276            debug!(
277                "Tried to remove entity {entity:?} from the per-instance entity id index but it was not there. This may be expected if you're in a shared instance."
278            );
279        }
280
281        // remove it from every client's EntityIdIndex
282        for mut entity_id_index in entity_id_index_query.iter_mut() {
283            entity_id_index.remove_by_ecs_entity(entity);
284        }
285
286        // and now remove the entity from the ecs
287        commands.entity(entity).despawn();
288        debug!("Despawned entity {entity:?} because it was not loaded by anything.");
289    }
290}
291
292pub fn add_entity_to_indexes(
293    entity_id: MinecraftEntityId,
294    ecs_entity: Entity,
295    entity_uuid: Option<Uuid>,
296    entity_id_index: &mut EntityIdIndex,
297    entity_uuid_index: &mut EntityUuidIndex,
298    instance: &mut Instance,
299) {
300    // per-client id index
301    entity_id_index.insert(entity_id, ecs_entity);
302
303    // per-instance id index
304    instance.entity_by_id.insert(entity_id, ecs_entity);
305
306    if let Some(uuid) = entity_uuid {
307        // per-instance uuid index
308        entity_uuid_index.insert(uuid, ecs_entity);
309    }
310}