Add memory extensions for LRUCache, hashbrown

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"]

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<T: MallocSizeOf> DerefMut for Measurable<T> {
 		&mut self.0
 	}
 }
+
+#[cfg(feature = "hashbrown")]
+impl<K, V, S> MallocShallowSizeOf for hashbrown::HashMap<K, V, S> {
+	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::<V>() + size_of::<K>() + size_of::<usize>())
+		}
+	}
+}
+
+#[cfg(feature = "hashbrown")]
+impl<K, V, S> MallocSizeOf for hashbrown::HashMap<K, V, S>
+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<K, V, S> MallocSizeOf for lru::LruCache<K, V, S>
+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
+	}
+}

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 <http://www.gnu.org/licenses/>.
 
+#![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<u64, Vec<u8>>,
+		cache: lru::LruCache<u128, Vec<u8>>,
+	}
+
+	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);
+}