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 Box<T> {
43    fn heap_usage(&self) -> usize {
44        mem::size_of_val::<T>(self) + T::heap_usage(self)
45    }
46}
47
48impl<T: MemoryUsage + ?Sized> MemoryUsage for std::sync::Arc<T> {
49    fn heap_usage(&self) -> usize {
50        let refcounts = mem::size_of::<usize>() * 2;
51        refcounts + mem::size_of_val::<T>(self) + T::heap_usage(self)
52    }
53}
54
55impl<T: MemoryUsage + ?Sized> MemoryUsage for std::rc::Rc<T> {
56    fn heap_usage(&self) -> usize {
57        let refcounts = mem::size_of::<usize>() * 2;
58        refcounts + mem::size_of_val::<T>(self) + T::heap_usage(self)
59    }
60}
61
62impl<T: MemoryUsage> MemoryUsage for [T] {
63    fn heap_usage(&self) -> usize {
64        self.iter().map(T::heap_usage).sum()
65    }
66}
67
68impl MemoryUsage for str {}
69
70impl<T: MemoryUsage> MemoryUsage for Option<T> {
71    fn heap_usage(&self) -> usize {
72        self.as_ref().map_or(0, T::heap_usage)
73    }
74}
75
76impl<A: MemoryUsage, B: MemoryUsage> MemoryUsage for (A, B) {
77    fn heap_usage(&self) -> usize {
78        self.0.heap_usage() + self.1.heap_usage()
79    }
80}
81
82impl MemoryUsage for String {
83    fn heap_usage(&self) -> usize {
84        self.capacity()
85    }
86}
87
88impl<T: MemoryUsage> MemoryUsage for Vec<T> {
89    fn heap_usage(&self) -> usize {
90        self.capacity() * mem::size_of::<T>() + self.iter().map(T::heap_usage).sum::<usize>()
91    }
92}
93
94impl<K: MemoryUsage + Eq + Hash, V: MemoryUsage, S: BuildHasher> MemoryUsage
95    for spacetimedb_data_structures::map::HashMap<K, V, S>
96{
97    fn heap_usage(&self) -> usize {
98        self.allocation_size() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
99    }
100}
101
102impl<K: MemoryUsage, V: MemoryUsage> MemoryUsage for std::collections::BTreeMap<K, V> {
103    fn heap_usage(&self) -> usize {
104        // NB: this is best-effort, since we don't have a `capacity()` method on `BTreeMap`.
105        self.len() * mem::size_of::<(K, V)>() + self.iter().map(|(k, v)| k.heap_usage() + v.heap_usage()).sum::<usize>()
106    }
107}
108
109impl<A: smallvec::Array> MemoryUsage for smallvec::SmallVec<A>
110where
111    A::Item: MemoryUsage,
112{
113    fn heap_usage(&self) -> usize {
114        self.as_slice().heap_usage()
115            + if self.spilled() {
116                self.capacity() * mem::size_of::<A::Item>()
117            } else {
118                0
119            }
120    }
121}
122
123impl MemoryUsage for spacetimedb_primitives::TableId {}
124impl MemoryUsage for spacetimedb_primitives::SequenceId {}
125impl MemoryUsage for spacetimedb_primitives::ConstraintId {}
126impl MemoryUsage for spacetimedb_primitives::IndexId {}
127impl MemoryUsage for spacetimedb_primitives::ColId {}
128impl MemoryUsage for spacetimedb_primitives::ColList {
129    fn heap_usage(&self) -> usize {
130        self.heap_size()
131    }
132}
133
134impl MemoryUsage for AlgebraicValue {
135    fn heap_usage(&self) -> usize {
136        match self {
137            AlgebraicValue::Sum(x) => x.heap_usage(),
138            AlgebraicValue::Product(x) => x.heap_usage(),
139            AlgebraicValue::Array(x) => x.heap_usage(),
140            AlgebraicValue::String(x) => x.heap_usage(),
141            _ => 0,
142        }
143    }
144}
145impl MemoryUsage for SumValue {
146    fn heap_usage(&self) -> usize {
147        self.value.heap_usage()
148    }
149}
150impl MemoryUsage for ProductValue {
151    fn heap_usage(&self) -> usize {
152        self.elements.heap_usage()
153    }
154}
155impl MemoryUsage for ArrayValue {
156    fn heap_usage(&self) -> usize {
157        match self {
158            ArrayValue::Sum(v) => v.heap_usage(),
159            ArrayValue::Product(v) => v.heap_usage(),
160            ArrayValue::Bool(v) => v.heap_usage(),
161            ArrayValue::I8(v) => v.heap_usage(),
162            ArrayValue::U8(v) => v.heap_usage(),
163            ArrayValue::I16(v) => v.heap_usage(),
164            ArrayValue::U16(v) => v.heap_usage(),
165            ArrayValue::I32(v) => v.heap_usage(),
166            ArrayValue::U32(v) => v.heap_usage(),
167            ArrayValue::I64(v) => v.heap_usage(),
168            ArrayValue::U64(v) => v.heap_usage(),
169            ArrayValue::I128(v) => v.heap_usage(),
170            ArrayValue::U128(v) => v.heap_usage(),
171            ArrayValue::I256(v) => v.heap_usage(),
172            ArrayValue::U256(v) => v.heap_usage(),
173            ArrayValue::F32(v) => v.heap_usage(),
174            ArrayValue::F64(v) => v.heap_usage(),
175            ArrayValue::String(v) => v.heap_usage(),
176            ArrayValue::Array(v) => v.heap_usage(),
177        }
178    }
179}
180impl MemoryUsage for AlgebraicType {
181    fn heap_usage(&self) -> usize {
182        match self {
183            AlgebraicType::Ref(_) => 0,
184            AlgebraicType::Sum(x) => x.heap_usage(),
185            AlgebraicType::Product(x) => x.heap_usage(),
186            AlgebraicType::Array(x) => x.heap_usage(),
187            AlgebraicType::String
188            | AlgebraicType::Bool
189            | AlgebraicType::I8
190            | AlgebraicType::U8
191            | AlgebraicType::I16
192            | AlgebraicType::U16
193            | AlgebraicType::I32
194            | AlgebraicType::U32
195            | AlgebraicType::I64
196            | AlgebraicType::U64
197            | AlgebraicType::I128
198            | AlgebraicType::U128
199            | AlgebraicType::I256
200            | AlgebraicType::U256
201            | AlgebraicType::F32
202            | AlgebraicType::F64 => 0,
203        }
204    }
205}
206impl MemoryUsage for SumType {
207    fn heap_usage(&self) -> usize {
208        self.variants.heap_usage()
209    }
210}
211impl MemoryUsage for SumTypeVariant {
212    fn heap_usage(&self) -> usize {
213        self.name.heap_usage() + self.algebraic_type.heap_usage()
214    }
215}
216impl MemoryUsage for ProductType {
217    fn heap_usage(&self) -> usize {
218        self.elements.heap_usage()
219    }
220}
221impl MemoryUsage for ProductTypeElement {
222    fn heap_usage(&self) -> usize {
223        self.name.heap_usage() + self.algebraic_type.heap_usage()
224    }
225}
226impl MemoryUsage for ArrayType {
227    fn heap_usage(&self) -> usize {
228        self.elem_ty.heap_usage()
229    }
230}
231
232impl<T: MemoryUsage + Copy> MemoryUsage for Packed<T> {
233    fn heap_usage(&self) -> usize {
234        { self.0 }.heap_usage()
235    }
236}
237
238impl MemoryUsage for spacetimedb_lib::ConnectionId {}
239impl MemoryUsage for spacetimedb_lib::Identity {}