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        // BigUint has internal Vec<u8>
206        self.to_bytes_le().capacity()
207    }
208}
209
210impl HeapSize for num_bigint::BigInt {
211    fn heap_size(&self) -> usize {
212        self.to_bytes_le().1.capacity()
213    }
214}
215
216// Implementations for tuples
217impl HeapSize for () {
218    fn heap_size(&self) -> usize {
219        0
220    }
221}
222
223impl<T0: HeapSize> HeapSize for (T0,) {
224    fn heap_size(&self) -> usize {
225        self.0.heap_size()
226    }
227}
228
229impl<T0: HeapSize, T1: HeapSize> HeapSize for (T0, T1) {
230    fn heap_size(&self) -> usize {
231        self.0.heap_size() + self.1.heap_size()
232    }
233}
234
235impl<T0: HeapSize, T1: HeapSize, T2: HeapSize> HeapSize for (T0, T1, T2) {
236    fn heap_size(&self) -> usize {
237        self.0.heap_size() + self.1.heap_size() + self.2.heap_size()
238    }
239}
240
241impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize> HeapSize for (T0, T1, T2, T3) {
242    fn heap_size(&self) -> usize {
243        self.0.heap_size() + self.1.heap_size() + self.2.heap_size() + self.3.heap_size()
244    }
245}
246
247impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize> HeapSize
248    for (T0, T1, T2, T3, T4)
249{
250    fn heap_size(&self) -> usize {
251        self.0.heap_size()
252            + self.1.heap_size()
253            + self.2.heap_size()
254            + self.3.heap_size()
255            + self.4.heap_size()
256    }
257}
258
259impl<T0: HeapSize, T1: HeapSize, T2: HeapSize, T3: HeapSize, T4: HeapSize, T5: HeapSize> HeapSize
260    for (T0, T1, T2, T3, T4, T5)
261{
262    fn heap_size(&self) -> usize {
263        self.0.heap_size()
264            + self.1.heap_size()
265            + self.2.heap_size()
266            + self.3.heap_size()
267            + self.4.heap_size()
268            + self.5.heap_size()
269    }
270}
271
272impl<
273    T0: HeapSize,
274    T1: HeapSize,
275    T2: HeapSize,
276    T3: HeapSize,
277    T4: HeapSize,
278    T5: HeapSize,
279    T6: HeapSize,
280> HeapSize for (T0, T1, T2, T3, T4, T5, T6)
281{
282    fn heap_size(&self) -> usize {
283        self.0.heap_size()
284            + self.1.heap_size()
285            + self.2.heap_size()
286            + self.3.heap_size()
287            + self.4.heap_size()
288            + self.5.heap_size()
289            + self.6.heap_size()
290    }
291}
292
293impl<
294    T0: HeapSize,
295    T1: HeapSize,
296    T2: HeapSize,
297    T3: HeapSize,
298    T4: HeapSize,
299    T5: HeapSize,
300    T6: HeapSize,
301    T7: HeapSize,
302> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7)
303{
304    fn heap_size(&self) -> usize {
305        self.0.heap_size()
306            + self.1.heap_size()
307            + self.2.heap_size()
308            + self.3.heap_size()
309            + self.4.heap_size()
310            + self.5.heap_size()
311            + self.6.heap_size()
312            + self.7.heap_size()
313    }
314}
315
316impl<
317    T0: HeapSize,
318    T1: HeapSize,
319    T2: HeapSize,
320    T3: HeapSize,
321    T4: HeapSize,
322    T5: HeapSize,
323    T6: HeapSize,
324    T7: HeapSize,
325    T8: HeapSize,
326> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8)
327{
328    fn heap_size(&self) -> usize {
329        self.0.heap_size()
330            + self.1.heap_size()
331            + self.2.heap_size()
332            + self.3.heap_size()
333            + self.4.heap_size()
334            + self.5.heap_size()
335            + self.6.heap_size()
336            + self.7.heap_size()
337            + self.8.heap_size()
338    }
339}
340
341impl<
342    T0: HeapSize,
343    T1: HeapSize,
344    T2: HeapSize,
345    T3: HeapSize,
346    T4: HeapSize,
347    T5: HeapSize,
348    T6: HeapSize,
349    T7: HeapSize,
350    T8: HeapSize,
351    T9: HeapSize,
352> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9)
353{
354    fn heap_size(&self) -> usize {
355        self.0.heap_size()
356            + self.1.heap_size()
357            + self.2.heap_size()
358            + self.3.heap_size()
359            + self.4.heap_size()
360            + self.5.heap_size()
361            + self.6.heap_size()
362            + self.7.heap_size()
363            + self.8.heap_size()
364            + self.9.heap_size()
365    }
366}
367
368impl<
369    T0: HeapSize,
370    T1: HeapSize,
371    T2: HeapSize,
372    T3: HeapSize,
373    T4: HeapSize,
374    T5: HeapSize,
375    T6: HeapSize,
376    T7: HeapSize,
377    T8: HeapSize,
378    T9: HeapSize,
379    T10: HeapSize,
380> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10)
381{
382    fn heap_size(&self) -> usize {
383        self.0.heap_size()
384            + self.1.heap_size()
385            + self.2.heap_size()
386            + self.3.heap_size()
387            + self.4.heap_size()
388            + self.5.heap_size()
389            + self.6.heap_size()
390            + self.7.heap_size()
391            + self.8.heap_size()
392            + self.9.heap_size()
393            + self.10.heap_size()
394    }
395}
396
397impl<
398    T0: HeapSize,
399    T1: HeapSize,
400    T2: HeapSize,
401    T3: HeapSize,
402    T4: HeapSize,
403    T5: HeapSize,
404    T6: HeapSize,
405    T7: HeapSize,
406    T8: HeapSize,
407    T9: HeapSize,
408    T10: HeapSize,
409    T11: HeapSize,
410> HeapSize for (T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11)
411{
412    fn heap_size(&self) -> usize {
413        self.0.heap_size()
414            + self.1.heap_size()
415            + self.2.heap_size()
416            + self.3.heap_size()
417            + self.4.heap_size()
418            + self.5.heap_size()
419            + self.6.heap_size()
420            + self.7.heap_size()
421            + self.8.heap_size()
422            + self.9.heap_size()
423            + self.10.heap_size()
424            + self.11.heap_size()
425    }
426}
cairo_lang_utils/heap_size.rs

cairo_lang_utils/
heap_size.rs
