Implement primitive translation arguments (#263)

* implement primitive translation arguments

* add distinct number types to PrimitiveOrComponent

* fix deserializing PrimitiveOrComponent::Boolean from nbt

* improve nbt serialization for translatable component arguments

1 files changed, 146 insertions, 66 deletions

diff --git a/azalea-chat/src/translatable_component.rs b/azalea-chat/src/translatable_component.rs
index b8e93245..13afe6b7 100644
--- a/azalea-chat/src/translatable_component.rs
+++ b/azalea-chat/src/translatable_component.rs
@@ -1,24 +1,36 @@
 use std::fmt::{self, Display};
 
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 #[cfg(feature = "simdnbt")]
-use simdnbt::Serialize as _;
+use simdnbt::ToNbtTag;
 
 use crate::{FormattedText, base_component::BaseComponent, text_component::TextComponent};
 
-#[derive(Clone, Debug, PartialEq, Serialize)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 #[serde(untagged)]
-pub enum StringOrComponent {
+pub enum PrimitiveOrComponent {
+    Boolean(bool),
+    Short(i16),
+    Integer(i32),
+    Long(i64),
+    Float(f32),
+    Double(f64),
     String(String),
     FormattedText(FormattedText),
 }
 
 #[cfg(feature = "simdnbt")]
-impl simdnbt::ToNbtTag for StringOrComponent {
+impl simdnbt::ToNbtTag for PrimitiveOrComponent {
     fn to_nbt_tag(self) -> simdnbt::owned::NbtTag {
         match self {
-            StringOrComponent::String(s) => s.to_nbt_tag(),
-            StringOrComponent::FormattedText(c) => c.to_nbt_tag(),
+            PrimitiveOrComponent::Boolean(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::Short(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::Integer(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::Long(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::Float(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::Double(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::String(value) => value.to_nbt_tag(),
+            PrimitiveOrComponent::FormattedText(value) => value.to_nbt_tag(),
         }
     }
 }
@@ -33,43 +45,97 @@ pub struct TranslatableComponent {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub fallback: Option<String>,
     #[serde(rename = "with")]
-    pub args: Vec<StringOrComponent>,
+    pub args: Vec<PrimitiveOrComponent>,
 }
 
 #[cfg(feature = "simdnbt")]
-fn serialize_args_as_nbt(args: &[StringOrComponent]) -> simdnbt::owned::NbtList {
-    // if it's all strings then make it a string list
-    // if it's all components then make it a compound list
-    // if it's a mix then return an error
-
-    use tracing::debug;
-
-    let mut string_list = Vec::new();
-    let mut compound_list = Vec::new();
-
-    for arg in args {
-        match arg {
-            StringOrComponent::String(s) => {
-                string_list.push(s.clone());
-            }
-            StringOrComponent::FormattedText(c) => {
-                compound_list.push(c.clone().to_compound());
-            }
+fn serialize_args_as_nbt(args: Vec<PrimitiveOrComponent>) -> simdnbt::owned::NbtList {
+    let tags: Vec<simdnbt::owned::NbtTag> = args.into_iter().map(|arg| arg.to_nbt_tag()).collect();
+
+    if let Some(first_element) = tags.first() {
+        let tags = tags.clone();
+        let homogenous_list = match first_element {
+            simdnbt::owned::NbtTag::Byte(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_byte())
+                .collect::<Option<Vec<i8>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Short(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_short())
+                .collect::<Option<Vec<i16>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Int(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_int())
+                .collect::<Option<Vec<i32>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Long(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_long())
+                .collect::<Option<Vec<i64>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Float(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_float())
+                .collect::<Option<Vec<f32>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Double(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_double())
+                .collect::<Option<Vec<f64>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::ByteArray(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_byte_array())
+                .collect::<Option<Vec<Vec<u8>>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::String(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_string())
+                .collect::<Option<Vec<simdnbt::Mutf8String>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::List(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_list())
+                .collect::<Option<Vec<simdnbt::owned::NbtList>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::Compound(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_compound())
+                .collect::<Option<Vec<simdnbt::owned::NbtCompound>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::IntArray(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_int_array())
+                .collect::<Option<Vec<Vec<i32>>>>()
+                .map(simdnbt::owned::NbtList::from),
+            simdnbt::owned::NbtTag::LongArray(_) => tags
+                .into_iter()
+                .map(|tag| tag.into_long_array())
+                .collect::<Option<Vec<Vec<i64>>>>()
+                .map(simdnbt::owned::NbtList::from),
+        };
+
+        if let Some(homogenous_list) = homogenous_list {
+            return homogenous_list;
         }
     }
 
-    if !string_list.is_empty() && !compound_list.is_empty() {
-        // i'm actually not sure what vanilla does here, so i just made it return the
-        // string list
-        debug!("Tried to serialize a TranslatableComponent with a mix of strings and components.");
-        return string_list.into();
-    }
-
-    if !string_list.is_empty() {
-        return string_list.into();
+    let mut compounds = Vec::with_capacity(tags.len());
+    for tag in tags {
+        let compound = if let simdnbt::owned::NbtTag::Compound(compound) = tag {
+            compound
+        } else {
+            let mut compound = simdnbt::owned::NbtCompound::new();
+            compound.insert("", tag);
+            compound
+        };
+
+        compounds.push(compound);
     }
 
-    compound_list.into()
+    compounds.into()
 }
 
 #[cfg(feature = "simdnbt")]
@@ -79,13 +145,13 @@ impl simdnbt::Serialize for TranslatableComponent {
         compound.insert("translate", self.key);
         compound.extend(self.base.style.to_compound());
 
-        compound.insert("with", serialize_args_as_nbt(&self.args));
+        compound.insert("with", serialize_args_as_nbt(self.args));
         compound
     }
 }
 
 impl TranslatableComponent {
-    pub fn new(key: String, args: Vec<StringOrComponent>) -> Self {
+    pub fn new(key: String, args: Vec<PrimitiveOrComponent>) -> Self {
         Self {
             base: BaseComponent::new(),
             key,
@@ -97,7 +163,7 @@ impl TranslatableComponent {
     pub fn with_fallback(
         key: String,
         fallback: Option<String>,
-        args: Vec<StringOrComponent>,
+        args: Vec<PrimitiveOrComponent>,
     ) -> Self {
         Self {
             base: BaseComponent::new(),
@@ -142,7 +208,7 @@ impl TranslatableComponent {
                             .args
                             .get(matched)
                             .cloned()
-                            .unwrap_or_else(|| StringOrComponent::String("".to_string()));
+                            .unwrap_or_else(|| PrimitiveOrComponent::String("".to_string()));
 
                         components.push(TextComponent::new(built_text.clone()));
                         built_text.clear();
@@ -160,7 +226,9 @@ impl TranslatableComponent {
                                         &self
                                             .args
                                             .get((d - 1) as usize)
-                                            .unwrap_or(&StringOrComponent::String("".to_string()))
+                                            .unwrap_or(&PrimitiveOrComponent::String(
+                                                "".to_string(),
+                                            ))
                                             .to_string(),
                                     );
                                 } else {
@@ -217,20 +285,32 @@ impl Display for TranslatableComponent {
     }
 }
 
-impl Display for StringOrComponent {
+impl Display for PrimitiveOrComponent {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
         match self {
-            StringOrComponent::String(s) => write!(f, "{s}"),
-            StringOrComponent::FormattedText(c) => write!(f, "{c}"),
+            PrimitiveOrComponent::Boolean(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::Short(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::Integer(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::Long(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::Float(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::Double(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::String(value) => write!(f, "{value}"),
+            PrimitiveOrComponent::FormattedText(value) => write!(f, "{value}"),
         }
     }
 }
 
-impl From<StringOrComponent> for TextComponent {
-    fn from(soc: StringOrComponent) -> Self {
+impl From<PrimitiveOrComponent> for TextComponent {
+    fn from(soc: PrimitiveOrComponent) -> Self {
         match soc {
-            StringOrComponent::String(s) => TextComponent::new(s),
-            StringOrComponent::FormattedText(c) => TextComponent::new(c.to_string()),
+            PrimitiveOrComponent::String(value) => TextComponent::new(value),
+            PrimitiveOrComponent::Boolean(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::Short(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::Integer(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::Long(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::Float(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::Double(value) => TextComponent::new(value.to_string()),
+            PrimitiveOrComponent::FormattedText(value) => TextComponent::new(value.to_string()),
         }
     }
 }
@@ -254,10 +334,10 @@ mod tests {
         let c = TranslatableComponent::new(
             "translation.test.complex".to_string(),
             vec![
-                StringOrComponent::String("a".to_string()),
-                StringOrComponent::String("b".to_string()),
-                StringOrComponent::String("c".to_string()),
-                StringOrComponent::String("d".to_string()),
+                PrimitiveOrComponent::String("a".to_string()),
+                PrimitiveOrComponent::String("b".to_string()),
+                PrimitiveOrComponent::String("c".to_string()),
+                PrimitiveOrComponent::String("d".to_string()),
             ],
         );
         // so true mojang
@@ -271,10 +351,10 @@ mod tests {
         let c = TranslatableComponent::new(
             "translation.test.escape".to_string(),
             vec![
-                StringOrComponent::String("a".to_string()),
-                StringOrComponent::String("b".to_string()),
-                StringOrComponent::String("c".to_string()),
-                StringOrComponent::String("d".to_string()),
+                PrimitiveOrComponent::String("a".to_string()),
+                PrimitiveOrComponent::String("b".to_string()),
+                PrimitiveOrComponent::String("c".to_string()),
+                PrimitiveOrComponent::String("d".to_string()),
             ],
         );
         assert_eq!(c.read().unwrap().to_string(), "%s %a %%s %%b".to_string());
@@ -284,10 +364,10 @@ mod tests {
         let c = TranslatableComponent::new(
             "translation.test.invalid".to_string(),
             vec![
-                StringOrComponent::String("a".to_string()),
-                StringOrComponent::String("b".to_string()),
-                StringOrComponent::String("c".to_string()),
-                StringOrComponent::String("d".to_string()),
+                PrimitiveOrComponent::String("a".to_string()),
+                PrimitiveOrComponent::String("b".to_string()),
+                PrimitiveOrComponent::String("c".to_string()),
+                PrimitiveOrComponent::String("d".to_string()),
             ],
         );
         assert_eq!(c.read().unwrap().to_string(), "hi %".to_string());
@@ -297,10 +377,10 @@ mod tests {
         let c = TranslatableComponent::new(
             "translation.test.invalid2".to_string(),
             vec![
-                StringOrComponent::String("a".to_string()),
-                StringOrComponent::String("b".to_string()),
-                StringOrComponent::String("c".to_string()),
-                StringOrComponent::String("d".to_string()),
+                PrimitiveOrComponent::String("a".to_string()),
+                PrimitiveOrComponent::String("b".to_string()),
+                PrimitiveOrComponent::String("c".to_string()),
+                PrimitiveOrComponent::String("d".to_string()),
             ],
         );
         assert_eq!(c.read().unwrap().to_string(), "hi %  s".to_string());
@@ -310,7 +390,7 @@ mod tests {
     fn test_undefined() {
         let c = TranslatableComponent::new(
             "translation.test.undefined".to_string(),
-            vec![StringOrComponent::String("a".to_string())],
+            vec![PrimitiveOrComponent::String("a".to_string())],
         );
         assert_eq!(
             c.read().unwrap().to_string(),
@@ -323,7 +403,7 @@ mod tests {
         let c = TranslatableComponent::with_fallback(
             "translation.test.undefined".to_string(),
             Some("translation fallback: %s".to_string()),
-            vec![StringOrComponent::String("a".to_string())],
+            vec![PrimitiveOrComponent::String("a".to_string())],
         );
         assert_eq!(
             c.read().unwrap().to_string(),