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
161impl<K: HeapSize, V: HeapSize> HeapSize for std::collections::HashMap<K, V> {
162    fn heap_size(&self) -> usize {
163        // Approximate: capacity * size_of key/value + heap of keys/values
164        self.capacity() * (std::mem::size_of::<K>() + std::mem::size_of::<V>())
165            + self.iter().map(|(k, v)| k.heap_size() + v.heap_size()).sum::<usize>()
166    }
167}
168
169impl<T: HeapSize> HeapSize for std::collections::HashSet<T> {
170    fn heap_size(&self) -> usize {
171        self.capacity() * std::mem::size_of::<T>()
172            + self.iter().map(|x| x.heap_size()).sum::<usize>()
173    }
174}
175
176impl<T: HeapSize> HeapSize for std::collections::BTreeSet<T> {
177    fn heap_size(&self) -> usize {
178        // BTreeSet has internal structure, approximate with len * size + heap
179        self.len() * std::mem::size_of::<T>() + self.iter().map(|x| x.heap_size()).sum::<usize>()
180    }
181}
182
183impl<K: HeapSize, V: HeapSize, BH> HeapSize for OrderedHashMap<K, V, BH> {
184    fn heap_size(&self) -> usize {
185        self.iter().map(|(k, v)| k.heap_size() + v.heap_size()).sum()
186    }
187}
188
189// For smol_str::SmolStr
190impl HeapSize for smol_str::SmolStr {
191    fn heap_size(&self) -> usize {
192        if *COUNT_SHARED_ALLOCATIONS && self.is_heap_allocated() {
193            // SmolStr stores short strings inline, but longer strings on the heap
194            // Note: This may double-count if the same SmolStr is cloned and shares the allocation
195            self.len()
196        } else {
197            0
198        }
199    }
200}
201
202// For num-bigint
203impl HeapSize for num_bigint::BigUint {
204    fn heap_size(&self) -> usize {
205        // Approximate the number of bytes required to store the magnitude.
206        let bits = self.bits() as usize;
207        bits.div_ceil(8)
208    }
209}
210
211impl HeapSize for num_bigint::BigInt {
212    fn heap_size(&self) -> usize {
213        self.magnitude().heap_size()
214    }
215}
216
217// Implementations for tuples
218impl HeapSize for () {
219    fn heap_size(&self) -> usize {
220        0
221    }
222}
223
224impl<T0: HeapSize> HeapSize for (T0,) {
225    fn heap_size(&self) -> usize {
226        self.0.heap_size()
227    }
228}
229
230impl<T0: HeapSize, T1: HeapSize> HeapSize for (T0, T1) {
231    fn heap_size(&self) -> usize {
232        self.0.heap_size() + self.1.heap_size()
233    }
234}
235
236impl<T0: HeapSize, T1: HeapSize, T2: HeapSize> HeapSize for (T0, T1, T2) {
237    fn heap_size(&self) -> usize {
238        self.0.heap_size() + self.1.heap_size() + self.2.heap_size()
239    }
240}
241
242impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize> HeapSize for (T0, T1, T2, T3) {
243    fn heap_size(&self) -> usize {
244        self.0.heap_size() + self.1.heap_size() + self.2.heap_size() + self.3.heap_size()
245    }
246}
247
248impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize> HeapSize
249    for (T0, T1, T2, T3, T4)
250{
251    fn heap_size(&self) -> usize {
252        self.0.heap_size()
253            + self.1.heap_size()
254            + self.2.heap_size()
255            + self.3.heap_size()
256            + self.4.heap_size()
257    }
258}
259
260impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize, T5: HeapSize> HeapSize
261    for (T0, T1, T2, T3, T4, T5)
262{
263    fn heap_size(&self) -> usize {
264        self.0.heap_size()
265            + self.1.heap_size()
266            + self.2.heap_size()
267            + self.3.heap_size()
268            + self.4.heap_size()
269            + self.5.heap_size()
270    }
271}
272
273impl<
274    T0: HeapSize,
275    T1: HeapSize,
276    T2: HeapSize,
277    T3: HeapSize,
278    T4: HeapSize,
279    T5: HeapSize,
280    T6: HeapSize,
281> HeapSize for (T0, T1, T2, T3, T4, T5, T6)
282{
283    fn heap_size(&self) -> usize {
284        self.0.heap_size()
285            + self.1.heap_size()
286            + self.2.heap_size()
287            + self.3.heap_size()
288            + self.4.heap_size()
289            + self.5.heap_size()
290            + self.6.heap_size()
291    }
292}
293
294impl<
295    T0: HeapSize,
296    T1: HeapSize,
297    T2: HeapSize,
298    T3: HeapSize,
299    T4: HeapSize,
300    T5: HeapSize,
301    T6: HeapSize,
302    T7: HeapSize,
303> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7)
304{
305    fn heap_size(&self) -> usize {
306        self.0.heap_size()
307            + self.1.heap_size()
308            + self.2.heap_size()
309            + self.3.heap_size()
310            + self.4.heap_size()
311            + self.5.heap_size()
312            + self.6.heap_size()
313            + self.7.heap_size()
314    }
315}
316
317impl<
318    T0: HeapSize,
319    T1: HeapSize,
320    T2: HeapSize,
321    T3: HeapSize,
322    T4: HeapSize,
323    T5: HeapSize,
324    T6: HeapSize,
325    T7: HeapSize,
326    T8: HeapSize,
327> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8)
328{
329    fn heap_size(&self) -> usize {
330        self.0.heap_size()
331            + self.1.heap_size()
332            + self.2.heap_size()
333            + self.3.heap_size()
334            + self.4.heap_size()
335            + self.5.heap_size()
336            + self.6.heap_size()
337            + self.7.heap_size()
338            + self.8.heap_size()
339    }
340}
341
342impl<
343    T0: HeapSize,
344    T1: HeapSize,
345    T2: HeapSize,
346    T3: HeapSize,
347    T4: HeapSize,
348    T5: HeapSize,
349    T6: HeapSize,
350    T7: HeapSize,
351    T8: HeapSize,
352    T9: HeapSize,
353> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9)
354{
355    fn heap_size(&self) -> usize {
356        self.0.heap_size()
357            + self.1.heap_size()
358            + self.2.heap_size()
359            + self.3.heap_size()
360            + self.4.heap_size()
361            + self.5.heap_size()
362            + self.6.heap_size()
363            + self.7.heap_size()
364            + self.8.heap_size()
365            + self.9.heap_size()
366    }
367}
368
369impl<
370    T0: HeapSize,
371    T1: HeapSize,
372    T2: HeapSize,
373    T3: HeapSize,
374    T4: HeapSize,
375    T5: HeapSize,
376    T6: HeapSize,
377    T7: HeapSize,
378    T8: HeapSize,
379    T9: HeapSize,
380    T10: HeapSize,
381> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10)
382{
383    fn heap_size(&self) -> usize {
384        self.0.heap_size()
385            + self.1.heap_size()
386            + self.2.heap_size()
387            + self.3.heap_size()
388            + self.4.heap_size()
389            + self.5.heap_size()
390            + self.6.heap_size()
391            + self.7.heap_size()
392            + self.8.heap_size()
393            + self.9.heap_size()
394            + self.10.heap_size()
395    }
396}
397
398impl<
399    T0: HeapSize,
400    T1: HeapSize,
401    T2: HeapSize,
402    T3: HeapSize,
403    T4: HeapSize,
404    T5: HeapSize,
405    T6: HeapSize,
406    T7: HeapSize,
407    T8: HeapSize,
408    T9: HeapSize,
409    T10: HeapSize,
410    T11: HeapSize,
411> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11)
412{
413    fn heap_size(&self) -> usize {
414        self.0.heap_size()
415            + self.1.heap_size()
416            + self.2.heap_size()
417            + self.3.heap_size()
418            + self.4.heap_size()
419            + self.5.heap_size()
420            + self.6.heap_size()
421            + self.7.heap_size()
422            + self.8.heap_size()
423            + self.9.heap_size()
424            + self.10.heap_size()
425            + self.11.heap_size()
426    }
427}
cairo_lang_utils/heap_size.rs

cairo_lang_utils/
heap_size.rs
