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::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 std::fmt::Formatter<'_>) -> std::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, world_name, mut entity_chunk_pos) in query.iter_mut() {
135        let instance_lock = instance_container.get(world_name).unwrap();
136        let mut instance = instance_lock.write();
137
138        let old_chunk = **entity_chunk_pos;
139        let new_chunk = ChunkPos::from(*pos);
140        if old_chunk != new_chunk {
141            **entity_chunk_pos = new_chunk;
142
143            if old_chunk != new_chunk {
144                // move the entity from the old chunk to the new one
145                if let Some(entities) = instance.entities_by_chunk.get_mut(&old_chunk) {
146                    entities.remove(&entity);
147                }
148                instance
149                    .entities_by_chunk
150                    .entry(new_chunk)
151                    .or_default()
152                    .insert(entity);
153                trace!("Entity {entity:?} moved from {old_chunk:?} to {new_chunk:?}");
154            }
155        }
156    }
157}
158
159/// Insert new entities into [`Instance::entities_by_chunk`].
160pub fn insert_entity_chunk_position(
161    query: Query<(Entity, &Position, &InstanceName), Added<EntityChunkPos>>,
162    instance_container: Res<InstanceContainer>,
163) {
164    for (entity, pos, world_name) in query.iter() {
165        let instance_lock = instance_container.get(world_name).unwrap();
166        let mut instance = instance_lock.write();
167
168        let chunk = ChunkPos::from(*pos);
169        instance
170            .entities_by_chunk
171            .entry(chunk)
172            .or_default()
173            .insert(entity);
174    }
175}
176
177/// Despawn entities that aren't being loaded by anything.
178#[allow(clippy::type_complexity)]
179pub fn remove_despawned_entities_from_indexes(
180    mut commands: Commands,
181    mut entity_uuid_index: ResMut<EntityUuidIndex>,
182    instance_container: Res<InstanceContainer>,
183    query: Query<
184        (
185            Entity,
186            &EntityUuid,
187            &MinecraftEntityId,
188            &Position,
189            &InstanceName,
190            &LoadedBy,
191        ),
192        (Changed<LoadedBy>, Without<LocalEntity>),
193    >,
194    mut entity_id_index_query: Query<&mut EntityIdIndex>,
195) {
196    for (entity, uuid, minecraft_id, position, instance_name, loaded_by) in &query {
197        let Some(instance_lock) = instance_container.get(instance_name) else {
198            // the instance isn't even loaded by us, so we can safely delete the entity
199            debug!(
200                "Despawned entity {entity:?} because it's in an instance that isn't loaded anymore"
201            );
202            if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
203                warn!(
204                    "Tried to remove entity {entity:?} from the uuid index but it was not there."
205                );
206            }
207            // and now remove the entity from the ecs
208            commands.entity(entity).despawn();
209
210            continue;
211        };
212
213        let mut instance = instance_lock.write();
214
215        // if the entity has no references left, despawn it
216        if !loaded_by.is_empty() {
217            continue;
218        }
219
220        // remove the entity from the chunk index
221        let chunk = ChunkPos::from(*position);
222        match instance.entities_by_chunk.get_mut(&chunk) {
223            Some(entities_in_chunk) => {
224                if entities_in_chunk.remove(&entity) {
225                    // remove the chunk if there's no entities in it anymore
226                    if entities_in_chunk.is_empty() {
227                        instance.entities_by_chunk.remove(&chunk);
228                    }
229                } else {
230                    // search all the other chunks for it :(
231                    let mut found_in_other_chunks = HashSet::new();
232                    for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
233                        if entities_in_other_chunk.remove(&entity) {
234                            found_in_other_chunks.insert(other_chunk);
235                        }
236                    }
237                    if found_in_other_chunks.is_empty() {
238                        warn!(
239                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there or in any other chunks."
240                        );
241                    } else {
242                        warn!(
243                            "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:?}"
244                        );
245                    }
246                }
247            }
248            _ => {
249                debug!(
250                    "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found."
251                );
252            }
253        }
254        // remove it from the uuid index
255        if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
256            warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
257        }
258        if instance.entity_by_id.remove(minecraft_id).is_none() {
259            debug!(
260                "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."
261            );
262        }
263
264        // remove it from every client's EntityIdIndex
265        for mut entity_id_index in entity_id_index_query.iter_mut() {
266            entity_id_index.remove_by_ecs_entity(entity);
267        }
268
269        // and now remove the entity from the ecs
270        commands.entity(entity).despawn();
271        debug!("Despawned entity {entity:?} because it was not loaded by anything.");
272    }
273}
274
275pub fn add_entity_to_indexes(
276    entity_id: MinecraftEntityId,
277    ecs_entity: Entity,
278    entity_uuid: Option<Uuid>,
279    entity_id_index: &mut EntityIdIndex,
280    entity_uuid_index: &mut EntityUuidIndex,
281    instance: &mut Instance,
282) {
283    // per-client id index
284    entity_id_index.insert(entity_id, ecs_entity);
285
286    // per-instance id index
287    instance.entity_by_id.insert(entity_id, ecs_entity);
288
289    if let Some(uuid) = entity_uuid {
290        // per-instance uuid index
291        entity_uuid_index.insert(uuid, ecs_entity);
292    }
293}