Merge pull request #6 from mat-1/chunk-decoding

Chunk decoding

3 files changed, 585 insertions, 0 deletions

diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
new file mode 100644
index 00000000..aba52aef
--- /dev/null
+++ b/azalea-world/src/bit_storage.rs
@@ -0,0 +1,209 @@
+use std::{error::Error, fmt};
+
+// this is from minecraft's code
+// yeah idk either
+const MAGIC: [(i32, i32, i32); 64] = [
+    (-1, -1, 0),
+    (-2147483648, 0, 0),
+    (1431655765, 1431655765, 0),
+    (-2147483648, 0, 1),
+    (858993459, 858993459, 0),
+    (715827882, 715827882, 0),
+    (613566756, 613566756, 0),
+    (-2147483648, 0, 2),
+    (477218588, 477218588, 0),
+    (429496729, 429496729, 0),
+    (390451572, 390451572, 0),
+    (357913941, 357913941, 0),
+    (330382099, 330382099, 0),
+    (306783378, 306783378, 0),
+    (286331153, 286331153, 0),
+    (-2147483648, 0, 3),
+    (252645135, 252645135, 0),
+    (238609294, 238609294, 0),
+    (226050910, 226050910, 0),
+    (214748364, 214748364, 0),
+    (204522252, 204522252, 0),
+    (195225786, 195225786, 0),
+    (186737708, 186737708, 0),
+    (178956970, 178956970, 0),
+    (171798691, 171798691, 0),
+    (165191049, 165191049, 0),
+    (159072862, 159072862, 0),
+    (153391689, 153391689, 0),
+    (148102320, 148102320, 0),
+    (143165576, 143165576, 0),
+    (138547332, 138547332, 0),
+    (-2147483648, 0, 4),
+    (130150524, 130150524, 0),
+    (126322567, 126322567, 0),
+    (122713351, 122713351, 0),
+    (119304647, 119304647, 0),
+    (116080197, 116080197, 0),
+    (113025455, 113025455, 0),
+    (110127366, 110127366, 0),
+    (107374182, 107374182, 0),
+    (104755299, 104755299, 0),
+    (102261126, 102261126, 0),
+    (99882960, 99882960, 0),
+    (97612893, 97612893, 0),
+    (95443717, 95443717, 0),
+    (93368854, 93368854, 0),
+    (91382282, 91382282, 0),
+    (89478485, 89478485, 0),
+    (87652393, 87652393, 0),
+    (85899345, 85899345, 0),
+    (84215045, 84215045, 0),
+    (82595524, 82595524, 0),
+    (81037118, 81037118, 0),
+    (79536431, 79536431, 0),
+    (78090314, 78090314, 0),
+    (76695844, 76695844, 0),
+    (75350303, 75350303, 0),
+    (74051160, 74051160, 0),
+    (72796055, 72796055, 0),
+    (71582788, 71582788, 0),
+    (70409299, 70409299, 0),
+    (69273666, 69273666, 0),
+    (68174084, 68174084, 0),
+    (-2147483648, 0, 5),
+];
+
+/// A compact list of integers with the given number of bits per entry.
+#[derive(Clone, Debug, Default)]
+pub struct BitStorage {
+    pub data: Vec<u64>,
+    bits: usize,
+    mask: u64,
+    size: usize,
+    values_per_long: u8,
+    divide_mul: i32,
+    divide_add: i32,
+    divide_shift: i32,
+}
+
+#[derive(Debug)]
+pub enum BitStorageError {
+    InvalidLength { got: usize, expected: usize },
+}
+impl fmt::Display for BitStorageError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            BitStorageError::InvalidLength { got, expected } => write!(
+                f,
+                "Invalid length given for storage, got: {}, but expected: {}",
+                got, expected
+            ),
+        }
+    }
+}
+impl Error for BitStorageError {}
+
+impl BitStorage {
+    /// Create a new BitStorage with the given number of bits per entry.
+    /// `size` is the number of entries in the BitStorage.
+    pub fn new(bits: usize, size: usize, data: Option<Vec<u64>>) -> Result<Self, BitStorageError> {
+        // vanilla has this assert but it's not always true for some reason??
+        // assert!(bits >= 1 && bits <= 32);
+
+        if let Some(data) = &data {
+            if data.len() == 0 {
+                // TODO: make 0 bit storage actually work
+                return Ok(BitStorage::default());
+            }
+        }
+
+        let values_per_long = 64 / bits;
+        let magic_index = values_per_long - 1;
+        let (divide_mul, divide_add, divide_shift) = MAGIC[magic_index as usize];
+        let calculated_length = (size + values_per_long - 1) / values_per_long;
+
+        let mask = (1 << bits) - 1;
+
+        let using_data = if let Some(data) = data {
+            if data.len() != calculated_length as usize {
+                return Err(BitStorageError::InvalidLength {
+                    got: data.len(),
+                    expected: calculated_length as usize,
+                });
+            }
+            data
+        } else {
+            vec![0; calculated_length as usize]
+        };
+
+        Ok(BitStorage {
+            data: using_data,
+            bits,
+            mask,
+            size,
+            values_per_long: values_per_long as u8,
+            divide_mul,
+            divide_add,
+            divide_shift,
+        })
+    }
+
+    pub fn cell_index(&self, index: u64) -> usize {
+        // as unsigned wrap
+        let first = self.divide_mul as u32 as u64;
+        let second = self.divide_add as u64;
+
+        (((index * first) + second) >> 32 >> self.divide_shift)
+            .try_into()
+            .unwrap()
+    }
+
+    pub fn get(&self, index: usize) -> u64 {
+        // Validate.inclusiveBetween(0L, (long)(this.size - 1), (long)var1);
+        // int var2 = this.cellIndex(var1);
+        // long var3 = this.data[var2];
+        // int var5 = (var1 - var2 * this.valuesPerLong) * this.bits;
+        // return (int)(var3 >> var5 & this.mask);
+
+        assert!(
+            index <= self.size - 1,
+            "Index {} out of bounds (max is {})",
+            index,
+            self.size - 1
+        );
+        let cell_index = self.cell_index(index as u64);
+        let cell = &self.data[cell_index as usize];
+        let bit_index = (index - cell_index * self.values_per_long as usize) * self.bits;
+        cell >> bit_index & self.mask
+    }
+
+    pub fn set(&mut self, index: usize, value: u64) {
+        // Validate.inclusiveBetween(0L, (long)(this.size - 1), (long)var1);
+        // Validate.inclusiveBetween(0L, this.mask, (long)var2);
+        // int var3 = this.cellIndex(var1);
+        // long var4 = this.data[var3];
+        // int var6 = (var1 - var3 * this.valuesPerLong) * this.bits;
+        // this.data[var3] = var4 & ~(this.mask << var6) | ((long)var2 & this.mask) << var6;
+
+        assert!(index <= self.size - 1);
+        assert!(value <= self.mask);
+        let cell_index = self.cell_index(index as u64);
+        let cell = &mut self.data[cell_index as usize];
+        let bit_index = (index - cell_index * self.values_per_long as usize) * self.bits;
+        *cell = *cell & !(self.mask << bit_index) | (value & self.mask) << bit_index;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_wikivg_example() {
+        let data = [
+            1, 2, 2, 3, 4, 4, 5, 6, 6, 4, 8, 0, 7, 4, 3, 13, 15, 16, 9, 14, 10, 12, 0, 2,
+        ];
+        let compact_data: [u64; 2] = [0x0020863148418841, 0x01018A7260F68C87];
+        let storage = BitStorage::new(5, data.len(), Some(compact_data.to_vec())).unwrap();
+
+        for (i, expected) in data.iter().enumerate() {
+            assert_eq!(storage.get(i), *expected);
+        }
+    }
+}
diff --git a/azalea-world/src/lib.rs b/azalea-world/src/lib.rs
new file mode 100644
index 00000000..766c61f0
--- /dev/null
+++ b/azalea-world/src/lib.rs
@@ -0,0 +1,237 @@
+#![feature(int_roundings)]
+
+mod bit_storage;
+mod palette;
+
+use crate::palette::PalettedContainerType;
+use azalea_core::{BlockPos, ChunkBlockPos, ChunkPos, ChunkSectionBlockPos};
+use azalea_protocol::mc_buf::{McBufReadable, McBufWritable};
+pub use bit_storage::BitStorage;
+use palette::PalettedContainer;
+use std::{
+    io::{Read, Write},
+    ops::{Index, IndexMut},
+    sync::{Arc, Mutex},
+};
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn it_works() {
+        let result = 2 + 2;
+        assert_eq!(result, 4);
+    }
+}
+
+const SECTION_HEIGHT: u32 = 16;
+
+pub struct World {
+    pub storage: ChunkStorage,
+    pub height: u32,
+    pub min_y: i32,
+}
+
+impl World {
+    pub fn replace_with_packet_data(
+        &mut self,
+        pos: &ChunkPos,
+        data: &mut impl Read,
+    ) -> Result<(), String> {
+        if !self.storage.in_range(pos) {
+            println!(
+                "Ignoring chunk since it's not in the view range: {}, {}",
+                pos.x, pos.z
+            );
+            return Ok(());
+        }
+        // let existing_chunk = &self.storage[pos];
+
+        let chunk = Arc::new(Mutex::new(Chunk::read_with_world(data, self)?));
+        println!("Loaded chunk {:?}", pos);
+        self.storage[pos] = Some(chunk);
+
+        Ok(())
+    }
+
+    pub fn update_view_center(&mut self, pos: &ChunkPos) {
+        self.storage.view_center = *pos;
+    }
+
+    pub fn get_block_state(&self, pos: &BlockPos) -> Option<u32> {
+        self.storage.get_block_state(pos, self.min_y)
+    }
+}
+impl Index<&ChunkPos> for World {
+    type Output = Option<Arc<Mutex<Chunk>>>;
+
+    fn index(&self, pos: &ChunkPos) -> &Self::Output {
+        &self.storage[pos]
+    }
+}
+impl IndexMut<&ChunkPos> for World {
+    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
+        &mut self.storage[pos]
+    }
+}
+// impl Index<&BlockPos> for World {
+//     type Output = Option<Arc<Mutex<Chunk>>>;
+
+//     fn index(&self, pos: &BlockPos) -> &Self::Output {
+//         let chunk = &self[ChunkPos::from(pos)];
+//         // chunk.
+
+//     }
+// }
+
+pub struct ChunkStorage {
+    view_center: ChunkPos,
+    chunk_radius: u32,
+    view_range: u32,
+    // chunks is a list of size chunk_radius * chunk_radius
+    chunks: Vec<Option<Arc<Mutex<Chunk>>>>,
+}
+
+// java moment
+// it might be possible to replace this with just a modulo, but i copied java's floorMod just in case
+fn floor_mod(x: i32, y: u32) -> u32 {
+    if x < 0 {
+        y - ((-x) as u32 % y)
+    } else {
+        x as u32 % y
+    }
+}
+
+impl ChunkStorage {
+    pub fn new(chunk_radius: u32) -> Self {
+        let view_range = chunk_radius * 2 + 1;
+        ChunkStorage {
+            view_center: ChunkPos::new(0, 0),
+            chunk_radius,
+            view_range,
+            chunks: vec![None; (view_range * view_range) as usize],
+        }
+    }
+
+    fn get_index(&self, chunk_pos: &ChunkPos) -> usize {
+        (floor_mod(chunk_pos.x, self.view_range) * self.view_range
+            + floor_mod(chunk_pos.z, self.view_range)) as usize
+    }
+
+    pub fn in_range(&self, chunk_pos: &ChunkPos) -> bool {
+        (chunk_pos.x - self.view_center.x).unsigned_abs() <= self.chunk_radius
+            && (chunk_pos.z - self.view_center.z).unsigned_abs() <= self.chunk_radius
+    }
+
+    pub fn get_block_state(&self, pos: &BlockPos, min_y: i32) -> Option<u32> {
+        let chunk_pos = ChunkPos::from(pos);
+        println!("chunk_pos {:?} block_pos {:?}", chunk_pos, pos);
+        let chunk = &self[&chunk_pos];
+        match chunk {
+            Some(chunk) => Some(chunk.lock().unwrap().get(&ChunkBlockPos::from(pos), min_y)),
+            None => None,
+        }
+    }
+}
+
+impl Index<&ChunkPos> for ChunkStorage {
+    type Output = Option<Arc<Mutex<Chunk>>>;
+
+    fn index(&self, pos: &ChunkPos) -> &Self::Output {
+        &self.chunks[self.get_index(pos)]
+    }
+}
+impl IndexMut<&ChunkPos> for ChunkStorage {
+    fn index_mut<'a>(&'a mut self, pos: &ChunkPos) -> &'a mut Self::Output {
+        let index = self.get_index(pos);
+        &mut self.chunks[index]
+    }
+}
+
+#[derive(Debug)]
+pub struct Chunk {
+    pub sections: Vec<Section>,
+}
+
+impl Chunk {
+    pub fn read_with_world(buf: &mut impl Read, data: &World) -> Result<Self, String> {
+        Self::read_with_world_height(buf, data.height)
+    }
+
+    pub fn read_with_world_height(buf: &mut impl Read, world_height: u32) -> Result<Self, String> {
+        let section_count = world_height / SECTION_HEIGHT;
+        let mut sections = Vec::with_capacity(section_count as usize);
+        for _ in 0..section_count {
+            let section = Section::read_into(buf)?;
+            sections.push(section);
+        }
+        Ok(Chunk { sections })
+    }
+
+    pub fn section_index(&self, y: i32, min_y: i32) -> u32 {
+        // TODO: check the build height and stuff, this code will be broken if the min build height is 0
+        // (LevelHeightAccessor.getMinSection in vanilla code)
+        assert!(y >= 0);
+        let min_section_index = min_y.div_floor(16);
+        (y.div_floor(16) - min_section_index) as u32
+    }
+
+    pub fn get(&self, pos: &ChunkBlockPos, min_y: i32) -> u32 {
+        let section_index = self.section_index(pos.y, min_y);
+        // TODO: make sure the section exists
+        let section = &self.sections[section_index as usize];
+        let chunk_section_pos = ChunkSectionBlockPos::from(pos);
+        let block_state = section.get(chunk_section_pos);
+        block_state
+    }
+}
+
+impl McBufWritable for Chunk {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        for section in &self.sections {
+            section.write_into(buf)?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Section {
+    pub block_count: u16,
+    pub states: PalettedContainer,
+    pub biomes: PalettedContainer,
+}
+
+impl McBufReadable for Section {
+    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
+        let block_count = u16::read_into(buf)?;
+        // this is commented out because the vanilla server is wrong
+        // assert!(
+        //     block_count <= 16 * 16 * 16,
+        //     "A section has more blocks than what should be possible. This is a bug!"
+        // );
+        let states = PalettedContainer::read_with_type(buf, &PalettedContainerType::BlockStates)?;
+        let biomes = PalettedContainer::read_with_type(buf, &PalettedContainerType::Biomes)?;
+        Ok(Section {
+            block_count,
+            states,
+            biomes,
+        })
+    }
+}
+
+impl McBufWritable for Section {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        self.block_count.write_into(buf)?;
+        self.states.write_into(buf)?;
+        self.biomes.write_into(buf)?;
+        Ok(())
+    }
+}
+
+impl Section {
+    // TODO: return a BlockState instead of a u32
+    fn get(&self, pos: ChunkSectionBlockPos) -> u32 {
+        self.states
+            .get(pos.x as usize, pos.y as usize, pos.z as usize)
+    }
+}
diff --git a/azalea-world/src/palette.rs b/azalea-world/src/palette.rs
new file mode 100644
index 00000000..be6f022c
--- /dev/null
+++ b/azalea-world/src/palette.rs
@@ -0,0 +1,139 @@
+use azalea_protocol::mc_buf::{McBufReadable, McBufVarReadable, McBufWritable, Readable, Writable};
+use std::io::{Read, Write};
+
+use crate::BitStorage;
+
+#[derive(Clone, Debug, Copy)]
+pub enum PalettedContainerType {
+    Biomes,
+    BlockStates,
+}
+
+#[derive(Clone, Debug)]
+pub struct PalettedContainer {
+    pub bits_per_entry: u8,
+    pub palette: Palette,
+    /// Compacted list of indices pointing to entry IDs in the Palette.
+    pub storage: BitStorage,
+    pub container_type: PalettedContainerType,
+}
+
+impl PalettedContainer {
+    pub fn read_with_type(
+        buf: &mut impl Read,
+        type_: &'static PalettedContainerType,
+    ) -> Result<Self, String> {
+        let bits_per_entry = buf.read_byte()?;
+        let palette = match type_ {
+            PalettedContainerType::BlockStates => {
+                Palette::block_states_read_with_bits_per_entry(buf, bits_per_entry)?
+            }
+            PalettedContainerType::Biomes => {
+                Palette::biomes_read_with_bits_per_entry(buf, bits_per_entry)?
+            }
+        };
+        let size = match type_ {
+            PalettedContainerType::BlockStates => 4096,
+            PalettedContainerType::Biomes => 64,
+        };
+
+        let data = Vec::<u64>::read_into(buf)?;
+        debug_assert!(
+            bits_per_entry != 0 || data.is_empty(),
+            "Bits per entry is 0 but data is not empty."
+        );
+        let storage = BitStorage::new(bits_per_entry.into(), size, Some(data)).unwrap();
+
+        Ok(PalettedContainer {
+            bits_per_entry,
+            palette,
+            storage,
+            container_type: *type_,
+        })
+    }
+
+    pub fn get_index(&self, x: usize, y: usize, z: usize) -> usize {
+        let size_bits = match self.container_type {
+            PalettedContainerType::BlockStates => 4,
+            PalettedContainerType::Biomes => 2,
+        };
+
+        (((y << size_bits) | z) << size_bits) | x
+    }
+
+    pub fn get(&self, x: usize, y: usize, z: usize) -> u32 {
+        let paletted_value = self.storage.get(self.get_index(x, y, z));
+        println!("palette: {:?}", self.palette);
+        self.palette.value_for(paletted_value as usize)
+    }
+}
+
+impl McBufWritable for PalettedContainer {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        buf.write_byte(self.bits_per_entry)?;
+        self.palette.write_into(buf)?;
+        self.storage.data.write_into(buf)?;
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum Palette {
+    /// ID of the corresponding entry in its global palette
+    SingleValue(u32),
+    Linear(Vec<u32>),
+    Hashmap(Vec<u32>),
+    Global,
+}
+
+impl Palette {
+    pub fn block_states_read_with_bits_per_entry(
+        buf: &mut impl Read,
+        bits_per_entry: u8,
+    ) -> Result<Palette, String> {
+        Ok(match bits_per_entry {
+            0 => Palette::SingleValue(u32::var_read_into(buf)?),
+            1..=4 => Palette::Linear(Vec::<u32>::var_read_into(buf)?),
+            5..=8 => Palette::Hashmap(Vec::<u32>::var_read_into(buf)?),
+            _ => Palette::Global,
+        })
+    }
+
+    pub fn biomes_read_with_bits_per_entry(
+        buf: &mut impl Read,
+        bits_per_entry: u8,
+    ) -> Result<Palette, String> {
+        Ok(match bits_per_entry {
+            0 => Palette::SingleValue(u32::var_read_into(buf)?),
+            1..=3 => Palette::Linear(Vec::<u32>::var_read_into(buf)?),
+            _ => Palette::Global,
+        })
+    }
+
+    pub fn value_for(&self, value: usize) -> u32 {
+        match self {
+            Palette::SingleValue(v) => *v,
+            Palette::Linear(v) => v[value],
+            Palette::Hashmap(v) => v[value],
+            Palette::Global => value as u32,
+        }
+    }
+}
+
+impl McBufWritable for Palette {
+    fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        match self {
+            Palette::SingleValue(value) => {
+                value.write_into(buf)?;
+            }
+            Palette::Linear(values) => {
+                values.write_into(buf)?;
+            }
+            Palette::Hashmap(values) => {
+                values.write_into(buf)?;
+            }
+            Palette::Global => {}
+        }
+        Ok(())
+    }
+}