1use 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 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 fmt::Formatter<'_>) -> 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, instance_name, mut entity_chunk_pos) in query.iter_mut() {
135 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 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
160pub 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#[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 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 commands.entity(entity).despawn();
210
211 continue;
212 };
213
214 let mut instance = instance_lock.write();
215
216 if !loaded_by.is_empty() {
218 continue;
219 }
220
221 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 if entities_in_chunk.is_empty() {
228 instance.entities_by_chunk.remove(&chunk);
229 }
230 } else {
231 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 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 for mut entity_id_index in entity_id_index_query.iter_mut() {
267 entity_id_index.remove_by_ecs_entity(entity);
268 }
269
270 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 entity_id_index.insert(entity_id, ecs_entity);
286
287 instance.entity_by_id.insert(entity_id, ecs_entity);
289
290 if let Some(uuid) = entity_uuid {
291 entity_uuid_index.insert(uuid, ecs_entity);
293 }
294}