From 38863e83dbc84fc289a7f93c89e091ab454498c3 Mon Sep 17 00:00:00 2001
From: Amanda McCann <amanda@technomancy.org>
Date: Tue, 16 Jan 2024 22:14:51 +0100
Subject: [PATCH] Impl get_size::GetSize (behind feature flag)

Shows memory usage of the struct

cf. issue #331
---
 Cargo.toml   |  1 +
 src/lib.rs   | 26 ++++++++++++++++++++
 src/tests.rs | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 94 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index d728ba3..0efd9d2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,6 +28,7 @@ debugger_visualizer = []
 [dependencies]
 serde = { version = "1", optional = true, default-features = false }
 arbitrary = { version = "1", optional = true }
+get-size = { version = "0.1", optional = true, default-features = false }
 
 [dev_dependencies]
 bincode = "1.0.1"
diff --git a/src/lib.rs b/src/lib.rs
index cadb5d8..c34e12e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -88,6 +88,10 @@
 //! [Rustonomicon](https://doc.rust-lang.org/1.42.0/nomicon/dropck.html#an-escape-hatch).
 //!
 //! Tracking issue: [rust-lang/rust#34761](https://github.com/rust-lang/rust/issues/34761)
+//!
+//! ### `get-size`
+//!
+//! When this optional dependency is enabled, `SmallVec` implements the `get_size::GetSize` trait.
 
 #![no_std]
 #![cfg_attr(docsrs, feature(doc_cfg))]
@@ -141,6 +145,9 @@ use std::io;
 #[cfg(feature = "drain_keep_rest")]
 use core::mem::ManuallyDrop;
 
+#[cfg(feature = "get-size")]
+use get_size::GetSize;
+
 /// Creates a [`SmallVec`] containing the arguments.
 ///
 /// `smallvec!` allows `SmallVec`s to be defined with the same syntax as array expressions.
@@ -2469,3 +2476,22 @@ impl<T> Clone for ConstNonNull<T> {
 }
 
 impl<T> Copy for ConstNonNull<T> {}
+
+#[cfg(feature = "get-size")]
+impl<A: Array> GetSize for SmallVec<A>
+where
+    A::Item: GetSize,
+{
+    fn get_heap_size(&self) -> usize {
+        let mut total = 0;
+        if self.spilled() {
+            total += self.capacity * A::Item::get_stack_size();
+        }
+
+        for v in self.iter() {
+            total += v.get_heap_size();
+        }
+
+        total
+    }
+}
diff --git a/src/tests.rs b/src/tests.rs
index 3eab846..6e3bcce 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1023,3 +1023,70 @@ fn drain_keep_rest() {
 
     assert_eq!(a, SmallVec::<[i32; 3]>::from_slice(&[1i32, 3, 5, 6, 7, 8]));
 }
+
+#[cfg(all(feature = "get-size", target_pointer_width = "64"))]
+mod get_size {
+    use super::*;
+
+    #[test]
+    fn end_to_end1() {
+        use ::get_size::GetSize;
+
+        let mut a: SmallVec<[i32; 2]> = smallvec![];
+        assert!(!a.spilled());
+        assert_eq!(a.len(), 0);
+        assert_eq!(a.get_size(), 24);
+        assert_eq!(a.get_heap_size(), 0);
+
+        a.push(0);
+        assert_eq!(a.len(), 1);
+        assert!(!a.spilled());
+        assert_eq!(a.get_size(), 24);
+        assert_eq!(a.get_heap_size(), 0);
+
+        a.push(1);
+        assert_eq!(a.len(), 2);
+        assert!(!a.spilled());
+        assert_eq!(a.get_size(), 24);
+        assert_eq!(a.get_heap_size(), 0);
+
+        a.push(2);
+        assert_eq!(a.len(), 3);
+        assert!(a.spilled());
+        assert_eq!(a.get_size(), 40);
+        assert_eq!(a.get_heap_size(), 16);
+
+        a.push(3);
+        assert_eq!(a.len(), 4);
+        assert!(a.spilled());
+        assert_eq!(a.get_size(), 40);
+        assert_eq!(a.get_heap_size(), 16);
+
+        a.push(4);
+        assert_eq!(a.len(), 5);
+        assert!(a.spilled());
+        assert_eq!(a.get_size(), 56);
+        assert_eq!(a.get_heap_size(), 32);
+
+    }
+
+    #[cfg(not(feature = "union"))]
+    #[test]
+    fn stack_size_no_union1() {
+        use ::get_size::GetSize;
+
+        assert_eq!(SmallVec::<[i32; 2]>::get_stack_size(), 24);
+        assert_eq!(SmallVec::<[i32; 10]>::get_stack_size(), 56);
+    }
+
+    #[cfg(feature="union")]
+    #[test]
+    fn stack_size_union1() {
+        use ::get_size::GetSize;
+
+        assert_eq!(SmallVec::<[i32; 2]>::get_stack_size(), 24);
+        assert_eq!(SmallVec::<[i32; 10]>::get_stack_size(), 48);
+    }
+
+}
+