1use 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 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#[derive(Component, Default)]
60pub struct EntityIdIndex {
61 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 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#[derive(Component, Debug, Deref, DerefMut)]
127pub struct EntityChunkPos(pub ChunkPos);
128
129pub 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 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 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#[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 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 commands.entity(entity).despawn();
211
212 continue;
213 };
214
215 let mut instance = instance_lock.write();
216
217 if !loaded_by.is_empty() {
219 continue;
220 }
221
222 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 if entities_in_chunk.is_empty() {
229 instance.entities_by_chunk.remove(&chunk);
230 }
231 } else {
232 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 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 for mut entity_id_index in entity_id_index_query.iter_mut() {
268 entity_id_index.remove_by_ecs_entity(entity);
269 }
270
271 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 entity_id_index.insert(entity_id, ecs_entity);
287
288 instance.entity_by_id.insert(entity_id, ecs_entity);
290
291 if let Some(uuid) = entity_uuid {
292 entity_uuid_index.insert(uuid, ecs_entity);
294 }
295}