From d99ba0da5500a1a4917bf9c8c9e93e2caa4ada99 Mon Sep 17 00:00:00 2001 From: mat Date: Sun, 23 Jul 2023 22:17:17 -0500 Subject: clean up azalea-entity a little --- azalea-entity/src/plugin/indexing.rs | 243 +++++++++++++++++++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 azalea-entity/src/plugin/indexing.rs (limited to 'azalea-entity/src/plugin/indexing.rs') diff --git a/azalea-entity/src/plugin/indexing.rs b/azalea-entity/src/plugin/indexing.rs new file mode 100644 index 00000000..f7dfe0fa --- /dev/null +++ b/azalea-entity/src/plugin/indexing.rs @@ -0,0 +1,243 @@ +//! Stuff related to entity indexes and keeping track of entities in the world. + +use azalea_core::ChunkPos; +use azalea_world::{InstanceContainer, InstanceName, MinecraftEntityId}; +use bevy_ecs::{ + entity::Entity, + query::{Changed, With, Without}, + system::{Commands, Query, Res, ResMut, Resource}, +}; +use log::{debug, error, info, warn}; +use std::{collections::HashMap, fmt::Debug}; +use uuid::Uuid; + +use crate::{EntityUuid, LastSentPosition, Local, Position}; + +use super::LoadedBy; + +#[derive(Resource, Default)] +pub struct EntityUuidIndex { + /// An index of entities by their UUIDs + entity_by_uuid: HashMap, +} + +impl EntityUuidIndex { + pub fn new() -> Self { + Self { + entity_by_uuid: HashMap::default(), + } + } + + pub fn get(&self, uuid: &Uuid) -> Option { + self.entity_by_uuid.get(uuid).copied() + } + + pub fn contains_key(&self, uuid: &Uuid) -> bool { + self.entity_by_uuid.contains_key(uuid) + } + + pub fn insert(&mut self, uuid: Uuid, entity: Entity) { + self.entity_by_uuid.insert(uuid, entity); + } +} + +impl Debug for EntityUuidIndex { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EntityUuidIndex").finish() + } +} + +/// Remove new entities that have the same id as an existing entity, and +/// increase the reference counts. +/// +/// This is the reason why spawning entities into the ECS when you get a spawn +/// entity packet is okay. This system will make sure the new entity gets +/// combined into the old one. +#[allow(clippy::type_complexity)] +pub fn deduplicate_entities( + mut commands: Commands, + mut query: Query< + (Entity, &MinecraftEntityId, &InstanceName), + (Changed, Without), + >, + mut loaded_by_query: Query<&mut LoadedBy>, + instance_container: Res, +) { + // if this entity already exists, remove it + for (new_entity, id, world_name) in query.iter_mut() { + if let Some(world_lock) = instance_container.get(world_name) { + let world = world_lock.write(); + if let Some(old_entity) = world.entity_by_id.get(id) { + if old_entity == &new_entity { + continue; + } + + // this entity already exists!!! remove the one we just added but increase + // the reference count + let new_loaded_by = loaded_by_query + .get(new_entity) + .unwrap_or_else(|_| panic!( + "Entities should always have the LoadedBy component ({new_entity:?} did not)" + )) + .clone(); + let old_loaded_by = loaded_by_query.get_mut(*old_entity); + // merge them if possible + if let Ok(mut old_loaded_by) = old_loaded_by { + old_loaded_by.extend(new_loaded_by.iter()); + } + commands.entity(new_entity).despawn(); + info!( + "Entity with id {id:?} / {new_entity:?} already existed in the world, merging it with {old_entity:?}" + ); + break; + } + } else { + error!("Entity was inserted into a world that doesn't exist."); + } + } +} + +// when a local entity is added, if there was already an entity with the same id +// then delete the old entity +#[allow(clippy::type_complexity)] +pub fn deduplicate_local_entities( + mut commands: Commands, + mut query: Query< + (Entity, &MinecraftEntityId, &InstanceName), + (Changed, With), + >, + instance_container: Res, +) { + // if this entity already exists, remove the old one + for (new_entity, id, world_name) in query.iter_mut() { + if let Some(world_lock) = instance_container.get(world_name) { + let world = world_lock.write(); + if let Some(old_entity) = world.entity_by_id.get(id) { + if old_entity == &new_entity { + // lol + continue; + } + + commands.entity(*old_entity).despawn(); + debug!( + "Added local entity {id:?} / {new_entity:?} but already existed in world as {old_entity:?}, despawning {old_entity:?}" + ); + break; + } + } else { + error!("Entity was inserted into a world that doesn't exist."); + } + } +} + +pub fn update_uuid_index( + mut entity_infos: ResMut, + query: Query<(Entity, &EntityUuid, Option<&Local>), Changed>, +) { + for (entity, &uuid, local) in query.iter() { + // only add it if it doesn't already exist in + // entity_infos.entity_by_uuid + if local.is_none() { + if let Some(old_entity) = entity_infos.entity_by_uuid.get(&uuid) { + debug!( + "Entity with UUID {uuid:?} already existed in the world, not adding to + index (old ecs id: {old_entity:?} / new ecs id: {entity:?})" + ); + continue; + } + } + entity_infos.entity_by_uuid.insert(*uuid, entity); + } +} + +/// System to keep the entity_by_id index up-to-date. +pub fn update_entity_by_id_index( + mut query: Query< + (Entity, &MinecraftEntityId, &InstanceName, Option<&Local>), + Changed, + >, + instance_container: Res, +) { + for (entity, id, world_name, local) in query.iter_mut() { + let world_lock = instance_container.get(world_name).unwrap(); + let mut world = world_lock.write(); + if local.is_none() { + if let Some(old_entity) = world.entity_by_id.get(id) { + debug!( + "Entity with ID {id:?} already existed in the world, not adding to + index (old ecs id: {old_entity:?} / new ecs id: {entity:?})" + ); + continue; + } + } + world.entity_by_id.insert(*id, entity); + debug!("Added {entity:?} to {world_name:?} with {id:?}."); + } +} + +/// Update the chunk position indexes in [`EntityUuidIndex`]. +pub fn update_entity_chunk_positions( + mut query: Query<(Entity, &Position, &mut LastSentPosition, &InstanceName), Changed>, + instance_container: Res, +) { + for (entity, pos, last_pos, world_name) in query.iter_mut() { + let world_lock = instance_container.get(world_name).unwrap(); + let mut world = world_lock.write(); + + let old_chunk = ChunkPos::from(*last_pos); + let new_chunk = ChunkPos::from(*pos); + + if old_chunk != new_chunk { + // move the entity from the old chunk to the new one + if let Some(entities) = world.entities_by_chunk.get_mut(&old_chunk) { + entities.remove(&entity); + } + world + .entities_by_chunk + .entry(new_chunk) + .or_default() + .insert(entity); + } + } +} + +/// Despawn entities that aren't being loaded by anything. +pub fn remove_despawned_entities_from_indexes( + mut commands: Commands, + mut entity_infos: ResMut, + instance_container: Res, + query: Query<(Entity, &EntityUuid, &Position, &InstanceName, &LoadedBy), Changed>, +) { + for (entity, uuid, position, world_name, loaded_by) in &query { + let world_lock = instance_container.get(world_name).unwrap(); + let mut world = world_lock.write(); + + // if the entity has no references left, despawn it + if !loaded_by.is_empty() { + continue; + } + + // remove the entity from the chunk index + let chunk = ChunkPos::from(*position); + if let Some(entities_in_chunk) = world.entities_by_chunk.get_mut(&chunk) { + if entities_in_chunk.remove(&entity) { + // remove the chunk if there's no entities in it anymore + if entities_in_chunk.is_empty() { + world.entities_by_chunk.remove(&chunk); + } + } else { + warn!("Tried to remove entity from chunk {chunk:?} but the entity was not there."); + } + } else { + warn!("Tried to remove entity from chunk {chunk:?} but the chunk was not found."); + } + // remove it from the uuid index + if entity_infos.entity_by_uuid.remove(uuid).is_none() { + warn!("Tried to remove entity {entity:?} from the uuid index but it was not there."); + } + // and now remove the entity from the ecs + commands.entity(entity).despawn(); + debug!("Despawned entity {entity:?} because it was not loaded by anything."); + return; + } +} -- cgit v1.2.3