allocative/
size_of.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under both the MIT license found in the
5 * LICENSE-MIT file in the root directory of this source tree and the Apache
6 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
7 * of this source tree.
8 */
9
10use crate::visitor::NodeKind;
11use crate::visitor::VisitorImpl;
12use crate::Allocative;
13use crate::Key;
14use crate::Visitor;
15
16/// Size of data allocated in unique pointers in the struct.
17///
18/// * Exclude self
19/// * Exclude shared pointers
20/// * For unique pointers, include the size of the pointee plus this function recursively
21///
22/// # Example
23///
24/// ```
25/// use allocative::Allocative;
26///
27/// #[derive(Allocative)]
28/// struct Foo {
29///     data: Vec<u8>,
30/// }
31///
32/// assert_eq!(
33///     3,
34///     allocative::size_of_unique_allocated_data(&Foo {
35///         data: vec![10, 20, 30]
36///     })
37/// );
38/// ```
39pub fn size_of_unique_allocated_data(root: &dyn Allocative) -> usize {
40    struct SizeOfUniqueAllocatedDataVisitor {
41        /// Size we return.
42        size: usize,
43    }
44
45    impl VisitorImpl for SizeOfUniqueAllocatedDataVisitor {
46        fn enter_inline_impl(&mut self, _name: Key, size: usize, parent: NodeKind) {
47            if let NodeKind::Unique = parent {
48                self.size += size;
49            }
50        }
51
52        fn enter_unique_impl(&mut self, _name: Key, _size: usize, _parent: NodeKind) {}
53
54        fn enter_shared_impl(
55            &mut self,
56            _name: Key,
57            _size: usize,
58            _ptr: *const (),
59            _parent: NodeKind,
60        ) -> bool {
61            false
62        }
63
64        fn exit_inline_impl(&mut self) {}
65
66        fn exit_unique_impl(&mut self) {}
67
68        fn exit_shared_impl(&mut self) {
69            unreachable!("shared pointers are not visited")
70        }
71
72        fn exit_root_impl(&mut self) {}
73    }
74
75    let mut visitor_impl = SizeOfUniqueAllocatedDataVisitor { size: 0 };
76    let mut visitor = Visitor {
77        visitor: &mut visitor_impl,
78        node_kind: NodeKind::Root,
79    };
80    root.visit(&mut visitor);
81    visitor.exit();
82    visitor_impl.size
83}
84
85/// Size of a piece of data and data allocated in unique pointers in the struct.
86///
87/// * Excludes shared pointers
88///
89/// # Example
90///
91/// ```
92/// use allocative::Allocative;
93///
94/// #[derive(Allocative)]
95/// struct Foo {
96///     data: Vec<u8>,
97/// }
98///
99/// assert_eq!(
100///     3 + std::mem::size_of::<Vec<u8>>(),
101///     allocative::size_of_unique(&Foo {
102///         data: vec![10, 20, 30]
103///     })
104/// );
105/// ```
106pub fn size_of_unique<T>(root: &T) -> usize
107where
108    T: Allocative,
109{
110    std::mem::size_of::<T>() + size_of_unique_allocated_data(root)
111}
112
113#[cfg(test)]
114mod tests {
115    use std::mem;
116
117    use allocative_derive::Allocative;
118
119    use crate as allocative;
120    use crate::size_of_unique;
121    use crate::size_of_unique_allocated_data;
122
123    #[test]
124    fn test_box() {
125        #[derive(Allocative)]
126        struct Boxed {
127            data: Box<u32>,
128        }
129
130        let boxed = Boxed { data: Box::new(17) };
131
132        assert_eq!(mem::size_of::<u32>(), size_of_unique_allocated_data(&boxed));
133
134        assert_eq!(
135            mem::size_of::<u32>() + mem::size_of::<Boxed>(),
136            size_of_unique(&boxed)
137        );
138    }
139
140    #[test]
141    fn test_box_slice() {
142        #[derive(Allocative)]
143        struct Boxed {
144            data: Box<[u32]>,
145        }
146
147        let boxed = Boxed {
148            data: vec![1, 2, 3].into_boxed_slice(),
149        };
150
151        assert_eq!(
152            mem::size_of::<u32>() * 3,
153            size_of_unique_allocated_data(&boxed)
154        );
155
156        assert_eq!(
157            mem::size_of::<Boxed>() + mem::size_of::<u32>() * 3,
158            size_of_unique(&boxed)
159        );
160    }
161
162    #[test]
163    fn test_struct_in_box() {
164        #[derive(Allocative)]
165        struct Data {
166            a: u8,
167            b: Box<u32>,
168        }
169
170        #[derive(Allocative)]
171        struct Boxed {
172            data: Box<Data>,
173        }
174
175        let boxed = Boxed {
176            data: Box::new(Data {
177                a: 1,
178                b: Box::new(2),
179            }),
180        };
181
182        assert_eq!(
183            mem::size_of::<Data>() + mem::size_of::<u32>(),
184            size_of_unique_allocated_data(&boxed)
185        );
186
187        assert_eq!(
188            mem::size_of::<Boxed>() + mem::size_of::<Data>() + mem::size_of::<u32>(),
189            size_of_unique(&boxed)
190        );
191    }
192}
allocative/size_of.rs

allocative/
size_of.rs
