clean up azalea-entity a little

1 files changed, 243 insertions, 0 deletions

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<Uuid, Entity>,
+}
+
+impl EntityUuidIndex {
+    pub fn new() -> Self {
+        Self {
+            entity_by_uuid: HashMap::default(),
+        }
+    }
+
+    pub fn get(&self, uuid: &Uuid) -> Option<Entity> {
+        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<MinecraftEntityId>, Without<Local>),
+    >,
+    mut loaded_by_query: Query<&mut LoadedBy>,
+    instance_container: Res<InstanceContainer>,
+) {
+    // 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<MinecraftEntityId>, With<Local>),
+    >,
+    instance_container: Res<InstanceContainer>,
+) {
+    // 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<EntityUuidIndex>,
+    query: Query<(Entity, &EntityUuid, Option<&Local>), Changed<EntityUuid>>,
+) {
+    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<MinecraftEntityId>,
+    >,
+    instance_container: Res<InstanceContainer>,
+) {
+    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<Position>>,
+    instance_container: Res<InstanceContainer>,
+) {
+    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<EntityUuidIndex>,
+    instance_container: Res<InstanceContainer>,
+    query: Query<(Entity, &EntityUuid, &Position, &InstanceName, &LoadedBy), Changed<LoadedBy>>,
+) {
+    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;
+    }
+}