Extensible ChunkStorage

Co-authored-by: sdwhw <191973436+sdwhw@users.noreply.github.com>

19 files changed, 861 insertions, 722 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8ddb1c76..89d7f9e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,7 +12,7 @@ is breaking anyways, semantic versioning is not followed.
 
 - Add `SimulationPathfinderExecutionPlugin`, an alternative execution engine for the pathfinder with smoother movement.
 - The pathfinder can now traverse on the surface of water.
-- `AccountTrait` was implemented, which allows for custom refresh and join behavior for `Account`s.
+- `Account`s can now have custom refresh and join behavior using `AccountTrait`.
 - Add `Account::microsoft_with_opts` to make it easier to create accounts with custom cache files. (@ElijahBare)
 - Add an `EntityRef` type to simplify interactions with entities.
 - Implement speed/swiftness.
@@ -21,6 +21,7 @@ is breaking anyways, semantic versioning is not followed.
 - Re-implement `Client::map_component` and `map_get_component`.
 - Add `Client::exit` and `Swarm::exit` to make it easier to return from `ClientBuilder::start` or `SwarmBuilder::start`.
 - Add `Event::ConnectionFailed` for when the client failed to create its initial connection to the server.
+- `ChunkStorage` can now have custom implementations using `ChunkStorageTrait`.
 - Setting blocks now updates `Section::block_count`.
 
 ### Changed
diff --git a/azalea-client/src/local_player.rs b/azalea-client/src/local_player.rs
index e19ffd1f..274f9fcd 100644
--- a/azalea-client/src/local_player.rs
+++ b/azalea-client/src/local_player.rs
@@ -137,7 +137,7 @@ impl WorldHolder {
         WorldHolder {
             shared,
             partial: Arc::new(RwLock::new(PartialWorld::new(
-                azalea_world::chunk_storage::calculate_chunk_storage_range(
+                azalea_world::chunk::calculate_chunk_storage_range(
                     client_information.view_distance.into(),
                 ),
                 Some(entity),
diff --git a/azalea-client/src/plugins/chunks.rs b/azalea-client/src/plugins/chunks.rs
index 0937c2ef..50f9ddf8 100644
--- a/azalea-client/src/plugins/chunks.rs
+++ b/azalea-client/src/plugins/chunks.rs
@@ -99,7 +99,7 @@ pub fn handle_receive_chunk_event(
         ) {
             error!(
                 "Couldn't set chunk data: {e}. World height: {}",
-                world.chunks.height
+                world.chunks.height()
             );
         }
     }
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index e7e709e8..ecc1bc2f 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -281,7 +281,7 @@ impl GamePacketHandler<'_> {
                 // will be in the `worlds`)
 
                 *world_holder.partial.write() = PartialWorld::new(
-                    azalea_world::chunk_storage::calculate_chunk_storage_range(
+                    azalea_world::chunk::calculate_chunk_storage_range(
                         client_information.view_distance.into(),
                     ),
                     // this argument makes it so other clients don't update this player entity in a
@@ -1422,7 +1422,7 @@ impl GamePacketHandler<'_> {
                 // those will be in the `worlds`)
 
                 *world_holder.partial.write() = PartialWorld::new(
-                    azalea_world::chunk_storage::calculate_chunk_storage_range(
+                    azalea_world::chunk::calculate_chunk_storage_range(
                         client_information.view_distance.into(),
                     ),
                     Some(self.player),
diff --git a/azalea-client/tests/simulation/login_to_dimension_with_same_name.rs b/azalea-client/tests/simulation/login_to_dimension_with_same_name.rs
index ef67130f..ecd6369b 100644
--- a/azalea-client/tests/simulation/login_to_dimension_with_same_name.rs
+++ b/azalea-client/tests/simulation/login_to_dimension_with_same_name.rs
@@ -117,7 +117,7 @@ fn generic_test_login_to_dimension_with_same_name(using_respawn: bool) {
             .shared
             .read()
             .chunks
-            .height,
+            .height(),
         256
     );
 
diff --git a/azalea-physics/src/collision/world_collisions.rs b/azalea-physics/src/collision/world_collisions.rs
index 06c11a15..470f5aa8 100644
--- a/azalea-physics/src/collision/world_collisions.rs
+++ b/azalea-physics/src/collision/world_collisions.rs
@@ -86,7 +86,7 @@ impl<'a> BlockCollisionsState<'a> {
         let block_state: BlockState = if item_chunk_pos == initial_chunk_pos {
             initial_chunk
                 .and_then(|chunk| {
-                    chunk.get_block_state(&ChunkBlockPos::from(item.pos), self.world.chunks.min_y)
+                    chunk.get_block_state(&ChunkBlockPos::from(item.pos), self.world.chunks.min_y())
                 })
                 .unwrap_or(BlockState::AIR)
         } else {
@@ -174,7 +174,7 @@ impl<'a> BlockCollisionsState<'a> {
     }
 
     fn get_block_state(&mut self, block_pos: BlockPos) -> BlockState {
-        if block_pos.y < self.world.chunks.min_y {
+        if block_pos.y < self.world.chunks.min_y() {
             // below the world
             return BlockState::AIR;
         }
@@ -196,8 +196,7 @@ impl<'a> BlockCollisionsState<'a> {
 
         let sections = &chunk.sections;
         let section_index =
-            azalea_world::chunk_storage::section_index(block_pos.y, self.world.chunks.min_y)
-                as usize;
+            azalea_world::chunk::section_index(block_pos.y, self.world.chunks.min_y()) as usize;
 
         let Some(section) = sections.get(section_index) else {
             return BlockState::AIR;
diff --git a/azalea-world/benches/chunks.rs b/azalea-world/benches/chunks.rs
index e1ca94a6..44533412 100644
--- a/azalea-world/benches/chunks.rs
+++ b/azalea-world/benches/chunks.rs
@@ -1,8 +1,8 @@
 use std::hint::black_box;
 
-use azalea_core::position::ChunkBlockPos;
+use azalea_core::position::{BlockPos, ChunkBlockPos, ChunkPos};
 use azalea_registry::builtin::BlockKind;
-use azalea_world::{BitStorage, Chunk};
+use azalea_world::{BitStorage, Chunk, ChunkStorage};
 use criterion::{BatchSize, Bencher, Criterion, criterion_group, criterion_main};
 
 fn bench_chunks(c: &mut Criterion) {
@@ -23,6 +23,24 @@ fn bench_chunks(c: &mut Criterion) {
             black_box(chunk);
         });
     });
+
+    c.bench_function("ChunkStorage::get_block_state", |b| {
+        b.iter(|| {
+            let mut storage = ChunkStorage::default();
+
+            let chunk = storage.upsert(ChunkPos::new(0, 0), Chunk::default());
+
+            for x in 0..16 {
+                for z in 0..16 {
+                    for y in 0..16 {
+                        black_box(storage.get_block_state(BlockPos::new(x, y, z)));
+                    }
+                }
+            }
+
+            black_box(chunk);
+        });
+    });
 }
 
 fn bench_bitstorage_with(b: &mut Bencher, bits: usize, size: usize) {
diff --git a/azalea-world/src/chunk/mod.rs b/azalea-world/src/chunk/mod.rs
new file mode 100644
index 00000000..cd52026a
--- /dev/null
+++ b/azalea-world/src/chunk/mod.rs
@@ -0,0 +1,305 @@
+pub mod partial;
+pub mod storage;
+
+use std::{
+    collections::HashMap,
+    fmt::Debug,
+    io,
+    io::{Cursor, Write},
+};
+
+use azalea_block::block_state::{BlockState, BlockStateIntegerRepr};
+use azalea_buf::{AzBuf, BufReadError};
+use azalea_core::{
+    heightmap_kind::HeightmapKind,
+    position::{ChunkBiomePos, ChunkBlockPos, ChunkSectionBiomePos, ChunkSectionBlockPos},
+};
+use azalea_registry::data::Biome;
+use tracing::warn;
+
+use crate::{heightmap::Heightmap, palette::PalettedContainer};
+
+const SECTION_HEIGHT: u32 = 16;
+
+/// A single chunk in a world (16*?*16 blocks).
+///
+/// This only contains blocks and biomes. You can derive the height of the chunk
+/// from the number of sections, but you need a [`ChunkStorage`] to get the
+/// minimum Y coordinate.
+#[derive(Debug)]
+pub struct Chunk {
+    pub sections: Box<[Section]>,
+    /// Heightmaps are used for identifying the surface blocks in a chunk.
+    /// Usually for clients only `WorldSurface` and `MotionBlocking` are
+    /// present.
+    pub heightmaps: HashMap<HeightmapKind, Heightmap>,
+}
+
+/// A section of a chunk, i.e. a 16*16*16 block area.
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct Section {
+    /// The number of non-empty blocks in the section, which is initialized
+    /// based on a value sent to us by the server.
+    ///
+    /// This may be updated every time [`Self::get_and_set_block_state`] is
+    /// called.
+    pub block_count: u16,
+    pub states: PalettedContainer<BlockState>,
+    pub biomes: PalettedContainer<Biome>,
+}
+
+/// Get the actual stored view distance for the selected view distance.
+///
+/// For some reason, Minecraft stores an extra 3 chunks.
+pub fn calculate_chunk_storage_range(view_distance: u32) -> u32 {
+    u32::max(view_distance, 2) + 3
+}
+
+impl Default for Chunk {
+    fn default() -> Self {
+        Chunk {
+            sections: vec![Section::default(); (384 / 16) as usize].into(),
+            heightmaps: HashMap::new(),
+        }
+    }
+}
+
+impl Chunk {
+    pub fn read_with_dimension_height(
+        buf: &mut Cursor<&[u8]>,
+        dimension_height: u32,
+        min_y: i32,
+        heightmaps_data: &[(HeightmapKind, Box<[u64]>)],
+    ) -> Result<Self, BufReadError> {
+        let section_count = dimension_height / SECTION_HEIGHT;
+        let mut sections = Vec::with_capacity(section_count as usize);
+        for _ in 0..section_count {
+            let section = Section::azalea_read(buf)?;
+            sections.push(section);
+        }
+        let sections = sections.into_boxed_slice();
+
+        let mut heightmaps = HashMap::new();
+        for (kind, data) in heightmaps_data {
+            let data = data.clone();
+            let heightmap = Heightmap::new(*kind, dimension_height, min_y, data);
+            heightmaps.insert(*kind, heightmap);
+        }
+
+        Ok(Chunk {
+            sections,
+            heightmaps,
+        })
+    }
+
+    pub fn get_block_state(&self, pos: &ChunkBlockPos, min_y: i32) -> Option<BlockState> {
+        get_block_state_from_sections(&self.sections, pos, min_y)
+    }
+
+    #[must_use = "Use Chunk::set_block_state instead if you don't need the previous state"]
+    pub fn get_and_set_block_state(
+        &mut self,
+        pos: &ChunkBlockPos,
+        state: BlockState,
+        min_y: i32,
+    ) -> BlockState {
+        let section_index = section_index(pos.y, min_y);
+        let Some(section) = self.sections.get_mut(section_index as usize) else {
+            warn!(
+                "Tried to get and set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
+            );
+            return BlockState::AIR;
+        };
+        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+        let previous_state = section.get_and_set_block_state(chunk_section_pos, state);
+
+        for heightmap in self.heightmaps.values_mut() {
+            heightmap.update(pos, state, &self.sections);
+        }
+
+        previous_state
+    }
+
+    pub fn set_block_state(&mut self, pos: &ChunkBlockPos, state: BlockState, min_y: i32) {
+        let section_index = section_index(pos.y, min_y);
+        let Some(section) = self.sections.get_mut(section_index as usize) else {
+            warn!(
+                "Tried to set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
+            );
+            return;
+        };
+        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+        section.get_and_set_block_state(chunk_section_pos, state);
+
+        for heightmap in self.heightmaps.values_mut() {
+            heightmap.update(pos, state, &self.sections);
+        }
+    }
+
+    /// Get the biome at the given position, or `None` if it's out of bounds.
+    pub fn get_biome(&self, pos: ChunkBiomePos, min_y: i32) -> Option<Biome> {
+        if pos.y < min_y {
+            // y position is out of bounds
+            return None;
+        }
+        let section_index = section_index(pos.y, min_y);
+        let Some(section) = self.sections.get(section_index as usize) else {
+            warn!("Tried to get biome at out-of-bounds relative chunk position {pos:?}",);
+            return None;
+        };
+        let chunk_section_pos = ChunkSectionBiomePos::from(pos);
+        Some(section.get_biome(chunk_section_pos))
+    }
+}
+
+/// Get the block state at the given position from a list of sections. Returns
+/// `None` if the position is out of bounds.
+#[inline]
+pub fn get_block_state_from_sections(
+    sections: &[Section],
+    pos: &ChunkBlockPos,
+    min_y: i32,
+) -> Option<BlockState> {
+    if pos.y < min_y {
+        // y position is out of bounds
+        return None;
+    }
+    let section_index = section_index(pos.y, min_y) as usize;
+    if section_index >= sections.len() {
+        // y position is out of bounds
+        return None;
+    };
+    let section = &sections[section_index];
+
+    let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+    Some(section.get_block_state(chunk_section_pos))
+}
+
+impl AzBuf for Section {
+    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
+        let block_count = u16::azalea_read(buf)?;
+
+        // this is commented out because the vanilla server is wrong
+        // TODO: ^ this comment was written ages ago. needs more investigation.
+        // assert!(
+        //     block_count <= 16 * 16 * 16,
+        //     "A section has more blocks than what should be possible. This is a bug!"
+        // );
+
+        let states = PalettedContainer::<BlockState>::read(buf)?;
+
+        for i in 0..states.storage.size() {
+            if !BlockState::is_valid_state(states.storage.get(i) as BlockStateIntegerRepr) {
+                return Err(BufReadError::Custom(format!(
+                    "Invalid block state {} (index {i}) found in section.",
+                    states.storage.get(i)
+                )));
+            }
+        }
+
+        let biomes = PalettedContainer::<Biome>::read(buf)?;
+        Ok(Section {
+            block_count,
+            states,
+            biomes,
+        })
+    }
+    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
+        self.block_count.azalea_write(buf)?;
+        self.states.write(buf)?;
+        self.biomes.write(buf)?;
+        Ok(())
+    }
+}
+
+impl Section {
+    pub fn get_block_state(&self, pos: ChunkSectionBlockPos) -> BlockState {
+        self.states.get(pos)
+    }
+    pub fn get_and_set_block_state(
+        &mut self,
+        pos: ChunkSectionBlockPos,
+        state: BlockState,
+    ) -> BlockState {
+        let previous_state = self.states.get_and_set(pos, state);
+
+        if previous_state.is_air() && !state.is_air() {
+            self.block_count += 1;
+        } else if !previous_state.is_air() && state.is_air() {
+            self.block_count -= 1;
+        }
+
+        previous_state
+    }
+
+    pub fn get_biome(&self, pos: ChunkSectionBiomePos) -> Biome {
+        self.biomes.get(pos)
+    }
+    pub fn set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) {
+        self.biomes.set(pos, biome);
+    }
+    pub fn get_and_set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) -> Biome {
+        self.biomes.get_and_set(pos, biome)
+    }
+}
+
+/// Get the index of where a section is in a chunk based on its y coordinate
+/// and the minimum y coordinate of the world.
+#[inline]
+pub fn section_index(y: i32, min_y: i32) -> u32 {
+    if y < min_y {
+        #[cfg(debug_assertions)]
+        tracing::warn!("y ({y}) must be at least {min_y}");
+        #[cfg(not(debug_assertions))]
+        tracing::trace!("y ({y}) must be at least {min_y}")
+    };
+    let min_section_index = min_y >> 4;
+    ((y >> 4) - min_section_index) as u32
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::palette::SectionPos;
+
+    #[test]
+    fn test_section_index() {
+        assert_eq!(section_index(0, 0), 0);
+        assert_eq!(section_index(128, 0), 8);
+        assert_eq!(section_index(127, 0), 7);
+        assert_eq!(section_index(0, -64), 4);
+        assert_eq!(section_index(-64, -64), 0);
+        assert_eq!(section_index(-49, -64), 0);
+        assert_eq!(section_index(-48, -64), 1);
+        assert_eq!(section_index(128, -64), 12);
+    }
+
+    #[test]
+    fn serialize_and_deserialize_section() {
+        let mut states = PalettedContainer::new();
+
+        states.set(
+            SectionPos::new(1, 2, 3),
+            BlockState::try_from(BlockState::MAX_STATE).unwrap(),
+        );
+        states.set(
+            SectionPos::new(4, 5, 6),
+            BlockState::try_from(BlockState::MAX_STATE).unwrap(),
+        );
+        let biomes = PalettedContainer::new();
+        let section = Section {
+            block_count: 2,
+            states,
+            biomes,
+        };
+
+        let mut buf = Vec::new();
+        section.azalea_write(&mut buf).unwrap();
+
+        let mut cur = Cursor::new(buf.as_slice());
+        let deserialized_section = Section::azalea_read(&mut cur).unwrap();
+        assert_eq!(cur.position(), buf.len() as u64);
+
+        assert_eq!(section, deserialized_section);
+    }
+}
diff --git a/azalea-world/src/chunk/partial.rs b/azalea-world/src/chunk/partial.rs
new file mode 100644
index 00000000..e0f4912b
--- /dev/null
+++ b/azalea-world/src/chunk/partial.rs
@@ -0,0 +1,246 @@
+use std::{
+    fmt::{self, Debug},
+    io::Cursor,
+    sync::Arc,
+};
+
+use azalea_block::BlockState;
+use azalea_buf::BufReadError;
+use azalea_core::{
+    heightmap_kind::HeightmapKind,
+    position::{BlockPos, ChunkBlockPos, ChunkPos},
+};
+use parking_lot::RwLock;
+use tracing::{debug, trace, warn};
+
+use crate::{Chunk, chunk::storage::ChunkStorage};
+
+/// An efficient storage of chunks for a client that has a limited render
+/// distance.
+///
+/// This has support for using a shared [`ChunkStorage`].
+pub struct PartialChunkStorage {
+    /// The center of the view, i.e. the chunk the player is currently in.
+    view_center: ChunkPos,
+    pub(crate) chunk_radius: u32,
+    view_range: u32,
+    // chunks is a list of size chunk_radius * chunk_radius
+    chunks: Box<[Option<Arc<RwLock<Chunk>>>]>,
+}
+
+impl PartialChunkStorage {
+    pub fn new(chunk_radius: u32) -> Self {
+        let view_range = chunk_radius * 2 + 1;
+        PartialChunkStorage {
+            view_center: ChunkPos::new(0, 0),
+            chunk_radius,
+            view_range,
+            chunks: vec![None; (view_range * view_range) as usize].into(),
+        }
+    }
+
+    /// Update the chunk to center the view on.
+    ///
+    /// This should be called when the client receives a `SetChunkCacheCenter`
+    /// packet.
+    pub fn update_view_center(&mut self, view_center: ChunkPos) {
+        // this code block makes it force unload the chunks that are out of range after
+        // updating the view center. it's usually fine without it but the commented code
+        // is there in case you want to temporarily uncomment to test something
+
+        // ```
+        // for index in 0..self.chunks.len() {
+        //     let chunk_pos = self.chunk_pos_from_index(index);
+        //     if !in_range_for_view_center_and_radius(&chunk_pos, view_center, self.chunk_radius) {
+        //         self.chunks[index] = None;
+        //     }
+        // }
+        // ```
+
+        self.view_center = view_center;
+    }
+
+    /// Get the center of the view. This is usually the chunk that the player is
+    /// in.
+    pub fn view_center(&self) -> ChunkPos {
+        self.view_center
+    }
+
+    pub fn view_range(&self) -> u32 {
+        self.view_range
+    }
+
+    pub fn index_from_chunk_pos(&self, chunk_pos: &ChunkPos) -> usize {
+        let view_range = self.view_range as i32;
+
+        let x = i32::rem_euclid(chunk_pos.x, view_range) * view_range;
+        let z = i32::rem_euclid(chunk_pos.z, view_range);
+        (x + z) as usize
+    }
+
+    pub fn chunk_pos_from_index(&self, index: usize) -> ChunkPos {
+        let view_range = self.view_range as i32;
+
+        // find the base from the view center
+        let base_x = self.view_center.x.div_euclid(view_range) * view_range;
+        let base_z = self.view_center.z.div_euclid(view_range) * view_range;
+
+        // add the offset from the base
+        let offset_x = index as i32 / view_range;
+        let offset_z = index as i32 % view_range;
+
+        ChunkPos::new(base_x + offset_x, base_z + offset_z)
+    }
+
+    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
+        in_range_for_view_center_and_radius(chunk_pos, self.view_center, self.chunk_radius)
+    }
+
+    pub fn set_block_state(
+        &self,
+        pos: BlockPos,
+        state: BlockState,
+        chunk_storage: &ChunkStorage,
+    ) -> Option<BlockState> {
+        if pos.y < chunk_storage.min_y()
+            || pos.y >= (chunk_storage.min_y() + chunk_storage.height() as i32)
+        {
+            return None;
+        }
+        let chunk_pos = ChunkPos::from(pos);
+        let chunk_lock = chunk_storage.get(&chunk_pos)?;
+        let mut chunk = chunk_lock.write();
+        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, chunk_storage.min_y()))
+    }
+
+    pub fn replace_with_packet_data(
+        &mut self,
+        pos: &ChunkPos,
+        data: &mut Cursor<&[u8]>,
+        heightmaps: &[(HeightmapKind, Box<[u64]>)],
+        chunk_storage: &mut ChunkStorage,
+    ) -> Result<(), BufReadError> {
+        debug!("Replacing chunk at {:?}", pos);
+        if !self.in_range(pos) {
+            warn!("Ignoring chunk since it's not in the view range: {pos:?}");
+            return Ok(());
+        }
+
+        let chunk = Chunk::read_with_dimension_height(
+            data,
+            chunk_storage.height(),
+            chunk_storage.min_y(),
+            heightmaps,
+        )?;
+
+        self.set(pos, Some(chunk), chunk_storage);
+        trace!("Loaded chunk {pos:?}");
+
+        Ok(())
+    }
+
+    /// Get a [`Chunk`] within render distance, or `None` if it's not loaded.
+    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
+    pub fn limited_get(&self, pos: &ChunkPos) -> Option<&Arc<RwLock<Chunk>>> {
+        if !self.in_range(pos) {
+            warn!(
+                "Chunk at {:?} is not in the render distance (center: {:?}, {} chunks)",
+                pos, self.view_center, self.chunk_radius,
+            );
+            return None;
+        }
+
+        let index = self.index_from_chunk_pos(pos);
+        self.chunks[index].as_ref()
+    }
+    /// Get a mutable reference to a [`Chunk`] within render distance, or
+    /// `None` if it's not loaded.
+    ///
+    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
+    pub fn limited_get_mut(&mut self, pos: &ChunkPos) -> Option<&mut Option<Arc<RwLock<Chunk>>>> {
+        if !self.in_range(pos) {
+            return None;
+        }
+
+        let index = self.index_from_chunk_pos(pos);
+
+        Some(&mut self.chunks[index])
+    }
+
+    /// Set a chunk in the shared storage and reference it from the limited
+    /// storage.
+    ///
+    /// Use [`Self::limited_set`] if you already have an `Arc<RwLock<Chunk>>`.
+    ///
+    /// # Panics
+    /// If the chunk is not in the render distance.
+    pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Chunk>, chunk_storage: &mut ChunkStorage) {
+        let new_chunk = chunk.map(|c| chunk_storage.upsert(*pos, c));
+        self.limited_set(pos, new_chunk);
+    }
+
+    /// Set a chunk in our limited storage, useful if your chunk is already
+    /// referenced somewhere else and you want to make it also be referenced by
+    /// this storage.
+    ///
+    /// Use [`Self::set`] if you don't already have an `Arc<RwLock<Chunk>>`.
+    ///
+    /// # Panics
+    /// If the chunk is not in the render distance.
+    pub fn limited_set(&mut self, pos: &ChunkPos, chunk: Option<Arc<RwLock<Chunk>>>) {
+        if let Some(chunk_mut) = self.limited_get_mut(pos) {
+            *chunk_mut = chunk;
+        }
+    }
+
+    /// Get an iterator over all the chunks in the storage.
+    pub fn chunks(&self) -> impl Iterator<Item = &Option<Arc<RwLock<Chunk>>>> {
+        self.chunks.iter()
+    }
+}
+
+impl Debug for PartialChunkStorage {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("PartialChunkStorage")
+            .field("view_center", &self.view_center)
+            .field("chunk_radius", &self.chunk_radius)
+            .field("view_range", &self.view_range)
+            // .field("chunks", &self.chunks)
+            .field("chunks", &format_args!("{} items", self.chunks.len()))
+            .finish()
+    }
+}
+
+impl Default for PartialChunkStorage {
+    fn default() -> Self {
+        Self::new(8)
+    }
+}
+
+pub fn in_range_for_view_center_and_radius(
+    chunk_pos: &ChunkPos,
+    view_center: ChunkPos,
+    chunk_radius: u32,
+) -> bool {
+    (chunk_pos.x - view_center.x).unsigned_abs() <= chunk_radius
+        && (chunk_pos.z - view_center.z).unsigned_abs() <= chunk_radius
+}
+
+#[cfg(test)]
+mod tests {
+    use azalea_core::position::ChunkPos;
+
+    use crate::chunk::partial::PartialChunkStorage;
+
+    #[test]
+    fn test_chunk_pos_from_index() {
+        let mut partial_chunk_storage = PartialChunkStorage::new(5);
+        partial_chunk_storage.update_view_center(ChunkPos::new(0, -1));
+        assert_eq!(
+            partial_chunk_storage.chunk_pos_from_index(
+                partial_chunk_storage.index_from_chunk_pos(&ChunkPos::new(2, -1))
+            ),
+            ChunkPos::new(2, -1),
+        );
+    }
+}
diff --git a/azalea-world/src/chunk/storage.rs b/azalea-world/src/chunk/storage.rs
new file mode 100644
index 00000000..926c0b44
--- /dev/null
+++ b/azalea-world/src/chunk/storage.rs
@@ -0,0 +1,242 @@
+use std::{
+    any::Any,
+    collections::hash_map::Entry,
+    fmt::{self, Debug},
+    ops::{Deref, DerefMut},
+    sync::{Arc, Weak},
+};
+
+use azalea_block::{BlockState, fluid_state::FluidState};
+use azalea_core::position::{BlockPos, ChunkBiomePos, ChunkBlockPos, ChunkPos};
+use azalea_registry::data::Biome;
+use nohash_hasher::IntMap;
+use parking_lot::RwLock;
+
+use crate::Chunk;
+
+/// An abstract chunk storage backed by a [`ChunkStorageTrait`] implementation.
+///
+/// By default, this wraps a [`WeakChunkStorage`].
+pub struct ChunkStorage(pub Box<dyn ChunkStorageTrait>);
+
+pub trait ChunkStorageTrait: Send + Sync + Any {
+    /// Return the lowest y coordinate in the world, usually `-64`.
+    fn min_y(&self) -> i32;
+    /// Return the height of the world in blocks, usually `384`.
+    fn height(&self) -> u32;
+    /// Return a reference to the chunk from the storage.
+    #[must_use]
+    fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>>;
+    /// Insert the chunk into the storage and return a reference to it.
+    ///
+    /// Since the storage may be a [`WeakChunkStorage`], you must immediately
+    /// put the returned `Arc<RwLock<Chunk>>` somewhere (probably a
+    /// [`PartialChunkStorage`](crate::PartialChunkStorage)).
+    #[must_use]
+    fn upsert(&mut self, pos: ChunkPos, chunk: Chunk) -> Arc<RwLock<Chunk>>;
+    fn chunks(&self) -> Box<[&ChunkPos]>;
+    fn clone_box(&self) -> Box<dyn ChunkStorageTrait>;
+
+    // these impls are here instead of in the `impl ChunkStorage` so rust is able to
+    // inline more and thus optimize them a lot better.
+
+    /// Returns the [`BlockState`] at the given position.
+    ///
+    /// If the block is outside of the world, then `None` is returned.
+    fn get_block_state(&self, pos: BlockPos) -> Option<BlockState> {
+        let chunk = self.get(&ChunkPos::from(pos))?;
+        let chunk = chunk.read();
+        chunk.get_block_state(&ChunkBlockPos::from(pos), self.min_y())
+    }
+    /// Set a [`BlockState`] at the given position.
+    ///
+    /// Returns the block that was previously there, or `None` if the position
+    /// is outside of the world.
+    fn set_block_state(&self, pos: BlockPos, state: BlockState) -> Option<BlockState> {
+        if pos.y < self.min_y() || pos.y >= (self.min_y() + self.height() as i32) {
+            return None;
+        }
+        let chunk = self.get(&ChunkPos::from(pos))?;
+        let mut chunk = chunk.write();
+        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, self.min_y()))
+    }
+    fn get_fluid_state(&self, pos: BlockPos) -> Option<FluidState> {
+        let block_state = self.get_block_state(pos)?;
+        Some(FluidState::from(block_state))
+    }
+    fn get_biome(&self, pos: BlockPos) -> Option<Biome> {
+        let chunk = self.get(&ChunkPos::from(pos))?;
+        let chunk = chunk.read();
+        chunk.get_biome(ChunkBiomePos::from(pos), self.min_y())
+    }
+}
+impl ChunkStorage {
+    /// Create a storage backed by a [`WeakChunkStorage`] with the given world
+    /// dimensions.
+    pub fn new(height: u32, min_y: i32) -> Self {
+        Self(Box::new(WeakChunkStorage::new(height, min_y)))
+    }
+
+    /// Create a storage backed by a custom [`ChunkStorageTrait`]
+    /// implementation.
+    pub fn new_with(inner: Box<dyn ChunkStorageTrait>) -> Self {
+        Self(inner)
+    }
+}
+
+impl Deref for ChunkStorage {
+    type Target = dyn ChunkStorageTrait;
+
+    fn deref(&self) -> &Self::Target {
+        &*self.0
+    }
+}
+impl DerefMut for ChunkStorage {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut *self.0
+    }
+}
+impl Clone for ChunkStorage {
+    fn clone(&self) -> Self {
+        Self(self.0.clone_box())
+    }
+}
+
+/// A storage for chunks where they're only stored weakly, so if they're not
+/// actively being used somewhere else they'll be forgotten.
+///
+/// This is used for shared worlds.
+///
+/// This is relatively cheap to clone since it's just an `IntMap` with `Weak`
+/// pointers.
+#[derive(Clone, Debug)]
+pub struct WeakChunkStorage {
+    /// The height of the world.
+    ///
+    /// To get the maximum y position (exclusive), you have to combine this with
+    /// [`Self::min_y`].
+    pub height: u32,
+    /// The lowest y position in the world that can still have blocks placed on
+    /// it.
+    ///
+    /// This exists because in modern Minecraft versions, worlds can extend
+    /// below y=0.
+    pub min_y: i32,
+    pub map: IntMap<ChunkPos, Weak<RwLock<Chunk>>>,
+}
+
+impl WeakChunkStorage {
+    pub fn new(height: u32, min_y: i32) -> Self {
+        WeakChunkStorage {
+            height,
+            min_y,
+            map: IntMap::default(),
+        }
+    }
+}
+impl ChunkStorageTrait for WeakChunkStorage {
+    fn min_y(&self) -> i32 {
+        self.min_y
+    }
+    fn height(&self) -> u32 {
+        self.height
+    }
+    fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
+        self.map.get(pos).and_then(|chunk| chunk.upgrade())
+    }
+
+    fn upsert(&mut self, pos: ChunkPos, chunk: Chunk) -> Arc<RwLock<Chunk>> {
+        match self.map.entry(pos) {
+            Entry::Occupied(mut e) => {
+                if let Some(existing) = e.get_mut().upgrade() {
+                    *existing.write() = chunk;
+                    existing
+                } else {
+                    let arc = Arc::new(RwLock::new(chunk));
+                    e.insert(Arc::downgrade(&arc));
+                    arc
+                }
+            }
+            Entry::Vacant(e) => {
+                let arc = Arc::new(RwLock::new(chunk));
+                e.insert(Arc::downgrade(&arc));
+                arc
+            }
+        }
+    }
+
+    fn chunks(&self) -> Box<[&ChunkPos]> {
+        self.map.keys().collect::<Vec<_>>().into_boxed_slice()
+    }
+
+    fn clone_box(&self) -> Box<dyn ChunkStorageTrait> {
+        Box::new(self.clone())
+    }
+}
+
+impl Default for WeakChunkStorage {
+    fn default() -> Self {
+        Self::new(384, -64)
+    }
+}
+impl Default for ChunkStorage {
+    fn default() -> Self {
+        Self::new(384, -64)
+    }
+}
+
+impl Debug for ChunkStorage {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ChunkStorage")
+            .field("min_y", &self.0.min_y())
+            .field("height", &self.0.height())
+            .field("chunk_count", &self.0.chunks().len())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use azalea_core::position::{BlockPos, ChunkPos};
+
+    use crate::{
+        Chunk,
+        chunk::{partial::PartialChunkStorage, storage::ChunkStorage},
+    };
+
+    #[test]
+    fn test_out_of_bounds_y() {
+        let mut chunk_storage = ChunkStorage::default();
+        let mut partial_chunk_storage = PartialChunkStorage::default();
+        partial_chunk_storage.set(
+            &ChunkPos { x: 0, z: 0 },
+            Some(Chunk::default()),
+            &mut chunk_storage,
+        );
+        assert!(
+            chunk_storage
+                .get_block_state(BlockPos { x: 0, y: 319, z: 0 })
+                .is_some()
+        );
+        assert!(
+            chunk_storage
+                .get_block_state(BlockPos { x: 0, y: 320, z: 0 })
+                .is_none()
+        );
+        assert!(
+            chunk_storage
+                .get_block_state(BlockPos { x: 0, y: 338, z: 0 })
+                .is_none()
+        );
+        assert!(
+            chunk_storage
+                .get_block_state(BlockPos { x: 0, y: -64, z: 0 })
+                .is_some()
+        );
+        assert!(
+            chunk_storage
+                .get_block_state(BlockPos { x: 0, y: -65, z: 0 })
+                .is_none()
+        );
+    }
+}
diff --git a/azalea-world/src/chunk_storage.rs b/azalea-world/src/chunk_storage.rs
deleted file mode 100644
index 282c7d51..00000000
--- a/azalea-world/src/chunk_storage.rs
+++ /dev/null
@@ -1,677 +0,0 @@
-use std::{
-    collections::{HashMap, hash_map::Entry},
-    fmt,
-    fmt::Debug,
-    io,
-    io::{Cursor, Write},
-    sync::{Arc, Weak},
-};
-
-use azalea_block::{
-    block_state::{BlockState, BlockStateIntegerRepr},
-    fluid_state::FluidState,
-};
-use azalea_buf::{AzBuf, BufReadError};
-use azalea_core::{
-    heightmap_kind::HeightmapKind,
-    position::{
-        BlockPos, ChunkBiomePos, ChunkBlockPos, ChunkPos, ChunkSectionBiomePos,
-        ChunkSectionBlockPos,
-    },
-};
-use azalea_registry::data::Biome;
-use nohash_hasher::IntMap;
-use parking_lot::RwLock;
-use tracing::{debug, trace, warn};
-
-use crate::{heightmap::Heightmap, palette::PalettedContainer};
-
-const SECTION_HEIGHT: u32 = 16;
-
-/// An efficient storage of chunks for a client that has a limited render
-/// distance.
-///
-/// This has support for using a shared [`ChunkStorage`].
-pub struct PartialChunkStorage {
-    /// The center of the view, i.e. the chunk the player is currently in.
-    view_center: ChunkPos,
-    pub(crate) chunk_radius: u32,
-    view_range: u32,
-    // chunks is a list of size chunk_radius * chunk_radius
-    chunks: Box<[Option<Arc<RwLock<Chunk>>>]>,
-}
-
-/// A storage for chunks where they're only stored weakly, so if they're not
-/// actively being used somewhere else they'll be forgotten.
-///
-/// This is used for shared worlds.
-///
-/// This is relatively cheap to clone since it's just an `IntMap` with `Weak`
-/// pointers.
-#[derive(Clone, Debug)]
-pub struct ChunkStorage {
-    /// The height of the world.
-    ///
-    /// To get the maximum y position (exclusive), you have to combine this with
-    /// [`Self::min_y`].
-    pub height: u32,
-    /// The lowest y position in the world that can still have blocks placed on
-    /// it.
-    ///
-    /// This exists because in modern Minecraft versions, worlds can extend
-    /// below y=0.
-    pub min_y: i32,
-    pub map: IntMap<ChunkPos, Weak<RwLock<Chunk>>>,
-}
-
-/// A single chunk in a world (16*?*16 blocks).
-///
-/// This only contains blocks and biomes. You can derive the height of the chunk
-/// from the number of sections, but you need a [`ChunkStorage`] to get the
-/// minimum Y coordinate.
-#[derive(Debug)]
-pub struct Chunk {
-    pub sections: Box<[Section]>,
-    /// Heightmaps are used for identifying the surface blocks in a chunk.
-    /// Usually for clients only `WorldSurface` and `MotionBlocking` are
-    /// present.
-    pub heightmaps: HashMap<HeightmapKind, Heightmap>,
-}
-
-/// A section of a chunk, i.e. a 16*16*16 block area.
-#[derive(Clone, Debug, Default, PartialEq)]
-pub struct Section {
-    /// The number of non-empty blocks in the section, which is initialized
-    /// based on a value sent to us by the server.
-    ///
-    /// This may be updated every time [`Self::get_and_set_block_state`] is
-    /// called.
-    pub block_count: u16,
-    pub states: PalettedContainer<BlockState>,
-    pub biomes: PalettedContainer<Biome>,
-}
-
-/// Get the actual stored view distance for the selected view distance.
-///
-/// For some reason, Minecraft stores an extra 3 chunks.
-pub fn calculate_chunk_storage_range(view_distance: u32) -> u32 {
-    u32::max(view_distance, 2) + 3
-}
-
-impl Default for Chunk {
-    fn default() -> Self {
-        Chunk {
-            sections: vec![Section::default(); (384 / 16) as usize].into(),
-            heightmaps: HashMap::new(),
-        }
-    }
-}
-
-impl PartialChunkStorage {
-    pub fn new(chunk_radius: u32) -> Self {
-        let view_range = chunk_radius * 2 + 1;
-        PartialChunkStorage {
-            view_center: ChunkPos::new(0, 0),
-            chunk_radius,
-            view_range,
-            chunks: vec![None; (view_range * view_range) as usize].into(),
-        }
-    }
-
-    /// Update the chunk to center the view on.
-    ///
-    /// This should be called when the client receives a `SetChunkCacheCenter`
-    /// packet.
-    pub fn update_view_center(&mut self, view_center: ChunkPos) {
-        // this code block makes it force unload the chunks that are out of range after
-        // updating the view center. it's usually fine without it but the commented code
-        // is there in case you want to temporarily uncomment to test something
-
-        // ```
-        // for index in 0..self.chunks.len() {
-        //     let chunk_pos = self.chunk_pos_from_index(index);
-        //     if !in_range_for_view_center_and_radius(&chunk_pos, view_center, self.chunk_radius) {
-        //         self.chunks[index] = None;
-        //     }
-        // }
-        // ```
-
-        self.view_center = view_center;
-    }
-
-    /// Get the center of the view. This is usually the chunk that the player is
-    /// in.
-    pub fn view_center(&self) -> ChunkPos {
-        self.view_center
-    }
-
-    pub fn view_range(&self) -> u32 {
-        self.view_range
-    }
-
-    pub fn index_from_chunk_pos(&self, chunk_pos: &ChunkPos) -> usize {
-        let view_range = self.view_range as i32;
-
-        let x = i32::rem_euclid(chunk_pos.x, view_range) * view_range;
-        let z = i32::rem_euclid(chunk_pos.z, view_range);
-        (x + z) as usize
-    }
-
-    pub fn chunk_pos_from_index(&self, index: usize) -> ChunkPos {
-        let view_range = self.view_range as i32;
-
-        // find the base from the view center
-        let base_x = self.view_center.x.div_euclid(view_range) * view_range;
-        let base_z = self.view_center.z.div_euclid(view_range) * view_range;
-
-        // add the offset from the base
-        let offset_x = index as i32 / view_range;
-        let offset_z = index as i32 % view_range;
-
-        ChunkPos::new(base_x + offset_x, base_z + offset_z)
-    }
-
-    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
-        in_range_for_view_center_and_radius(chunk_pos, self.view_center, self.chunk_radius)
-    }
-
-    pub fn set_block_state(
-        &self,
-        pos: BlockPos,
-        state: BlockState,
-        chunk_storage: &ChunkStorage,
-    ) -> Option<BlockState> {
-        if pos.y < chunk_storage.min_y
-            || pos.y >= (chunk_storage.min_y + chunk_storage.height as i32)
-        {
-            return None;
-        }
-        let chunk_pos = ChunkPos::from(pos);
-        let chunk_lock = chunk_storage.get(&chunk_pos)?;
-        let mut chunk = chunk_lock.write();
-        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, chunk_storage.min_y))
-    }
-
-    pub fn replace_with_packet_data(
-        &mut self,
-        pos: &ChunkPos,
-        data: &mut Cursor<&[u8]>,
-        heightmaps: &[(HeightmapKind, Box<[u64]>)],
-        chunk_storage: &mut ChunkStorage,
-    ) -> Result<(), BufReadError> {
-        debug!("Replacing chunk at {:?}", pos);
-        if !self.in_range(pos) {
-            warn!("Ignoring chunk since it's not in the view range: {pos:?}");
-            return Ok(());
-        }
-
-        let chunk = Chunk::read_with_dimension_height(
-            data,
-            chunk_storage.height,
-            chunk_storage.min_y,
-            heightmaps,
-        )?;
-
-        self.set(pos, Some(chunk), chunk_storage);
-        trace!("Loaded chunk {pos:?}");
-
-        Ok(())
-    }
-
-    /// Get a [`Chunk`] within render distance, or `None` if it's not loaded.
-    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
-    pub fn limited_get(&self, pos: &ChunkPos) -> Option<&Arc<RwLock<Chunk>>> {
-        if !self.in_range(pos) {
-            warn!(
-                "Chunk at {:?} is not in the render distance (center: {:?}, {} chunks)",
-                pos, self.view_center, self.chunk_radius,
-            );
-            return None;
-        }
-
-        let index = self.index_from_chunk_pos(pos);
-        self.chunks[index].as_ref()
-    }
-    /// Get a mutable reference to a [`Chunk`] within render distance, or
-    /// `None` if it's not loaded.
-    ///
-    /// Use [`ChunkStorage::get`] to get a chunk from the shared storage.
-    pub fn limited_get_mut(&mut self, pos: &ChunkPos) -> Option<&mut Option<Arc<RwLock<Chunk>>>> {
-        if !self.in_range(pos) {
-            return None;
-        }
-
-        let index = self.index_from_chunk_pos(pos);
-
-        Some(&mut self.chunks[index])
-    }
-
-    /// Set a chunk in the shared storage and reference it from the limited
-    /// storage.
-    ///
-    /// Use [`Self::limited_set`] if you already have an `Arc<RwLock<Chunk>>`.
-    ///
-    /// # Panics
-    /// If the chunk is not in the render distance.
-    pub fn set(&mut self, pos: &ChunkPos, chunk: Option<Chunk>, chunk_storage: &mut ChunkStorage) {
-        let new_chunk;
-
-        // add the chunk to the shared storage
-        if let Some(chunk) = chunk {
-            match chunk_storage.map.entry(*pos) {
-                Entry::Occupied(mut e) => {
-                    if let Some(old_chunk) = e.get_mut().upgrade() {
-                        *old_chunk.write() = chunk;
-                        new_chunk = Some(old_chunk);
-                    } else {
-                        let chunk_lock = Arc::new(RwLock::new(chunk));
-                        e.insert(Arc::downgrade(&chunk_lock));
-                        new_chunk = Some(chunk_lock);
-                    }
-                }
-                Entry::Vacant(e) => {
-                    let chunk_lock = Arc::new(RwLock::new(chunk));
-                    e.insert(Arc::downgrade(&chunk_lock));
-                    new_chunk = Some(chunk_lock);
-                }
-            }
-        } else {
-            // don't remove it from the shared storage, since it'll be removed
-            // automatically if this was the last reference
-
-            new_chunk = None;
-        }
-
-        self.limited_set(pos, new_chunk);
-    }
-
-    /// Set a chunk in our limited storage, useful if your chunk is already
-    /// referenced somewhere else and you want to make it also be referenced by
-    /// this storage.
-    ///
-    /// Use [`Self::set`] if you don't already have an `Arc<RwLock<Chunk>>`.
-    ///
-    /// # Panics
-    /// If the chunk is not in the render distance.
-    pub fn limited_set(&mut self, pos: &ChunkPos, chunk: Option<Arc<RwLock<Chunk>>>) {
-        if let Some(chunk_mut) = self.limited_get_mut(pos) {
-            *chunk_mut = chunk;
-        }
-    }
-
-    /// Get an iterator over all the chunks in the storage.
-    pub fn chunks(&self) -> impl Iterator<Item = &Option<Arc<RwLock<Chunk>>>> {
-        self.chunks.iter()
-    }
-}
-impl ChunkStorage {
-    pub fn new(height: u32, min_y: i32) -> Self {
-        ChunkStorage {
-            height,
-            min_y,
-            map: IntMap::default(),
-        }
-    }
-
-    pub fn get(&self, pos: &ChunkPos) -> Option<Arc<RwLock<Chunk>>> {
-        self.map.get(pos).and_then(|chunk| chunk.upgrade())
-    }
-
-    pub fn get_block_state(&self, pos: BlockPos) -> Option<BlockState> {
-        let chunk_pos = ChunkPos::from(pos);
-        let chunk = self.get(&chunk_pos)?;
-        let chunk = chunk.read();
-        chunk.get_block_state(&ChunkBlockPos::from(pos), self.min_y)
-    }
-
-    pub fn get_fluid_state(&self, pos: BlockPos) -> Option<FluidState> {
-        let block_state = self.get_block_state(pos)?;
-        Some(FluidState::from(block_state))
-    }
-
-    pub fn get_biome(&self, pos: BlockPos) -> Option<Biome> {
-        let chunk_pos = ChunkPos::from(pos);
-        let chunk = self.get(&chunk_pos)?;
-        let chunk = chunk.read();
-        chunk.get_biome(ChunkBiomePos::from(pos), self.min_y)
-    }
-
-    pub fn set_block_state(&self, pos: BlockPos, state: BlockState) -> Option<BlockState> {
-        if pos.y < self.min_y || pos.y >= (self.min_y + self.height as i32) {
-            return None;
-        }
-        let chunk_pos = ChunkPos::from(pos);
-        let chunk = self.get(&chunk_pos)?;
-        let mut chunk = chunk.write();
-        Some(chunk.get_and_set_block_state(&ChunkBlockPos::from(pos), state, self.min_y))
-    }
-}
-
-pub fn in_range_for_view_center_and_radius(
-    chunk_pos: &ChunkPos,
-    view_center: ChunkPos,
-    chunk_radius: u32,
-) -> bool {
-    (chunk_pos.x - view_center.x).unsigned_abs() <= chunk_radius
-        && (chunk_pos.z - view_center.z).unsigned_abs() <= chunk_radius
-}
-
-impl Chunk {
-    pub fn read_with_dimension_height(
-        buf: &mut Cursor<&[u8]>,
-        dimension_height: u32,
-        min_y: i32,
-        heightmaps_data: &[(HeightmapKind, Box<[u64]>)],
-    ) -> Result<Self, BufReadError> {
-        let section_count = dimension_height / SECTION_HEIGHT;
-        let mut sections = Vec::with_capacity(section_count as usize);
-        for _ in 0..section_count {
-            let section = Section::azalea_read(buf)?;
-            sections.push(section);
-        }
-        let sections = sections.into_boxed_slice();
-
-        let mut heightmaps = HashMap::new();
-        for (kind, data) in heightmaps_data {
-            let data = data.clone();
-            let heightmap = Heightmap::new(*kind, dimension_height, min_y, data);
-            heightmaps.insert(*kind, heightmap);
-        }
-
-        Ok(Chunk {
-            sections,
-            heightmaps,
-        })
-    }
-
-    pub fn get_block_state(&self, pos: &ChunkBlockPos, min_y: i32) -> Option<BlockState> {
-        get_block_state_from_sections(&self.sections, pos, min_y)
-    }
-
-    #[must_use = "Use Chunk::set_block_state instead if you don't need the previous state"]
-    pub fn get_and_set_block_state(
-        &mut self,
-        pos: &ChunkBlockPos,
-        state: BlockState,
-        min_y: i32,
-    ) -> BlockState {
-        let section_index = section_index(pos.y, min_y);
-        let Some(section) = self.sections.get_mut(section_index as usize) else {
-            warn!(
-                "Tried to get and set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
-            );
-            return BlockState::AIR;
-        };
-        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-        let previous_state = section.get_and_set_block_state(chunk_section_pos, state);
-
-        for heightmap in self.heightmaps.values_mut() {
-            heightmap.update(pos, state, &self.sections);
-        }
-
-        previous_state
-    }
-
-    pub fn set_block_state(&mut self, pos: &ChunkBlockPos, state: BlockState, min_y: i32) {
-        let section_index = section_index(pos.y, min_y);
-        let Some(section) = self.sections.get_mut(section_index as usize) else {
-            warn!(
-                "Tried to set block state {state:?} at out-of-bounds relative chunk position {pos:?}",
-            );
-            return;
-        };
-        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-        section.get_and_set_block_state(chunk_section_pos, state);
-
-        for heightmap in self.heightmaps.values_mut() {
-            heightmap.update(pos, state, &self.sections);
-        }
-    }
-
-    /// Get the biome at the given position, or `None` if it's out of bounds.
-    pub fn get_biome(&self, pos: ChunkBiomePos, min_y: i32) -> Option<Biome> {
-        if pos.y < min_y {
-            // y position is out of bounds
-            return None;
-        }
-        let section_index = section_index(pos.y, min_y);
-        let Some(section) = self.sections.get(section_index as usize) else {
-            warn!("Tried to get biome at out-of-bounds relative chunk position {pos:?}",);
-            return None;
-        };
-        let chunk_section_pos = ChunkSectionBiomePos::from(pos);
-        Some(section.get_biome(chunk_section_pos))
-    }
-}
-
-/// Get the block state at the given position from a list of sections. Returns
-/// `None` if the position is out of bounds.
-#[inline]
-pub fn get_block_state_from_sections(
-    sections: &[Section],
-    pos: &ChunkBlockPos,
-    min_y: i32,
-) -> Option<BlockState> {
-    if pos.y < min_y {
-        // y position is out of bounds
-        return None;
-    }
-    let section_index = section_index(pos.y, min_y) as usize;
-    if section_index >= sections.len() {
-        // y position is out of bounds
-        return None;
-    };
-    let section = &sections[section_index];
-
-    let chunk_section_pos = ChunkSectionBlockPos::from(pos);
-    Some(section.get_block_state(chunk_section_pos))
-}
-
-// impl AzBuf for Chunk {
-//     fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
-//         for section in &self.sections {
-//             section.azalea_write(buf)?;
-//         }
-//         Ok(())
-//     }
-// }
-
-impl Debug for PartialChunkStorage {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("PartialChunkStorage")
-            .field("view_center", &self.view_center)
-            .field("chunk_radius", &self.chunk_radius)
-            .field("view_range", &self.view_range)
-            // .field("chunks", &self.chunks)
-            .field("chunks", &format_args!("{} items", self.chunks.len()))
-            .finish()
-    }
-}
-
-impl AzBuf for Section {
-    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
-        let block_count = u16::azalea_read(buf)?;
-
-        // this is commented out because the vanilla server is wrong
-        // TODO: ^ this comment was written ages ago. needs more investigation.
-        // assert!(
-        //     block_count <= 16 * 16 * 16,
-        //     "A section has more blocks than what should be possible. This is a bug!"
-        // );
-
-        let states = PalettedContainer::<BlockState>::read(buf)?;
-
-        for i in 0..states.storage.size() {
-            if !BlockState::is_valid_state(states.storage.get(i) as BlockStateIntegerRepr) {
-                return Err(BufReadError::Custom(format!(
-                    "Invalid block state {} (index {i}) found in section.",
-                    states.storage.get(i)
-                )));
-            }
-        }
-
-        let biomes = PalettedContainer::<Biome>::read(buf)?;
-        Ok(Section {
-            block_count,
-            states,
-            biomes,
-        })
-    }
-    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
-        self.block_count.azalea_write(buf)?;
-        self.states.write(buf)?;
-        self.biomes.write(buf)?;
-        Ok(())
-    }
-}
-
-impl Section {
-    pub fn get_block_state(&self, pos: ChunkSectionBlockPos) -> BlockState {
-        self.states.get(pos)
-    }
-    pub fn get_and_set_block_state(
-        &mut self,
-        pos: ChunkSectionBlockPos,
-        state: BlockState,
-    ) -> BlockState {
-        let previous_state = self.states.get_and_set(pos, state);
-
-        if previous_state.is_air() && !state.is_air() {
-            self.block_count += 1;
-        } else if !previous_state.is_air() && state.is_air() {
-            self.block_count -= 1;
-        }
-
-        previous_state
-    }
-
-    pub fn get_biome(&self, pos: ChunkSectionBiomePos) -> Biome {
-        self.biomes.get(pos)
-    }
-    pub fn set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) {
-        self.biomes.set(pos, biome);
-    }
-    pub fn get_and_set_biome(&mut self, pos: ChunkSectionBiomePos, biome: Biome) -> Biome {
-        self.biomes.get_and_set(pos, biome)
-    }
-}
-
-impl Default for PartialChunkStorage {
-    fn default() -> Self {
-        Self::new(8)
-    }
-}
-impl Default for ChunkStorage {
-    fn default() -> Self {
-        Self::new(384, -64)
-    }
-}
-
-/// Get the index of where a section is in a chunk based on its y coordinate
-/// and the minimum y coordinate of the world.
-#[inline]
-pub fn section_index(y: i32, min_y: i32) -> u32 {
-    if y < min_y {
-        #[cfg(debug_assertions)]
-        warn!("y ({y}) must be at least {min_y}");
-        #[cfg(not(debug_assertions))]
-        trace!("y ({y}) must be at least {min_y}")
-    };
-    let min_section_index = min_y >> 4;
-    ((y >> 4) - min_section_index) as u32
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::palette::SectionPos;
-
-    #[test]
-    fn test_section_index() {
-        assert_eq!(section_index(0, 0), 0);
-        assert_eq!(section_index(128, 0), 8);
-        assert_eq!(section_index(127, 0), 7);
-        assert_eq!(section_index(0, -64), 4);
-        assert_eq!(section_index(-64, -64), 0);
-        assert_eq!(section_index(-49, -64), 0);
-        assert_eq!(section_index(-48, -64), 1);
-        assert_eq!(section_index(128, -64), 12);
-    }
-
-    #[test]
-    fn test_out_of_bounds_y() {
-        let mut chunk_storage = ChunkStorage::default();
-        let mut partial_chunk_storage = PartialChunkStorage::default();
-        partial_chunk_storage.set(
-            &ChunkPos { x: 0, z: 0 },
-            Some(Chunk::default()),
-            &mut chunk_storage,
-        );
-        assert!(
-            chunk_storage
-                .get_block_state(BlockPos { x: 0, y: 319, z: 0 })
-                .is_some()
-        );
-        assert!(
-            chunk_storage
-                .get_block_state(BlockPos { x: 0, y: 320, z: 0 })
-                .is_none()
-        );
-        assert!(
-            chunk_storage
-                .get_block_state(BlockPos { x: 0, y: 338, z: 0 })
-                .is_none()
-        );
-        assert!(
-            chunk_storage
-                .get_block_state(BlockPos { x: 0, y: -64, z: 0 })
-                .is_some()
-        );
-        assert!(
-            chunk_storage
-                .get_block_state(BlockPos { x: 0, y: -65, z: 0 })
-                .is_none()
-        );
-    }
-
-    #[test]
-    fn test_chunk_pos_from_index() {
-        let mut partial_chunk_storage = PartialChunkStorage::new(5);
-        partial_chunk_storage.update_view_center(ChunkPos::new(0, -1));
-        assert_eq!(
-            partial_chunk_storage.chunk_pos_from_index(
-                partial_chunk_storage.index_from_chunk_pos(&ChunkPos::new(2, -1))
-            ),
-            ChunkPos::new(2, -1),
-        );
-    }
-
-    #[test]
-    fn serialize_and_deserialize_section() {
-        let mut states = PalettedContainer::new();
-
-        states.set(
-            SectionPos::new(1, 2, 3),
-            BlockState::try_from(BlockState::MAX_STATE).unwrap(),
-        );
-        states.set(
-            SectionPos::new(4, 5, 6),
-            BlockState::try_from(BlockState::MAX_STATE).unwrap(),
-        );
-        let biomes = PalettedContainer::new();
-        let section = Section {
-            block_count: 2,
-            states,
-            biomes,
-        };
-
-        let mut buf = Vec::new();
-        section.azalea_write(&mut buf).unwrap();
-
-        let mut cur = Cursor::new(buf.as_slice());
-        let deserialized_section = Section::azalea_read(&mut cur).unwrap();
-        assert_eq!(cur.position(), buf.len() as u64);
-
-        assert_eq!(section, deserialized_section);
-    }
-}
diff --git a/azalea-world/src/container.rs b/azalea-world/src/container.rs
index c0fa2633..9b9aeed9 100644
--- a/azalea-world/src/container.rs
+++ b/azalea-world/src/container.rs
@@ -60,16 +60,16 @@ impl Worlds {
         match self.map.get(&name).and_then(|world| world.upgrade()) {
             Some(existing_lock) => {
                 let existing = existing_lock.read();
-                if existing.chunks.height != height {
+                if existing.chunks.height() != height {
                     error!(
                         "Shared world height mismatch: {} != {height}",
-                        existing.chunks.height
+                        existing.chunks.height()
                     );
                 }
-                if existing.chunks.min_y != min_y {
+                if existing.chunks.min_y() != min_y {
                     error!(
                         "Shared world min_y mismatch: {} != {min_y}",
-                        existing.chunks.min_y
+                        existing.chunks.min_y()
                     );
                 }
                 existing_lock.clone()
diff --git a/azalea-world/src/find_blocks.rs b/azalea-world/src/find_blocks.rs
index 753809ef..7de7a91f 100644
--- a/azalea-world/src/find_blocks.rs
+++ b/azalea-world/src/find_blocks.rs
@@ -1,7 +1,9 @@
 use azalea_block::{BlockState, BlockStates};
 use azalea_core::position::{BlockPos, ChunkPos};
 
-use crate::{Chunk, ChunkStorage, World, iterators::ChunkIterator, palette::Palette};
+use crate::{
+    Chunk, World, chunk::storage::ChunkStorage, iterators::ChunkIterator, palette::Palette,
+};
 
 impl World {
     /// Find the coordinates of a block in the world.
@@ -45,7 +47,7 @@ impl World {
                 block_states,
                 chunk_pos,
                 &chunk.read(),
-                self.chunks.min_y,
+                self.chunks.min_y(),
                 |this_block_pos| {
                     let this_block_distance = (nearest_to - this_block_pos).length_manhattan();
                     // only update if it's closer
@@ -157,7 +159,7 @@ impl Iterator for FindBlocks<'_> {
                 self.block_states,
                 chunk_pos,
                 &chunk.read(),
-                self.chunks.min_y,
+                self.chunks.min_y(),
                 |this_block_pos| {
                     let this_block_distance = (self.nearest_to - this_block_pos).length_manhattan();
 
@@ -250,7 +252,7 @@ mod tests {
     use azalea_registry::builtin::BlockKind;
 
     use super::*;
-    use crate::{Chunk, PartialChunkStorage};
+    use crate::{Chunk, chunk::partial::PartialChunkStorage};
 
     #[test]
     fn find_block() {
diff --git a/azalea-world/src/heightmap.rs b/azalea-world/src/heightmap.rs
index 6d99ba9d..789f83eb 100644
--- a/azalea-world/src/heightmap.rs
+++ b/azalea-world/src/heightmap.rs
@@ -3,7 +3,7 @@ use azalea_core::{heightmap_kind::HeightmapKind as HeightmapKind_, math, positio
 use azalea_registry::tags::blocks::LEAVES;
 use tracing::warn;
 
-use crate::{BitStorage, Section, chunk_storage::get_block_state_from_sections};
+use crate::{BitStorage, Section, chunk::get_block_state_from_sections};
 
 // TODO: when removing this deprecated marker, also rename `HeightmapKind_` back
 // to `HeightmapKind`.
diff --git a/azalea-world/src/lib.rs b/azalea-world/src/lib.rs
index 4f9307c4..e9d78594 100644
--- a/azalea-world/src/lib.rs
+++ b/azalea-world/src/lib.rs
@@ -1,7 +1,7 @@
 #![doc = include_str!("../README.md")]
 
 mod bit_storage;
-pub mod chunk_storage;
+pub mod chunk;
 mod container;
 pub mod find_blocks;
 pub mod heightmap;
@@ -10,7 +10,7 @@ pub mod palette;
 mod world;
 
 pub use bit_storage::BitStorage;
-pub use chunk_storage::{Chunk, ChunkStorage, PartialChunkStorage, Section};
+pub use chunk::{Chunk, Section, partial::PartialChunkStorage, storage::ChunkStorage};
 pub use container::{WorldName, Worlds};
 pub use world::*;
 
diff --git a/azalea/benches/pathfinder.rs b/azalea/benches/pathfinder.rs
index defd7c18..3adb9903 100644
--- a/azalea/benches/pathfinder.rs
+++ b/azalea/benches/pathfinder.rs
@@ -46,13 +46,13 @@ fn generate_bedrock_world(
                     chunk.set_block_state(
                         &ChunkBlockPos::new(x, 1, z),
                         BlockKind::Bedrock.into(),
-                        chunks.min_y,
+                        chunks.min_y(),
                     );
                     if rng.random_bool(0.5) {
                         chunk.set_block_state(
                             &ChunkBlockPos::new(x, 2, z),
                             BlockKind::Bedrock.into(),
-                            chunks.min_y,
+                            chunks.min_y(),
                         );
                     }
                 }
@@ -98,13 +98,13 @@ fn generate_mining_world(
             let chunk_pos = ChunkPos::new(chunk_x, chunk_z);
             let chunk = chunks.get(&chunk_pos).unwrap();
             let mut chunk = chunk.write();
-            for y in chunks.min_y..(chunks.min_y + chunks.height as i32) {
+            for y in chunks.min_y()..(chunks.min_y() + chunks.height() as i32) {
                 for x in 0..16_u8 {
                     for z in 0..16_u8 {
                         chunk.set_block_state(
                             &ChunkBlockPos::new(x, y, z),
                             BlockKind::Stone.into(),
-                            chunks.min_y,
+                            chunks.min_y(),
                         );
                     }
                 }
diff --git a/azalea/benches/physics.rs b/azalea/benches/physics.rs
index 8b15cc6c..430809cc 100644
--- a/azalea/benches/physics.rs
+++ b/azalea/benches/physics.rs
@@ -33,7 +33,7 @@ fn generate_world(partial_chunks: &mut PartialChunkStorage, size: u32) -> ChunkS
                     chunk.set_block_state(
                         &ChunkBlockPos::new(x, 1, z),
                         BlockKind::Stone.into(),
-                        chunks.min_y,
+                        chunks.min_y(),
                     );
                 }
             }
diff --git a/azalea/examples/testbot/commands/debug.rs b/azalea/examples/testbot/commands/debug.rs
index cd487abb..62baef1d 100644
--- a/azalea/examples/testbot/commands/debug.rs
+++ b/azalea/examples/testbot/commands/debug.rs
@@ -1,6 +1,6 @@
 //! Commands for debugging and getting the current state of the bot.
 
-use std::{env, fs::File, io::Write, thread, time::Duration};
+use std::{any::Any, env, fs::File, io::Write, thread, time::Duration};
 
 use azalea::{
     BlockPos,
@@ -16,7 +16,7 @@ use azalea::{
 use azalea_core::hit_result::HitResult;
 use azalea_entity::{EntityKindComponent, metadata};
 use azalea_inventory::{Menu, components::MaxStackSize};
-use azalea_world::Worlds;
+use azalea_world::{Worlds, chunk::storage::WeakChunkStorage};
 use bevy_app::AppExit;
 use bevy_ecs::{message::Messages, query::With, world::EntityRef};
 use parking_lot::Mutex;
@@ -349,19 +349,22 @@ pub fn register(commands: &mut CommandDispatcher<Mutex<CommandSource>>) {
                                 .unwrap();
                             if let Some(world) = world.upgrade() {
                                 let world = world.read();
-                                let strong_chunks = world
-                                    .chunks
-                                    .map
-                                    .iter()
-                                    .filter(|(_, v)| v.strong_count() > 0)
-                                    .count();
-                                writeln!(
-                                    report,
-                                    "- Chunks: {} strongly referenced, {} in map",
-                                    strong_chunks,
-                                    world.chunks.map.len()
-                                )
-                                .unwrap();
+                                let chunks = &world.chunks;
+                                let chunks = (chunks as &dyn Any).downcast_ref::<WeakChunkStorage>();
+                                if let Some(chunks) = chunks {
+                                    let strong_chunks = chunks
+                                        .map
+                                        .iter()
+                                        .filter(|(_, v)| v.strong_count() > 0)
+                                        .count();
+                                    writeln!(
+                                        report,
+                                        "- Chunks: {} strongly referenced, {} in map",
+                                        strong_chunks,
+                                        chunks.map.len()
+                                    )
+                                    .unwrap();
+                                }
                                 writeln!(
                                     report,
                                     "- Entities: {}",
diff --git a/azalea/src/pathfinder/world.rs b/azalea/src/pathfinder/world.rs
index 98b22d70..6c5d84d2 100644
--- a/azalea/src/pathfinder/world.rs
+++ b/azalea/src/pathfinder/world.rs
@@ -110,7 +110,7 @@ pub struct SectionBitsets {
 
 impl CachedWorld {
     pub fn new(world_lock: Arc<RwLock<World>>, origin: BlockPos) -> Self {
-        let min_y = world_lock.read().chunks.min_y;
+        let min_y = world_lock.read().chunks.min_y();
         Self {
             origin,
             min_y,
@@ -145,7 +145,7 @@ impl CachedWorld {
 
         let chunk_pos = ChunkPos::new(section_pos.x as i32, section_pos.z as i32);
         let section_index =
-            azalea_world::chunk_storage::section_index(section_pos.y * 16, self.min_y) as usize;
+            azalea_world::chunk::section_index(section_pos.y * 16, self.min_y) as usize;
 
         let mut cache_idx = 0;