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(&self, id: MinecraftEntityId) -> Option<Entity> {
68        self.entity_by_id.get(&id).copied()
69    }
70
71    pub fn contains_key(&self, id: MinecraftEntityId) -> bool {
72        self.entity_by_id.contains_key(&id)
73    }
74    pub fn contains_ecs_entity(&self, id: Entity) -> bool {
75        self.id_by_entity.contains_key(&id)
76    }
77
78    pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
79        self.entity_by_id.insert(id, entity);
80        self.id_by_entity.insert(entity, id);
81        trace!("Inserted {id} -> {entity:?} into a client's EntityIdIndex");
82    }
83
84    pub fn remove(&mut self, id: MinecraftEntityId) -> Option<Entity> {
85        if let Some(entity) = self.entity_by_id.remove(&id) {
86            trace!(
87                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove)"
88            );
89            self.id_by_entity.remove(&entity);
90            Some(entity)
91        } else {
92            trace!(
93                "Failed to remove {id} from a client's EntityIdIndex (using EntityIdIndex::remove)"
94            );
95            None
96        }
97    }
98
99    pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
100        if let Some(id) = self.id_by_entity.remove(&entity) {
101            trace!(
102                "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity)."
103            );
104            self.entity_by_id.remove(&id);
105            Some(id)
106        } else {
107            // this is expected to happen when despawning entities if it was already
108            // despawned for another reason (like because the client received a
109            // remove_entities packet, or if we're in a shared instance where entity ids are
110            // different for each client)
111            trace!(
112                "Failed to remove {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity). This may be expected behavior."
113            );
114            None
115        }
116    }
117}
118
119impl Debug for EntityUuidIndex {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        f.debug_struct("EntityUuidIndex").finish()
122    }
123}
124
125/// The chunk position that an entity is currently in.
126#[derive(Component, Debug, Deref, DerefMut)]
127pub struct EntityChunkPos(pub ChunkPos);
128
129/// Update the chunk position indexes in [`Instance::entities_by_chunk`].
130///
131/// [`Instance::entities_by_chunk`]: azalea_world::Instance::entities_by_chunk
132pub fn update_entity_chunk_positions(
133    mut query: Query<(Entity, &Position, &InstanceName, &mut EntityChunkPos), Changed<Position>>,
134    instance_container: Res<InstanceContainer>,
135) {
136    for (entity, pos, world_name, mut entity_chunk_pos) in query.iter_mut() {
137        let instance_lock = instance_container.get(world_name).unwrap();
138        let mut instance = instance_lock.write();
139
140        let old_chunk = **entity_chunk_pos;
141        let new_chunk = ChunkPos::from(*pos);
142        if old_chunk != new_chunk {
143            **entity_chunk_pos = new_chunk;
144
145            if old_chunk != new_chunk {
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 instance_lock = instance_container.get(world_name).unwrap();
168        let mut instance = instance_lock.write();
169
170        let chunk = ChunkPos::from(*pos);
171        instance
172            .entities_by_chunk
173            .entry(chunk)
174            .or_default()
175            .insert(entity);
176    }
177}
178
179/// Despawn entities that aren't being loaded by anything.
180#[allow(clippy::type_complexity)]
181pub fn remove_despawned_entities_from_indexes(
182    mut commands: Commands,
183    mut entity_uuid_index: ResMut<EntityUuidIndex>,
184    instance_container: Res<InstanceContainer>,
185    query: Query<
186        (
187            Entity,
188            &EntityUuid,
189            &MinecraftEntityId,
190            &Position,
191            &InstanceName,
192            &LoadedBy,
193        ),
194        (Changed<LoadedBy>, Without<LocalEntity>),
195    >,
196    mut entity_id_index_query: Query<&mut EntityIdIndex>,
197) {
198    for (entity, uuid, minecraft_id, position, instance_name, loaded_by) in &query {
199        let Some(instance_lock) = instance_container.get(instance_name) else {
200            // the instance isn't even loaded by us, so we can safely delete the entity
201            debug!(
202                "Despawned entity {entity:?} because it's in an instance that isn't loaded anymore"
203            );
204            if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
205                warn!(
206                    "Tried to remove entity {entity:?} from the uuid index but it was not there."
207                );
208            }
209            // and now remove the entity from the ecs
210            commands.entity(entity).despawn();
211
212            continue;
213        };
214
215        let mut instance = instance_lock.write();
216
217        // if the entity has no references left, despawn it
218        if !loaded_by.is_empty() {
219            continue;
220        }
221
222        // remove the entity from the chunk index
223        let chunk = ChunkPos::from(*position);
224        match instance.entities_by_chunk.get_mut(&chunk) {
225            Some(entities_in_chunk) => {
226                if entities_in_chunk.remove(&entity) {
227                    // remove the chunk if there's no entities in it anymore
228                    if entities_in_chunk.is_empty() {
229                        instance.entities_by_chunk.remove(&chunk);
230                    }
231                } else {
232                    // search all the other chunks for it :(
233                    let mut found_in_other_chunks = HashSet::new();
234                    for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
235                        if entities_in_other_chunk.remove(&entity) {
236                            found_in_other_chunks.insert(other_chunk);
237                        }
238                    }
239                    if found_in_other_chunks.is_empty() {
240                        warn!(
241                            "Tried to remove entity {entity:?} from chunk {chunk:?} but the entity was not there or in any other chunks."
242                        );
243                    } else {
244                        warn!(
245                            "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:?}"
246                        );
247                    }
248                }
249            }
250            _ => {
251                debug!(
252                    "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found."
253                );
254            }
255        }
256        // remove it from the uuid index
257        if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
258            warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
259        }
260        if instance.entity_by_id.remove(minecraft_id).is_none() {
261            debug!(
262                "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."
263            );
264        }
265
266        // remove it from every client's EntityIdIndex
267        for mut entity_id_index in entity_id_index_query.iter_mut() {
268            entity_id_index.remove_by_ecs_entity(entity);
269        }
270
271        // and now remove the entity from the ecs
272        commands.entity(entity).despawn();
273        debug!("Despawned entity {entity:?} because it was not loaded by anything.");
274    }
275}
276
277pub fn add_entity_to_indexes(
278    entity_id: MinecraftEntityId,
279    ecs_entity: Entity,
280    entity_uuid: Option<Uuid>,
281    entity_id_index: &mut EntityIdIndex,
282    entity_uuid_index: &mut EntityUuidIndex,
283    instance: &mut Instance,
284) {
285    // per-client id index
286    entity_id_index.insert(entity_id, ecs_entity);
287
288    // per-instance id index
289    instance.entity_by_id.insert(entity_id, ecs_entity);
290
291    if let Some(uuid) = entity_uuid {
292        // per-instance uuid index
293        entity_uuid_index.insert(uuid, ecs_entity);
294    }
295}