Merge pull request #10 from mat-1/azalea-entity

azalea-entity

6 files changed, 465 insertions, 196 deletions

diff --git a/azalea-world/Cargo.toml b/azalea-world/Cargo.toml
index 79e6155d..ca890d82 100644
--- a/azalea-world/Cargo.toml
+++ b/azalea-world/Cargo.toml
@@ -6,7 +6,14 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+azalea-block = {path = "../azalea-block"}
+azalea-buf = {path = "../azalea-buf"}
 azalea-core = {path = "../azalea-core"}
+azalea-entity = {path = "../azalea-entity"}
 azalea-nbt = {path = "../azalea-nbt"}
-azalea-protocol = {path = "../azalea-protocol"}
-azalea-block = {path = "../azalea-block"}
+log = "0.4.17"
+nohash-hasher = "0.2.0"
+uuid = "1.1.2"
+
+[profile.release]
+lto = true
diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
index 0dc81f9a..fcb3f8f9 100644
--- a/azalea-world/src/bit_storage.rs
+++ b/azalea-world/src/bit_storage.rs
@@ -107,7 +107,7 @@ impl BitStorage {
         // assert!(bits >= 1 && bits <= 32);
 
         if let Some(data) = &data {
-            if data.len() == 0 {
+            if data.is_empty() {
                 // TODO: make 0 bit storage actually work
                 return Ok(BitStorage::default());
             }
@@ -162,7 +162,7 @@ impl BitStorage {
         // return (int)(var3 >> var5 & this.mask);
 
         assert!(
-            index <= self.size - 1,
+            index < self.size,
             "Index {} out of bounds (max is {})",
             index,
             self.size - 1
@@ -181,7 +181,7 @@ impl BitStorage {
         // int var6 = (var1 - var3 * this.valuesPerLong) * this.bits;
         // this.data[var3] = var4 & ~(this.mask << var6) | ((long)var2 & this.mask) << var6;
 
-        assert!(index <= self.size - 1);
+        assert!(index < self.size);
         assert!(value <= self.mask);
         let cell_index = self.cell_index(index as u64);
         let cell = &mut self.data[cell_index as usize];
diff --git a/azalea-world/src/chunk.rs b/azalea-world/src/chunk.rs
new file mode 100644
index 00000000..27ef1ca7
--- /dev/null
+++ b/azalea-world/src/chunk.rs
@@ -0,0 +1,221 @@
+use crate::palette::PalettedContainer;
+use crate::palette::PalettedContainerType;
+use crate::World;
+use azalea_block::BlockState;
+use azalea_buf::{McBufReadable, McBufWritable};
+use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
+use std::fmt::Debug;
+use std::{
+    io::{Read, Write},
+    ops::{Index, IndexMut},
+    sync::{Arc, Mutex},
+};
+
+const SECTION_HEIGHT: u32 = 16;
+
+pub struct ChunkStorage {
+    pub view_center: ChunkPos,
+    chunk_radius: u32,
+    view_range: u32,
+    pub height: u32,
+    pub min_y: i32,
+    // chunks is a list of size chunk_radius * chunk_radius
+    chunks: Vec<Option<Arc<Mutex<Chunk>>>>,
+}
+
+// java moment
+// it might be possible to replace this with just a modulo, but i copied java's floorMod just in case
+fn floor_mod(x: i32, y: u32) -> u32 {
+    if x < 0 {
+        y - ((-x) as u32 % y)
+    } else {
+        x as u32 % y
+    }
+}
+
+impl ChunkStorage {
+    pub fn new(chunk_radius: u32, height: u32, min_y: i32) -> Self {
+        let view_range = chunk_radius * 2 + 1;
+        ChunkStorage {
+            view_center: ChunkPos::new(0, 0),
+            chunk_radius,
+            view_range,
+            height,
+            min_y,
+            chunks: vec![None; (view_range * view_range) as usize],
+        }
+    }
+
+    fn get_index(&self, chunk_pos: &ChunkPos) -> usize {
+        (floor_mod(chunk_pos.x, self.view_range) * self.view_range
+            + floor_mod(chunk_pos.z, self.view_range)) as usize
+    }
+
+    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
+        (chunk_pos.x - self.view_center.x).unsigned_abs() <= self.chunk_radius
+            && (chunk_pos.z - self.view_center.z).unsigned_abs() <= self.chunk_radius
+    }
+
+    pub fn get_block_state(&self, pos: &BlockPos, min_y: i32) -> Option<BlockState> {
+        let chunk_pos = ChunkPos::from(pos);
+        println!("chunk_pos {:?} block_pos {:?}", chunk_pos, pos);
+        let chunk = &self[&chunk_pos];
+        chunk
+            .as_ref()
+            .map(|chunk| chunk.lock().unwrap().get(&ChunkBlockPos::from(pos), min_y))
+    }
+
+    pub fn replace_with_packet_data(
+        &mut self,
+        pos: &ChunkPos,
+        data: &mut impl Read,
+    ) -> Result<(), String> {
+        if !self.in_range(pos) {
+            println!(
+                "Ignoring chunk since it's not in the view range: {}, {}",
+                pos.x, pos.z
+            );
+            return Ok(());
+        }
+
+        let chunk = Arc::new(Mutex::new(Chunk::read_with_world_height(
+            data,
+            self.height,
+        )?));
+        println!("Loaded chunk {:?}", pos);
+        self[pos] = Some(chunk);
+
+        Ok(())
+    }
+}
+
+impl Index<&ChunkPos> for ChunkStorage {
+    type Output = Option<Arc<Mutex<Chunk>>>;
+
+    fn index(&self, pos: &ChunkPos) -> &Self::Output {
+        &self.chunks[self.get_index(pos)]
+    }
+}
+impl IndexMut<&ChunkPos> for ChunkStorage {
+    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
+        let index = self.get_index(pos);
+        &mut self.chunks[index]
+    }
+}
+
+#[derive(Debug)]
+pub struct Chunk {
+    pub sections: Vec<Section>,
+}
+
+impl Chunk {
+    pub fn read_with_world(buf: &mut impl Read, data: &World) -> Result<Self, String> {
+        Self::read_with_world_height(buf, data.height())
+    }
+
+    pub fn read_with_world_height(buf: &mut impl Read, world_height: u32) -> Result<Self, String> {
+        let section_count = world_height / SECTION_HEIGHT;
+        let mut sections = Vec::with_capacity(section_count as usize);
+        for _ in 0..section_count {
+            let section = Section::read_into(buf)?;
+            sections.push(section);
+        }
+        Ok(Chunk { sections })
+    }
+
+    pub fn section_index(&self, y: i32, min_y: i32) -> u32 {
+        // TODO: check the build height and stuff, this code will be broken if the min build height is 0
+        // (LevelHeightAccessor.getMinSection in vanilla code)
+        assert!(y >= 0);
+        let min_section_index = min_y.div_floor(16);
+        (y.div_floor(16) - min_section_index) as u32
+    }
+
+    pub fn get(&self, pos: &ChunkBlockPos, min_y: i32) -> BlockState {
+        let section_index = self.section_index(pos.y, min_y);
+        // TODO: make sure the section exists
+        let section = &self.sections[section_index as usize];
+        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+        section.get(chunk_section_pos)
+    }
+}
+
+impl McBufWritable for Chunk {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        for section in &self.sections {
+            section.write_into(buf)?;
+        }
+        Ok(())
+    }
+}
+
+impl Debug for ChunkStorage {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ChunkStorage")
+            .field("view_center", &self.view_center)
+            .field("chunk_radius", &self.chunk_radius)
+            .field("view_range", &self.view_range)
+            .field("height", &self.height)
+            .field("min_y", &self.min_y)
+            // .field("chunks", &self.chunks)
+            .field("chunks", &format_args!("{} items", self.chunks.len()))
+            .finish()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Section {
+    pub block_count: u16,
+    pub states: PalettedContainer,
+    pub biomes: PalettedContainer,
+}
+
+impl McBufReadable for Section {
+    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
+        let block_count = u16::read_into(buf)?;
+
+        // this is commented out because the vanilla server is wrong
+        // assert!(
+        //     block_count <= 16 * 16 * 16,
+        //     "A section has more blocks than what should be possible. This is a bug!"
+        // );
+
+        let states = PalettedContainer::read_with_type(buf, &PalettedContainerType::BlockStates)?;
+
+        for i in 0..states.storage.size() {
+            if !BlockState::is_valid_state(states.storage.get(i) as u32) {
+                return Err(format!(
+                    "Invalid block state {} (index {}) found in section.",
+                    states.storage.get(i),
+                    i
+                ));
+            }
+        }
+
+        let biomes = PalettedContainer::read_with_type(buf, &PalettedContainerType::Biomes)?;
+        Ok(Section {
+            block_count,
+            states,
+            biomes,
+        })
+    }
+}
+
+impl McBufWritable for Section {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        self.block_count.write_into(buf)?;
+        self.states.write_into(buf)?;
+        self.biomes.write_into(buf)?;
+        Ok(())
+    }
+}
+
+impl Section {
+    fn get(&self, pos: ChunkSectionBlockPos) -> BlockState {
+        // TODO: use the unsafe method and do the check earlier
+        self.states
+            .get(pos.x as usize, pos.y as usize, pos.z as usize)
+            .try_into()
+            .expect("Invalid block state.")
+    }
+}
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()
+    }
+}
diff --git a/azalea-world/src/lib.rs b/azalea-world/src/lib.rs
index e651e455..3afa4fee 100644
--- a/azalea-world/src/lib.rs
+++ b/azalea-world/src/lib.rs
@@ -1,19 +1,22 @@
 #![feature(int_roundings)]
 
 mod bit_storage;
+mod chunk;
+mod entity;
 mod palette;
 
-use crate::palette::PalettedContainerType;
 use azalea_block::BlockState;
-use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
-use azalea_protocol::mc_buf::{McBufReadable, McBufWritable};
+use azalea_core::{BlockPos, ChunkPos, EntityPos, PositionDelta8};
+use azalea_entity::Entity;
 pub use bit_storage::BitStorage;
-use palette::PalettedContainer;
+pub use chunk::{Chunk, ChunkStorage};
+pub use entity::EntityStorage;
 use std::{
-    io::{Read, Write},
+    io::Read,
     ops::{Index, IndexMut},
     sync::{Arc, Mutex},
 };
+use uuid::Uuid;
 
 #[cfg(test)]
 mod tests {
@@ -24,230 +27,123 @@ mod tests {
     }
 }
 
-const SECTION_HEIGHT: u32 = 16;
-
+#[derive(Debug)]
 pub struct World {
-    pub storage: ChunkStorage,
-    pub height: u32,
-    pub min_y: i32,
+    chunk_storage: ChunkStorage,
+    entity_storage: EntityStorage,
 }
 
 impl World {
+    pub fn new(chunk_radius: u32, height: u32, min_y: i32) -> Self {
+        World {
+            chunk_storage: ChunkStorage::new(chunk_radius, height, min_y),
+            entity_storage: EntityStorage::new(),
+        }
+    }
+
     pub fn replace_with_packet_data(
         &mut self,
         pos: &ChunkPos,
         data: &mut impl Read,
     ) -> Result<(), String> {
-        if !self.storage.in_range(pos) {
-            println!(
-                "Ignoring chunk since it's not in the view range: {}, {}",
-                pos.x, pos.z
-            );
-            return Ok(());
-        }
-        // let existing_chunk = &self.storage[pos];
-
-        let chunk = Arc::new(Mutex::new(Chunk::read_with_world(data, self)?));
-        println!("Loaded chunk {:?}", pos);
-        self.storage[pos] = Some(chunk);
-
-        Ok(())
+        self.chunk_storage.replace_with_packet_data(pos, data)
     }
 
     pub fn update_view_center(&mut self, pos: &ChunkPos) {
-        self.storage.view_center = *pos;
+        self.chunk_storage.view_center = *pos;
     }
 
     pub fn get_block_state(&self, pos: &BlockPos) -> Option<BlockState> {
-        self.storage.get_block_state(pos, self.min_y)
-    }
-}
-impl Index<&ChunkPos> for World {
-    type Output = Option<Arc<Mutex<Chunk>>>;
-
-    fn index(&self, pos: &ChunkPos) -> &Self::Output {
-        &self.storage[pos]
-    }
-}
-impl IndexMut<&ChunkPos> for World {
-    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
-        &mut self.storage[pos]
-    }
-}
-// impl Index<&BlockPos> for World {
-//     type Output = Option<Arc<Mutex<Chunk>>>;
-
-//     fn index(&self, pos: &BlockPos) -> &Self::Output {
-//         let chunk = &self[ChunkPos::from(pos)];
-//         // chunk.
-
-//     }
-// }
-
-pub struct ChunkStorage {
-    view_center: ChunkPos,
-    chunk_radius: u32,
-    view_range: u32,
-    // chunks is a list of size chunk_radius * chunk_radius
-    chunks: Vec<Option<Arc<Mutex<Chunk>>>>,
-}
-
-// java moment
-// it might be possible to replace this with just a modulo, but i copied java's floorMod just in case
-fn floor_mod(x: i32, y: u32) -> u32 {
-    if x < 0 {
-        y - ((-x) as u32 % y)
-    } else {
-        x as u32 % y
-    }
-}
-
-impl ChunkStorage {
-    pub fn new(chunk_radius: u32) -> Self {
-        let view_range = chunk_radius * 2 + 1;
-        ChunkStorage {
-            view_center: ChunkPos::new(0, 0),
-            chunk_radius,
-            view_range,
-            chunks: vec![None; (view_range * view_range) as usize],
+        self.chunk_storage.get_block_state(pos, self.min_y())
+    }
+
+    pub fn move_entity(&mut self, entity_id: u32, new_pos: EntityPos) -> Result<(), String> {
+        let entity = self
+            .entity_storage
+            .get_mut_by_id(entity_id)
+            .ok_or_else(|| "Moving entity that doesn't exist".to_string())?;
+
+        let old_chunk = ChunkPos::from(entity.pos());
+        let new_chunk = ChunkPos::from(&new_pos);
+        // this is fine because we update the chunk below
+        entity.unsafe_move(new_pos);
+        if old_chunk != new_chunk {
+            self.entity_storage
+                .update_entity_chunk(entity_id, &old_chunk, &new_chunk);
         }
+        Ok(())
     }
 
-    fn get_index(&self, chunk_pos: &ChunkPos) -> usize {
-        (floor_mod(chunk_pos.x, self.view_range) * self.view_range
-            + floor_mod(chunk_pos.z, self.view_range)) as usize
+    pub fn move_entity_with_delta(
+        &mut self,
+        entity_id: u32,
+        delta: &PositionDelta8,
+    ) -> Result<(), String> {
+        let entity = self
+            .entity_storage
+            .get_mut_by_id(entity_id)
+            .ok_or_else(|| "Moving entity that doesn't exist".to_string())?;
+        let new_pos = entity.pos().with_delta(delta);
+
+        let old_chunk = ChunkPos::from(entity.pos());
+        let new_chunk = ChunkPos::from(&new_pos);
+        // this is fine because we update the chunk below
+
+        entity.unsafe_move(new_pos);
+        if old_chunk != new_chunk {
+            self.entity_storage
+                .update_entity_chunk(entity_id, &old_chunk, &new_chunk);
+        }
+        Ok(())
     }
 
-    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
-        (chunk_pos.x - self.view_center.x).unsigned_abs() <= self.chunk_radius
-            && (chunk_pos.z - self.view_center.z).unsigned_abs() <= self.chunk_radius
+    pub fn add_entity(&mut self, entity: Entity) {
+        self.entity_storage.insert(entity);
     }
 
-    pub fn get_block_state(&self, pos: &BlockPos, min_y: i32) -> Option<BlockState> {
-        let chunk_pos = ChunkPos::from(pos);
-        println!("chunk_pos {:?} block_pos {:?}", chunk_pos, pos);
-        let chunk = &self[&chunk_pos];
-        match chunk {
-            Some(chunk) => Some(chunk.lock().unwrap().get(&ChunkBlockPos::from(pos), min_y)),
-            None => None,
-        }
+    pub fn height(&self) -> u32 {
+        self.chunk_storage.height
     }
-}
-
-impl Index<&ChunkPos> for ChunkStorage {
-    type Output = Option<Arc<Mutex<Chunk>>>;
 
-    fn index(&self, pos: &ChunkPos) -> &Self::Output {
-        &self.chunks[self.get_index(pos)]
+    pub fn min_y(&self) -> i32 {
+        self.chunk_storage.min_y
     }
-}
-impl IndexMut<&ChunkPos> for ChunkStorage {
-    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
-        let index = self.get_index(pos);
-        &mut self.chunks[index]
-    }
-}
 
-#[derive(Debug)]
-pub struct Chunk {
-    pub sections: Vec<Section>,
-}
-
-impl Chunk {
-    pub fn read_with_world(buf: &mut impl Read, data: &World) -> Result<Self, String> {
-        Self::read_with_world_height(buf, data.height)
+    pub fn entity_by_id(&self, id: u32) -> Option<&Entity> {
+        self.entity_storage.get_by_id(id)
     }
 
-    pub fn read_with_world_height(buf: &mut impl Read, world_height: u32) -> Result<Self, String> {
-        let section_count = world_height / SECTION_HEIGHT;
-        let mut sections = Vec::with_capacity(section_count as usize);
-        for _ in 0..section_count {
-            let section = Section::read_into(buf)?;
-            sections.push(section);
-        }
-        Ok(Chunk { sections })
+    pub fn mut_entity_by_id(&mut self, id: u32) -> Option<&mut Entity> {
+        self.entity_storage.get_mut_by_id(id)
     }
 
-    pub fn section_index(&self, y: i32, min_y: i32) -> u32 {
-        // TODO: check the build height and stuff, this code will be broken if the min build height is 0
-        // (LevelHeightAccessor.getMinSection in vanilla code)
-        assert!(y >= 0);
-        let min_section_index = min_y.div_floor(16);
-        (y.div_floor(16) - min_section_index) as u32
+    pub fn entity_by_uuid(&self, uuid: &Uuid) -> Option<&Entity> {
+        self.entity_storage.get_by_uuid(uuid)
     }
 
-    pub fn get(&self, pos: &ChunkBlockPos, min_y: i32) -> BlockState {
-        let section_index = self.section_index(pos.y, min_y);
-        // TODO: make sure the section exists
-        let section = &self.sections[section_index as usize];
-        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-        let block_state = section.get(chunk_section_pos);
-        block_state
+    /// Get an iterator over all entities.
+    #[inline]
+    pub fn entities(&self) -> std::collections::hash_map::Values<'_, u32, Entity> {
+        self.entity_storage.entities()
     }
-}
 
-impl McBufWritable for Chunk {
-    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
-        for section in &self.sections {
-            section.write_into(buf)?;
-        }
-        Ok(())
+    pub fn find_one_entity<F>(&self, mut f: F) -> Option<&Entity>
+    where
+        F: FnMut(&Entity) -> bool,
+    {
+        self.entity_storage.find_one_entity(|entity| f(entity))
     }
 }
 
-#[derive(Clone, Debug)]
-pub struct Section {
-    pub block_count: u16,
-    pub states: PalettedContainer,
-    pub biomes: PalettedContainer,
-}
-
-impl McBufReadable for Section {
-    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
-        let block_count = u16::read_into(buf)?;
-
-        // this is commented out because the vanilla server is wrong
-        // assert!(
-        //     block_count <= 16 * 16 * 16,
-        //     "A section has more blocks than what should be possible. This is a bug!"
-        // );
-
-        let states = PalettedContainer::read_with_type(buf, &PalettedContainerType::BlockStates)?;
-
-        for i in 0..states.storage.size() {
-            if !BlockState::is_valid_state(states.storage.get(i) as u32) {
-                return Err(format!(
-                    "Invalid block state {} (index {}) found in section.",
-                    states.storage.get(i),
-                    i
-                ));
-            }
-        }
-
-        let biomes = PalettedContainer::read_with_type(buf, &PalettedContainerType::Biomes)?;
-        Ok(Section {
-            block_count,
-            states,
-            biomes,
-        })
-    }
-}
+impl Index<&ChunkPos> for World {
+    type Output = Option<Arc<Mutex<Chunk>>>;
 
-impl McBufWritable for Section {
-    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
-        self.block_count.write_into(buf)?;
-        self.states.write_into(buf)?;
-        self.biomes.write_into(buf)?;
-        Ok(())
+    fn index(&self, pos: &ChunkPos) -> &Self::Output {
+        &self.chunk_storage[pos]
     }
 }
-
-impl Section {
-    fn get(&self, pos: ChunkSectionBlockPos) -> BlockState {
-        // TODO: use the unsafe method and do the check earlier
-        self.states
-            .get(pos.x as usize, pos.y as usize, pos.z as usize)
-            .try_into()
-            .expect("Invalid block state.")
+impl IndexMut<&ChunkPos> for World {
+    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
+        &mut self.chunk_storage[pos]
     }
 }
diff --git a/azalea-world/src/palette.rs b/azalea-world/src/palette.rs
index be6f022c..4779dc4b 100644
--- a/azalea-world/src/palette.rs
+++ b/azalea-world/src/palette.rs
@@ -1,4 +1,4 @@
-use azalea_protocol::mc_buf::{McBufReadable, McBufVarReadable, McBufWritable, Readable, Writable};
+use azalea_buf::{McBufReadable, McBufVarReadable, McBufWritable, Readable, Writable};
 use std::io::{Read, Write};
 
 use crate::BitStorage;