Extensible ChunkStorage

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

1 files changed, 242 insertions, 0 deletions

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()
+        );
+    }
+}