From acdb01604aead70d0a0702a8e9ddf2c0504989c4 Mon Sep 17 00:00:00 2001
From: Meadow Liu <meadowsys@kiwin.gay>
Date: Mon, 19 Aug 2024 14:52:51 -0700
Subject: [PATCH] impl some shrink_extra, str/String, Box, Vec

---
 src/memory_usage/mod.rs | 106 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 106 insertions(+)

diff --git a/src/memory_usage/mod.rs b/src/memory_usage/mod.rs
index f4f888eca..33e3ff561 100644
--- a/src/memory_usage/mod.rs
+++ b/src/memory_usage/mod.rs
@@ -153,6 +153,8 @@ impl<'h, T: ?Sized + MemoryUsage> MemoryUsage for &'h T {
 	fn mem_use_heap_excl_extra_capacity(&self) -> usize {
 		T::mem_use_excl_extra_capacity(self)
 	}
+
+	// what to do about shrink extra? should we panic? or is no op fine?
 }
 
 impl<'h, T: ?Sized + MemoryUsageStatic> MemoryUsageStatic for &'h T {
@@ -178,6 +180,11 @@ impl<'h, T: ?Sized + MemoryUsage> MemoryUsage for &'h mut T {
 	fn mem_use_heap_excl_extra_capacity(&self) -> usize {
 		T::mem_use_excl_extra_capacity(self)
 	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		T::shrink_extra(self)
+	}
 }
 
 impl<'h, T: ?Sized + MemoryUsageStatic> MemoryUsageStatic for &'h mut T {
@@ -227,6 +234,12 @@ impl<T: MemoryUsage> MemoryUsage for [T] {
 			.map(T::mem_use_excl_extra_capacity)
 			.sum()
 	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		self.iter_mut()
+			.for_each(T::shrink_extra)
+	}
 }
 
 impl<T: MemoryUsageStatic> MemoryUsageStatic for [T] {
@@ -244,6 +257,22 @@ impl<T: MemoryUsageStatic> MemoryUsageStatic for [T] {
 	}
 }
 
+impl MemoryUsage for str {
+	#[inline]
+	fn mem_use_stack(&self) -> usize {
+		self.len()
+	}
+
+	mem_use_heap_zero_impl!();
+}
+
+impl MemoryUsageStatic for str {
+	#[inline]
+	fn mem_use_static(&self) -> usize {
+		self.len()
+	}
+}
+
 impl<T: MemoryUsage, const N: usize> MemoryUsage for [T; N] {
 	mem_use_stack_size_of_impl!();
 
@@ -266,6 +295,11 @@ impl<T: MemoryUsage, const N: usize> MemoryUsage for [T; N] {
 	fn mem_use_excl_extra_capacity(&self) -> usize {
 		<[T]>::mem_use_excl_extra_capacity(self)
 	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		<[T]>::shrink_extra(self)
+	}
 }
 
 impl<T: MemoryUsageStatic, const N: usize> MemoryUsageStatic for [T; N] {
@@ -279,6 +313,78 @@ impl<T: MemoryUsageConst, const N: usize> MemoryUsageConst for [T; N] {
 	const MEM_USE_CONST: usize = T::MEM_USE_CONST * N;
 }
 
+impl<T: ?Sized + MemoryUsage> MemoryUsage for Box<T> {
+	mem_use_stack_size_of_impl!();
+
+	#[inline]
+	fn mem_use_heap(&self) -> usize {
+		size_of::<Box<T>>() + T::mem_use_heap(self)
+	}
+
+	#[inline]
+	fn mem_use_heap_excl_extra_capacity(&self) -> usize {
+		size_of::<Box<T>>() + T::mem_use_heap_excl_extra_capacity(self)
+	}
+
+	#[inline]
+	fn mem_use(&self) -> usize {
+		size_of::<Box<T>>() + T::mem_use(self)
+	}
+
+	#[inline]
+	fn mem_use_excl_extra_capacity(&self) -> usize {
+		size_of::<Box<T>>() + T::mem_use_excl_extra_capacity(self)
+	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		T::shrink_extra(self)
+	}
+}
+
+impl<T: MemoryUsage> MemoryUsage for Vec<T> {
+	mem_use_stack_size_of_impl!();
+
+	#[inline]
+	fn mem_use_heap(&self) -> usize {
+		let full = self.mem_use_excl_extra_capacity();
+		let empty = (self.capacity() + self.len()) * size_of::<T>();
+
+		full + empty
+	}
+
+	#[inline]
+	fn mem_use_heap_excl_extra_capacity(&self) -> usize {
+		<[T]>::mem_use_heap_excl_extra_capacity(self)
+	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		self.iter_mut()
+			.for_each(T::shrink_extra);
+		self.shrink_to_fit();
+	}
+}
+
+impl MemoryUsage for String {
+	mem_use_stack_size_of_impl!();
+
+	#[inline]
+	fn mem_use_heap(&self) -> usize {
+		self.capacity()
+	}
+
+	#[inline]
+	fn mem_use_heap_excl_extra_capacity(&self) -> usize {
+		self.len()
+	}
+
+	#[inline]
+	fn shrink_extra(&mut self) {
+		self.shrink_to_fit()
+	}
+}
+
 // /// Trait for types that can calculate their actual total memory usage, not
 // /// just the stack usage (ie. not just [`size_of`])
 // ///