From 1089b93a8e9ba360f7e6ef8d212f5bd923329c20 Mon Sep 17 00:00:00 2001
From: Amanieu d'Antras <amanieu@gmail.com>
Date: Thu, 12 Sep 2024 11:32:01 -0400
Subject: [PATCH] Re-introduce a way to get the allocation size of a table

This was previously removed from `RawTable` in #546. This is now added
as a public API on `HashMap`, `HashSet` and `HashTable`.
---
 src/map.rs     | 19 +++++++++++++++++++
 src/raw/mod.rs | 39 +++++++++++++++++++++++++++++++++++++++
 src/set.rs     | 17 +++++++++++++++++
 src/table.rs   | 22 ++++++++++++++++++++++
 4 files changed, 97 insertions(+)

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::<Q, S>(&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<K, V, S, A> PartialEq for HashMap<K, V, S, A>
@@ -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::<u32, u32>::new().allocation_size(), 0);
+        assert!(
+            HashMap::<u32, u32>::with_capacity(1).allocation_size() > core::mem::size_of::<u32>()
+        );
+    }
 }
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<T, A: Allocator> RawTable<T, A> {
         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<T>) -> 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<T, S, A> PartialEq for HashSet<T, S, A>
@@ -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::<u32>::new().allocation_size(), 0);
+        assert!(HashSet::<u32>::with_capacity(1).allocation_size() > core::mem::size_of::<u32>());
+    }
 }
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<T, A> IntoIterator for HashTable<T, A>
@@ -2208,3 +2218,15 @@ where
 }
 
 impl<T, F, A: Allocator> 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::<u32>::new().allocation_size(), 0);
+        assert!(HashTable::<u32>::with_capacity(1).allocation_size() > core::mem::size_of::<u32>());
+    }
+}