shrink_to_fit IdentityMap before storing it

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.
@@ -137,6 +137,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/execute.rs

@@ -214,6 +214,8 @@ where
                 );
             }
 
+            revisions.tracked_struct_ids.shrink_to_fit();
+
             return self.insert_memo(
                 zalsa,
                 id,

src/tracked_struct.rs

@@ -1,11 +1,11 @@
 #![allow(clippy::undocumented_unsafe_blocks)] // TODO(#697) document safety
 
 use std::any::TypeId;
-use std::fmt;
 use std::hash::Hash;
 use std::marker::PhantomData;
 use std::ops::DerefMut;
 use std::sync::Arc;
+use std::{fmt, mem};
 
 use crossbeam_queue::SegQueue;
 use tracked_field::FieldIngredientImpl;
@@ -213,55 +213,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

@@ -268,7 +268,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)
         })
     }
 
@@ -278,7 +278,7 @@ impl ZalsaLocal {
             let top_query = stack
                 .last_mut()
                 .expect("cannot store a tracked struct ID outside of a tracked function");
-            let old_id = top_query.tracked_struct_ids.insert(identity, id);
+            let old_id = top_query.tracked_struct_ids_mut().insert(identity, id);
             assert!(
                 old_id.is_none(),
                 "overwrote a previous id for `{identity:?}`"
@@ -485,8 +485,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);
         })
     }