binary search map

4 files changed, 78 insertions, 14 deletions

diff --git a/azalea-nbt/src/decode.rs b/azalea-nbt/src/decode.rs
index 635e2f6d..e2e9c7c9 100755
--- a/azalea-nbt/src/decode.rs
+++ b/azalea-nbt/src/decode.rs
@@ -1,6 +1,5 @@
 use crate::tag::*;
 use crate::Error;
-use ahash::AHashMap;
 use azalea_buf::{BufReadError, McBufReadable};
 use byteorder::{ReadBytesExt, BE};
 use flate2::read::{GzDecoder, ZlibDecoder};
@@ -265,7 +264,7 @@ impl Tag {
         }
         let name = read_string(stream)?;
         let tag = Tag::read_known(stream, tag_id)?;
-        let mut map = AHashMap::with_capacity(1);
+        let mut map = NbtCompound::with_capacity(1);
         map.insert(name, tag);
 
         Ok(Tag::Compound(map))
diff --git a/azalea-nbt/src/encode.rs b/azalea-nbt/src/encode.rs
index ba2b77a5..03715691 100755
--- a/azalea-nbt/src/encode.rs
+++ b/azalea-nbt/src/encode.rs
@@ -15,7 +15,7 @@ fn write_string(writer: &mut dyn Write, string: &str) -> Result<(), Error> {
 
 #[inline]
 fn write_compound(writer: &mut dyn Write, value: &NbtCompound, end_tag: bool) -> Result<(), Error> {
-    for (key, tag) in value {
+    for (key, tag) in value.inner() {
         match tag {
             Tag::End => {}
             Tag::Byte(value) => {
diff --git a/azalea-nbt/src/lib.rs b/azalea-nbt/src/lib.rs
index 048e466e..3c8aa2cd 100755
--- a/azalea-nbt/src/lib.rs
+++ b/azalea-nbt/src/lib.rs
@@ -6,23 +6,23 @@ mod error;
 mod tag;
 
 pub use error::Error;
-pub use tag::NbtList;
-pub use tag::Tag;
+pub use tag::{NbtCompound, NbtList, Tag};
 
 #[cfg(test)]
 mod tests {
     use std::io::Cursor;
 
+    use crate::tag::NbtCompound;
+
     use super::*;
-    use ahash::AHashMap;
     use azalea_buf::{McBufReadable, McBufWritable};
 
     #[test]
     fn mcbuf_nbt() {
         let mut buf = Vec::new();
-        let tag = Tag::Compound(AHashMap::from_iter(vec![(
+        let tag = Tag::Compound(NbtCompound::from_iter(vec![(
             "hello world".into(),
-            Tag::Compound(AHashMap::from_iter(vec![(
+            Tag::Compound(NbtCompound::from_iter(vec![(
                 "name".into(),
                 Tag::String("Bananrama".into()),
             )])),
@@ -34,9 +34,9 @@ mod tests {
         let result = Tag::read_from(&mut buf).unwrap();
         assert_eq!(
             result,
-            Tag::Compound(AHashMap::from_iter(vec![(
+            Tag::Compound(NbtCompound::from_iter(vec![(
                 "hello world".into(),
-                Tag::Compound(AHashMap::from_iter(vec![(
+                Tag::Compound(NbtCompound::from_iter(vec![(
                     "name".into(),
                     Tag::String("Bananrama".into()),
                 )])),
diff --git a/azalea-nbt/src/tag.rs b/azalea-nbt/src/tag.rs
index e1aecb91..e23db913 100755
--- a/azalea-nbt/src/tag.rs
+++ b/azalea-nbt/src/tag.rs
@@ -1,9 +1,7 @@
-use ahash::AHashMap;
-
 use compact_str::CompactString;
 use enum_as_inner::EnumAsInner;
 #[cfg(feature = "serde")]
-use serde::{Deserialize, Serialize};
+use serde::{ser::SerializeMap, Deserialize, Serialize};
 
 pub type NbtByte = i8;
 pub type NbtShort = i16;
@@ -13,7 +11,6 @@ pub type NbtFloat = f32;
 pub type NbtDouble = f64;
 pub type NbtByteArray = Vec<u8>;
 pub type NbtString = CompactString;
-pub type NbtCompound = AHashMap<CompactString, Tag>;
 pub type NbtIntArray = Vec<i32>;
 pub type NbtLongArray = Vec<i64>;
 
@@ -94,3 +91,71 @@ impl NbtList {
         unsafe { *<*const _>::from(self).cast::<u8>() }
     }
 }
+
+// thanks to Moulberry/Graphite for the idea to use a vec and binary search
+#[derive(Debug, Clone, PartialEq)]
+pub struct NbtCompound {
+    inner: Vec<(NbtString, Tag)>,
+}
+impl NbtCompound {
+    #[inline]
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            inner: Vec::with_capacity(capacity),
+        }
+    }
+
+    #[inline]
+    fn binary_search(&self, key: &NbtString) -> Result<usize, usize> {
+        self.inner.binary_search_by(|(k, _)| k.cmp(key))
+    }
+
+    #[inline]
+    pub fn get(&self, key: &NbtString) -> Option<&Tag> {
+        self.binary_search(key).ok().map(|i| &self.inner[i].1)
+    }
+
+    #[inline]
+    pub fn insert(&mut self, key: NbtString, value: Tag) -> Option<Tag> {
+        match self.binary_search(&key) {
+            Ok(i) => Some(std::mem::replace(&mut self.inner[i].1, value)),
+            Err(i) => {
+                self.inner.insert(i, (key, value));
+                None
+            }
+        }
+    }
+
+    #[inline]
+    pub fn inner(&self) -> &Vec<(NbtString, Tag)> {
+        &self.inner
+    }
+}
+#[cfg(feature = "serde")]
+impl Serialize for NbtCompound {
+    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
+        let mut map = serializer.serialize_map(Some(self.inner.len()))?;
+        for (key, value) in &self.inner {
+            map.serialize_entry(key, value)?;
+        }
+        map.end()
+    }
+}
+#[cfg(feature = "serde")]
+impl<'de> Deserialize<'de> for NbtCompound {
+    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+        use std::collections::BTreeMap;
+        let map = <BTreeMap<NbtString, Tag> as Deserialize>::deserialize(deserializer)?;
+        Ok(Self {
+            inner: map.into_iter().collect(),
+        })
+    }
+}
+
+impl FromIterator<(NbtString, Tag)> for NbtCompound {
+    fn from_iter<T: IntoIterator<Item = (NbtString, Tag)>>(iter: T) -> Self {
+        let mut inner = iter.into_iter().collect::<Vec<_>>();
+        inner.sort_unstable_by_key(|(k, _)| k.clone());
+        Self { inner }
+    }
+}