Skip to main content

forest/utils/get_size/
mod.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use cid::Cid;
5use derive_more::{From, Into};
6// re-exports the trait
7pub use get_size2::GetSize;
8use num_bigint::BigInt;
9
10#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, From, Into)]
11pub struct CidWrapper(pub Cid);
12impl GetSize for CidWrapper {}
13
14impl quick_cache::Equivalent<CidWrapper> for Cid {
15    fn equivalent(&self, other: &CidWrapper) -> bool {
16        self == &other.0
17    }
18}
19
20macro_rules! impl_vec_alike_heap_size_with_fn_helper {
21    ($name:ident, $tracker:ident, $t:ty, $get_stack_size: expr, $get_heap_size_with_tracker: expr) => {{
22        let mut heap_size = 0;
23        let mut tr = $tracker;
24        // use `____v` to avoid naming conflict
25        for ____v in $name.iter() {
26            let (_s, _tr) = $get_heap_size_with_tracker(____v, tr);
27            heap_size += $get_stack_size() + _s;
28            tr = _tr;
29        }
30        let additional = usize::from($name.capacity()) - usize::from($name.len());
31        heap_size += additional * $get_stack_size();
32        (heap_size, tr)
33    }};
34}
35
36macro_rules! impl_vec_alike_heap_size_helper {
37    ($name:ident, $tracker:ident, $t:ty) => {
38        impl_vec_alike_heap_size_with_fn_helper!(
39            $name,
40            $tracker,
41            $t,
42            <$t>::get_stack_size,
43            GetSize::get_heap_size_with_tracker
44        )
45    };
46}
47
48pub fn vec_heap_size_with_fn_helper<T, Tr: get_size2::GetSizeTracker>(
49    v: &Vec<T>,
50    tracker: Tr,
51    get_heap_size_with_tracker: impl Fn(&T, Tr) -> (usize, Tr),
52) -> (usize, Tr) {
53    impl_vec_alike_heap_size_with_fn_helper!(
54        v,
55        tracker,
56        T,
57        std::mem::size_of::<T>,
58        get_heap_size_with_tracker
59    )
60}
61
62pub fn nunny_vec_heap_size_helper<T: GetSize, Tr: get_size2::GetSizeTracker>(
63    v: &nunny::Vec<T>,
64    tracker: Tr,
65) -> (usize, Tr) {
66    impl_vec_alike_heap_size_helper!(v, tracker, T)
67}
68
69// This is a rough estimation. Use `b.allocation_size()`
70// once https://github.com/rust-num/num-bigint/pull/333 is accepted and released.
71pub fn big_int_heap_size_helper(b: &BigInt) -> usize {
72    b.bits().div_ceil(8) as usize
73}
74
75pub fn raw_bytes_heap_size_helper(b: &fvm_ipld_encoding::RawBytes) -> usize {
76    // Note: this is a cheap but inaccurate estimation,
77    // the correct implementation should be `Vec<u8>.from(b.clone()).get_heap_size()`,
78    // or `b.bytes.get_heap_size()` if `bytes` is made public.
79    b.bytes().get_heap_size()
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use crate::utils::multihash::MultihashCode;
86    use fvm_ipld_encoding::DAG_CBOR;
87    use multihash_derive::MultihashDigest as _;
88
89    #[test]
90    fn test_cid() {
91        let cid = Cid::new_v1(DAG_CBOR, MultihashCode::Blake2b256.digest(&[0, 1, 2, 3]));
92        let wrapper = CidWrapper(cid);
93        assert_eq!(std::mem::size_of_val(&cid), wrapper.get_size());
94    }
95
96    #[test]
97    fn test_heap_size_helper() {
98        let keys: nunny::Vec<CidWrapper> = nunny::vec![Cid::default().into(); 3];
99        // It's likely > 3 (4 on my laptop)
100        println!("keys.capacity() = {}", keys.capacity());
101        let tracker = get_size2::StandardTracker::new();
102        assert_eq!(
103            nunny_vec_heap_size_helper(&keys, tracker).0,
104            CidWrapper::get_stack_size() * usize::from(keys.capacity())
105        );
106    }
107
108    #[test]
109    fn test_derive_macro() {
110        #[derive(GetSize)]
111        struct A {
112            #[get_size(ignore)]
113            _cid: Cid,
114        }
115
116        #[derive(GetSize)]
117        struct B {
118            #[get_size(size = 0)]
119            _cid: Cid,
120        }
121
122        #[derive(GetSize)]
123        struct C {
124            #[get_size(size = 8)]
125            _cid: Cid,
126        }
127
128        let _cid = Cid::default();
129        let a = vec![A { _cid }];
130        assert_eq!(
131            a.get_heap_size(),
132            std::mem::size_of_val(&_cid) * a.capacity()
133        );
134
135        let b = vec![B { _cid }];
136        assert_eq!(
137            b.get_heap_size(),
138            std::mem::size_of_val(&_cid) * b.capacity()
139        );
140
141        let c = vec![C { _cid }];
142        assert_eq!(
143            c.get_heap_size(),
144            (std::mem::size_of_val(&_cid) + 8) * c.capacity()
145        );
146    }
147}