Skip to content

Commit

Permalink
Add typesize::TypeSize implementation for DashMap/DashSet (#308)
Browse files Browse the repository at this point in the history
  • Loading branch information
GnomedDev authored Sep 5, 2024
1 parent 633aadb commit da6ac5e
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 4 deletions.
29 changes: 25 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ categories = ["concurrency", "algorithms", "data-structures"]

[features]
raw-api = []
typesize = ["dep:typesize"]
inline = ["hashbrown/inline-more"]

[dependencies]
Expand All @@ -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"]
39 changes: 39 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
17 changes: 17 additions & 0 deletions src/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit da6ac5e

Please sign in to comment.