From a41a489c933c5c0711188d23f540a947bb788cb9 Mon Sep 17 00:00:00 2001
From: Nico Burns <nico@nicoburns.com>
Date: Sat, 15 Feb 2025 16:57:41 +1300
Subject: [PATCH] Implement MallocSizeOf for SmallVec (v2)

Signed-off-by: Nico Burns <nico@nicoburns.com>
---
 .github/workflows/main.yml |  3 +++
 Cargo.toml                 |  1 +
 src/lib.rs                 | 24 ++++++++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 149116b..b17cda8 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -39,6 +39,9 @@ jobs:
       - name: Cargo test w/ serde
         run: cargo test --verbose --features serde
 
+      - name: Cargo test w/ malloc_size_of
+        run: cargo test --verbose --features malloc_size_of
+
       - name: Cargo check w/o default features
         if: matrix.toolchain == 'nightly'
         run: cargo check --verbose --no-default-features
diff --git a/Cargo.toml b/Cargo.toml
index f77070f..a90e634 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,6 +20,7 @@ extract_if = []
 
 [dependencies]
 serde = { version = "1", optional = true, default-features = false }
+malloc_size_of = { 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 37c2123..3459da5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -92,6 +92,8 @@ use core::ptr::copy;
 use core::ptr::copy_nonoverlapping;
 use core::ptr::NonNull;
 
+#[cfg(feature = "malloc_size_of")]
+use malloc_size_of::{MallocShallowSizeOf, MallocSizeOf, MallocSizeOfOps};
 #[cfg(feature = "serde")]
 use serde::{
     de::{Deserialize, Deserializer, SeqAccess, Visitor},
@@ -2176,6 +2178,28 @@ where
     }
 }
 
+#[cfg(feature = "malloc_size_of")]
+impl<T, const N: usize> MallocShallowSizeOf for SmallVec<T, N> {
+    fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+        if self.spilled() {
+            unsafe { ops.malloc_size_of(self.as_ptr()) }
+        } else {
+            0
+        }
+    }
+}
+
+#[cfg(feature = "malloc_size_of")]
+impl<T: MallocSizeOf, const N: usize> MallocSizeOf for SmallVec<T, N> {
+    fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+        let mut n = self.shallow_size_of(ops);
+        for elem in self.iter() {
+            n += elem.size_of(ops);
+        }
+        n
+    }
+}
+
 #[cfg(feature = "std")]
 #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
 impl<const N: usize> io::Write for SmallVec<u8, N> {