Improvements to azalea-world for entities

1 files changed, 42 insertions, 5 deletions

diff --git a/azalea-world/src/entity.rs b/azalea-world/src/entity.rs
index 49e1ae73..7077d0c4 100644
--- a/azalea-world/src/entity.rs
+++ b/azalea-world/src/entity.rs
@@ -1,14 +1,14 @@
-use std::collections::HashMap;
-
 use azalea_core::ChunkPos;
 use azalea_entity::Entity;
-use nohash_hasher::IntMap;
+use log::warn;
+use nohash_hasher::{IntMap, IntSet};
+use std::collections::HashMap;
 
 #[derive(Debug)]
 pub struct EntityStorage {
     by_id: IntMap<u32, Entity>,
     // TODO: this doesn't work yet (should be updated in the set_pos method in azalea-entity)
-    by_chunk: HashMap<ChunkPos, u32>,
+    by_chunk: HashMap<ChunkPos, IntSet<u32>>,
 }
 
 impl EntityStorage {
@@ -22,13 +22,24 @@ impl EntityStorage {
     /// Add an entity to the storage.
     #[inline]
     pub fn insert(&mut self, entity: Entity) {
+        self.by_chunk
+            .entry(ChunkPos::from(entity.pos()))
+            .or_default()
+            .insert(entity.id);
         self.by_id.insert(entity.id, entity);
     }
 
     /// Remove an entity from the storage by its id.
     #[inline]
     pub fn remove_by_id(&mut self, id: u32) {
-        self.by_id.remove(&id);
+        if let Some(entity) = self.by_id.remove(&id) {
+            let entity_chunk = ChunkPos::from(entity.pos());
+            if let None = self.by_chunk.remove(&entity_chunk) {
+                warn!("Tried to remove entity with id {id} from chunk {entity_chunk:?} but it was not found.");
+            }
+        } else {
+            warn!("Tried to remove entity with id {id} but it was not found.")
+        }
     }
 
     /// Get a reference to an entity by its id.
@@ -42,4 +53,30 @@ impl EntityStorage {
     pub fn get_mut_by_id(&mut self, id: u32) -> Option<&mut Entity> {
         self.by_id.get_mut(&id)
     }
+
+    /// Clear all entities in a chunk.
+    pub fn clear_chunk(&mut self, chunk: &ChunkPos) {
+        if let Some(entities) = self.by_chunk.remove(chunk) {
+            for entity_id in entities {
+                self.by_id.remove(&entity_id);
+            }
+        }
+    }
+
+    /// Updates an entity from its old chunk.
+    #[inline]
+    pub fn update_entity_chunk(
+        &mut self,
+        entity_id: u32,
+        old_chunk: &ChunkPos,
+        new_chunk: &ChunkPos,
+    ) {
+        if let Some(entities) = self.by_chunk.get_mut(old_chunk) {
+            entities.remove(&entity_id);
+        }
+        self.by_chunk
+            .entry(*new_chunk)
+            .or_default()
+            .insert(entity_id);
+    }
 }