1 files changed, 145 insertions, 0 deletions
diff --git a/azalea-world/src/entity.rs b/azalea-world/src/entity.rs
new file mode 100644
index 00000000..5219e410
--- /dev/null
+++ b/azalea-world/src/entity.rs
@@ -0,0 +1,145 @@
+use azalea_core::ChunkPos;
+use azalea_entity::Entity;
+use log::warn;
+use nohash_hasher::{IntMap, IntSet};
+use std::collections::HashMap;
+use uuid::Uuid;
+
+#[derive(Debug)]
+pub struct EntityStorage {
+    by_id: IntMap<u32, Entity>,
+    by_chunk: HashMap<ChunkPos, IntSet<u32>>,
+    by_uuid: HashMap<Uuid, u32>,
+}
+
+impl EntityStorage {
+    pub fn new() -> Self {
+        Self {
+            by_id: IntMap::default(),
+            by_chunk: HashMap::default(),
+            by_uuid: HashMap::default(),
+        }
+    }
+
+    /// 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_uuid.insert(entity.uuid, 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) {
+        if let Some(entity) = self.by_id.remove(&id) {
+            let entity_chunk = ChunkPos::from(entity.pos());
+            let entity_uuid = entity.uuid;
+            if self.by_chunk.remove(&entity_chunk).is_none() {
+                warn!("Tried to remove entity with id {id} from chunk {entity_chunk:?} but it was not found.");
+            }
+            if self.by_uuid.remove(&entity_uuid).is_none() {
+                warn!("Tried to remove entity with id {id} from uuid {entity_uuid:?} 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.
+    #[inline]
+    pub fn get_by_id(&self, id: u32) -> Option<&Entity> {
+        self.by_id.get(&id)
+    }
+
+    /// Get a mutable reference to an entity by its id.
+    #[inline]
+    pub fn get_mut_by_id(&mut self, id: u32) -> Option<&mut Entity> {
+        self.by_id.get_mut(&id)
+    }
+
+    /// Get a reference to an entity by its uuid.
+    #[inline]
+    pub fn get_by_uuid(&self, uuid: &Uuid) -> Option<&Entity> {
+        self.by_uuid.get(uuid).and_then(|id| self.by_id.get(id))
+    }
+
+    /// Get a mutable reference to an entity by its uuid.
+    #[inline]
+    pub fn get_mut_by_uuid(&mut self, uuid: &Uuid) -> Option<&mut Entity> {
+        self.by_uuid.get(uuid).and_then(|id| 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 {
+                if let Some(entity) = self.by_id.remove(&entity_id) {
+                    self.by_uuid.remove(&entity.uuid);
+                } else {
+                    warn!("While clearing chunk {chunk:?}, found an entity that isn't in by_id {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);
+    }
+
+    /// Get an iterator over all entities.
+    #[inline]
+    pub fn entities(&self) -> std::collections::hash_map::Values<'_, u32, Entity> {
+        self.by_id.values()
+    }
+
+    pub fn find_one_entity<F>(&self, mut f: F) -> Option<&Entity>
+    where
+        F: FnMut(&Entity) -> bool,
+    {
+        for entity in self.entities() {
+            if f(entity) {
+                return Some(entity);
+            }
+        }
+        None
+    }
+
+    pub fn find_one_entity_in_chunk<F>(&self, chunk: &ChunkPos, mut f: F) -> Option<&Entity>
+    where
+        F: FnMut(&Entity) -> bool,
+    {
+        if let Some(entities) = self.by_chunk.get(chunk) {
+            for entity_id in entities {
+                if let Some(entity) = self.by_id.get(entity_id) {
+                    if f(entity) {
+                        return Some(entity);
+                    }
+                }
+            }
+        }
+        None
+    }
+}
+
+impl Default for EntityStorage {
+    fn default() -> Self {
+        Self::new()
+    }
+}