feat!: proper heterogeneous lists

Cargo.toml

@@ -24,6 +24,7 @@ simd_cesu8 = "1.0.1"
 derive_more = { version = "2.0.1", features = ["into", "from"] }
 thiserror = "2.0.11"
 serde = { version = "1.0.218", optional = true, features = ["derive"] }
+polonius-the-crab = "0.5.0"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }

src/error.rs

@@ -19,6 +19,8 @@ pub enum Error {
     SerdeError(String),
     #[error("NBT doesn't support this type {0}")]
     UnsupportedType(String),
+    #[error("List contains element of different type")]
+    ListDifferentType,
     #[error(transparent)]
     Io(#[from] io::Error),
 }

src/nbt/compound.rs

@@ -20,6 +20,16 @@ impl NbtCompound {
         }
     }
 
+    pub fn wrap(child: impl Into<NbtTag>) -> Self {
+        Self {
+            child_tags: vec![(String::new(), child.into())],
+        }
+    }
+
+    pub fn is_wrapper(&self) -> bool {
+        self.child_tags.len() == 1 && self.child_tags[0].0.is_empty()
+    }
+
     pub fn deserialize_content(bytes: &mut impl Buf) -> Result<NbtCompound, Error> {
         let mut compound = NbtCompound::new();
 

src/nbt/tag.rs

@@ -60,10 +60,39 @@ impl NbtTag {
                 bytes.put_slice(&java_string);
             }
             NbtTag::List(list) => {
-                bytes.put_u8(list.first().unwrap_or(&NbtTag::End).get_type_id());
-                bytes.put_i32(list.len() as i32);
-                for nbt_tag in list {
-                    bytes.put(nbt_tag.serialize_data())
+                let ty = list
+                    .first()
+                    .map(|e| e.get_type_id())
+                    .unwrap_or(NbtTag::End.get_type_id());
+
+                let mut homogeneous = true;
+                for element in list {
+                    if element.get_type_id() != ty {
+                        homogeneous = false;
+                    }
+                }
+
+                if homogeneous {
+                    bytes.put_u8(ty);
+                    bytes.put_i32(list.len() as i32);
+                    for tag in list {
+                        bytes.put(tag.serialize_data())
+                    }
+                } else {
+                    bytes.put_u8(COMPOUND_ID);
+                    bytes.put_i32(list.len() as i32);
+                    for tag in list {
+                        match tag {
+                            NbtTag::Compound(_) => bytes.put(tag.serialize_data()),
+                            value => {
+                                // Manually serialize value as a Compound, to avoid a Clone
+                                bytes.put_u8(value.get_type_id());
+                                bytes.put_u16(0); // Empty string key
+                                bytes.put(value.serialize_data());
+                                bytes.put_u8(END_ID);
+                            }
+                        }
+                    }
                 }
             }
             NbtTag::Compound(compound) => {
@@ -131,11 +160,36 @@ impl NbtTag {
                 let tag_type_id = bytes.get_u8();
                 let len = bytes.get_i32();
                 let mut list = Vec::with_capacity(len as usize);
+                let mut all_compounds = len != 0;
+                let mut any_singletons = false;
                 for _ in 0..len {
                     let tag = NbtTag::deserialize_data(bytes, tag_type_id)?;
-                    assert_eq!(tag.get_type_id(), tag_type_id);
+                    if tag.get_type_id() != tag_type_id {
+                        return Err(Error::ListDifferentType);
+                    }
+
+                    if let NbtTag::Compound(compound) = &tag {
+                        if !any_singletons && compound.is_wrapper() {
+                            any_singletons = true;
+                        }
+                    } else {
+                        all_compounds = false;
+                    }
+
                     list.push(tag);
                 }
+
+                if all_compounds && any_singletons {
+                    for tag in &mut list {
+                        match tag {
+                            NbtTag::Compound(ref mut compound) if compound.is_wrapper() => {
+                                *tag = compound.child_tags.remove(0).1;
+                            }
+                            _ => (),
+                        }
+                    }
+                }
+
                 Ok(NbtTag::List(list))
             }
             COMPOUND_ID => Ok(NbtTag::Compound(NbtCompound::deserialize_content(bytes)?)),

tests/list.rs

@@ -0,0 +1,55 @@
+use crab_nbt::{NbtCompound, NbtTag};
+
+#[allow(clippy::perf)]
+#[test]
+fn list_types() {
+    let empty: Vec<NbtTag> = Vec::new();
+    let mut iter = empty.clone().into_iter();
+    assert_eq!(iter.next(), None);
+
+    let mut homogeneous = Vec::new();
+    homogeneous.push(1.into());
+    homogeneous.push(2.into());
+    homogeneous.push(3.into());
+    let mut iter = homogeneous.iter();
+    assert_eq!(iter.next(), Some(&1.into()));
+    assert_eq!(iter.next(), Some(&2.into()));
+    assert_eq!(iter.next(), Some(&3.into()));
+    assert_eq!(iter.next(), None);
+
+    let mut serialized = NbtTag::List(homogeneous).serialize();
+    let Ok(NbtTag::List(homogeneous)) = NbtTag::deserialize(&mut serialized) else {
+        panic!("homogeneous list did not survive round trip");
+    };
+
+    let mut heterogeneous = homogeneous.clone();
+    heterogeneous.push("four".into());
+    heterogeneous.push(vec![5].into());
+    heterogeneous.push(
+        NbtCompound {
+            child_tags: vec![("six".into(), 7.into())],
+        }
+        .into(),
+    );
+    let mut iter = heterogeneous.iter();
+    assert_eq!(iter.next(), Some(&1.into()));
+    assert_eq!(iter.next(), Some(&2.into()));
+    assert_eq!(iter.next(), Some(&3.into()));
+    assert_eq!(iter.next(), Some(&"four".into()));
+    assert_eq!(iter.next(), Some(&vec![5].into()));
+    assert_eq!(
+        iter.next(),
+        Some(
+            &NbtCompound {
+                child_tags: vec![("six".into(), 7.into())]
+            }
+            .into()
+        )
+    );
+    assert_eq!(iter.next(), None);
+
+    let mut serialized = NbtTag::List(heterogeneous).serialize();
+    let Ok(NbtTag::List(_heterogeneous)) = NbtTag::deserialize(&mut serialized) else {
+        panic!("heterogeneous list did not survive round trip");
+    };
+}