spacetimedb_memory_usage/
lib.rs

1use core::mem;
2use core::sync::atomic::AtomicUsize;
3
4/// For inspecting how much memory a value is using.
5///
6/// This trait specifically measures heap memory. If you want to measure stack memory too, add
7/// `mem::size_of_val()` to it. (This only really matters for the outermost type in a hierarchy.)
8pub trait MemoryUsage {
9    /// The **heap** memory usage of this type. The default implementation returns 0.
10    #[inline(always)]
11    fn heap_usage(&self) -> usize {
12        0
13    }
14}
15
16impl MemoryUsage for () {}
17impl MemoryUsage for bool {}
18impl MemoryUsage for u8 {}
19impl MemoryUsage for u16 {}
20impl MemoryUsage for u32 {}
21impl MemoryUsage for u64 {}
22impl MemoryUsage for u128 {}
23#[cfg(feature = "ethnum")]
24impl MemoryUsage for ethnum::u256 {}
25impl MemoryUsage for usize {}
26impl MemoryUsage for AtomicUsize {}
27impl MemoryUsage for i8 {}
28impl MemoryUsage for i16 {}
29impl MemoryUsage for i32 {}
30impl MemoryUsage for i64 {}
31impl MemoryUsage for i128 {}
32#[cfg(feature = "ethnum")]
33impl MemoryUsage for ethnum::i256 {}
34impl MemoryUsage for isize {}
35impl MemoryUsage for f32 {}
36impl MemoryUsage for f64 {}
37#[cfg(feature = "decorum")]
38impl MemoryUsage for decorum::Total<f32> {}
39#[cfg(feature = "decorum")]
40impl MemoryUsage for decorum::Total<f64> {}
41
42impl<T: MemoryUsage + ?Sized> MemoryUsage for &T {
43    fn heap_usage(&self) -> usize {
44        (*self).heap_usage()
45    }
46}
47
48impl<T: MemoryUsage + ?Sized> MemoryUsage for Box<T> {
49    fn heap_usage(&self) -> usize {
50        mem::size_of_val::<T>(self) + T::heap_usage(self)
51    }
52}
53
54impl<T: MemoryUsage + ?Sized> MemoryUsage for std::sync::Arc<T> {
55    fn heap_usage(&self) -> usize {
56        let refcounts = mem::size_of::<usize>() * 2;
57        refcounts + mem::size_of_val::<T>(self) + T::heap_usage(self)
58    }
59}
60
61impl<T: MemoryUsage + ?Sized> MemoryUsage for std::rc::Rc<T> {
62    fn heap_usage(&self) -> usize {
63        let refcounts = mem::size_of::<usize>() * 2;
64        refcounts + mem::size_of_val::<T>(self) + T::heap_usage(self)
65    }
66}
67
68impl<T: MemoryUsage> MemoryUsage for [T] {
69    fn heap_usage(&self) -> usize {
70        self.iter().map(T::heap_usage).sum()
71    }
72}
73
74impl<T: MemoryUsage, const N: usize> MemoryUsage for [T; N] {
75    fn heap_usage(&self) -> usize {
76        self.iter().map(T::heap_usage).sum()
77    }
78}
79
80impl MemoryUsage for str {}
81
82impl<T: MemoryUsage> MemoryUsage for Option<T> {
83    fn heap_usage(&self) -> usize {
84        self.as_ref().map_or(0, T::heap_usage)
85    }
86}
87
88impl<A: MemoryUsage, B: MemoryUsage> MemoryUsage for (A, B) {
89    fn heap_usage(&self) -> usize {
90        self.0.heap_usage() + self.1.heap_usage()
91    }
92}
93
94impl MemoryUsage for String {
95    fn heap_usage(&self) -> usize {
96        self.capacity()
97    }
98}
99
100impl<T: MemoryUsage> MemoryUsage for Vec<T> {
101    fn heap_usage(&self) -> usize {
102        self.capacity() * mem::size_of::<T>() + self.iter().map(T::heap_usage).sum::<usize>()
103    }
104}
105
106#[cfg(feature = "hashbrown")]
107impl<K: MemoryUsage + Eq + core::hash::Hash, V: MemoryUsage, S: core::hash::BuildHasher> MemoryUsage
108    for hashbrown::HashMap<K, V, S>
109{
110    fn heap_usage(&self) -> usize {
111        self.allocation_size() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
112    }
113}
114
115#[cfg(feature = "hashbrown")]
116impl<K: MemoryUsage + Eq + core::hash::Hash, S: core::hash::BuildHasher> MemoryUsage for hashbrown::HashSet<K, S> {
117    fn heap_usage(&self) -> usize {
118        self.allocation_size() + self.iter().map(|k| k.heap_usage()).sum::<usize>()
119    }
120}
121
122impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for std::collections::BTreeMap<K, V> {
123    fn heap_usage(&self) -> usize {
124        // NB: this is best-effort, since we don't have a `capacity()` method on `BTreeMap`.
125        self.len() * mem::size_of::<(K, V)>() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
126    }
127}
128
129#[cfg(feature = "smallvec")]
130impl<A: smallvec::Array> MemoryUsage for smallvec::SmallVec<A>
131where
132    A::Item: MemoryUsage,
133{
134    fn heap_usage(&self) -> usize {
135        self.as_slice().heap_usage()
136            + if self.spilled() {
137                self.capacity() * mem::size_of::<A::Item>()
138            } else {
139                0
140            }
141    }
142}