diff --git a/parity-util-mem/Cargo.toml b/parity-util-mem/Cargo.toml index 10a90c3cc..e15e460c1 100644 --- a/parity-util-mem/Cargo.toml +++ b/parity-util-mem/Cargo.toml @@ -11,6 +11,8 @@ edition = "2018" cfg-if = "0.1.10" dlmalloc = { version = "0.1.3", features = ["global"], optional = true } wee_alloc = { version = "0.4.5", optional = true } +lru = { version = "0.4", optional = true } +hashbrown = { version = "0.6", optional = true } # from https://github.com/microsoft/mimalloc: # mimalloc can be built in secure mode, # adding guard pages, randomized allocation, encrypted free lists, etc. @@ -32,7 +34,7 @@ version = "0.3.2" optional = true [features] -default = ["std", "ethereum-impls"] +default = ["std", "ethereum-impls", "lru", "hashbrown"] std = ["parking_lot"] # use dlmalloc as global allocator dlmalloc-global = ["dlmalloc", "estimate-heapsize"] diff --git a/parity-util-mem/src/malloc_size.rs b/parity-util-mem/src/malloc_size.rs index 49f9c9bce..3c26a917a 100644 --- a/parity-util-mem/src/malloc_size.rs +++ b/parity-util-mem/src/malloc_size.rs @@ -68,6 +68,7 @@ use std::sync::Arc; pub use alloc::boxed::Box; #[cfg(not(feature = "std"))] use core::ffi::c_void; +#[cfg(feature = "std")] use rstd::hash::Hash; use rstd::mem::size_of; use rstd::ops::Range; @@ -623,3 +624,48 @@ impl DerefMut for Measurable { &mut self.0 } } + +#[cfg(feature = "hashbrown")] +impl MallocShallowSizeOf for hashbrown::HashMap { + fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize { + // See the implementation for std::collections::HashSet for details. + if ops.has_malloc_enclosing_size_of() { + self.values().next().map_or(0, |v| unsafe { ops.malloc_enclosing_size_of(v) }) + } else { + self.capacity() * (size_of::() + size_of::() + size_of::()) + } + } +} + +#[cfg(feature = "hashbrown")] +impl MallocSizeOf for hashbrown::HashMap +where + K: MallocSizeOf, + V: MallocSizeOf, +{ + fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize { + let mut n = self.shallow_size_of(ops); + for (k, v) in self.iter() { + n += k.size_of(ops); + n += v.size_of(ops); + } + n + } +} + +#[cfg(feature = "lru")] +impl MallocSizeOf for lru::LruCache +where + K: MallocSizeOf + rstd::cmp::Eq + rstd::hash::Hash, + V: MallocSizeOf, + S: rstd::hash::BuildHasher, +{ + fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize { + let mut n = 0; + for (k, v) in self.iter() { + n += k.size_of(ops); + n += v.size_of(ops); + } + n + } +} diff --git a/parity-util-mem/tests/derive.rs b/parity-util-mem/tests/derive.rs index 6338e2cc8..34674fd5f 100644 --- a/parity-util-mem/tests/derive.rs +++ b/parity-util-mem/tests/derive.rs @@ -14,10 +14,11 @@ // You should have received a copy of the GNU General Public License // along with Parity. If not, see . +#![cfg(feature = "std")] + use parity_util_mem::{MallocSizeOf, MallocSizeOfExt}; #[test] -#[cfg(feature = "std")] fn derive_vec() { #[derive(MallocSizeOf)] struct Trivia { @@ -29,7 +30,6 @@ fn derive_vec() { assert!(t.malloc_size_of() > 1000); } -#[cfg(feature = "std")] #[test] fn derive_hashmap() { #[derive(MallocSizeOf, Default)] @@ -45,7 +45,6 @@ fn derive_hashmap() { } #[test] -#[cfg(feature = "std")] fn derive_ignore() { #[derive(MallocSizeOf, Default)] struct Trivia { @@ -60,3 +59,20 @@ fn derive_ignore() { t.v = vec![0u8; 1024]; assert!(t.malloc_size_of() < 3000); } + +#[test] +fn derive_morecomplex() { + #[derive(MallocSizeOf)] + struct Trivia { + hm: hashbrown::HashMap>, + cache: lru::LruCache>, + } + + let mut t = Trivia { hm: hashbrown::HashMap::new(), cache: lru::LruCache::unbounded() }; + + t.hm.insert(1, vec![0u8; 2048]); + t.cache.put(1, vec![0u8; 2048]); + t.cache.put(2, vec![0u8; 4096]); + + assert!(t.malloc_size_of() > 8000); +}