Add typesize::TypeSize implementation for DashMap/DashSet (#308)

Cargo.toml

@@ -15,6 +15,7 @@ categories = ["concurrency", "algorithms", "data-structures"]
 
 [features]
 raw-api = []
+typesize = ["dep:typesize"]
 inline = ["hashbrown/inline-more"]
 
 [dependencies]
@@ -27,6 +28,7 @@ rayon = { version = "1.7.0", optional = true }
 once_cell = "1.18.0"
 arbitrary = { version = "1.3.0", optional = true }
 crossbeam-utils = "0.8"
+typesize = { version = "0.1.8", default-features = false, optional = true }
 
 [package.metadata.docs.rs]
 features = ["rayon", "raw-api", "serde"]

src/lib.rs

@@ -1346,6 +1346,45 @@ impl<K: Eq + Hash, V, S: BuildHasher + Clone + Default> FromIterator<(K, V)> for
     }
 }
 
+#[cfg(feature = "typesize")]
+impl<K, V, S> typesize::TypeSize for DashMap<K, V, S>
+where
+    K: typesize::TypeSize + Eq + Hash,
+    V: typesize::TypeSize,
+    S: typesize::TypeSize + Clone + BuildHasher,
+{
+    fn extra_size(&self) -> usize {
+        let shards_extra_size: usize = self
+            .shards
+            .iter()
+            .map(|shard_lock| {
+                let shard = shard_lock.read();
+                let hashtable_size = shard.allocation_info().1.size();
+
+                // Safety: The iterator is dropped before the HashTable
+                let iter = unsafe { shard.iter() };
+                let entry_size_iter = iter.map(|bucket| {
+                    // Safety: The iterator returns buckets with valid pointers to entries
+                    let (key, value) = unsafe { bucket.as_ref() };
+                    key.extra_size() + value.get().extra_size()
+                });
+
+                core::mem::size_of::<CachePadded<RwLock<HashMap<K, V>>>>()
+                    + hashtable_size
+                    + entry_size_iter.sum::<usize>()
+            })
+            .sum();
+
+        self.hasher.extra_size() + shards_extra_size
+    }
+
+    typesize::if_typesize_details! {
+        fn get_collection_item_count(&self) -> Option<usize> {
+            Some(self.len())
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::DashMap;

src/set.rs

@@ -415,6 +415,23 @@ impl<K: Eq + Hash, S: BuildHasher + Clone + Default> FromIterator<K> for DashSet
     }
 }
 
+#[cfg(feature = "typesize")]
+impl<K, S> typesize::TypeSize for DashSet<K, S>
+where
+    K: typesize::TypeSize + Eq + Hash,
+    S: typesize::TypeSize + Clone + BuildHasher,
+{
+    fn extra_size(&self) -> usize {
+        self.inner.extra_size()
+    }
+
+    typesize::if_typesize_details! {
+        fn get_collection_item_count(&self) -> Option<usize> {
+            Some(self.len())
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::DashSet;