size_of/
collections.rs

1use crate::{Context, SizeOf};
2use alloc::{
3    collections::{BinaryHeap, LinkedList, VecDeque},
4    ffi::CString,
5    string::String,
6    vec::Vec,
7};
8use core::mem::size_of;
9
10impl SizeOf for String {
11    #[inline]
12    fn size_of_children(&self, context: &mut Context) {
13        if self.capacity() != 0 {
14            context
15                .add_vectorlike(self.len(), self.capacity(), size_of::<u8>())
16                .add_distinct_allocation();
17        }
18    }
19}
20
21impl SizeOf for CString {
22    fn size_of_children(&self, context: &mut Context) {
23        let length = self.to_bytes_with_nul().len();
24        if length != 0 {
25            context
26                .add_arraylike(length, size_of::<u8>())
27                .add_distinct_allocation();
28        }
29    }
30}
31
32impl<T> SizeOf for Vec<T>
33where
34    T: SizeOf,
35{
36    #[inline]
37    fn size_of_children(&self, context: &mut Context) {
38        if self.capacity() != 0 {
39            if size_of::<T>() != 0 {
40                context
41                    .add_vectorlike(self.len(), self.capacity(), size_of::<T>())
42                    .add_distinct_allocation();
43            }
44
45            self.as_slice().size_of_children(context);
46        }
47    }
48}
49
50impl<T> SizeOf for VecDeque<T>
51where
52    T: SizeOf,
53{
54    fn size_of_children(&self, context: &mut Context) {
55        if self.capacity() != 0 {
56            if size_of::<T>() != 0 {
57                context
58                    .add_vectorlike(self.len(), self.capacity(), size_of::<T>())
59                    .add_distinct_allocation();
60            }
61
62            let (left, right) = self.as_slices();
63            left.size_of_children(context);
64            right.size_of_children(context);
65        }
66    }
67}
68
69impl<T> SizeOf for BinaryHeap<T>
70where
71    T: SizeOf,
72{
73    fn size_of_children(&self, context: &mut Context) {
74        if self.capacity() != 0 {
75            if size_of::<T>() != 0 {
76                context
77                    .add_vectorlike(self.len(), self.capacity(), size_of::<T>())
78                    .add_distinct_allocation();
79            }
80
81            self.iter()
82                .for_each(|element| element.size_of_children(context));
83        }
84    }
85}
86
87impl<T> SizeOf for LinkedList<T>
88where
89    T: SizeOf,
90{
91    fn size_of_children(&self, context: &mut Context) {
92        let length = self.len();
93
94        if length != 0 {
95            // Record each node as a `{ T, *const (), *const () }`
96            context
97                .add_arraylike(length, size_of::<T>() + (size_of::<*const ()>() * 2))
98                .add_distinct_allocations(length);
99
100            self.iter()
101                .for_each(|element| element.size_of_children(context));
102        }
103    }
104}
105
106// A btree node has 2*B - 1 (K,V) pairs and (usize, u16, u16)
107// overhead, and an internal btree node additionally has 2*B
108// `usize` overhead.
109// A node can contain between B - 1 and 2*B - 1 elements, so
110// we assume it has the midpoint 3/2*B - 1.
111pub(crate) mod btree {
112    use crate::{Context, SizeOf};
113    use alloc::collections::{BTreeMap, BTreeSet};
114    use core::mem::size_of;
115
116    // Constants from rust's source:
117    // https://doc.rust-lang.org/src/alloc/collections/btree/node.rs.html#43-45
118    const B: usize = 6;
119    const BTREE_MAX: usize = 2 * B - 1;
120    const BTREE_MIN: usize = B - 1;
121
122    // A fake btree node, this isn't 100% accurate since the real btree node can
123    // have a different layout, but it should be close enough
124    #[allow(dead_code)]
125    struct FakeNode<K, V> {
126        parent: *const (),
127        parent_idx: u16,
128        len: u16,
129        keys: [K; BTREE_MAX],
130        values: [V; BTREE_MAX],
131    }
132
133    // TODO: Figure out the unused capacity as well
134    // TODO: Estimate the number of allocated buckets each btree makes
135    pub(crate) const fn estimate_btree_size<K, V>(length: usize) -> usize {
136        length * size_of::<FakeNode<K, V>>() * 2 / (BTREE_MAX + BTREE_MIN)
137    }
138
139    impl<K> SizeOf for BTreeSet<K>
140    where
141        K: SizeOf,
142    {
143        fn size_of_children(&self, context: &mut Context) {
144            if !self.is_empty() {
145                context
146                    .add(estimate_btree_size::<K, ()>(self.len()))
147                    // FIXME: Estimate the number of allocated buckets
148                    .add_distinct_allocation();
149
150                self.iter().for_each(|key| key.size_of_children(context));
151            }
152        }
153    }
154
155    impl<K, V> SizeOf for BTreeMap<K, V>
156    where
157        K: SizeOf,
158        V: SizeOf,
159    {
160        fn size_of_children(&self, context: &mut Context) {
161            if !self.is_empty() {
162                context
163                    .add(estimate_btree_size::<K, V>(self.len()))
164                    // FIXME: Estimate the number of allocated buckets
165                    .add_distinct_allocation();
166
167                self.iter().for_each(|(key, value)| {
168                    key.size_of_children(context);
169                    value.size_of_children(context);
170                });
171            }
172        }
173    }
174}
size_of/collections.rs

size_of/
collections.rs
