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 old_chunk = **entity_chunk_pos;
136 let new_chunk = ChunkPos::from(*pos);
137 if old_chunk != new_chunk {
138 **entity_chunk_pos = new_chunk;
139
140 if old_chunk != new_chunk {
141 let Some(instance_lock) = instance_container.get(instance_name) else {
142 continue;
143 };
144 let mut instance = instance_lock.write();
145
146 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
161pub 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 Some(instance_lock) = instance_container.get(world_name) else {
168 continue;
170 };
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#[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 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 commands.entity(entity).despawn();
214
215 continue;
216 };
217
218 let mut instance = instance_lock.write();
219
220 if !loaded_by.is_empty() {
222 continue;
223 }
224
225 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 if entities_in_chunk.is_empty() {
232 instance.entities_by_chunk.remove(&chunk);
233 }
234 } else {
235 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 let mut found_in_other_chunks = HashSet::new();
255 for (other_chunk, entities_in_other_chunk) in &mut instance.entities_by_chunk {
256 if entities_in_other_chunk.remove(&entity) {
257 found_in_other_chunks.insert(other_chunk);
258 }
259 }
260 if found_in_other_chunks.is_empty() {
261 warn!(
262 "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found and the entity wasn't in any other chunks."
263 );
264 } else {
265 warn!(
266 "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found. Entity found in and removed from other chunk(s): {found_in_other_chunks:?}"
267 );
268 }
269 }
270 }
271 if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
273 warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
274 }
275 if instance.entity_by_id.remove(minecraft_id).is_none() {
276 debug!(
277 "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."
278 );
279 }
280
281 for mut entity_id_index in entity_id_index_query.iter_mut() {
283 entity_id_index.remove_by_ecs_entity(entity);
284 }
285
286 commands.entity(entity).despawn();
288 debug!("Despawned entity {entity:?} because it was not loaded by anything.");
289 }
290}
291
292pub fn add_entity_to_indexes(
293 entity_id: MinecraftEntityId,
294 ecs_entity: Entity,
295 entity_uuid: Option<Uuid>,
296 entity_id_index: &mut EntityIdIndex,
297 entity_uuid_index: &mut EntityUuidIndex,
298 instance: &mut Instance,
299) {
300 entity_id_index.insert(entity_id, ecs_entity);
302
303 instance.entity_by_id.insert(entity_id, ecs_entity);
305
306 if let Some(uuid) = entity_uuid {
307 entity_uuid_index.insert(uuid, ecs_entity);
309 }
310}