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_by_minecraft_entity(&self, id: MinecraftEntityId) -> Option<Entity> {
68 self.entity_by_id.get(&id).copied()
69 }
70 pub fn get_by_ecs_entity(&self, entity: Entity) -> Option<MinecraftEntityId> {
71 self.id_by_entity.get(&entity).copied()
72 }
73
74 pub fn contains_minecraft_entity(&self, id: MinecraftEntityId) -> bool {
75 self.entity_by_id.contains_key(&id)
76 }
77 pub fn contains_ecs_entity(&self, id: Entity) -> bool {
78 self.id_by_entity.contains_key(&id)
79 }
80
81 pub fn insert(&mut self, id: MinecraftEntityId, entity: Entity) {
82 self.entity_by_id.insert(id, entity);
83 self.id_by_entity.insert(entity, id);
84 trace!("Inserted {id} -> {entity:?} into a client's EntityIdIndex");
85 }
86
87 pub fn remove_by_minecraft_entity(&mut self, id: MinecraftEntityId) -> Option<Entity> {
88 if let Some(entity) = self.entity_by_id.remove(&id) {
89 trace!(
90 "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove)"
91 );
92 self.id_by_entity.remove(&entity);
93 Some(entity)
94 } else {
95 trace!(
96 "Failed to remove {id} from a client's EntityIdIndex (using EntityIdIndex::remove)"
97 );
98 None
99 }
100 }
101
102 pub fn remove_by_ecs_entity(&mut self, entity: Entity) -> Option<MinecraftEntityId> {
103 if let Some(id) = self.id_by_entity.remove(&entity) {
104 trace!(
105 "Removed {id} -> {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity)."
106 );
107 self.entity_by_id.remove(&id);
108 Some(id)
109 } else {
110 trace!(
115 "Failed to remove {entity:?} from a client's EntityIdIndex (using EntityIdIndex::remove_by_ecs_entity). This may be expected behavior."
116 );
117 None
118 }
119 }
120}
121
122impl Debug for EntityUuidIndex {
123 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124 f.debug_struct("EntityUuidIndex").finish()
125 }
126}
127
128#[derive(Component, Debug, Deref, DerefMut)]
130pub struct EntityChunkPos(pub ChunkPos);
131
132pub fn update_entity_chunk_positions(
136 mut query: Query<(Entity, &Position, &InstanceName, &mut EntityChunkPos), Changed<Position>>,
137 instance_container: Res<InstanceContainer>,
138) {
139 for (entity, pos, world_name, mut entity_chunk_pos) in query.iter_mut() {
140 let instance_lock = instance_container.get(world_name).unwrap();
141 let mut instance = instance_lock.write();
142
143 let old_chunk = **entity_chunk_pos;
144 let new_chunk = ChunkPos::from(*pos);
145 if old_chunk != new_chunk {
146 **entity_chunk_pos = new_chunk;
147
148 if old_chunk != new_chunk {
149 if let Some(entities) = instance.entities_by_chunk.get_mut(&old_chunk) {
151 entities.remove(&entity);
152 }
153 instance
154 .entities_by_chunk
155 .entry(new_chunk)
156 .or_default()
157 .insert(entity);
158 trace!("Entity {entity:?} moved from {old_chunk:?} to {new_chunk:?}");
159 }
160 }
161 }
162}
163
164pub fn insert_entity_chunk_position(
166 query: Query<(Entity, &Position, &InstanceName), Added<EntityChunkPos>>,
167 instance_container: Res<InstanceContainer>,
168) {
169 for (entity, pos, world_name) in query.iter() {
170 let instance_lock = instance_container.get(world_name).unwrap();
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 debug!(
255 "Tried to remove entity {entity:?} from chunk {chunk:?} but the chunk was not found."
256 );
257 }
258 }
259 if entity_uuid_index.entity_by_uuid.remove(uuid).is_none() {
261 warn!("Tried to remove entity {entity:?} from the uuid index but it was not there.");
262 }
263 if instance.entity_by_id.remove(minecraft_id).is_none() {
264 debug!(
265 "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."
266 );
267 }
268
269 for mut entity_id_index in entity_id_index_query.iter_mut() {
271 entity_id_index.remove_by_ecs_entity(entity);
272 }
273
274 commands.entity(entity).despawn();
276 debug!("Despawned entity {entity:?} because it was not loaded by anything.");
277 }
278}
279
280pub fn add_entity_to_indexes(
281 entity_id: MinecraftEntityId,
282 ecs_entity: Entity,
283 entity_uuid: Option<Uuid>,
284 entity_id_index: &mut EntityIdIndex,
285 entity_uuid_index: &mut EntityUuidIndex,
286 instance: &mut Instance,
287) {
288 entity_id_index.insert(entity_id, ecs_entity);
290
291 instance.entity_by_id.insert(entity_id, ecs_entity);
293
294 if let Some(uuid) = entity_uuid {
295 entity_uuid_index.insert(uuid, ecs_entity);
297 }
298}