Skip to main content

s2_common/
deep_size.rs

1pub trait DeepSize {
2    /// - size_of(primitive)
3    /// - length for chunks of data like strings and bytes (so not including the container overhead)
4    /// - deep size of all struct fields
5    /// - deep size of actual variant for enums
6    fn deep_size(&self) -> usize;
7}
8
9impl<X: DeepSize, Y: DeepSize> DeepSize for (X, Y) {
10    fn deep_size(&self) -> usize {
11        self.0.deep_size() + self.1.deep_size()
12    }
13}
14
15impl<T: DeepSize> DeepSize for &[T] {
16    fn deep_size(&self) -> usize {
17        self.iter().map(DeepSize::deep_size).sum::<usize>()
18    }
19}
20
21impl<T: DeepSize> DeepSize for Vec<T> {
22    fn deep_size(&self) -> usize {
23        self.iter().map(DeepSize::deep_size).sum::<usize>()
24    }
25}
26
27impl<T: DeepSize> DeepSize for Option<T> {
28    fn deep_size(&self) -> usize {
29        match self {
30            Some(v) => v.deep_size(),
31            None => 1,
32        }
33    }
34}
35
36impl DeepSize for String {
37    fn deep_size(&self) -> usize {
38        self.len()
39    }
40}
41
42impl DeepSize for bytes::Bytes {
43    fn deep_size(&self) -> usize {
44        self.len()
45    }
46}
47
48impl<T: DeepSize> DeepSize for std::ops::Bound<T> {
49    fn deep_size(&self) -> usize {
50        match self {
51            std::ops::Bound::Included(x) => x.deep_size(),
52            std::ops::Bound::Excluded(x) => x.deep_size(),
53            std::ops::Bound::Unbounded => 1,
54        }
55    }
56}
57
58macro_rules! impl_deep_size_prim {
59    ($($t:ty),+) => {
60        $(
61            impl DeepSize for $t {
62                fn deep_size(&self) -> usize {
63                    size_of_val(self)
64                }
65            }
66        )+
67    };
68}
69
70impl_deep_size_prim!(bool, u64, usize, std::num::NonZeroU64);
71
72#[cfg(test)]
73mod tests {
74    use proptest::prelude::*;
75
76    use super::*;
77
78    fn string_strategy(max_chars: usize) -> impl Strategy<Value = String> {
79        prop::collection::vec(any::<char>(), 0..=max_chars)
80            .prop_map(|chars| chars.into_iter().collect())
81    }
82
83    #[test]
84    fn primitives() {
85        assert_eq!(42u64.deep_size(), size_of::<u64>());
86        assert_eq!(true.deep_size(), size_of::<bool>());
87        assert_eq!(0usize.deep_size(), size_of::<usize>());
88        let nz = std::num::NonZeroU64::new(1).unwrap();
89        assert_eq!(nz.deep_size(), size_of::<std::num::NonZeroU64>());
90    }
91
92    #[test]
93    fn option_some() {
94        let o: Option<u64> = Some(42);
95        assert_eq!(o.deep_size(), size_of::<u64>());
96    }
97
98    #[test]
99    fn option_none() {
100        let o: Option<u64> = None;
101        assert_eq!(o.deep_size(), 1);
102    }
103
104    #[test]
105    fn tuple_deep_size() {
106        let t = (42u64, String::from("hi"));
107        assert_eq!(t.deep_size(), size_of::<u64>() + 2);
108    }
109
110    #[test]
111    fn bound_included() {
112        let b = std::ops::Bound::Included(100u64);
113        assert_eq!(b.deep_size(), size_of::<u64>());
114    }
115
116    #[test]
117    fn bound_excluded() {
118        let b = std::ops::Bound::Excluded(100u64);
119        assert_eq!(b.deep_size(), size_of::<u64>());
120    }
121
122    #[test]
123    fn bound_unbounded() {
124        let b: std::ops::Bound<u64> = std::ops::Bound::Unbounded;
125        assert_eq!(b.deep_size(), 1);
126    }
127
128    proptest! {
129        #[test]
130        fn string_deep_size_matches_byte_len(s in string_strategy(256)) {
131            prop_assert_eq!(s.deep_size(), s.len());
132        }
133
134        #[test]
135        fn bytes_deep_size_matches_len(bytes in prop::collection::vec(any::<u8>(), 0..=1024)) {
136            let bytes = bytes::Bytes::from(bytes);
137            prop_assert_eq!(bytes.deep_size(), bytes.len());
138        }
139
140        #[test]
141        fn vec_and_slice_deep_size_sum_elements(values in prop::collection::vec(any::<u64>(), 0..=256)) {
142            let expected = values.len() * size_of::<u64>();
143            prop_assert_eq!(values.deep_size(), expected);
144            prop_assert_eq!(values.as_slice().deep_size(), expected);
145        }
146
147        #[test]
148        fn nested_vec_of_strings_sums_string_lengths(
149            values in prop::collection::vec(string_strategy(64), 0..=32),
150        ) {
151            let expected = values.iter().map(String::len).sum::<usize>();
152            prop_assert_eq!(values.deep_size(), expected);
153        }
154    }
155}