cairo_lang_utils/
heap_size.rs

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

cairo_lang_utils/
heap_size.rs
