azalea_entity/plugin/
indexing.rs1use 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 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#[derive(Component, Default)]
55pub struct EntityIdIndex {
56 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 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#[derive(Component, Debug, Deref, DerefMut)]
125pub struct EntityChunkPos(pub ChunkPos);
126
127pub 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 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
159pub 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#[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 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 commands.entity(entity).despawn();
209
210 continue;
211 };
212
213 let mut instance = instance_lock.write();
214
215 if !loaded_by.is_empty() {
217 continue;
218 }
219
220 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 if entities_in_chunk.is_empty() {
227 instance.entities_by_chunk.remove(&chunk);
228 }
229 } else {
230 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 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 for mut entity_id_index in entity_id_index_query.iter_mut() {
266 entity_id_index.remove_by_ecs_entity(entity);
267 }
268
269 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 entity_id_index.insert(entity_id, ecs_entity);
285
286 instance.entity_by_id.insert(entity_id, ecs_entity);
288
289 if let Some(uuid) = entity_uuid {
290 entity_uuid_index.insert(uuid, ecs_entity);
292 }
293}