Fix chunk decoding

7 files changed, 85 insertions, 19 deletions

diff --git a/azalea-protocol/src/mc_buf/mod.rs b/azalea-protocol/src/mc_buf/mod.rs
index a82334fb..debf2991 100755
--- a/azalea-protocol/src/mc_buf/mod.rs
+++ b/azalea-protocol/src/mc_buf/mod.rs
@@ -12,6 +12,8 @@ pub use write::{McBufVarintWritable, McBufWritable, Writable};
 const MAX_STRING_LENGTH: u16 = 32767;
 // const MAX_COMPONENT_STRING_LENGTH: u32 = 262144;
 
+// TODO: maybe get rid of the readable/writable traits so there's not two ways to do the same thing and improve McBufReadable/McBufWritable
+
 // TODO: have a definitions.rs in mc_buf that contains UnsizedByteArray and BitSet
 
 /// A Vec<u8> that isn't prefixed by a VarInt with the size.
diff --git a/azalea-protocol/src/mc_buf/read.rs b/azalea-protocol/src/mc_buf/read.rs
index e1ae321c..a8c44bdf 100755
--- a/azalea-protocol/src/mc_buf/read.rs
+++ b/azalea-protocol/src/mc_buf/read.rs
@@ -255,6 +255,17 @@ impl McBufVarintReadable for i32 {
     }
 }
 
+impl<T: McBufVarintReadable> McBufVarintReadable for Vec<T> {
+    fn varint_read_into(buf: &mut impl Read) -> Result<Self, String> {
+        let length = u32::varint_read_into(buf)?;
+        let mut vec = Vec::with_capacity(length as usize);
+        for _ in 0..length {
+            vec.push(T::varint_read_into(buf)?);
+        }
+        Ok(vec)
+    }
+}
+
 impl McBufReadable for UnsizedByteArray {
     fn read_into(buf: &mut impl Read) -> Result<Self, String> {
         Ok(UnsizedByteArray(buf.read_bytes()?))
diff --git a/azalea-protocol/src/mc_buf/write.rs b/azalea-protocol/src/mc_buf/write.rs
index 3a4a02f8..cc5f2284 100755
--- a/azalea-protocol/src/mc_buf/write.rs
+++ b/azalea-protocol/src/mc_buf/write.rs
@@ -210,6 +210,13 @@ impl McBufVarintWritable for u32 {
     }
 }
 
+// Vec<T> varint
+impl<T: McBufVarintWritable> McBufVarintWritable for Vec<T> {
+    fn varint_write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
+        buf.write_list(self, |buf, i| i.varint_write_into(buf))
+    }
+}
+
 // u16
 impl McBufWritable for u16 {
     fn write_into(&self, buf: &mut impl Write) -> Result<(), std::io::Error> {
diff --git a/azalea-world/src/bit_storage.rs b/azalea-world/src/bit_storage.rs
index 5ced67fe..9cc3a053 100644
--- a/azalea-world/src/bit_storage.rs
+++ b/azalea-world/src/bit_storage.rs
@@ -1,4 +1,4 @@
-use std::{error::Error, fmt, ops::Index};
+use std::{error::Error, fmt};
 
 // this is from minecraft's code
 // yeah idk either
@@ -70,6 +70,7 @@ const MAGIC: [(i32, i32, i32); 64] = [
 ];
 
 /// A compact list of integers with the given number of bits per entry.
+#[derive(Clone)]
 pub struct BitStorage {
     data: Vec<u64>,
     bits: usize,
@@ -135,7 +136,7 @@ impl BitStorage {
 
     pub fn cell_index(&self, index: u64) -> usize {
         let first = self.divide_mul as u64;
-        let second = self.divide_mul as u64;
+        let second = self.divide_add as u64;
 
         (index * first + second >> 32 >> self.divide_shift)
             .try_into()
@@ -182,8 +183,8 @@ mod tests {
         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 expected_compact: [u64; 2] = [0x0020863148418841, 0x01018A7260F68C87];
-        let storage = BitStorage::new(5, data.len(), Some(expected_compact.to_vec())).unwrap();
+        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
index 54961401..abeac181 100644
--- a/azalea-world/src/lib.rs
+++ b/azalea-world/src/lib.rs
@@ -11,6 +11,8 @@ use std::{
     sync::{Arc, Mutex},
 };
 
+use crate::palette::PalettedContainerType;
+
 #[cfg(test)]
 mod tests {
     #[test]
@@ -43,7 +45,7 @@ impl World {
         // let existing_chunk = &self.storage[pos];
 
         let chunk = Arc::new(Mutex::new(Chunk::read_with_world(data, self)?));
-        println!("Loaded chunk {:?}", chunk);
+        println!("Loaded chunk {:?}", pos);
         self.storage[pos] = Some(chunk);
 
         Ok(())
@@ -121,7 +123,9 @@ impl Chunk {
     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 {
+        println!("\n\nreading {} sections", section_count);
+        for i in 0..section_count {
+            println!("reading section #{}", i);
             let section = Section::read_into(buf)?;
             sections.push(section);
         }
@@ -148,8 +152,15 @@ pub struct Section {
 impl McBufReadable for Section {
     fn read_into(buf: &mut impl Read) -> Result<Self, String> {
         let block_count = u16::read_into(buf)?;
-        let states = PalettedContainer::read_into(buf)?;
-        let biomes = PalettedContainer::read_into(buf)?;
+        println!("block count: {}\n", block_count);
+        // 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)?;
+        println!("! read states, reading biomes next");
+        let biomes = PalettedContainer::read_with_type(buf, &PalettedContainerType::Biomes)?;
+        println!();
         Ok(Section {
             block_count,
             states,
diff --git a/azalea-world/src/palette.rs b/azalea-world/src/palette.rs
index 69900fe6..47354c94 100644
--- a/azalea-world/src/palette.rs
+++ b/azalea-world/src/palette.rs
@@ -1,20 +1,43 @@
+use azalea_protocol::mc_buf::{
+    McBufReadable, McBufVarintReadable, McBufWritable, Readable, Writable,
+};
 use std::io::{Read, Write};
 
-use azalea_protocol::mc_buf::{McBufReadable, McBufWritable, Readable, Writable};
+#[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 data: Vec<i64>,
+    pub data: Vec<u64>,
 }
 
-impl McBufReadable for PalettedContainer {
-    fn read_into(buf: &mut impl Read) -> Result<Self, String> {
+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 = Palette::read_with_bits_per_entry(buf, bits_per_entry)?;
-        let data = Vec::<i64>::read_into(buf)?;
+        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 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."
+        );
+
         Ok(PalettedContainer {
             bits_per_entry,
             palette,
@@ -41,14 +64,25 @@ pub enum Palette {
 }
 
 impl Palette {
-    pub fn read_with_bits_per_entry(
+    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::varint_read_into(buf)?),
+            1..=4 => Palette::Linear(Vec::<u32>::varint_read_into(buf)?),
+            5..=8 => Palette::Hashmap(Vec::<u32>::varint_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::read_into(buf)?),
-            1..=4 => Palette::Linear(Vec::<u32>::read_into(buf)?),
-            5..=8 => Palette::Hashmap(Vec::<u32>::read_into(buf)?),
+            0 => Palette::SingleValue(u32::varint_read_into(buf)?),
+            1..=3 => Palette::Linear(Vec::<u32>::varint_read_into(buf)?),
             _ => Palette::Global,
         })
     }
diff --git a/bot/src/main.rs b/bot/src/main.rs
index 167f4edc..7cf056bf 100644
--- a/bot/src/main.rs
+++ b/bot/src/main.rs
@@ -5,7 +5,7 @@ async fn main() {
     println!("Hello, world!");
 
     // let address = "95.111.249.143:10000";
-    let address = "172.23.192.1:65111";
+    let address = "172.23.192.1:62522";
     // let response = azalea_client::ping::ping_server(&address.try_into().unwrap())
     //     .await
     //     .unwrap();