diff --git a/src/map.rs b/src/map.rs index e9cd44aad..d0a3d7224 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1929,6 +1929,16 @@ where let hash = make_hash::(&self.hash_builder, k); self.table.remove_entry(hash, equivalent_key(k)) } + + /// Returns the total amount of memory allocated internally by the hash + /// set, in bytes. + /// + /// The returned number is informational only. It is intended to be + /// primarily used for memory profiling. + #[inline] + pub fn allocation_size(&self) -> usize { + self.table.allocation_size() + } } impl PartialEq for HashMap @@ -6416,4 +6426,13 @@ mod test_map { // All allocator clones should already be dropped. assert_eq!(dropped.load(Ordering::SeqCst), 0); } + + #[test] + fn test_allocation_info() { + assert_eq!(HashMap::<(), ()>::new().allocation_size(), 0); + assert_eq!(HashMap::::new().allocation_size(), 0); + assert!( + HashMap::::with_capacity(1).allocation_size() > core::mem::size_of::() + ); + } } diff --git a/src/raw/mod.rs b/src/raw/mod.rs index 72004d7a0..0b4a58f14 100644 --- a/src/raw/mod.rs +++ b/src/raw/mod.rs @@ -771,6 +771,18 @@ impl RawTable { NonNull::new_unchecked(self.data_end().as_ptr().wrapping_sub(self.buckets())) } + /// Returns the total amount of memory allocated internally by the hash + /// table, in bytes. + /// + /// The returned number is informational only. It is intended to be + /// primarily used for memory profiling. + #[inline] + pub fn allocation_size(&self) -> usize { + // SAFETY: We use the same `table_layout` that was used to allocate + // this table. + unsafe { self.table.allocation_size_or_zero(Self::TABLE_LAYOUT) } + } + /// Returns the index of a bucket from a `Bucket`. #[inline] pub unsafe fn bucket_index(&self, bucket: &Bucket) -> usize { @@ -3056,6 +3068,33 @@ impl RawTableInner { ) } + /// Returns the total amount of memory allocated internally by the hash + /// table, in bytes. + /// + /// The returned number is informational only. It is intended to be + /// primarily used for memory profiling. + /// + /// # Safety + /// + /// The `table_layout` must be the same [`TableLayout`] as the `TableLayout` + /// that was used to allocate this table. Failure to comply with this condition + /// may result in [`undefined behavior`]. + /// + /// + /// [`undefined behavior`]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html + #[inline] + unsafe fn allocation_size_or_zero(&self, table_layout: TableLayout) -> usize { + if self.is_empty_singleton() { + 0 + } else { + // SAFETY: + // 1. We have checked that our table is allocated. + // 2. The caller ensures that `table_layout` matches the [`TableLayout`] + // that was used to allocate this table. + unsafe { self.allocation_info(table_layout).1.size() } + } + } + /// Marks all table buckets as empty without dropping their contents. #[inline] fn clear_no_drop(&mut self) { diff --git a/src/set.rs b/src/set.rs index e5e9496ae..15a8489ee 100644 --- a/src/set.rs +++ b/src/set.rs @@ -1230,6 +1230,16 @@ where None => None, } } + + /// Returns the total amount of memory allocated internally by the hash + /// set, in bytes. + /// + /// The returned number is informational only. It is intended to be + /// primarily used for memory profiling. + #[inline] + pub fn allocation_size(&self) -> usize { + self.map.allocation_size() + } } impl PartialEq for HashSet @@ -3055,4 +3065,11 @@ mod test_set { // (and without the `map`, it does not). let mut _set: HashSet<_> = (0..3).map(|_| ()).collect(); } + + #[test] + fn test_allocation_info() { + assert_eq!(HashSet::<()>::new().allocation_size(), 0); + assert_eq!(HashSet::::new().allocation_size(), 0); + assert!(HashSet::::with_capacity(1).allocation_size() > core::mem::size_of::()); + } } diff --git a/src/table.rs b/src/table.rs index 421f52ca3..216d88b65 100644 --- a/src/table.rs +++ b/src/table.rs @@ -1088,6 +1088,16 @@ where ) -> Option<[&'_ mut T; N]> { self.raw.get_many_unchecked_mut(hashes, eq) } + + /// Returns the total amount of memory allocated internally by the hash + /// table, in bytes. + /// + /// The returned number is informational only. It is intended to be + /// primarily used for memory profiling. + #[inline] + pub fn allocation_size(&self) -> usize { + self.raw.allocation_size() + } } impl IntoIterator for HashTable @@ -2208,3 +2218,15 @@ where } impl FusedIterator for ExtractIf<'_, T, F, A> where F: FnMut(&mut T) -> bool {} + +#[cfg(test)] +mod tests { + use super::HashTable; + + #[test] + fn test_allocation_info() { + assert_eq!(HashTable::<()>::new().allocation_size(), 0); + assert_eq!(HashTable::::new().allocation_size(), 0); + assert!(HashTable::::with_capacity(1).allocation_size() > core::mem::size_of::()); + } +}