heapsize && squeeze

Cargo.toml

@@ -20,6 +20,7 @@ lazy_static = "0.1.*"
 secp256k1 = "0.5.1"
 rust-crypto = "0.2.34"
 elastic-array = "0.4"
+heapsize = "0.2"
 
 [dev-dependencies]
 json-tests = { path = "json-tests" }

src/heapsizeof.rs

@@ -0,0 +1,5 @@
+use uint::*;
+use hash::*;
+
+known_heap_size!(0, H32, H64, H128, Address, H256, H264, H512, H520, H1024, H2048);
+known_heap_size!(0, U128, U256);

src/lib.rs

@@ -32,10 +32,14 @@ extern crate mio;
 extern crate rand;
 extern crate rocksdb;
 extern crate tiny_keccak;
+
+#[macro_use]
+extern crate heapsize;
 #[macro_use]
 extern crate log;
 #[macro_use]
 extern crate lazy_static;
+
 extern crate env_logger;
 
 extern crate time;
@@ -60,5 +64,7 @@ pub mod crypto;
 pub mod triehash;
 pub mod trie;
 pub mod nibbleslice;
+pub mod heapsizeof;
+pub mod squeeze;
 
 //pub mod network;

src/squeeze.rs

@@ -0,0 +1,33 @@
+use std::collections::HashMap;
+use std::hash::Hash;
+use heapsize::HeapSizeOf;
+
+/// Should be used to squeeze collections to certain size in bytes
+trait Squeeze {
+	fn squeeze(&mut self, size: usize);
+}
+
+impl<K, T> Squeeze for HashMap<K, T> where K: Eq + Hash + Clone + HeapSizeOf, T: HeapSizeOf {
+	fn squeeze(&mut self, size: usize) {
+		if self.len() == 0 {
+			return
+		}
+		
+		let size_of_entry = self.heap_size_of_children() / self.capacity();
+		let mut shrinked_size = size_of_entry * self.len();
+
+		while self.len() > 0 || shrinked_size > size {
+			// could be optimized
+			let key = self.keys().next().unwrap().clone();
+			self.remove(&key);
+			shrinked_size -= size_of_entry;
+		}
+
+		self.shrink_to_fit();
+
+		// if we havent shrinked enough, squeeze again
+		if self.heap_size_of_children() > size {
+			self.squeeze(size);
+		}
+	}
+}