nbt lookup optimization

7 files changed, 70 insertions, 113 deletions

diff --git a/azalea-nbt/README.md b/azalea-nbt/README.md
index d051af87..e13bcf22 100755
--- a/azalea-nbt/README.md
+++ b/azalea-nbt/README.md
@@ -31,3 +31,5 @@ assert_eq!(
 At the time of writing, Azalea NBT is the fastest NBT decoder (approximately twice as fast as Graphite NBT, the second fastest) and on-par with the fastest NBT encoders (sometimes the fastest, depending on the data).
 
 You can run the benchmarks to compare against other NBT libraries with `cargo bench --bench compare` and the normal benchmarks with `cargo bench --bench nbt`.
+
+Note: For best performance, use `RUSTFLAGS='-C target-cpu=native'`.
diff --git a/azalea-nbt/benches/compare.rs b/azalea-nbt/benches/compare.rs
index 1634b45b..f0fa17c0 100755
--- a/azalea-nbt/benches/compare.rs
+++ b/azalea-nbt/benches/compare.rs
@@ -22,21 +22,21 @@ pub fn bench_read_file(filename: &str, c: &mut Criterion) {
     let mut group = c.benchmark_group(filename);
     group.throughput(Throughput::Bytes(input.len() as u64));
 
-    group.bench_function("azalea_parse", |b| {
-        b.iter(|| {
-            let input = black_box(input);
-            let nbt = azalea_nbt::Nbt::read(&mut Cursor::new(&input)).unwrap();
-            black_box(nbt);
-        })
-    });
+    // group.bench_function("azalea_parse", |b| {
+    //     b.iter(|| {
+    //         let input = black_box(input);
+    //         let nbt = azalea_nbt::Nbt::read(&mut Cursor::new(&input)).unwrap();
+    //         black_box(nbt);
+    //     })
+    // });
 
-    group.bench_function("graphite_parse", |b| {
-        b.iter(|| {
-            let input = black_box(input);
-            let nbt = graphite_binary::nbt::decode::read(&mut &input[..]).unwrap();
-            black_box(nbt);
-        })
-    });
+    // group.bench_function("graphite_parse", |b| {
+    //     b.iter(|| {
+    //         let input = black_box(input);
+    //         let nbt = graphite_binary::nbt::decode::read(&mut
+    // &input[..]).unwrap();         black_box(nbt);
+    //     })
+    // });
 
     // group.bench_function("valence_parse", |b| {
     //     b.iter(|| {
@@ -81,7 +81,7 @@ pub fn bench_read_file(filename: &str, c: &mut Criterion) {
 fn bench(c: &mut Criterion) {
     bench_read_file("tests/bigtest.nbt", c);
     // bench_read_file("tests/simple_player.dat", c);
-    bench_read_file("tests/complex_player.dat", c);
+    // bench_read_file("tests/complex_player.dat", c);
     // bench_read_file("tests/level.dat", c);
     // bench_read_file("tests/stringtest.nbt", c);
     // bench_read_file("tests/inttest.nbt", c);
diff --git a/azalea-nbt/benches/nbt.rs b/azalea-nbt/benches/nbt.rs
index 9d017990..0d413103 100755
--- a/azalea-nbt/benches/nbt.rs
+++ b/azalea-nbt/benches/nbt.rs
@@ -33,9 +33,15 @@ fn bench_file(filename: &str, c: &mut Criterion) {
         })
     });
 
-    group.bench_function("Encode", |b| {
+    // group.bench_function("Encode", |b| {
+    //     b.iter(|| {
+    //         nbt.write(&mut black_box(Vec::new()));
+    //     })
+    // });
+
+    group.bench_function("Get", |b| {
         b.iter(|| {
-            nbt.write(&mut black_box(Vec::new()));
+            black_box(nbt.as_compound().unwrap().get(""));
         })
     });
     group.finish();
@@ -43,11 +49,11 @@ fn bench_file(filename: &str, c: &mut Criterion) {
 
 fn bench(c: &mut Criterion) {
     bench_file("tests/bigtest.nbt", c);
-    bench_file("tests/simple_player.dat", c);
-    bench_file("tests/complex_player.dat", c);
-    bench_file("tests/level.dat", c);
-    bench_file("tests/stringtest.nbt", c);
-    bench_file("tests/inttest.nbt", c);
+    // bench_file("tests/simple_player.dat", c);
+    // bench_file("tests/complex_player.dat", c);
+    // bench_file("tests/level.dat", c);
+    // bench_file("tests/stringtest.nbt", c);
+    // bench_file("tests/inttest.nbt", c);
 }
 
 criterion_group!(benches, bench);
diff --git a/azalea-nbt/src/decode.rs b/azalea-nbt/src/decode.rs
index 5421e1c5..35392b15 100755
--- a/azalea-nbt/src/decode.rs
+++ b/azalea-nbt/src/decode.rs
@@ -170,8 +170,9 @@ fn read_compound(stream: &mut Cursor<&[u8]>) -> Result<NbtCompound, Error> {
         }
         let name = read_string(stream)?;
         let tag = Nbt::read_known(stream, tag_id)?;
-        map.insert(name, tag);
+        map.insert_unsorted(name, tag);
     }
+    map.sort();
     Ok(map)
 }
 
@@ -265,7 +266,7 @@ impl Nbt {
         let name = read_string(stream)?;
         let tag = Nbt::read_known(stream, tag_id)?;
         let mut map = NbtCompound::with_capacity(1);
-        map.insert(name, tag);
+        map.insert_unsorted(name, tag);
 
         Ok(Nbt::Compound(map))
     }
diff --git a/azalea-nbt/src/encode.rs b/azalea-nbt/src/encode.rs
index 279cab03..bf56c223 100755
--- a/azalea-nbt/src/encode.rs
+++ b/azalea-nbt/src/encode.rs
@@ -5,74 +5,52 @@ use flate2::write::{GzEncoder, ZlibEncoder};
 use std::io::Write;
 
 #[inline(always)]
-fn write_string(writer: &mut impl Write, string: &str) {
+fn write_string(writer: &mut impl Write, string: &NbtString) {
     writer.write_u16::<BE>(string.len() as u16).unwrap();
     writer.write_all(string.as_bytes()).unwrap();
 }
 
-#[inline]
+#[inline(always)]
 fn write_compound(writer: &mut impl Write, value: &NbtCompound, end_tag: bool) {
     for (key, tag) in value.iter() {
+        writer.write_u8(tag.id()).unwrap();
+        write_string(writer, key);
         match tag {
             Nbt::End => {}
             Nbt::Byte(value) => {
-                writer.write_u8(BYTE_ID).unwrap();
-                write_string(writer, key);
                 writer.write_i8(*value).unwrap();
             }
             Nbt::Short(value) => {
-                writer.write_u8(SHORT_ID).unwrap();
-                write_string(writer, key);
                 writer.write_i16::<BE>(*value).unwrap();
             }
             Nbt::Int(value) => {
-                writer.write_u8(INT_ID).unwrap();
-                write_string(writer, key);
                 writer.write_i32::<BE>(*value).unwrap();
             }
             Nbt::Long(value) => {
-                writer.write_u8(LONG_ID).unwrap();
-                write_string(writer, key);
                 writer.write_i64::<BE>(*value).unwrap();
             }
             Nbt::Float(value) => {
-                writer.write_u8(FLOAT_ID).unwrap();
-                write_string(writer, key);
                 writer.write_f32::<BE>(*value).unwrap();
             }
             Nbt::Double(value) => {
-                writer.write_u8(DOUBLE_ID).unwrap();
-                write_string(writer, key);
                 writer.write_f64::<BE>(*value).unwrap();
             }
             Nbt::ByteArray(value) => {
-                writer.write_u8(BYTE_ARRAY_ID).unwrap();
-                write_string(writer, key);
                 write_byte_array(writer, value);
             }
             Nbt::String(value) => {
-                writer.write_u8(STRING_ID).unwrap();
-                write_string(writer, key);
                 write_string(writer, value);
             }
             Nbt::List(value) => {
-                writer.write_u8(LIST_ID).unwrap();
-                write_string(writer, key);
                 write_list(writer, value);
             }
             Nbt::Compound(value) => {
-                writer.write_u8(COMPOUND_ID).unwrap();
-                write_string(writer, key);
                 write_compound(writer, value, true);
             }
             Nbt::IntArray(value) => {
-                writer.write_u8(INT_ARRAY_ID).unwrap();
-                write_string(writer, key);
                 write_int_array(writer, value);
             }
             Nbt::LongArray(value) => {
-                writer.write_u8(LONG_ARRAY_ID).unwrap();
-                write_string(writer, key);
                 write_long_array(writer, value);
             }
         }
@@ -82,12 +60,12 @@ fn write_compound(writer: &mut impl Write, value: &NbtCompound, end_tag: bool) {
     }
 }
 
-#[inline]
+#[inline(always)]
 fn write_list(writer: &mut impl Write, value: &NbtList) {
+    writer.write_u8(value.id()).unwrap();
     match value {
-        NbtList::Empty => writer.write_all(&[0; 5]).unwrap(),
+        NbtList::Empty => writer.write_all(&[0; 4]).unwrap(),
         NbtList::Byte(l) => {
-            writer.write_u8(BYTE_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             let l = l.as_slice();
             writer
@@ -96,77 +74,66 @@ fn write_list(writer: &mut impl Write, value: &NbtList) {
                 .unwrap();
         }
         NbtList::Short(l) => {
-            writer.write_u8(SHORT_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for &v in l {
                 writer.write_i16::<BE>(v).unwrap();
             }
         }
         NbtList::Int(l) => {
-            writer.write_u8(INT_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for &v in l {
                 writer.write_i32::<BE>(v).unwrap();
             }
         }
         NbtList::Long(l) => {
-            writer.write_u8(LONG_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for &v in l {
                 writer.write_i64::<BE>(v).unwrap();
             }
         }
         NbtList::Float(l) => {
-            writer.write_u8(FLOAT_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for &v in l {
                 writer.write_f32::<BE>(v).unwrap();
             }
         }
         NbtList::Double(l) => {
-            writer.write_u8(DOUBLE_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for &v in l {
                 writer.write_f64::<BE>(v).unwrap();
             }
         }
         NbtList::ByteArray(l) => {
-            writer.write_u8(BYTE_ARRAY_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_byte_array(writer, v);
             }
         }
         NbtList::String(l) => {
-            writer.write_u8(STRING_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_string(writer, v);
             }
         }
         NbtList::List(l) => {
-            writer.write_u8(LIST_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_list(writer, v);
             }
         }
         NbtList::Compound(l) => {
-            writer.write_u8(COMPOUND_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_compound(writer, v, true);
             }
         }
         NbtList::IntArray(l) => {
-            writer.write_u8(INT_ARRAY_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_int_array(writer, v);
             }
         }
         NbtList::LongArray(l) => {
-            writer.write_u8(LONG_ARRAY_ID).unwrap();
             writer.write_i32::<BE>(l.len() as i32).unwrap();
             for v in l {
                 write_long_array(writer, v);
@@ -176,7 +143,7 @@ fn write_list(writer: &mut impl Write, value: &NbtList) {
 }
 
 #[inline]
-fn write_byte_array(writer: &mut impl Write, value: &Vec<u8>) {
+fn write_byte_array(writer: &mut impl Write, value: &[u8]) {
     writer.write_u32::<BE>(value.len() as u32).unwrap();
     writer.write_all(value).unwrap();
 }
diff --git a/azalea-nbt/src/lib.rs b/azalea-nbt/src/lib.rs
index b6a7a971..fe0c34cb 100755
--- a/azalea-nbt/src/lib.rs
+++ b/azalea-nbt/src/lib.rs
@@ -1,4 +1,5 @@
 #![doc = include_str!("../README.md")]
+#![feature(min_specialization)]
 
 mod decode;
 mod encode;
@@ -6,7 +7,7 @@ mod error;
 mod tag;
 
 pub use error::Error;
-pub use tag::{NbtCompound, NbtList, Nbt};
+pub use tag::{Nbt, NbtCompound, NbtList};
 
 #[cfg(test)]
 mod tests {
diff --git a/azalea-nbt/src/tag.rs b/azalea-nbt/src/tag.rs
index 573867ab..620d0577 100755
--- a/azalea-nbt/src/tag.rs
+++ b/azalea-nbt/src/tag.rs
@@ -93,16 +93,14 @@ impl NbtList {
 }
 
 // thanks to Moulberry/Graphite for the idea to use a vec and binary search
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone, Default, PartialEq)]
 pub struct NbtCompound {
-    sorted: bool,
     inner: Vec<(NbtString, Nbt)>,
 }
 impl NbtCompound {
     #[inline]
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
-            sorted: false,
             inner: Vec::with_capacity(capacity),
         }
     }
@@ -117,11 +115,23 @@ impl NbtCompound {
     /// If you previously used [`Self::insert_unsorted`] without [`Self::sort`],
     /// this function may return incorrect results.
     #[inline]
-    pub fn get(&mut self, key: &NbtString) -> Option<&Nbt> {
-        if !self.sorted {
-            self.sort()
+    pub fn get(&self, key: &str) -> Option<&Nbt> {
+        if self.is_worth_sorting() {
+            let key = NbtString::from(key);
+            self.binary_search(&key).ok().map(|i| &self.inner[i].1)
+        } else {
+            for (k, v) in &self.inner {
+                if &key == k {
+                    return Some(v);
+                }
+            }
+            None
         }
-        self.binary_search(key).ok().map(|i| &self.inner[i].1)
+    }
+
+    #[inline]
+    pub fn insert_unsorted(&mut self, key: NbtString, value: Nbt) {
+        self.inner.push((key, value));
     }
 
     /// Insert an item into the compound, returning the previous value if it
@@ -133,11 +143,14 @@ impl NbtCompound {
     #[inline]
     pub fn insert(&mut self, key: NbtString, value: Nbt) {
         self.inner.push((key, value));
+        self.sort()
     }
 
     #[inline]
     pub fn sort(&mut self) {
-        self.sorted = true;
+        if !self.is_worth_sorting() {
+            return;
+        }
         self.inner.sort_unstable_by(|(a, _), (b, _)| a.cmp(b));
     }
 
@@ -145,6 +158,11 @@ impl NbtCompound {
     pub fn iter(&self) -> std::slice::Iter<'_, (CompactString, Nbt)> {
         self.inner.iter()
     }
+
+    #[inline]
+    fn is_worth_sorting(&self) -> bool {
+        self.inner.len() > 128
+    }
 }
 #[cfg(feature = "serde")]
 impl Serialize for NbtCompound {
@@ -163,7 +181,6 @@ impl<'de> Deserialize<'de> for NbtCompound {
         let map = <BTreeMap<NbtString, Nbt> as Deserialize>::deserialize(deserializer)?;
         Ok(Self {
             inner: map.into_iter().collect(),
-            sorted: false,
         })
     }
 }
@@ -171,43 +188,6 @@ impl<'de> Deserialize<'de> for NbtCompound {
 impl FromIterator<(NbtString, Nbt)> for NbtCompound {
     fn from_iter<T: IntoIterator<Item = (NbtString, Nbt)>>(iter: T) -> Self {
         let inner = iter.into_iter().collect::<Vec<_>>();
-        Self {
-            inner,
-            sorted: false,
-        }
-    }
-}
-
-impl PartialEq for NbtCompound {
-    /// Compare two NBT compounds for equality, ignoring the order of the keys.
-    /// Note that this will execute fastest if the keys are already sorted with
-    /// [`Self::sort`].
-    fn eq(&self, other: &Self) -> bool {
-        if self.inner.len() != other.inner.len() {
-            return false;
-        }
-        if self.inner == other.inner {
-            return true;
-        }
-        if !self.sorted && !other.sorted {
-            // neither are sorted, so sort both
-            let mut a = self.clone();
-            let mut b = other.clone();
-            a.sort();
-            b.sort();
-            a == b
-        } else if !self.sorted {
-            // only self is sorted, so sort self
-            let mut a = self.clone();
-            a.sort();
-            a == *other
-        } else if !other.sorted {
-            // only other is sorted, so sort other
-            let mut b = other.clone();
-            b.sort();
-            *self == b
-        } else {
-            self.inner == other.inner
-        }
+        Self { inner }
     }
 }