add basic support for getting biome ids in chunks

3 files changed, 509 insertions, 0 deletions

diff --git a/azalea-world/src/palette/container.rs b/azalea-world/src/palette/container.rs
new file mode 100644
index 00000000..2f3d9238
--- /dev/null
+++ b/azalea-world/src/palette/container.rs
@@ -0,0 +1,314 @@
+use std::{
+    fmt::Debug,
+    io::{self, Cursor, Write},
+};
+
+use azalea_block::BlockState;
+use azalea_buf::{AzaleaRead, AzaleaWrite, BufReadError};
+use azalea_core::position::{ChunkSectionBiomePos, ChunkSectionBlockPos};
+use azalea_registry::Biome;
+use tracing::warn;
+
+use super::{Palette, PaletteKind};
+use crate::BitStorage;
+
+#[derive(Clone, Debug)]
+pub struct PalettedContainer<S: PalletedContainerKind> {
+    pub bits_per_entry: u8,
+    /// This is usually a list of unique values that appear in the container so
+    /// they can be indexed by the bit storage.
+    ///
+    /// Sometimes it doesn't contain anything if there's too many unique items
+    /// in the bit storage, though.
+    pub palette: Palette<S>,
+    /// Compacted list of indices pointing to entry IDs in the Palette.
+    pub storage: BitStorage,
+}
+
+pub trait PalletedContainerKind: Copy + Clone + Debug + Default + TryFrom<u32> + Into<u32> {
+    type SectionPos: SectionPos;
+
+    fn size_bits() -> usize;
+
+    fn size() -> usize {
+        1 << (Self::size_bits() * 3)
+    }
+
+    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind;
+}
+impl PalletedContainerKind for BlockState {
+    type SectionPos = ChunkSectionBlockPos;
+
+    fn size_bits() -> usize {
+        4
+    }
+
+    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind {
+        match bits_per_entry {
+            0 => PaletteKind::SingleValue,
+            1..=4 => PaletteKind::Linear,
+            5..=8 => PaletteKind::Hashmap,
+            _ => PaletteKind::Global,
+        }
+    }
+}
+impl PalletedContainerKind for Biome {
+    type SectionPos = ChunkSectionBiomePos;
+
+    fn size_bits() -> usize {
+        2
+    }
+
+    fn bits_per_entry_to_palette_kind(bits_per_entry: u8) -> PaletteKind {
+        match bits_per_entry {
+            0 => PaletteKind::SingleValue,
+            1..=3 => PaletteKind::Linear,
+            _ => PaletteKind::Global,
+        }
+    }
+}
+
+/// A trait for position types that are sometimes valid ways to index into a
+/// chunk section.
+pub trait SectionPos {
+    fn coords(&self) -> (usize, usize, usize);
+    fn new(x: usize, y: usize, z: usize) -> Self;
+}
+impl SectionPos for ChunkSectionBlockPos {
+    fn coords(&self) -> (usize, usize, usize) {
+        (self.x as usize, self.y as usize, self.z as usize)
+    }
+
+    fn new(x: usize, y: usize, z: usize) -> Self {
+        ChunkSectionBlockPos {
+            x: x as u8,
+            y: y as u8,
+            z: z as u8,
+        }
+    }
+}
+impl SectionPos for ChunkSectionBiomePos {
+    fn coords(&self) -> (usize, usize, usize) {
+        (self.x as usize, self.y as usize, self.z as usize)
+    }
+
+    fn new(x: usize, y: usize, z: usize) -> Self {
+        ChunkSectionBiomePos {
+            x: x as u8,
+            y: y as u8,
+            z: z as u8,
+        }
+    }
+}
+
+impl<S: PalletedContainerKind> PalettedContainer<S> {
+    pub fn new() -> Self {
+        let palette = Palette::SingleValue(S::default());
+        let size = S::size();
+        let storage = BitStorage::new(0, size, Some(Box::new([]))).unwrap();
+
+        PalettedContainer {
+            bits_per_entry: 0,
+            palette,
+            storage,
+        }
+    }
+
+    pub fn read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
+        let bits_per_entry = u8::azalea_read(buf)?;
+        let palette_type = S::bits_per_entry_to_palette_kind(bits_per_entry);
+        let palette = palette_type.read(buf)?;
+        let size = S::size();
+
+        let mut storage = match BitStorage::new(
+            bits_per_entry as usize,
+            size,
+            if bits_per_entry == 0 {
+                Some(Box::new([]))
+            } else {
+                // we're going to update the data after creating the bitstorage
+                None
+            },
+        ) {
+            Ok(storage) => storage,
+            Err(e) => {
+                warn!("Failed to create bit storage: {:?}", e);
+                return Err(BufReadError::Custom(
+                    "Failed to create bit storage".to_string(),
+                ));
+            }
+        };
+
+        // now read the data
+        for i in 0..storage.data.len() {
+            storage.data[i] = u64::azalea_read(buf)?;
+        }
+
+        Ok(PalettedContainer {
+            bits_per_entry,
+            palette,
+            storage,
+        })
+    }
+
+    /// Calculates the index of the given coordinates.
+    pub fn index_from_coords(&self, pos: S::SectionPos) -> usize {
+        let size_bits = S::size_bits();
+        let (x, y, z) = pos.coords();
+        (((y << size_bits) | z) << size_bits) | x
+    }
+
+    pub fn coords_from_index(&self, index: usize) -> S::SectionPos {
+        let size_bits = S::size_bits();
+        let mask = (1 << size_bits) - 1;
+        S::SectionPos::new(
+            index & mask,
+            (index >> size_bits >> size_bits) & mask,
+            (index >> size_bits) & mask,
+        )
+    }
+
+    /// Returns the value at the given index.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if the index is greater than or equal to the number
+    /// of things in the storage. (So for block states, it must be less than
+    /// 4096).
+    pub fn get_at_index(&self, index: usize) -> S {
+        // first get the palette id
+        let paletted_value = self.storage.get(index);
+        // and then get the value from that id
+        self.palette.value_for(paletted_value as usize)
+    }
+
+    /// Returns the value at the given coordinates.
+    pub fn get(&self, pos: S::SectionPos) -> S {
+        // let paletted_value = self.storage.get(self.get_index(x, y, z));
+        // self.palette.value_for(paletted_value as usize)
+        self.get_at_index(self.index_from_coords(pos))
+    }
+
+    /// Sets the id at the given coordinates and return the previous id
+    pub fn get_and_set(&mut self, pos: S::SectionPos, value: S) -> S {
+        let paletted_value = self.id_for(value);
+        let block_state_id = self
+            .storage
+            .get_and_set(self.index_from_coords(pos), paletted_value as u64);
+        // error in debug mode
+        #[cfg(debug_assertions)]
+        if block_state_id > BlockState::MAX_STATE.into() {
+            warn!(
+                "Old block state from get_and_set {block_state_id} was greater than max state {}",
+                BlockState::MAX_STATE
+            );
+        }
+
+        S::try_from(block_state_id as u32).unwrap_or_default()
+    }
+
+    /// Sets the id at the given index and return the previous id. You probably
+    /// want `.set` instead.
+    pub fn set_at_index(&mut self, index: usize, value: S) {
+        let paletted_value = self.id_for(value);
+        self.storage.set(index, paletted_value as u64);
+    }
+
+    /// Sets the id at the given coordinates and return the previous id
+    pub fn set(&mut self, pos: S::SectionPos, value: S) {
+        self.set_at_index(self.index_from_coords(pos), value);
+    }
+
+    fn create_or_reuse_data(&self, bits_per_entry: u8) -> PalettedContainer<S> {
+        let new_palette_type = S::bits_per_entry_to_palette_kind(bits_per_entry);
+
+        let old_palette_type = (&self.palette).into();
+        if bits_per_entry == self.bits_per_entry && new_palette_type == old_palette_type {
+            return self.clone();
+        }
+        let storage = BitStorage::new(bits_per_entry as usize, S::size(), None).unwrap();
+
+        // sanity check
+        debug_assert_eq!(storage.size(), S::size());
+
+        // let palette = new_palette_type.as_empty_palette(1usize << (bits_per_entry as
+        // usize));
+        let palette = new_palette_type.as_empty_palette();
+        PalettedContainer {
+            bits_per_entry,
+            palette,
+            storage,
+        }
+    }
+
+    fn on_resize(&mut self, bits_per_entry: u8, value: S) -> usize {
+        // in vanilla this is always true, but it's sometimes false in purpur servers
+        // assert!(bits_per_entry <= 5, "bits_per_entry must be <= 5");
+        let mut new_data = self.create_or_reuse_data(bits_per_entry);
+        new_data.copy_from(&self.palette, &self.storage);
+        *self = new_data;
+        self.id_for(value)
+    }
+
+    fn copy_from(&mut self, palette: &Palette<S>, storage: &BitStorage) {
+        for i in 0..storage.size() {
+            let value = palette.value_for(storage.get(i) as usize);
+            let id = self.id_for(value) as u64;
+            self.storage.set(i, id);
+        }
+    }
+
+    pub fn id_for(&mut self, value: S) -> usize {
+        match &mut self.palette {
+            Palette::SingleValue(v) => {
+                if (*v).into() != value.into() {
+                    self.on_resize(1, value)
+                } else {
+                    0
+                }
+            }
+            Palette::Linear(palette) => {
+                if let Some(index) = palette.iter().position(|&v| v.into() == value.into()) {
+                    return index;
+                }
+                let capacity = 2usize.pow(self.bits_per_entry.into());
+                if capacity > palette.len() {
+                    palette.push(value);
+                    palette.len() - 1
+                } else {
+                    self.on_resize(self.bits_per_entry + 1, value)
+                }
+            }
+            Palette::Hashmap(palette) => {
+                // TODO? vanilla keeps this in memory as a hashmap, but it should be benchmarked
+                // before changing it
+                if let Some(index) = palette.iter().position(|v| (*v).into() == value.into()) {
+                    return index;
+                }
+                let capacity = 2usize.pow(self.bits_per_entry.into());
+                if capacity > palette.len() {
+                    palette.push(value);
+                    palette.len() - 1
+                } else {
+                    self.on_resize(self.bits_per_entry + 1, value)
+                }
+            }
+            Palette::Global => value.into() as usize,
+        }
+    }
+}
+
+impl<S: PalletedContainerKind> AzaleaWrite for PalettedContainer<S> {
+    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
+        self.bits_per_entry.azalea_write(buf)?;
+        self.palette.azalea_write(buf)?;
+        self.storage.data.azalea_write(buf)?;
+        Ok(())
+    }
+}
+
+impl<S: PalletedContainerKind> Default for PalettedContainer<S> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/azalea-world/src/palette/mod.rs b/azalea-world/src/palette/mod.rs
new file mode 100644
index 00000000..65a04f6a
--- /dev/null
+++ b/azalea-world/src/palette/mod.rs
@@ -0,0 +1,115 @@
+mod container;
+
+#[cfg(test)]
+mod tests;
+
+use std::{
+    fmt::Debug,
+    io::{self, Cursor, Write},
+};
+
+use azalea_buf::{AzaleaReadVar, AzaleaWrite, AzaleaWriteVar, BufReadError};
+pub use container::*;
+
+/// A representation of the different types of chunk palettes Minecraft uses.
+#[derive(Clone, Debug)]
+pub enum Palette<S: PalletedContainerKind> {
+    /// ID of the corresponding entry in its global palette
+    SingleValue(S),
+    // in vanilla this keeps a `size` field that might be less than the length, but i'm not sure
+    // it's actually needed?
+    Linear(Vec<S>),
+    Hashmap(Vec<S>),
+    Global,
+}
+
+impl<S: PalletedContainerKind> Palette<S> {
+    pub fn value_for(&self, id: usize) -> S {
+        match self {
+            Palette::SingleValue(v) => *v,
+            Palette::Linear(v) => v.get(id).copied().unwrap_or_default(),
+            Palette::Hashmap(v) => v.get(id).copied().unwrap_or_default(),
+            Palette::Global => S::try_from(id as u32).unwrap_or_default(),
+        }
+    }
+}
+
+impl<S: PalletedContainerKind> AzaleaWrite for Palette<S> {
+    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
+        match self {
+            Palette::SingleValue(value) => {
+                (*value).into().azalea_write_var(buf)?;
+            }
+            Palette::Linear(values) => {
+                (values.len() as u32).azalea_write_var(buf)?;
+                for value in values {
+                    (*value).into().azalea_write_var(buf)?;
+                }
+            }
+            Palette::Hashmap(values) => {
+                (values.len() as u32).azalea_write_var(buf)?;
+                for value in values {
+                    (*value).into().azalea_write_var(buf)?;
+                }
+            }
+            Palette::Global => {}
+        }
+        Ok(())
+    }
+}
+
+impl PaletteKind {
+    pub fn read<S: PalletedContainerKind>(
+        &self,
+        buf: &mut Cursor<&[u8]>,
+    ) -> Result<Palette<S>, BufReadError> {
+        Ok(match self {
+            // since they're read as varints it's actually fine to just use BlockStateIntegerRepr
+            // instead of the correct type (u32)
+            PaletteKind::SingleValue => {
+                Palette::SingleValue(S::try_from(u32::azalea_read_var(buf)?).unwrap_or_default())
+            }
+            PaletteKind::Linear => Palette::Linear(
+                Vec::<u32>::azalea_read_var(buf)?
+                    .into_iter()
+                    .map(|v| S::try_from(v).unwrap_or_default())
+                    .collect(),
+            ),
+            PaletteKind::Hashmap => Palette::Hashmap(
+                Vec::<u32>::azalea_read_var(buf)?
+                    .into_iter()
+                    .map(|v| S::try_from(v).unwrap_or_default())
+                    .collect(),
+            ),
+            PaletteKind::Global => Palette::Global,
+        })
+    }
+
+    pub fn as_empty_palette<S: PalletedContainerKind>(&self) -> Palette<S> {
+        match self {
+            PaletteKind::SingleValue => Palette::SingleValue(S::default()),
+            PaletteKind::Linear => Palette::Linear(Vec::new()),
+            PaletteKind::Hashmap => Palette::Hashmap(Vec::new()),
+            PaletteKind::Global => Palette::Global,
+        }
+    }
+}
+
+impl<S: PalletedContainerKind> From<&Palette<S>> for PaletteKind {
+    fn from(palette: &Palette<S>) -> Self {
+        match palette {
+            Palette::SingleValue(_) => PaletteKind::SingleValue,
+            Palette::Linear(_) => PaletteKind::Linear,
+            Palette::Hashmap(_) => PaletteKind::Hashmap,
+            Palette::Global => PaletteKind::Global,
+        }
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum PaletteKind {
+    SingleValue,
+    Linear,
+    Hashmap,
+    Global,
+}
diff --git a/azalea-world/src/palette/tests.rs b/azalea-world/src/palette/tests.rs
new file mode 100644
index 00000000..d1423306
--- /dev/null
+++ b/azalea-world/src/palette/tests.rs
@@ -0,0 +1,80 @@
+use azalea_block::BlockState;
+use azalea_core::position::ChunkSectionBlockPos;
+
+use super::*;
+
+#[test]
+fn test_resize_0_bits_to_1() {
+    let mut palette_container = PalettedContainer::<BlockState>::new();
+
+    assert_eq!(palette_container.bits_per_entry, 0);
+    assert_eq!(palette_container.get_at_index(0), BlockState::AIR);
+    assert_eq!(
+        PaletteKind::from(&palette_container.palette),
+        PaletteKind::SingleValue
+    );
+    let block_state_1 = BlockState::try_from(1_u32).unwrap();
+    palette_container.set_at_index(0, block_state_1);
+    assert_eq!(palette_container.get_at_index(0), block_state_1);
+    assert_eq!(
+        PaletteKind::from(&palette_container.palette),
+        PaletteKind::Linear
+    );
+}
+
+#[test]
+fn test_resize_0_bits_to_5() {
+    let mut palette_container = PalettedContainer::<BlockState>::new();
+
+    let set = |pc: &mut PalettedContainer<BlockState>, i, v: u32| {
+        pc.set_at_index(i, BlockState::try_from(v).unwrap());
+    };
+
+    set(&mut palette_container, 0, 0); // 0 bits
+    assert_eq!(palette_container.bits_per_entry, 0);
+
+    set(&mut palette_container, 1, 1); // 1 bit
+    assert_eq!(palette_container.bits_per_entry, 1);
+
+    set(&mut palette_container, 2, 2); // 2 bits
+    assert_eq!(palette_container.bits_per_entry, 2);
+    set(&mut palette_container, 3, 3);
+
+    set(&mut palette_container, 4, 4); // 3 bits
+    assert_eq!(palette_container.bits_per_entry, 3);
+    set(&mut palette_container, 5, 5);
+    set(&mut palette_container, 6, 6);
+    set(&mut palette_container, 7, 7);
+
+    set(&mut palette_container, 8, 8); // 4 bits
+    assert_eq!(palette_container.bits_per_entry, 4);
+    set(&mut palette_container, 9, 9);
+    set(&mut palette_container, 10, 10);
+    set(&mut palette_container, 11, 11);
+    set(&mut palette_container, 12, 12);
+    set(&mut palette_container, 13, 13);
+    set(&mut palette_container, 14, 14);
+    set(&mut palette_container, 15, 15);
+    assert_eq!(palette_container.bits_per_entry, 4);
+
+    set(&mut palette_container, 16, 16); // 5 bits
+    assert_eq!(palette_container.bits_per_entry, 5);
+}
+
+#[test]
+fn test_coords_from_index() {
+    let palette_container = PalettedContainer::<BlockState>::new();
+
+    for x in 0..15 {
+        for y in 0..15 {
+            for z in 0..15 {
+                assert_eq!(
+                    palette_container.coords_from_index(
+                        palette_container.index_from_coords(ChunkSectionBlockPos::new(x, y, z))
+                    ),
+                    ChunkSectionBlockPos::new(x, y, z)
+                );
+            }
+        }
+    }
+}