spacetimedb_table/
memory_usage.rs

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