cairo_lang_utils/
heap_size.rs

1use std::sync::LazyLock;
2
3use crate::ordered_hash_map::OrderedHashMap;
4use crate::ordered_hash_set::OrderedHashSet;
5
6/// Environment variable to control whether shared allocations (Arc, SmolStr) are counted.
7/// When set to "1" or "true", these types will include their heap allocations in the count.
8/// Note: This may lead to double-counting when the same Arc is referenced multiple times.
9static COUNT_SHARED_ALLOCATIONS: LazyLock<bool> = LazyLock::new(|| {
10    std::env::var("CAIRO_HEAPSIZE_COUNT_SHARED")
11        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
12        .unwrap_or(false)
13});
14
15/// Trait for calculating the heap size of an object's owned data.
16/// Arc's and references are not included in the sum by default.
17/// Set CAIRO_HEAPSIZE_COUNT_SHARED=1 to include shared allocations (may double-count).
18pub trait HeapSize {
19    /// Returns the size of the heap-allocated, owned data in bytes.
20    /// It does not include the size of the object itself (stack size), nor references.
21    fn heap_size(&self) -> usize;
22}
23
24// Implementations for standard library types
25impl<T: HeapSize> HeapSize for Vec<T> {
26    fn heap_size(&self) -> usize {
27        // Vec's heap size is the capacity * size_of::<T>() + heap_size of elements
28        self.capacity() * std::mem::size_of::<T>()
29            + self.iter().map(|x| x.heap_size()).sum::<usize>()
30    }
31}
32
33impl HeapSize for String {
34    fn heap_size(&self) -> usize {
35        self.capacity()
36    }
37}
38
39impl<T: HeapSize> HeapSize for Option<T> {
40    fn heap_size(&self) -> usize {
41        match self {
42            Some(x) => x.heap_size(),
43            None => 0,
44        }
45    }
46}
47
48impl<T: HeapSize, E: HeapSize> HeapSize for Result<T, E> {
49    fn heap_size(&self) -> usize {
50        match self {
51            Ok(value) => value.heap_size(),
52            Err(err) => err.heap_size(),
53        }
54    }
55}
56
57impl<T: HeapSize> HeapSize for Box<T> {
58    fn heap_size(&self) -> usize {
59        std::mem::size_of::<T>() + self.as_ref().heap_size()
60    }
61}
62
63impl<T: HeapSize> HeapSize for std::sync::Arc<T> {
64    fn heap_size(&self) -> usize {
65        if *COUNT_SHARED_ALLOCATIONS {
66            // Count the Arc's allocation: the data plus Arc's metadata
67            // Note: This may double-count if the same Arc is referenced multiple times
68            std::mem::size_of::<T>() + self.as_ref().heap_size()
69        } else {
70            // We do not count "unowned" data by default
71            0
72        }
73    }
74}
75
76impl<T: HeapSize> HeapSize for std::rc::Rc<T> {
77    fn heap_size(&self) -> usize {
78        if *COUNT_SHARED_ALLOCATIONS {
79            // Count the Rc's allocation: the data plus Rc's metadata
80            // Note: This may double-count if the same Rc is referenced multiple times
81            std::mem::size_of::<T>() + self.as_ref().heap_size()
82        } else {
83            // We do not count "unowned" data by default
84            0
85        }
86    }
87}
88
89// For types with no heap allocation
90impl HeapSize for u8 {
91    fn heap_size(&self) -> usize {
92        0
93    }
94}
95
96impl HeapSize for i8 {
97    fn heap_size(&self) -> usize {
98        0
99    }
100}
101
102impl HeapSize for i32 {
103    fn heap_size(&self) -> usize {
104        0
105    }
106}
107
108impl HeapSize for u32 {
109    fn heap_size(&self) -> usize {
110        0
111    }
112}
113
114impl HeapSize for i64 {
115    fn heap_size(&self) -> usize {
116        0
117    }
118}
119
120impl HeapSize for u64 {
121    fn heap_size(&self) -> usize {
122        0
123    }
124}
125
126impl HeapSize for usize {
127    fn heap_size(&self) -> usize {
128        0
129    }
130}
131
132impl HeapSize for isize {
133    fn heap_size(&self) -> usize {
134        0
135    }
136}
137
138impl HeapSize for bool {
139    fn heap_size(&self) -> usize {
140        0
141    }
142}
143
144impl HeapSize for char {
145    fn heap_size(&self) -> usize {
146        0
147    }
148}
149
150impl<T> HeapSize for std::marker::PhantomData<T> {
151    fn heap_size(&self) -> usize {
152        0
153    }
154}
155
156impl HeapSize for std::path::PathBuf {
157    fn heap_size(&self) -> usize {
158        self.capacity()
159    }
160}
161
162#[expect(clippy::disallowed_types)]
163impl<K: HeapSize, V: HeapSize> HeapSize for std::collections::HashMap<K, V> {
164    fn heap_size(&self) -> usize {
165        // Approximate: capacity * size_of key/value + heap of keys/values
166        self.capacity() * (std::mem::size_of::<K>() + std::mem::size_of::<V>())
167            + self.iter().map(|(k, v)| k.heap_size() + v.heap_size()).sum::<usize>()
168    }
169}
170
171#[expect(clippy::disallowed_types)]
172impl<T: HeapSize> HeapSize for std::collections::HashSet<T> {
173    fn heap_size(&self) -> usize {
174        self.capacity() * std::mem::size_of::<T>()
175            + self.iter().map(|x| x.heap_size()).sum::<usize>()
176    }
177}
178
179impl<T: HeapSize> HeapSize for std::collections::BTreeSet<T> {
180    fn heap_size(&self) -> usize {
181        // BTreeSet has internal structure, approximate with len * size + heap
182        self.len() * std::mem::size_of::<T>() + self.iter().map(|x| x.heap_size()).sum::<usize>()
183    }
184}
185
186impl<K: HeapSize, V: HeapSize, BH> HeapSize for OrderedHashMap<K, V, BH> {
187    fn heap_size(&self) -> usize {
188        self.iter().map(|(k, v)| k.heap_size() + v.heap_size()).sum()
189    }
190}
191
192impl<K: HeapSize, BH> HeapSize for OrderedHashSet<K, BH> {
193    fn heap_size(&self) -> usize {
194        self.iter().map(|k| k.heap_size()).sum()
195    }
196}
197
198// For smol_str::SmolStr
199impl HeapSize for smol_str::SmolStr {
200    fn heap_size(&self) -> usize {
201        if *COUNT_SHARED_ALLOCATIONS && self.is_heap_allocated() {
202            // SmolStr stores short strings inline, but longer strings on the heap
203            // Note: This may double-count if the same SmolStr is cloned and shares the allocation
204            self.len()
205        } else {
206            0
207        }
208    }
209}
210
211// For num-bigint
212impl HeapSize for num_bigint::BigUint {
213    fn heap_size(&self) -> usize {
214        // Approximate the number of bytes required to store the magnitude.
215        let bits = self.bits() as usize;
216        bits.div_ceil(8)
217    }
218}
219
220impl HeapSize for num_bigint::BigInt {
221    fn heap_size(&self) -> usize {
222        self.magnitude().heap_size()
223    }
224}
225
226// Implementations for tuples
227impl HeapSize for () {
228    fn heap_size(&self) -> usize {
229        0
230    }
231}
232
233impl<T0: HeapSize> HeapSize for (T0,) {
234    fn heap_size(&self) -> usize {
235        self.0.heap_size()
236    }
237}
238
239impl<T0: HeapSize, T1: HeapSize> HeapSize for (T0, T1) {
240    fn heap_size(&self) -> usize {
241        self.0.heap_size() + self.1.heap_size()
242    }
243}
244
245impl<T0: HeapSize, T1: HeapSize, T2: HeapSize> HeapSize for (T0, T1, T2) {
246    fn heap_size(&self) -> usize {
247        self.0.heap_size() + self.1.heap_size() + self.2.heap_size()
248    }
249}
250
251impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize> HeapSize for (T0, T1, T2, T3) {
252    fn heap_size(&self) -> usize {
253        self.0.heap_size() + self.1.heap_size() + self.2.heap_size() + self.3.heap_size()
254    }
255}
256
257impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize> HeapSize
258    for (T0, T1, T2, T3, T4)
259{
260    fn heap_size(&self) -> usize {
261        self.0.heap_size()
262            + self.1.heap_size()
263            + self.2.heap_size()
264            + self.3.heap_size()
265            + self.4.heap_size()
266    }
267}
268
269impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize, T5: HeapSize> HeapSize
270    for (T0, T1, T2, T3, T4, T5)
271{
272    fn heap_size(&self) -> usize {
273        self.0.heap_size()
274            + self.1.heap_size()
275            + self.2.heap_size()
276            + self.3.heap_size()
277            + self.4.heap_size()
278            + self.5.heap_size()
279    }
280}
281
282impl<
283    T0: HeapSize,
284    T1: HeapSize,
285    T2: HeapSize,
286    T3: HeapSize,
287    T4: HeapSize,
288    T5: HeapSize,
289    T6: HeapSize,
290> HeapSize for (T0, T1, T2, T3, T4, T5, T6)
291{
292    fn heap_size(&self) -> usize {
293        self.0.heap_size()
294            + self.1.heap_size()
295            + self.2.heap_size()
296            + self.3.heap_size()
297            + self.4.heap_size()
298            + self.5.heap_size()
299            + self.6.heap_size()
300    }
301}
302
303impl<
304    T0: HeapSize,
305    T1: HeapSize,
306    T2: HeapSize,
307    T3: HeapSize,
308    T4: HeapSize,
309    T5: HeapSize,
310    T6: HeapSize,
311    T7: HeapSize,
312> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7)
313{
314    fn heap_size(&self) -> usize {
315        self.0.heap_size()
316            + self.1.heap_size()
317            + self.2.heap_size()
318            + self.3.heap_size()
319            + self.4.heap_size()
320            + self.5.heap_size()
321            + self.6.heap_size()
322            + self.7.heap_size()
323    }
324}
325
326impl<
327    T0: HeapSize,
328    T1: HeapSize,
329    T2: HeapSize,
330    T3: HeapSize,
331    T4: HeapSize,
332    T5: HeapSize,
333    T6: HeapSize,
334    T7: HeapSize,
335    T8: HeapSize,
336> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8)
337{
338    fn heap_size(&self) -> usize {
339        self.0.heap_size()
340            + self.1.heap_size()
341            + self.2.heap_size()
342            + self.3.heap_size()
343            + self.4.heap_size()
344            + self.5.heap_size()
345            + self.6.heap_size()
346            + self.7.heap_size()
347            + self.8.heap_size()
348    }
349}
350
351impl<
352    T0: HeapSize,
353    T1: HeapSize,
354    T2: HeapSize,
355    T3: HeapSize,
356    T4: HeapSize,
357    T5: HeapSize,
358    T6: HeapSize,
359    T7: HeapSize,
360    T8: HeapSize,
361    T9: HeapSize,
362> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9)
363{
364    fn heap_size(&self) -> usize {
365        self.0.heap_size()
366            + self.1.heap_size()
367            + self.2.heap_size()
368            + self.3.heap_size()
369            + self.4.heap_size()
370            + self.5.heap_size()
371            + self.6.heap_size()
372            + self.7.heap_size()
373            + self.8.heap_size()
374            + self.9.heap_size()
375    }
376}
377
378impl<
379    T0: HeapSize,
380    T1: HeapSize,
381    T2: HeapSize,
382    T3: HeapSize,
383    T4: HeapSize,
384    T5: HeapSize,
385    T6: HeapSize,
386    T7: HeapSize,
387    T8: HeapSize,
388    T9: HeapSize,
389    T10: HeapSize,
390> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10)
391{
392    fn heap_size(&self) -> usize {
393        self.0.heap_size()
394            + self.1.heap_size()
395            + self.2.heap_size()
396            + self.3.heap_size()
397            + self.4.heap_size()
398            + self.5.heap_size()
399            + self.6.heap_size()
400            + self.7.heap_size()
401            + self.8.heap_size()
402            + self.9.heap_size()
403            + self.10.heap_size()
404    }
405}
406
407impl<
408    T0: HeapSize,
409    T1: HeapSize,
410    T2: HeapSize,
411    T3: HeapSize,
412    T4: HeapSize,
413    T5: HeapSize,
414    T6: HeapSize,
415    T7: HeapSize,
416    T8: HeapSize,
417    T9: HeapSize,
418    T10: HeapSize,
419    T11: HeapSize,
420> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11)
421{
422    fn heap_size(&self) -> usize {
423        self.0.heap_size()
424            + self.1.heap_size()
425            + self.2.heap_size()
426            + self.3.heap_size()
427            + self.4.heap_size()
428            + self.5.heap_size()
429            + self.6.heap_size()
430            + self.7.heap_size()
431            + self.8.heap_size()
432            + self.9.heap_size()
433            + self.10.heap_size()
434            + self.11.heap_size()
435    }
436}
cairo_lang_utils/heap_size.rs

cairo_lang_utils/
heap_size.rs
