spacetimedb_table/
memory_usage.rs

1use std::hash::{BuildHasher, Hash};
2use std::mem;
3
4use spacetimedb_sats::{
5    algebraic_value::Packed, i256, u256, AlgebraicType, AlgebraicValue, ArrayType, ArrayValue, ProductType,
6    ProductTypeElement, ProductValue, SumType, SumTypeVariant, SumValue,
7};
8
9/// For inspecting how much memory a value is using.
10///
11/// This trait specifically measures heap memory. If you want to measure stack memory too, add
12/// `mem::size_of_val()` to it. (This only really matters for the outermost type in a hierarchy.)
13pub trait MemoryUsage {
14    /// The **heap** memory usage of this type. The default implementation returns 0.
15    #[inline(always)]
16    fn heap_usage(&self) -> usize {
17        0
18    }
19}
20
21impl MemoryUsage for bool {}
22impl MemoryUsage for u8 {}
23impl MemoryUsage for u16 {}
24impl MemoryUsage for u32 {}
25impl MemoryUsage for u64 {}
26impl MemoryUsage for u128 {}
27impl MemoryUsage for u256 {}
28impl MemoryUsage for usize {}
29impl MemoryUsage for i8 {}
30impl MemoryUsage for i16 {}
31impl MemoryUsage for i32 {}
32impl MemoryUsage for i64 {}
33impl MemoryUsage for i128 {}
34impl MemoryUsage for i256 {}
35impl MemoryUsage for isize {}
36impl MemoryUsage for f32 {}
37impl MemoryUsage for f64 {}
38
39impl MemoryUsage for spacetimedb_sats::F32 {}
40impl MemoryUsage for spacetimedb_sats::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
106impl<K: MemoryUsage + Eq + Hash, V: MemoryUsage, S: BuildHasher> MemoryUsage
107    for spacetimedb_data_structures::map::HashMap<K, V, S>
108{
109    fn heap_usage(&self) -> usize {
110        self.allocation_size() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
111    }
112}
113
114impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for std::collections::BTreeMap<K, V> {
115    fn heap_usage(&self) -> usize {
116        // NB: this is best-effort, since we don't have a `capacity()` method on `BTreeMap`.
117        self.len() * mem::size_of::<(K, V)>() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
118    }
119}
120
121impl<A: smallvec::Array> MemoryUsage for smallvec::SmallVec<A>
122where
123    A::Item: MemoryUsage,
124{
125    fn heap_usage(&self) -> usize {
126        self.as_slice().heap_usage()
127            + if self.spilled() {
128                self.capacity() * mem::size_of::<A::Item>()
129            } else {
130                0
131            }
132    }
133}
134
135impl MemoryUsage for spacetimedb_primitives::TableId {}
136impl MemoryUsage for spacetimedb_primitives::SequenceId {}
137impl MemoryUsage for spacetimedb_primitives::ConstraintId {}
138impl MemoryUsage for spacetimedb_primitives::IndexId {}
139impl MemoryUsage for spacetimedb_primitives::ColId {}
140impl MemoryUsage for spacetimedb_primitives::ColList {
141    fn heap_usage(&self) -> usize {
142        self.heap_size()
143    }
144}
145
146impl MemoryUsage for AlgebraicValue {
147    fn heap_usage(&self) -> usize {
148        match self {
149            AlgebraicValue::Sum(x) => x.heap_usage(),
150            AlgebraicValue::Product(x) => x.heap_usage(),
151            AlgebraicValue::Array(x) => x.heap_usage(),
152            AlgebraicValue::String(x) => x.heap_usage(),
153            _ => 0,
154        }
155    }
156}
157impl MemoryUsage for SumValue {
158    fn heap_usage(&self) -> usize {
159        self.value.heap_usage()
160    }
161}
162impl MemoryUsage for ProductValue {
163    fn heap_usage(&self) -> usize {
164        self.elements.heap_usage()
165    }
166}
167impl MemoryUsage for ArrayValue {
168    fn heap_usage(&self) -> usize {
169        match self {
170            ArrayValue::Sum(v) => v.heap_usage(),
171            ArrayValue::Product(v) => v.heap_usage(),
172            ArrayValue::Bool(v) => v.heap_usage(),
173            ArrayValue::I8(v) => v.heap_usage(),
174            ArrayValue::U8(v) => v.heap_usage(),
175            ArrayValue::I16(v) => v.heap_usage(),
176            ArrayValue::U16(v) => v.heap_usage(),
177            ArrayValue::I32(v) => v.heap_usage(),
178            ArrayValue::U32(v) => v.heap_usage(),
179            ArrayValue::I64(v) => v.heap_usage(),
180            ArrayValue::U64(v) => v.heap_usage(),
181            ArrayValue::I128(v) => v.heap_usage(),
182            ArrayValue::U128(v) => v.heap_usage(),
183            ArrayValue::I256(v) => v.heap_usage(),
184            ArrayValue::U256(v) => v.heap_usage(),
185            ArrayValue::F32(v) => v.heap_usage(),
186            ArrayValue::F64(v) => v.heap_usage(),
187            ArrayValue::String(v) => v.heap_usage(),
188            ArrayValue::Array(v) => v.heap_usage(),
189        }
190    }
191}
192impl MemoryUsage for AlgebraicType {
193    fn heap_usage(&self) -> usize {
194        match self {
195            AlgebraicType::Ref(_) => 0,
196            AlgebraicType::Sum(x) => x.heap_usage(),
197            AlgebraicType::Product(x) => x.heap_usage(),
198            AlgebraicType::Array(x) => x.heap_usage(),
199            AlgebraicType::String
200            | AlgebraicType::Bool
201            | AlgebraicType::I8
202            | AlgebraicType::U8
203            | AlgebraicType::I16
204            | AlgebraicType::U16
205            | AlgebraicType::I32
206            | AlgebraicType::U32
207            | AlgebraicType::I64
208            | AlgebraicType::U64
209            | AlgebraicType::I128
210            | AlgebraicType::U128
211            | AlgebraicType::I256
212            | AlgebraicType::U256
213            | AlgebraicType::F32
214            | AlgebraicType::F64 => 0,
215        }
216    }
217}
218impl MemoryUsage for SumType {
219    fn heap_usage(&self) -> usize {
220        self.variants.heap_usage()
221    }
222}
223impl MemoryUsage for SumTypeVariant {
224    fn heap_usage(&self) -> usize {
225        self.name.heap_usage() + self.algebraic_type.heap_usage()
226    }
227}
228impl MemoryUsage for ProductType {
229    fn heap_usage(&self) -> usize {
230        self.elements.heap_usage()
231    }
232}
233impl MemoryUsage for ProductTypeElement {
234    fn heap_usage(&self) -> usize {
235        self.name.heap_usage() + self.algebraic_type.heap_usage()
236    }
237}
238impl MemoryUsage for ArrayType {
239    fn heap_usage(&self) -> usize {
240        self.elem_ty.heap_usage()
241    }
242}
243
244impl<T: MemoryUsage + Copy> MemoryUsage for Packed<T> {
245    fn heap_usage(&self) -> usize {
246        { self.0 }.heap_usage()
247    }
248}
249
250impl MemoryUsage for spacetimedb_lib::ConnectionId {}
251impl MemoryUsage for spacetimedb_lib::Identity {}