NBT decoder optimizations (#17)

* replace HashMap with AHashMap

* faster read_string by just doing read_exact

* re-enable all the benchmarks

Cargo.lock

@@ -19,6 +19,17 @@ dependencies = [
  "cpufeatures",
 ]
 
+[[package]]
+name = "ahash"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47"
+dependencies = [
+ "getrandom",
+ "once_cell",
+ "version_check",
+]
+
 [[package]]
 name = "anyhow"
 version = "1.0.59"
@@ -183,6 +194,7 @@ dependencies = [
 name = "azalea-nbt"
 version = "0.1.0"
 dependencies = [
+ "ahash",
  "azalea-buf",
  "byteorder",
  "criterion",

azalea-nbt/Cargo.toml

@@ -6,6 +6,7 @@ version = "0.1.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+ahash = "0.7.6"
 azalea-buf = {path = "../azalea-buf"}
 byteorder = "1.4.3"
 flate2 = "1.0.23"

azalea-nbt/src/decode.rs

@@ -1,21 +1,17 @@
 use crate::Error;
 use crate::Tag;
-use azalea_buf::BufReadError;
-use azalea_buf::McBufReadable;
+use ahash::AHashMap;
+use azalea_buf::{BufReadError, McBufReadable};
 use byteorder::{ReadBytesExt, BE};
 use flate2::read::{GzDecoder, ZlibDecoder};
-use std::collections::HashMap;
-use std::io::BufRead;
-use std::io::Read;
+use std::io::{BufRead, Read};
 
 #[inline]
 fn read_string(stream: &mut impl Read) -> Result<String, Error> {
     let length = stream.read_u16::<BE>()?;
 
-    let mut buf = Vec::with_capacity(length as usize);
-    for _ in 0..length {
-        buf.push(stream.read_u8()?);
-    }
+    let mut buf = vec![0; length as usize];
+    stream.read_exact(&mut buf)?;
     Ok(String::from_utf8(buf)?)
 }
 
@@ -74,7 +70,7 @@ impl Tag {
             // Effectively a list of a named tags. Order is not guaranteed.
             10 => {
                 // we default to capacity 4 because it'll probably not be empty
-                let mut map = HashMap::with_capacity(4);
+                let mut map = AHashMap::with_capacity(4);
                 loop {
                     let tag_id = stream.read_u8().unwrap_or(0);
                     if tag_id == 0 {
@@ -122,7 +118,7 @@ impl Tag {
         }
         let name = read_string(stream)?;
         let tag = Tag::read_known(stream, tag_id)?;
-        let mut map = HashMap::with_capacity(1);
+        let mut map = AHashMap::with_capacity(1);
         map.insert(name, tag);
 
         Ok(Tag::Compound(map))

azalea-nbt/src/encode.rs

@@ -1,9 +1,9 @@
 use crate::Error;
 use crate::Tag;
+use ahash::AHashMap;
 use azalea_buf::McBufWritable;
 use byteorder::{WriteBytesExt, BE};
 use flate2::write::{GzEncoder, ZlibEncoder};
-use std::collections::HashMap;
 use std::io::Write;
 
 // who needs friends when you've got code that runs in nanoseconds?
@@ -19,7 +19,7 @@ fn write_string(writer: &mut dyn Write, string: &str) -> Result<(), Error> {
 #[inline]
 fn write_compound(
     writer: &mut dyn Write,
-    value: &HashMap<String, Tag>,
+    value: &AHashMap<String, Tag>,
     end_tag: bool,
 ) -> Result<(), Error> {
     for (key, tag) in value {

azalea-nbt/src/lib.rs

@@ -9,15 +9,16 @@ pub use tag::Tag;
 #[cfg(test)]
 mod tests {
     use super::*;
+    use ahash::AHashMap;
     use azalea_buf::{McBufReadable, McBufWritable};
-    use std::{collections::HashMap, io::Cursor};
+    use std::io::Cursor;
 
     #[test]
     fn mcbuf_nbt() {
         let mut buf = Vec::new();
-        let tag = Tag::Compound(HashMap::from_iter(vec![(
+        let tag = Tag::Compound(AHashMap::from_iter(vec![(
             "hello world".to_string(),
-            Tag::Compound(HashMap::from_iter(vec![(
+            Tag::Compound(AHashMap::from_iter(vec![(
                 "name".to_string(),
                 Tag::String("Bananrama".to_string()),
             )])),
@@ -29,9 +30,9 @@ mod tests {
         let result = Tag::read_from(&mut buf).unwrap();
         assert_eq!(
             result,
-            Tag::Compound(HashMap::from_iter(vec![(
+            Tag::Compound(AHashMap::from_iter(vec![(
                 "hello world".to_string(),
-                Tag::Compound(HashMap::from_iter(vec![(
+                Tag::Compound(AHashMap::from_iter(vec![(
                     "name".to_string(),
                     Tag::String("Bananrama".to_string()),
                 )])),

azalea-nbt/src/tag.rs

@@ -1,20 +1,20 @@
-use std::collections::HashMap;
+use ahash::AHashMap;
 
 #[derive(Clone, Debug, PartialEq)]
 pub enum Tag {
-    End,                            // 0
-    Byte(i8),                       // 1
-    Short(i16),                     // 2
-    Int(i32),                       // 3
-    Long(i64),                      // 4
-    Float(f32),                     // 5
-    Double(f64),                    // 6
-    ByteArray(Vec<i8>),             // 7
-    String(String),                 // 8
-    List(Vec<Tag>),                 // 9
-    Compound(HashMap<String, Tag>), // 10
-    IntArray(Vec<i32>),             // 11
-    LongArray(Vec<i64>),            // 12
+    End,                             // 0
+    Byte(i8),                        // 1
+    Short(i16),                      // 2
+    Int(i32),                        // 3
+    Long(i64),                       // 4
+    Float(f32),                      // 5
+    Double(f64),                     // 6
+    ByteArray(Vec<i8>),              // 7
+    String(String),                  // 8
+    List(Vec<Tag>),                  // 9
+    Compound(AHashMap<String, Tag>), // 10
+    IntArray(Vec<i32>),              // 11
+    LongArray(Vec<i64>),             // 12
 }
 
 impl Default for Tag {
@@ -107,7 +107,7 @@ impl Tag {
     }
 
     #[inline]
-    pub fn as_compound(&self) -> Option<&HashMap<String, Tag>> {
+    pub fn as_compound(&self) -> Option<&AHashMap<String, Tag>> {
         if let Tag::Compound(v) = self {
             Some(v)
         } else {

azalea-nbt/tests/tests.rs

@@ -1,6 +1,6 @@
+use ahash::AHashMap;
 use azalea_nbt::Tag;
 use std::{
-    collections::HashMap,
     fs::File,
     io::{Cursor, Read},
 };
@@ -12,9 +12,9 @@ fn test_decode_hello_world() {
     let tag = Tag::read(&mut file).unwrap();
     assert_eq!(
         tag,
-        Tag::Compound(HashMap::from_iter(vec![(
+        Tag::Compound(AHashMap::from_iter(vec![(
             "hello world".to_string(),
-            Tag::Compound(HashMap::from_iter(vec![(
+            Tag::Compound(AHashMap::from_iter(vec![(
                 "name".to_string(),
                 Tag::String("Bananrama".to_string()),
             )]))
@@ -58,7 +58,7 @@ fn test_bigtest() {
 
 #[test]
 fn test_stringtest() {
-    let correct_tag = Tag::Compound(HashMap::from_iter(vec![(
+    let correct_tag = Tag::Compound(AHashMap::from_iter(vec![(
         "😃".to_string(),
         Tag::List(vec![
             Tag::String("asdfkghasfjgihsdfogjsndfg".to_string()),