MallocSizeOf for BTreeSet

parity-util-mem/Cargo.toml

@@ -26,6 +26,7 @@ impl-trait-for-tuples = "0.1.3"
 smallvec = { version = "1.0.0", optional = true }
 ethereum-types = { version = "0.8.0", optional = true, path = "../ethereum-types" }
 parking_lot = { version = "0.9.0", optional = true }
+primitive-types = { version = "0.6", path = "../primitive-types", default-features = false, optional = true }
 
 [target.'cfg(target_os = "windows")'.dependencies]
 winapi = { version = "0.3.8", features = ["heapapi"] }
@@ -35,7 +36,7 @@ version = "0.3.2"
 optional = true
 
 [features]
-default = ["std", "ethereum-impls", "lru", "hashbrown", "smallvec"]
+default = ["std", "ethereum-impls", "lru", "hashbrown", "smallvec", "primitive-types"]
 std = ["parking_lot"]
 # use dlmalloc as global allocator
 dlmalloc-global = ["dlmalloc", "estimate-heapsize"]
@@ -46,6 +47,6 @@ jemalloc-global = ["jemallocator"]
 # use mimalloc as global allocator
 mimalloc-global = ["mimallocator", "mimalloc-sys"]
 # implement additional types
-ethereum-impls = ["ethereum-types"]
+ethereum-impls = ["ethereum-types", "primitive-types"]
 # Full estimate: no call to allocator
-estimate-heapsize = []
+estimate-heapsize = []

parity-util-mem/src/ethereum_impls.rs

@@ -17,6 +17,6 @@
 //! Implementation of `MallocSize` for common ethereum types: fixed hashes
 //! and uints.
 
-use ethereum_types::{Bloom, H128, H160, H256, H264, H32, H512, H520, H64, U128, U256, U512, U64};
+use ethereum_types::{Bloom, H128, H264, H32, H520, H64, U64};
 
-malloc_size_of_is_0!(U64, U128, U256, U512, H32, H64, H128, H160, H256, H264, H512, H520, Bloom);
+malloc_size_of_is_0!(U64, H32, H64, H128, H264, H520, Bloom);

parity-util-mem/src/lib.rs

@@ -68,6 +68,9 @@ mod malloc_size;
 #[cfg(feature = "ethereum-impls")]
 pub mod ethereum_impls;
 
+#[cfg(feature = "primitive-types")]
+pub mod primitives_impls;
+
 pub use allocators::MallocSizeOfExt;
 pub use malloc_size::{MallocSizeOf, MallocSizeOfOps};
 

parity-util-mem/src/malloc_size.rs

@@ -455,6 +455,31 @@ where
 	}
 }
 
+impl<T> MallocShallowSizeOf for rstd::collections::BTreeSet<T> {
+	fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+		if ops.has_malloc_enclosing_size_of() {
+			// See implementation for HashSet how this works.
+			self.iter().next().map_or(0, |t| unsafe { ops.malloc_enclosing_size_of(t) })
+		} else {
+			// An estimate.
+			self.len() * (size_of::<T>() + size_of::<usize>())
+		}
+	}
+}
+
+impl<T> MallocSizeOf for rstd::collections::BTreeSet<T>
+where
+	T: MallocSizeOf,
+{
+	fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+		let mut n = self.shallow_size_of(ops);
+		for k in self.iter() {
+			n += k.size_of(ops);
+		}
+		n
+	}
+}
+
 // PhantomData is always 0.
 impl<T> MallocSizeOf for rstd::marker::PhantomData<T> {
 	fn size_of(&self, _ops: &mut MallocSizeOfOps) -> usize {
@@ -676,6 +701,7 @@ malloc_size_of_is_0!(std::time::Duration);
 mod tests {
 	use crate::{allocators::new_malloc_size_ops, MallocSizeOf, MallocSizeOfOps};
 	use smallvec::SmallVec;
+	use std::collections::BTreeSet;
 	use std::mem;
 	impl_smallvec!(3);
 
@@ -727,4 +753,14 @@ mod tests {
 		let expected_min_allocs = mem::size_of::<String>() * 4 + "ÖWL".len() + "COW".len() + "PIG".len() + "DUCK".len();
 		assert!(v.size_of(&mut ops) >= expected_min_allocs);
 	}
+
+	#[test]
+	fn btree_set() {
+		let mut set = BTreeSet::new();
+		for t in 0..100 {
+			set.insert(vec![t]);
+		}
+		// ~36 per value
+		assert!(crate::malloc_size(&set) > 3000);
+	}
 }

parity-util-mem/src/primitives_impls.rs

@@ -0,0 +1,34 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Parity.
+
+// Parity is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Parity is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Parity.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Implementation of `MallocSize` primitive types.
+
+use primitive_types::{H160, H256, H512, U128, U256, U512};
+
+malloc_size_of_is_0!(U128, U256, U512, H160, H256, H512);
+
+#[cfg(test)]
+mod tests {
+
+	use primitive_types::H256;
+
+	#[test]
+	fn smoky() {
+		let v = vec![H256::zero(), H256::zero()];
+
+		assert!(crate::MallocSizeOfExt::malloc_size_of(&v) >= 64);
+	}
+}