feat: added unwrapping and wrapping for NBT lists to and from binary (#1742)

* added wrapping, unwrapping tags during serializing/deserializing NBT to/from bytes and added a serialization test

* fixed wrapping check in `wrap_tag_if_needed`

* added wrapped compound to test and renamed `put_component`

* fixed formatting

* fixed comment in test

* fixed test comments

Author: SomeYellowGuy

pumpkin-nbt/src/compound.rs

@@ -130,7 +130,7 @@ impl NbtCompound {
         self.put(name, NbtTag::Double(value));
     }
 
-    pub fn put_component(&mut self, name: &str, value: Self) {
+    pub fn put_compound(&mut self, name: &str, value: Self) {
         self.put(name, NbtTag::Compound(value));
     }
 

pumpkin-nbt/src/lib.rs

@@ -211,12 +211,14 @@ mod test {
     use std::io::Cursor;
 
     use crate::Error;
+    use crate::compound::NbtCompound;
     use crate::deserializer::from_bytes;
     use crate::nbt_byte_array;
     use crate::nbt_int_array;
     use crate::nbt_long_array;
-    use crate::serializer::to_bytes;
     use crate::serializer::to_bytes_named;
+    use crate::serializer::{WriteAdaptor, to_bytes};
+    use crate::tag::NbtTag;
     use crate::{deserializer::from_bytes_unnamed, serializer::to_bytes_unnamed};
     use serde::{Deserialize, Serialize};
 
@@ -406,6 +408,88 @@ mod test {
         assert_eq!(list_compound, recreated_struct);
     }
 
+    #[test]
+    fn wrapper_compound_lists() {
+        let mut vec: Vec<NbtTag> = Vec::new();
+
+        // These tags will be wrapped during serialization.
+        vec.push(NbtTag::Int(-1823));
+        vec.push(NbtTag::Int(123));
+        vec.push(NbtTag::String("Not an int".to_string()));
+        vec.push(NbtTag::Byte(2));
+
+        // This compound will not, since the list is already a list of compound tags.
+        // This compound cannot be unwrapped in any way, so it is preserved
+        // on deserialization.
+        vec.push(NbtTag::Compound({
+            let mut compound = NbtCompound::new();
+            compound.put_short("example", 1234);
+            compound
+        }));
+
+        // This wrapper compound will be wrapped because we want to preserve the
+        // original data during deserialization.
+        //
+        // Suppose we had {"": `tag`}. If we didn't wrap this, on deserialization,
+        // we would get `tag`, which doesn't match the serialized compound tag.
+        // Therefore, we must wrap it and serialize {"": {"": `tag`}}.
+        // Then on deserialization, we get {"": `tag`}, which matches what we wanted
+        // to serialize in the first place.
+        //
+        // This compound represents {"": 1L}.
+        vec.push(NbtTag::Compound({
+            let mut compound = NbtCompound::new();
+            compound.put_long("", 1);
+            compound
+        }));
+
+        let expected_bytes = [
+            0x09, // List type
+            0x0A, // This list is a compound tag list
+            0x00, 0x00, 0x00, 0x06, // This list has 6 elements.
+            // Now for parsing each compound tag:
+            0x03, // Int type
+            0x00, 0x00, // Empty key
+            0xFF, 0xFF, 0xF8, 0xE1, // -1823
+            0x00, // End
+            0x03, // Int type
+            0x00, 0x00, // Empty key
+            0x00, 0x00, 0x00, 0x7B, // 123
+            0x00, // End
+            0x08, // String type
+            0x00, 0x00, // Empty key
+            0x00, 0x0A, // The string is 10 characters long.
+            0x4E, 0x6F, 0x74, 0x20, 0x61, 0x6E, 0x20, 0x69, 0x6E, 0x74, // "Not an int"
+            0x00, // End
+            0x01, // Byte type
+            0x00, 0x00, // Empty key
+            0x02, // 2b
+            0x00, // End
+            // For the first (unwrapped) compound:
+            0x02, // Short type
+            0x00, 0x07, // The key is 7 characters long.
+            0x65, 0x78, 0x61, 0x6D, 0x70, 0x6C, 0x65, // "example"
+            0x04, 0xD2, // 1234
+            0x00, // End
+            // For the second (wrapped) wrapper compound:
+            0x0A, // Compound type
+            0x00, 0x00, // Empty key
+            0x04, // Long type
+            0x00, 0x00, // Empty key
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, // 1L
+            0x00, // End
+            0x00, // End
+        ];
+
+        let mut bytes = Vec::new();
+        let mut write_adaptor = WriteAdaptor::new(&mut bytes);
+        NbtTag::List(vec)
+            .serialize(&mut write_adaptor)
+            .expect("Expected serialization to succeed");
+
+        assert_eq!(bytes, expected_bytes);
+    }
+
     #[test]
     fn nbt_arrays() {
         #[derive(Serialize)]

pumpkin-nbt/src/nbt_compress.rs

@@ -131,7 +131,7 @@ mod tests {
         // Create a nested compound
         let mut nested = NbtCompound::new();
         nested.put_int("nested_int", 42);
-        compound.put_component("nested_compound", nested);
+        compound.put_compound("nested_compound", nested);
 
         // Write to GZip using streaming
         let mut buffer = Vec::new();

pumpkin-nbt/src/tag.rs

@@ -55,6 +55,70 @@ impl NbtTag {
         Ok(())
     }
 
+    /// Gets the element type of [`NbtTag::List`] the provided `Vec`
+    /// represents. If any elements in the `Vec` are found to be of
+    /// different types, this returns [`COMPOUND_ID`].
+    #[must_use]
+    fn get_list_element_type_id(vec: &Vec<Self>) -> u8 {
+        let mut element_id = END_ID;
+
+        for tag in vec {
+            let id = tag.get_type_id();
+            if element_id == END_ID {
+                element_id = id;
+            } else if element_id != id {
+                return COMPOUND_ID;
+            }
+        }
+
+        element_id
+    }
+
+    /// Tries to unwrap (flatten) a wrapped `NbtTag`. If there is a wrapped tag, it is returned.
+    /// If no unwrap is possible, this returns the given tag.
+    fn flatten(tag: Self) -> Self {
+        if let Self::Compound(mut compound) = tag {
+            // Try to get the wrapped tag, stored by "".
+            if Self::is_wrapper_compound(&compound) {
+                compound.child_tags.remove(0).1
+            } else {
+                Self::Compound(compound)
+            }
+        } else {
+            tag
+        }
+    }
+
+    /// Returns whether an [`NbtCompound`] is a wrapper compound.
+    ///
+    /// A *wrapper compound* is a compound that stores exactly one
+    /// key-value pair, an empty string key (`""`) and an `NbtTag`.
+    fn is_wrapper_compound(compound: &NbtCompound) -> bool {
+        compound.child_tags.len() == 1 && compound.child_tags[0].0.is_empty()
+    }
+
+    /// Wraps the provided tag if needed with the provided element type of list
+    /// the wrapped tag, if any, would be added to.
+    fn wrap_tag_if_needed(element_type: u8, tag: Self) -> Self {
+        if element_type == COMPOUND_ID {
+            if let Self::Compound(compound) = &tag
+                && !Self::is_wrapper_compound(compound)
+            {
+                tag
+            } else {
+                Self::wrap_tag(tag)
+            }
+        } else {
+            tag
+        }
+    }
+
+    fn wrap_tag(tag: Self) -> Self {
+        let mut compound = NbtCompound::new();
+        compound.put("", tag);
+        Self::Compound(compound)
+    }
+
     pub fn serialize_data<W: Write>(self, w: &mut WriteAdaptor<W>) -> serializer::Result<()> {
         match self {
             Self::End => {}
@@ -82,10 +146,15 @@ impl NbtTag {
                     return Err(Error::LargeLength(len));
                 }
 
-                w.write_u8_be(list.first().unwrap_or(&Self::End).get_type_id())?;
+                let list_element_id = Self::get_list_element_type_id(&list);
+
+                w.write_u8_be(list_element_id)?;
                 w.write_i32_be(len as i32)?;
                 for nbt_tag in list {
-                    nbt_tag.serialize_data(w)?;
+                    // Since tags in the same list tag must have the same type,
+                    // we need to handle those of different tag types by
+                    // wrapping them in `NbtCompound`s if needed.
+                    Self::wrap_tag_if_needed(list_element_id, nbt_tag).serialize_data(w)?;
                 }
             }
             Self::Compound(compound) => {
@@ -229,7 +298,8 @@ impl NbtTag {
                 for _ in 0..len {
                     let tag = Self::deserialize_data(reader, tag_type_id)?;
                     assert_eq!(tag.get_type_id(), tag_type_id);
-                    list.push(tag);
+                    // Try unwrapping the tag.
+                    list.push(Self::flatten(tag));
                 }
                 Ok(Self::List(list))
             }

pumpkin-world/src/block/entities/mob_spawner.rs

@@ -173,7 +173,7 @@ impl BlockEntity for MobSpawnerBlockEntity {
                 let mut entity_nbt = NbtCompound::new();
                 entity_nbt.put_string("id", format!("minecraft:{}", entity_type.resource_name));
 
-                nbt.put_component("entity", entity_nbt);
+                nbt.put_compound("entity", entity_nbt);
             }
         })
     }
@@ -186,9 +186,9 @@ impl BlockEntity for MobSpawnerBlockEntity {
             let mut entity_nbt = NbtCompound::new();
             entity_nbt.put_string("id", format!("minecraft:{}", entity_type.resource_name));
 
-            spawn_entry.put_component("entity", entity_nbt);
+            spawn_entry.put_compound("entity", entity_nbt);
 
-            final_nbt.put_component("SpawnData", spawn_entry);
+            final_nbt.put_compound("SpawnData", spawn_entry);
         }
         Some(final_nbt)
     }

pumpkin-world/src/item/mod.rs

@@ -436,7 +436,7 @@ impl ItemStack {
         }
 
         // Store custom data like enchantments, display name, etc. would go here
-        compound.put_component("components", tag);
+        compound.put_compound("components", tag);
     }
 
     #[must_use]

pumpkin/src/entity/player.rs

@@ -3019,27 +3019,27 @@ impl NBTStorage for PlayerInventory {
                     drop(stack);
                     match slot {
                         EquipmentSlot::OffHand(_) => {
-                            equipment_compound.put_component("offhand", item_compound);
+                            equipment_compound.put_compound("offhand", item_compound);
                         }
                         EquipmentSlot::Feet(_) => {
-                            equipment_compound.put_component("feet", item_compound);
+                            equipment_compound.put_compound("feet", item_compound);
                         }
                         EquipmentSlot::Legs(_) => {
-                            equipment_compound.put_component("legs", item_compound);
+                            equipment_compound.put_compound("legs", item_compound);
                         }
                         EquipmentSlot::Chest(_) => {
-                            equipment_compound.put_component("chest", item_compound);
+                            equipment_compound.put_compound("chest", item_compound);
                         }
                         EquipmentSlot::Head(_) => {
-                            equipment_compound.put_component("head", item_compound);
+                            equipment_compound.put_compound("head", item_compound);
                         }
                         _ => {
                             warn!("Invalid equipment slot for a player");
                         }
                     }
                 }
             }
-            nbt.put_component("equipment", equipment_compound);
+            nbt.put_compound("equipment", equipment_compound);
             nbt.put("Inventory", NbtTag::List(items));
         })
     }
@@ -3302,7 +3302,7 @@ impl NBTStorage for Abilities {
             component.put_bool("mayBuild", self.allow_modify_world);
             component.put_float("flySpeed", self.fly_speed);
             component.put_float("walkSpeed", self.walk_speed);
-            nbt.put_component("abilities", component);
+            nbt.put_compound("abilities", component);
         })
     }