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::{
11    component::Component,
12    entity::Entity,
13    query::{Added, Changed, Without},
14    system::{Commands, Query, Res, ResMut, Resource},
15};
16use derive_more::{Deref, DerefMut};
17use nohash_hasher::IntMap;
18use tracing::{debug, trace, warn};
19use uuid::Uuid;
20
21use super::LoadedBy;
22use crate::{EntityUuid, LocalEntity, Position};
23
24#[derive(Resource, Default)]
25pub struct EntityUuidIndex {
26    /// An index of entities by their UUIDs
27    entity_by_uuid: HashMap<Uuid, Entity>,
28}
29impl EntityUuidIndex {
30    pub fn new() -> Self {
31        Self {
32            entity_by_uuid: HashMap::default(),
33        }
34    }
35
36    pub fn get(&self, uuid: &Uuid) -> Option<Entity> {
37        self.entity_by_uuid.get(uuid).copied()
38    }
39
40    pub fn contains_key(&self, uuid: &Uuid) -> bool {
41        self.entity_by_uuid.contains_key(uuid)
42    }
43
44    pub fn insert(&mut self, uuid: Uuid, entity: Entity) {
45        self.entity_by_uuid.insert(uuid, entity);
46    }
47
48    pub fn remove(&mut self, uuid: &Uuid) -> Option<Entity> {
49        self.entity_by_uuid.remove(uuid)
50    }
51}
52
53/// An index of Minecraft entity IDs to Azalea ECS entities. This is a
54/// `Component` so local players can keep track of entity IDs independently from
55/// the instance.
56///
57/// If you need a per-instance instead of per-client version of this, you can
58/// use [`Instance::entity_by_id`].
59#[derive(Component, Default)]
60pub struct EntityIdIndex {
61    /// An index of entities by their MinecraftEntityId
62    entity_by_id: IntMap<MinecraftEntityId, Entity>,
63    id_by_entity: HashMap<Entity, MinecraftEntityId>,
64}
65
66impl EntityIdIndex {
67    pub fn get_by_minecraft_entity(&self, id: MinecraftEntityId) -> Option<Entity> {
68        self.entity_by_id.get(&id).copied()
69    }
70    pub fn get_by_ecs_entity(&self, entity: Entity) -> Option<MinecraftEntityId> {
71        self.id_by_entity.get(&entity).copied()
72    }
73
74    pub fn contains_minecraft_entity(&self, id: MinecraftEntityId) -> bool {
75        self.entity_by_id.contains_key(&id)
76    }
77    pub fn contains_ecs_entity(&self, id: Entity) -> bool {
78        self.id_by_entity.contains_key(&id)
79    }
80
81    pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
82        self.entity_by_id.insert(id, entity);
83        self.id_by_entity.insert(entity, id);
84        trace!("Inserted {id} -> {entity:?} into a client's EntityIdIndex");
85    }
86
87    pub fn remove_by_minecraft_entity(&mut self, id: MinecraftEntityId) -> Option<Entity> {
88        if let Some(entity) = self.entity_by_id.remove(&id) {
89            trace!(
90                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove)"
91            );
92            self.id_by_entity.remove(&entity);
93            Some(entity)
94        } else {
95            trace!(
96                "Failed to remove {id} from a client's EntityIdIndex (using EntityIdIndex::remove)"
97            );
98            None
99        }
100    }
101
102    pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
103        if let Some(id) = self.id_by_entity.remove(&entity) {
104            trace!(
105                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity)."
106            );
107            self.entity_by_id.remove(&id);
108            Some(id)
109        } else {
110            // this is expected to happen when despawning entities if it was already
111            // despawned for another reason (like because the client received a
112            // remove_entities packet, or if we're in a shared instance where entity ids are
113            // different for each client)
114            trace!(
115                "Failed to remove {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity). This may be expected behavior."
116            );
117            None
118        }
119    }
120}
121
122impl Debug for EntityUuidIndex {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("EntityUuidIndex").finish()
125    }
126}
127
128/// The chunk position that an entity is currently in.
129#[derive(Component, Debug, Deref, DerefMut)]
130pub struct EntityChunkPos(pub ChunkPos);
131
132/// Update the chunk position indexes in [`Instance::entities_by_chunk`].
133///
134/// [`Instance::entities_by_chunk`]: azalea_world::Instance::entities_by_chunk
135pub fn update_entity_chunk_positions(
136    mut query: Query<(Entity, &Position, &InstanceName, &mut EntityChunkPos), Changed<Position>>,
137    instance_container: Res<InstanceContainer>,
138) {
139    for (entity, pos, world_name, mut entity_chunk_pos) in query.iter_mut() {
140        let instance_lock = instance_container.get(world_name).unwrap();
141        let mut instance = instance_lock.write();
142
143        let old_chunk = **entity_chunk_pos;
144        let new_chunk = ChunkPos::from(*pos);
145        if old_chunk != new_chunk {
146            **entity_chunk_pos = new_chunk;
147
148            if old_chunk != new_chunk {
149                // move the entity from the old chunk to the new one
150                if let Some(entities) = instance.entities_by_chunk.get_mut(&old_chunk) {
151                    entities.remove(&entity);
152                }
153                instance
154                    .entities_by_chunk
155                    .entry(new_chunk)
156                    .or_default()
157                    .insert(entity);
158                trace!("Entity {entity:?} moved from {old_chunk:?} to {new_chunk:?}");
159            }
160        }
161    }
162}
163
164/// Insert new entities into [`Instance::entities_by_chunk`].
165pub fn insert_entity_chunk_position(
166    query: Query<(Entity, &Position, &InstanceName), Added<EntityChunkPos>>,
167    instance_container: Res<InstanceContainer>,
168) {
169    for (entity, pos, world_name) in query.iter() {
170        let instance_lock = instance_container.get(world_name).unwrap();
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 has no references left, 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                debug!(
255                    "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found."
256                );
257            }
258        }
259        // remove it from the uuid index
260        if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
261            warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
262        }
263        if instance.entity_by_id.remove(minecraft_id).is_none() {
264            debug!(
265                "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."
266            );
267        }
268
269        // remove it from every client's EntityIdIndex
270        for mut entity_id_index in entity_id_index_query.iter_mut() {
271            entity_id_index.remove_by_ecs_entity(entity);
272        }
273
274        // and now remove the entity from the ecs
275        commands.entity(entity).despawn();
276        debug!("Despawned entity {entity:?} because it was not loaded by anything.");
277    }
278}
279
280pub fn add_entity_to_indexes(
281    entity_id: MinecraftEntityId,
282    ecs_entity: Entity,
283    entity_uuid: Option<Uuid>,
284    entity_id_index: &mut EntityIdIndex,
285    entity_uuid_index: &mut EntityUuidIndex,
286    instance: &mut Instance,
287) {
288    // per-client id index
289    entity_id_index.insert(entity_id, ecs_entity);
290
291    // per-instance id index
292    instance.entity_by_id.insert(entity_id, ecs_entity);
293
294    if let Some(uuid) = entity_uuid {
295        // per-instance uuid index
296        entity_uuid_index.insert(uuid, ecs_entity);
297    }
298}