diff --git a/Cargo.toml b/Cargo.toml index 13295f766..4d23f49f8 100644 --- a/Cargo.toml +++ b/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" } diff --git a/src/heapsizeof.rs b/src/heapsizeof.rs new file mode 100644 index 000000000..c6d4cace4 --- /dev/null +++ b/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); diff --git a/src/lib.rs b/src/lib.rs index 24b404a88..9120f0977 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/squeeze.rs b/src/squeeze.rs new file mode 100644 index 000000000..220c61800 --- /dev/null +++ b/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 Squeeze for HashMap 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); + } + } +}