memuse 0.2.2

Traits for measuring dynamic memory usage of types
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! `DynamicUsage` impls for `HashMap` and `HashSet`.
//!
//! Rust's `HashMap` and `HashSet` are backed by the the `hashbrown` crate.

use alloc::alloc::Layout;
use core::mem;

#[cfg(feature = "std")]
use std::collections::{HashMap, HashSet};

use crate::DynamicUsage;

// The widths are sourced from here:
//   https://github.com/rust-lang/hashbrown/blob/dbd6dbe30a4076c0ea65ca5bd57036c27f3cc7c9/src/raw/mod.rs#L15-L36
//   https://github.com/rust-lang/hashbrown/blob/dbd6dbe30a4076c0ea65ca5bd57036c27f3cc7c9/src/raw/generic.rs#L5-L21
//   https://github.com/rust-lang/hashbrown/blob/dbd6dbe30a4076c0ea65ca5bd57036c27f3cc7c9/src/raw/sse2.rs#L14-L19

#[cfg(all(target_feature = "sse2", target_arch = "x86", not(miri)))]
pub(crate) const WIDTH: usize = mem::size_of::<core::arch::x86::__m128i>();
#[cfg(all(target_feature = "sse2", target_arch = "x86_64", not(miri)))]
pub(crate) const WIDTH: usize = mem::size_of::<core::arch::x86_64::__m128i>();
#[cfg(all(
    any(not(target_feature = "sse2"), miri),
    any(
        target_pointer_width = "64",
        target_arch = "aarch64",
        target_arch = "x86_64",
        target_arch = "wasm32",
    )
))]
pub(crate) const WIDTH: usize = mem::size_of::<u64>();
#[cfg(all(
    any(not(target_feature = "sse2"), miri),
    target_pointer_width = "32",
    not(target_arch = "aarch64"),
    not(target_arch = "x86_64"),
    not(target_arch = "wasm32"),
))]
pub(crate) const WIDTH: usize = mem::size_of::<u32>();

fn dynamic_usage_for_capacity<K, V>(cap: usize) -> usize {
    // The bucket calculation is sourced from here:
    //   https://github.com/rust-lang/hashbrown/blob/dbd6dbe30a4076c0ea65ca5bd57036c27f3cc7c9/src/raw/mod.rs#L187-L216
    //
    // hashbrown's RawTable::buckets is not accessible via the std HashMap, and
    // HashMap::capacity is a lower bound. However, hashbrown has an invariant that
    // the number of buckets is a power of two, so usually we'll calculate the correct
    // memory usage, and occasionally we'll undercount by around a factor of two.
    let buckets = {
        if cap < 8 {
            return if cap < 4 { 4 } else { 8 };
        }
        let adjusted_cap = (cap * 8) / 7;
        adjusted_cap.next_power_of_two()
    };

    // The memory usage calculation is sourced from here:
    //   https://github.com/rust-lang/hashbrown/blob/dbd6dbe30a4076c0ea65ca5bd57036c27f3cc7c9/src/raw/mod.rs#L240-L265

    let layout = Layout::new::<(K, V)>();
    let size = layout.size();
    let ctrl_align = usize::max(layout.align(), WIDTH);

    let ctrl_offset = (size * buckets + ctrl_align - 1) & !(ctrl_align - 1);
    ctrl_offset + buckets + WIDTH
}

impl<K: DynamicUsage, V: DynamicUsage> DynamicUsage for HashMap<K, V> {
    fn dynamic_usage(&self) -> usize {
        dynamic_usage_for_capacity::<K, V>(self.capacity())
            + self
                .iter()
                .map(|(k, v)| k.dynamic_usage() + v.dynamic_usage())
                .sum::<usize>()
    }

    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
        (
            dynamic_usage_for_capacity::<K, V>(self.capacity())
                + self
                    .iter()
                    .map(|(k, v)| k.dynamic_usage_bounds().0 + v.dynamic_usage_bounds().0)
                    .sum::<usize>(),
            None,
        )
    }
}

impl<T: DynamicUsage> DynamicUsage for HashSet<T> {
    fn dynamic_usage(&self) -> usize {
        // HashSet<T> is just HashMap<T, ()>
        dynamic_usage_for_capacity::<T, ()>(self.capacity())
            + self.iter().map(DynamicUsage::dynamic_usage).sum::<usize>()
    }

    fn dynamic_usage_bounds(&self) -> (usize, Option<usize>) {
        (
            dynamic_usage_for_capacity::<T, ()>(self.capacity())
                + self
                    .iter()
                    .map(|k| k.dynamic_usage_bounds().0)
                    .sum::<usize>(),
            None,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn hashmap() {
        let h: HashMap<u16, u32> = HashMap::with_capacity(12);

        // - Capacity of 12 -> 16 buckets
        // - Overhead is 1 byte per bucket
        // - Fixed overhead of WIDTH
        let lower = 16 * (mem::size_of::<(u16, u32)>() + 1) + WIDTH;
        assert_eq!(h.dynamic_usage(), lower);
        assert_eq!(h.dynamic_usage_bounds(), (lower, None));
    }

    #[test]
    fn hashset() {
        let h: HashSet<u16> = HashSet::with_capacity(17);

        // - Capacity of 17 -> 32 buckets
        // - Overhead is 1 byte per bucket
        // - Fixed overhead of WIDTH
        let lower = 32 * (mem::size_of::<u16>() + 1) + WIDTH;
        assert_eq!(h.dynamic_usage(), lower);
        assert_eq!(h.dynamic_usage_bounds(), (lower, None));
    }
}