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