From a0371d55068147f188548e7423331ae2cc712516 Mon Sep 17 00:00:00 2001
From: Micah <micah@uplift.games>
Date: Tue, 12 Nov 2024 15:57:05 -0800
Subject: [PATCH] Re-export AHashMap and give it a trait extension (#479)

---
 rbx_dom_weak/CHANGELOG.md |  5 +++--
 rbx_dom_weak/src/dom.rs   | 20 ++++++++++----------
 rbx_dom_weak/src/lib.rs   | 28 ++++++++++++++++++++++------
 3 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/rbx_dom_weak/CHANGELOG.md b/rbx_dom_weak/CHANGELOG.md
index 27d245cb..74f6ba2f 100644
--- a/rbx_dom_weak/CHANGELOG.md
+++ b/rbx_dom_weak/CHANGELOG.md
@@ -95,11 +95,12 @@ pub fn into_raw(self) -> (Ref, HashMap<Ref, Instance, RandomState>) {
 ```
 to
 ```rust
-pub fn into_raw(self) -> (Ref, HashMap<Ref, Instance, ahash::RandomState>) {
+pub fn into_raw(self) -> (Ref, AHashMap<Ref, Instance>) {
 ```
 
 ### Other changes
-* Added `UstrMapExt`, a helper trait providing convenience methods `UstrMap::new` and `UstrMap::with_capacity`.
+* Added `HashMapExt`, a helper trait providing convenience methods `UstrMap::new`, `UstrMap::with_capacity`, `AHashMap::new`, and `AHashMap::with_capacity`.
+* Added re-exports for `AHashMap`.
 * Added re-exports for `ustr` (a convenience function for creating `Ustr`s), `Ustr`, `UstrMap`, and `UstrSet`.
 * Added `InstanceBuilder::with_property_capacity`, which can preallocate an `InstanceBuilder`'s property table. [#464]
 * Added `WeakDom::reserve`, which can preallocate additional space for instances in the `WeakDom`. [#465]
diff --git a/rbx_dom_weak/src/dom.rs b/rbx_dom_weak/src/dom.rs
index ef93d495..262d12c3 100644
--- a/rbx_dom_weak/src/dom.rs
+++ b/rbx_dom_weak/src/dom.rs
@@ -1,6 +1,6 @@
 use std::collections::VecDeque;
 
-use ahash::{HashMap, HashMapExt, HashSet, HashSetExt};
+use ahash::{AHashMap, AHashSet};
 use rbx_types::{Ref, UniqueId, Variant};
 use ustr::ustr;
 
@@ -15,18 +15,18 @@ use crate::instance::{Instance, InstanceBuilder};
 /// objects and insert them into the tree.
 #[derive(Debug)]
 pub struct WeakDom {
-    instances: HashMap<Ref, Instance>,
+    instances: AHashMap<Ref, Instance>,
     root_ref: Ref,
-    unique_ids: HashSet<UniqueId>,
+    unique_ids: AHashSet<UniqueId>,
 }
 
 impl WeakDom {
     /// Construct a new `WeakDom` described by the given [`InstanceBuilder`].
     pub fn new(builder: InstanceBuilder) -> WeakDom {
         let mut dom = WeakDom {
-            instances: HashMap::new(),
+            instances: AHashMap::new(),
             root_ref: builder.referent,
-            unique_ids: HashSet::new(),
+            unique_ids: AHashSet::new(),
         };
 
         dom.insert(Ref::none(), builder);
@@ -42,7 +42,7 @@ impl WeakDom {
     /// Consumes the WeakDom, returning its underlying root ref and backing
     /// storage. This method is useful when tree-preserving operations are too
     /// slow.
-    pub fn into_raw(self) -> (Ref, HashMap<Ref, Instance>) {
+    pub fn into_raw(self) -> (Ref, AHashMap<Ref, Instance>) {
         (self.root_ref, self.instances)
     }
 
@@ -435,9 +435,9 @@ impl<'a> Iterator for WeakDomDescendants<'a> {
 impl Default for WeakDom {
     fn default() -> WeakDom {
         WeakDom {
-            instances: HashMap::new(),
+            instances: AHashMap::new(),
             root_ref: Ref::none(),
-            unique_ids: HashSet::new(),
+            unique_ids: AHashSet::new(),
         }
     }
 }
@@ -445,14 +445,14 @@ impl Default for WeakDom {
 #[derive(Debug, Default)]
 struct CloneContext {
     queue: VecDeque<(Ref, Ref)>,
-    ref_rewrites: HashMap<Ref, Ref>,
+    ref_rewrites: AHashMap<Ref, Ref>,
 }
 
 impl CloneContext {
     /// On any instances cloned during the operation, rewrite any Ref properties that
     /// point to instances that were also cloned.
     fn rewrite_refs(self, dest: &mut WeakDom) {
-        let mut existing_dest_refs = HashSet::new();
+        let mut existing_dest_refs = AHashSet::new();
 
         for (_, new_ref) in self.ref_rewrites.iter() {
             let instance = dest
diff --git a/rbx_dom_weak/src/lib.rs b/rbx_dom_weak/src/lib.rs
index 4500c619..250d0604 100644
--- a/rbx_dom_weak/src/lib.rs
+++ b/rbx_dom_weak/src/lib.rs
@@ -47,6 +47,7 @@ mod viewer;
 
 pub use rbx_types as types;
 
+pub use ahash::AHashMap;
 pub use ustr::{ustr, Ustr, UstrMap, UstrSet};
 
 pub use crate::{
@@ -55,22 +56,37 @@ pub use crate::{
     viewer::{DomViewer, ViewedInstance},
 };
 
-/// Helper trait that provides convenience methods for `UstrMap`.
-pub trait UstrMapExt {
-    /// Creates an empty `UstrMap` using the default value for its hasher.
+/// Helper trait that provides convenience methods for `AHashMap` and `UstrMap`.
+pub trait HashMapExt {
+    /// Constructs an empty map.
     fn new() -> Self;
 
-    /// Creates an empty `UstrMap` with at least the specified capacity using
-    /// the default value for its hasher.
+    /// Constructs an empty map with at least the specified capacity.
     fn with_capacity(capacity: usize) -> Self;
 }
 
-impl<V> UstrMapExt for UstrMap<V> {
+impl<V> HashMapExt for UstrMap<V> {
+    /// Creates an empty `UstrMap` using the default value for its hasher.
     fn new() -> Self {
         UstrMap::default()
     }
 
+    /// Creates an empty `UstrMap` with at least the specified capacity using
+    /// the default value for its hasher.
     fn with_capacity(capacity: usize) -> Self {
         UstrMap::with_capacity_and_hasher(capacity, Default::default())
     }
 }
+
+impl<K, V> HashMapExt for AHashMap<K, V> {
+    /// Creates an empty `AHashMap` using the default value for its hasher.
+    fn new() -> Self {
+        AHashMap::default()
+    }
+
+    /// Creates an empty `AHashMap` with at least the specified capacity using
+    /// the default value for its hasher.
+    fn with_capacity(capacity: usize) -> Self {
+        AHashMap::with_capacity_and_hasher(capacity, Default::default())
+    }
+}