Skip to main content

get_size2/impls/alloc_impls/
collections.rs

1use alloc::collections::{BTreeMap, BTreeSet, BinaryHeap, LinkedList, VecDeque};
2use alloc::vec::Vec;
3
4use crate::{GetSize, GetSizeTracker};
5
6macro_rules! impl_size_set {
7    ($name:ident) => {
8        impl<T> GetSize for $name<T>
9        where
10            T: GetSize,
11        {
12            fn get_heap_size_with_tracker<Tr: GetSizeTracker>(&self, tracker: Tr) -> (usize, Tr) {
13                let (size, tracker) = self.iter().fold((0, tracker), |(size, tracker), elem| {
14                    let (elem_size, tracker) = T::get_heap_size_with_tracker(elem, tracker);
15                    (size + elem_size, tracker)
16                });
17
18                let allocation_size = self.capacity() * T::get_stack_size();
19                (size + allocation_size, tracker)
20            }
21        }
22    };
23}
24
25macro_rules! impl_size_set_no_capacity {
26    ($name:ident) => {
27        impl<T> GetSize for $name<T>
28        where
29            T: GetSize,
30        {
31            fn get_heap_size_with_tracker<Tr: GetSizeTracker>(&self, tracker: Tr) -> (usize, Tr) {
32                let (size, tracker) = self.iter().fold((0, tracker), |(size, tracker), elem| {
33                    // We assume that values are held inside the heap.
34                    let (elem_size, tracker) = T::get_size_with_tracker(elem, tracker);
35                    (size + elem_size, tracker)
36                });
37
38                (size, tracker)
39            }
40        }
41    };
42}
43
44impl_size_set_no_capacity!(BTreeSet);
45impl_size_set!(BinaryHeap);
46impl_size_set_no_capacity!(LinkedList);
47impl_size_set!(VecDeque);
48impl_size_set!(Vec);
49
50impl<K, V> GetSize for BTreeMap<K, V>
51where
52    K: GetSize,
53    V: GetSize,
54{
55    fn get_heap_size_with_tracker<Tr: GetSizeTracker>(&self, tracker: Tr) -> (usize, Tr) {
56        self.iter()
57            .fold((0, tracker), |(size, tracker), (key, value)| {
58                let (key_size, tracker) = K::get_size_with_tracker(key, tracker);
59                let (value_size, tracker) = V::get_size_with_tracker(value, tracker);
60                (size + key_size + value_size, tracker)
61            })
62    }
63}