vroom vroom

2 files changed, 242 insertions, 69 deletions

diff --git a/azalea-nbt/src/encode.rs b/azalea-nbt/src/encode.rs
index d1126ee0..27cfe3af 100755
--- a/azalea-nbt/src/encode.rs
+++ b/azalea-nbt/src/encode.rs
@@ -12,96 +12,192 @@ fn write_string(writer: &mut dyn Write, string: &str) -> Result<(), Error> {
 
     Ok(())
 }
+
+// who needs friends when you've got code that runs in nanoseconds?
+
 #[inline]
-fn write_compound(writer: &mut dyn Write, value: &HashMap<String, Tag>) -> Result<(), Error> {
+fn write_compound(
+    writer: &mut dyn Write,
+    value: &HashMap<String, Tag>,
+    end_tag: bool,
+) -> Result<(), Error> {
     for (key, tag) in value {
-        writer.write_u8(tag.id())?;
-        write_string(writer, key)?;
-        tag.write_without_end(writer)?;
-    }
-    writer.write_u8(Tag::End.id())?;
-    Ok(())
-}
-
-impl Tag {
-    #[inline]
-    pub fn write_without_end(&self, writer: &mut dyn Write) -> Result<(), Error> {
-        match self {
+        match tag {
             Tag::End => {}
-            Tag::Byte(value) => writer.write_i8(*value)?,
-            Tag::Short(value) => writer.write_i16::<BE>(*value)?,
-            Tag::Int(value) => writer.write_i32::<BE>(*value)?,
-            Tag::Long(value) => writer.write_i64::<BE>(*value)?,
-            Tag::Float(value) => writer.write_f32::<BE>(*value)?,
-            Tag::Double(value) => writer.write_f64::<BE>(*value)?,
+            Tag::Byte(value) => {
+                writer.write_u8(1)?;
+                write_string(writer, key)?;
+                writer.write_i8(*value)?
+            }
+            Tag::Short(value) => {
+                writer.write_u8(2)?;
+                write_string(writer, key)?;
+                writer.write_i16::<BE>(*value)?
+            }
+            Tag::Int(value) => {
+                writer.write_u8(3)?;
+                write_string(writer, key)?;
+                writer.write_i32::<BE>(*value)?
+            }
+            Tag::Long(value) => {
+                writer.write_u8(4)?;
+                write_string(writer, key)?;
+                writer.write_i64::<BE>(*value)?
+            }
+            Tag::Float(value) => {
+                writer.write_u8(5)?;
+                write_string(writer, key)?;
+                writer.write_f32::<BE>(*value)?
+            }
+            Tag::Double(value) => {
+                writer.write_u8(6)?;
+                write_string(writer, key)?;
+                writer.write_f64::<BE>(*value)?
+            }
             Tag::ByteArray(value) => {
+                writer.write_u8(7)?;
+                write_string(writer, key)?;
                 writer.write_i32::<BE>(value.len() as i32)?;
                 for &byte in value {
                     writer.write_i8(byte)?;
                 }
             }
             Tag::String(value) => {
-                write_string(writer, value)?;
+                writer.write_u8(9)?;
+                write_string(writer, key)?;
+                write_string(writer, value)?
             }
             Tag::List(value) => {
-                // we just get the type from the first item, or default the type to END
-                if value.is_empty() {
-                    writer.write_all(&[0; 5])?;
-                } else {
-                    let first_tag = &value[0];
-                    writer.write_u8(first_tag.id())?;
-                    writer.write_i32::<BE>(value.len() as i32)?;
-                    match first_tag {
-                        Self::Int(_) => {
-                            for i in value {
-                                assert!(matches!(i, Tag::Int(_)));
-                                writer.write_i32::<BE>(
-                                    *i.as_int().expect("List of Int should only contains Int"),
-                                )?;
-                            }
-                        }
-                        Self::String(_) => {
-                            for i in value {
-                                assert!(matches!(i, Tag::String(_)));
-                                write_string(
-                                    writer,
-                                    i.as_string()
-                                        .expect("List of String should only contain String"),
-                                )?;
-                            }
-                        }
-                        Self::Compound(_) => {
-                            for i in value {
-                                assert!(matches!(i, Tag::Compound(_)));
-                                write_compound(
-                                    writer,
-                                    i.as_compound()
-                                        .expect("List of Compound should only contain Compound"),
-                                )?;
-                            }
-                        }
-                        _ => {
-                            for tag in value {
-                                tag.write_without_end(writer)?;
-                            }
-                        }
-                    }
-                }
+                writer.write_u8(10)?;
+                write_string(writer, key)?;
+                write_list(writer, value)?
+            }
+            Tag::Compound(value) => {
+                writer.write_u8(11)?;
+                write_string(writer, key)?;
+                write_compound(writer, value, true)?
             }
-            Tag::Compound(value) => write_compound(writer, value)?,
             Tag::IntArray(value) => {
+                writer.write_u8(12)?;
+                write_string(writer, key)?;
                 writer.write_i32::<BE>(value.len() as i32)?;
                 for &int in value {
                     writer.write_i32::<BE>(int)?;
                 }
             }
             Tag::LongArray(value) => {
+                writer.write_u8(13)?;
+                write_string(writer, key)?;
                 writer.write_i32::<BE>(value.len() as i32)?;
                 for &long in value {
                     writer.write_i64::<BE>(long)?;
                 }
             }
         }
+    }
+    if end_tag {
+        writer.write_u8(Tag::End.id())?;
+    }
+    return Ok(());
+}
+
+#[inline]
+fn write_list(writer: &mut dyn Write, value: &Vec<Tag>) -> Result<(), Error> {
+    // we just get the type from the first item, or default the type to END
+    if value.is_empty() {
+        writer.write_all(&[0; 5])?;
+    } else {
+        let first_tag = &value[0];
+        writer.write_u8(first_tag.id())?;
+        writer.write_i32::<BE>(value.len() as i32)?;
+        match first_tag {
+            Tag::Int(_) => {
+                // assert all items are the same variant
+                // assert_eq!(value.iter().all(|tag| tag.id() == 1), true);
+                for i in value {
+                    assert!(matches!(i, Tag::Int(_)));
+                    writer.write_i32::<BE>(
+                        *i.as_int().expect("List of Int should only contains Int"),
+                    )?;
+                }
+            }
+            Tag::String(_) => {
+                for i in value {
+                    assert!(matches!(i, Tag::String(_)));
+                    write_string(
+                        writer,
+                        i.as_string()
+                            .expect("List of String should only contain String"),
+                    )?;
+                }
+            }
+            Tag::Compound(_) => {
+                for i in value {
+                    assert!(matches!(i, Tag::Compound(_)));
+                    write_compound(
+                        writer,
+                        i.as_compound()
+                            .expect("List of Compound should only contain Compound"),
+                        true,
+                    )?;
+                }
+            }
+            _ => {
+                for i in value {
+                    i.write_without_end(writer)?;
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+#[inline]
+fn write_bytearray(writer: &mut dyn Write, value: &Vec<i8>) -> Result<(), Error> {
+    writer.write_i32::<BE>(value.len() as i32)?;
+    for &byte in value {
+        writer.write_i8(byte)?;
+    }
+    Ok(())
+}
+
+#[inline]
+fn write_intarray(writer: &mut dyn Write, value: &Vec<i32>) -> Result<(), Error> {
+    writer.write_i32::<BE>(value.len() as i32)?;
+    for &int in value {
+        writer.write_i32::<BE>(int)?;
+    }
+    Ok(())
+}
+
+#[inline]
+fn write_longarray(writer: &mut dyn Write, value: &Vec<i64>) -> Result<(), Error> {
+    writer.write_i32::<BE>(value.len() as i32)?;
+    for &long in value {
+        writer.write_i64::<BE>(long)?;
+    }
+    Ok(())
+}
+
+impl Tag {
+    #[inline]
+    pub fn write_without_end(&self, writer: &mut dyn Write) -> Result<(), Error> {
+        match self {
+            Tag::End => {}
+            Tag::Byte(value) => writer.write_i8(*value)?,
+            Tag::Short(value) => writer.write_i16::<BE>(*value)?,
+            Tag::Int(value) => writer.write_i32::<BE>(*value)?,
+            Tag::Long(value) => writer.write_i64::<BE>(*value)?,
+            Tag::Float(value) => writer.write_f32::<BE>(*value)?,
+            Tag::Double(value) => writer.write_f64::<BE>(*value)?,
+            Tag::ByteArray(value) => write_bytearray(writer, value)?,
+            Tag::String(value) => write_string(writer, value)?,
+            Tag::List(value) => write_list(writer, value)?,
+            Tag::Compound(value) => write_compound(writer, value, true)?,
+            Tag::IntArray(value) => write_intarray(writer, value)?,
+            Tag::LongArray(value) => write_longarray(writer, value)?,
+        }
 
         Ok(())
     }
@@ -109,11 +205,7 @@ impl Tag {
     pub fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
         match self {
             Tag::Compound(value) => {
-                for (key, tag) in value {
-                    writer.write_u8(tag.id())?;
-                    write_string(writer, key)?;
-                    tag.write_without_end(writer)?;
-                }
+                write_compound(writer, value, false)?;
                 Ok(())
             }
             _ => Err(Error::InvalidTag),
diff --git a/azalea-nbt/src/tag.rs b/azalea-nbt/src/tag.rs
index 705dff6a..f8960aa6 100755
--- a/azalea-nbt/src/tag.rs
+++ b/azalea-nbt/src/tag.rs
@@ -38,6 +38,24 @@ impl Tag {
     }
 
     #[inline]
+    pub fn as_byte(&self) -> Option<&i8> {
+        if let Tag::Byte(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_short(&self) -> Option<&i16> {
+        if let Tag::Short(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
     pub fn as_int(&self) -> Option<&i32> {
         if let Tag::Int(v) = self {
             Some(v)
@@ -47,6 +65,33 @@ impl Tag {
     }
 
     #[inline]
+    pub fn as_long(&self) -> Option<&i64> {
+        if let Tag::Long(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_float(&self) -> Option<&f32> {
+        if let Tag::Float(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_double(&self) -> Option<&f64> {
+        if let Tag::Double(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
     pub fn as_string(&self) -> Option<&str> {
         if let Tag::String(v) = self {
             Some(v)
@@ -63,4 +108,40 @@ impl Tag {
             None
         }
     }
+
+    #[inline]
+    pub fn as_bytearray(&self) -> Option<&Vec<i8>> {
+        if let Tag::ByteArray(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_intarray(&self) -> Option<&Vec<i32>> {
+        if let Tag::IntArray(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_longarray(&self) -> Option<&Vec<i64>> {
+        if let Tag::LongArray(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    pub fn as_list(&self) -> Option<&Vec<Tag>> {
+        if let Tag::List(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
 }