shrink_to_fit IdentityMap before storing it (#816)

src/active_query.rs

@@ -48,7 +48,7 @@ pub(crate) struct ActiveQuery {
 
     /// Map from tracked struct keys (which include the hash + disambiguator) to their
     /// final id.
-    pub(crate) tracked_struct_ids: IdentityMap,
+    tracked_struct_ids: IdentityMap,
 
     /// Stores the values accumulated to the given ingredient.
     /// The type of accumulated value is erased but known to the ingredient.
@@ -151,6 +151,14 @@ impl ActiveQuery {
     pub(super) fn iteration_count(&self) -> u32 {
         self.iteration_count
     }
+
+    pub(crate) fn tracked_struct_ids(&self) -> &IdentityMap {
+        &self.tracked_struct_ids
+    }
+
+    pub(crate) fn tracked_struct_ids_mut(&mut self) -> &mut IdentityMap {
+        &mut self.tracked_struct_ids
+    }
 }
 
 impl ActiveQuery {

src/function.rs

@@ -188,9 +188,10 @@ where
         &'db self,
         zalsa: &'db Zalsa,
         id: Id,
-        memo: memo::Memo<C::Output<'db>>,
+        mut memo: memo::Memo<C::Output<'db>>,
         memo_ingredient_index: MemoIngredientIndex,
     ) -> &'db memo::Memo<C::Output<'db>> {
+        memo.revisions.tracked_struct_ids.shrink_to_fit();
         // We convert to a `NonNull` here as soon as possible because we are going to alias
         // into the `Box`, which is a `noalias` type.
         // FIXME: Use `Box::into_non_null` once stable

src/tracked_struct.rs

@@ -205,55 +205,59 @@ pub struct IdentityHash {
     hash: u64,
 }
 
+/// A map from tracked struct keys (which include the hash + [Disambiguator]) to their
+/// final [Id].
 #[derive(Default, Debug)]
 pub(crate) struct IdentityMap {
-    // we use a non-hasher hashmap here as our key contains its own hash (`Identity::hash`)
-    // so we use the raw entry api instead to avoid the overhead of hashing unnecessarily
-    map: hashbrown::HashMap<Identity, Id, ()>,
+    // we use a hashtable here as our key contains its own hash (`Identity::hash`)
+    // so we do the hash wrangling ourselves
+    table: hashbrown::HashTable<(Identity, Id)>,
 }
 
 impl Clone for IdentityMap {
     fn clone(&self) -> Self {
         Self {
-            map: self.map.clone(),
+            table: self.table.clone(),
         }
     }
     fn clone_from(&mut self, source: &Self) {
-        self.map.clone_from(&source.map);
+        self.table.clone_from(&source.table);
     }
 }
 
 impl IdentityMap {
     pub(crate) fn insert(&mut self, key: Identity, id: Id) -> Option<Id> {
-        use hashbrown::hash_map::RawEntryMut;
-
-        let entry = self.map.raw_entry_mut().from_hash(key.hash, |k| *k == key);
+        let entry = self.table.find_mut(key.hash, |&(k, _)| k == key);
         match entry {
-            RawEntryMut::Occupied(mut occupied) => Some(occupied.insert(id)),
-            RawEntryMut::Vacant(vacant) => {
-                vacant.insert_with_hasher(key.hash, key, id, |k| k.hash);
+            Some(occupied) => Some(mem::replace(&mut occupied.1, id)),
+            None => {
+                self.table
+                    .insert_unique(key.hash, (key, id), |(k, _)| k.hash);
                 None
             }
         }
     }
 
     pub(crate) fn get(&self, key: &Identity) -> Option<Id> {
-        self.map
-            .raw_entry()
-            .from_hash(key.hash, |k| *k == *key)
-            .map(|(_, &v)| v)
+        self.table
+            .find(key.hash, |&(k, _)| k == *key)
+            .map(|&(_, v)| v)
     }
 
     pub(crate) fn is_empty(&self) -> bool {
-        self.map.is_empty()
+        self.table.is_empty()
     }
 
-    pub(crate) fn retain(&mut self, f: impl FnMut(&Identity, &mut Id) -> bool) {
-        self.map.retain(f);
+    pub(crate) fn retain(&mut self, mut f: impl FnMut(&Identity, &mut Id) -> bool) {
+        self.table.retain(|(k, v)| f(k, v));
     }
 
-    pub fn clear(&mut self) {
-        self.map.clear()
+    pub(crate) fn clear(&mut self) {
+        self.table.clear()
+    }
+
+    pub(crate) fn shrink_to_fit(&mut self) {
+        self.table.shrink_to_fit(|(k, _)| k.hash);
     }
 }
 

src/zalsa_local.rs

@@ -285,7 +285,7 @@ impl ZalsaLocal {
             let top_query = stack
                 .last()
                 .expect("cannot create a tracked struct ID outside of a tracked function");
-            top_query.tracked_struct_ids.get(identity)
+            top_query.tracked_struct_ids().get(identity)
         })
     }
 
@@ -295,7 +295,7 @@ impl ZalsaLocal {
             let top_query = stack
                 .last_mut()
                 .expect("cannot store a tracked struct ID outside of a tracked function");
-            top_query.tracked_struct_ids.insert(identity, id);
+            top_query.tracked_struct_ids_mut().insert(identity, id);
         })
     }
 
@@ -509,8 +509,10 @@ impl ActiveQueryGuard<'_> {
             #[cfg(debug_assertions)]
             assert_eq!(stack.len(), self.push_len);
             let frame = stack.last_mut().unwrap();
-            assert!(frame.tracked_struct_ids.is_empty());
-            frame.tracked_struct_ids.clone_from(tracked_struct_ids);
+            assert!(frame.tracked_struct_ids().is_empty());
+            frame
+                .tracked_struct_ids_mut()
+                .clone_from(tracked_struct_ids);
         })
     }