diff --git a/src/hashdb.rs b/src/hashdb.rs index dd0973bb3..f893c8df8 100644 --- a/src/hashdb.rs +++ b/src/hashdb.rs @@ -1,7 +1,11 @@ use hash::*; use bytes::*; +use std::collections::HashMap; pub trait HashDB { + /// Get the keys in the database together with number of underlying references. + fn keys(&self) -> HashMap; + /// Look up a given hash into the bytes that hash to it, returning None if the /// hash is not known. /// diff --git a/src/memorydb.rs b/src/memorydb.rs index 85fdc7c19..8c7eaff2c 100644 --- a/src/memorydb.rs +++ b/src/memorydb.rs @@ -116,6 +116,10 @@ impl MemoryDB { } self.data.get(key).unwrap() } + + pub fn raw_keys(&self) -> HashMap { + self.data.iter().filter_map(|(k, v)| if v.1 != 0 {Some((k.clone(), v.1))} else {None}).collect::>() + } } impl HashDB for MemoryDB { @@ -126,6 +130,10 @@ impl HashDB for MemoryDB { } } + fn keys(&self) -> HashMap { + self.data.iter().filter_map(|(k, v)| if v.1 > 0 {Some((k.clone(), v.1 as u32))} else {None} ).collect::>() + } + fn exists(&self, key: &H256) -> bool { match self.data.get(key) { Some(&(_, x)) if x > 0 => true, diff --git a/src/overlaydb.rs b/src/overlaydb.rs index d69afa0f1..78ca67d01 100644 --- a/src/overlaydb.rs +++ b/src/overlaydb.rs @@ -9,6 +9,7 @@ use memorydb::*; use std::ops::*; use std::sync::*; use std::env; +use std::collections::HashMap; use rocksdb::{DB, Writable}; #[derive(Clone)] @@ -135,6 +136,20 @@ impl OverlayDB { } impl HashDB for OverlayDB { + fn keys(&self) -> HashMap { + let mut ret: HashMap = HashMap::new(); + for (key, _) in self.backing.iterator().from_start() { + let h = H256::from_slice(key.deref()); + let r = self.payload(&h).unwrap().1; + ret.insert(h, r); + } + + for (key, refs) in self.overlay.raw_keys().into_iter() { + let refs = *ret.get(&key).unwrap_or(&0u32) as i32 + refs as i32; + ret.insert(key, refs as u32); + } + ret + } fn lookup(&self, key: &H256) -> Option<&[u8]> { // return ok if positive; if negative, check backing - might be enough references there to make // it positive again. diff --git a/src/rlp.rs b/src/rlp.rs index b055632b1..b875e61c5 100644 --- a/src/rlp.rs +++ b/src/rlp.rs @@ -125,6 +125,7 @@ impl<'a> From> for UntrustedRlp<'a> { } } +#[derive(Debug)] pub enum Prototype { Null, Data(usize), diff --git a/src/trie.rs b/src/trie.rs index 081fcfcff..3d825c707 100644 --- a/src/trie.rs +++ b/src/trie.rs @@ -8,6 +8,7 @@ use hash::*; use nibbleslice::*; use bytes::*; use rlp::*; +use std::collections::HashMap; //use log::*; @@ -238,6 +239,52 @@ impl TrieDB { } } + pub fn keys(&self) -> Vec { + let mut ret: Vec = Vec::new(); + ret.push(self.root.clone()); + self.accumulate_keys(self.root_node(), &mut ret); + ret + } + + fn accumulate_keys(&self, node: Node, acc: &mut Vec) { + let mut handle_payload = |payload| { + let p = Rlp::new(payload); + if p.is_data() && p.size() == 32 { + acc.push(H256::decode(&p)); + } + + self.accumulate_keys(self.get_node(payload), acc); + }; + + match node { + Node::Extension(_, payload) => handle_payload(payload), + Node::Branch(payloads, _) => for payload in payloads.iter() { handle_payload(payload) }, + _ => {}, + } + } + + fn to_map(hashes: Vec) -> HashMap { + let mut r: HashMap = HashMap::new(); + for h in hashes.into_iter() { + let c = *r.get(&h).unwrap_or(&0); + r.insert(h, c + 1); + } + r + } + + pub fn db_items_remaining(&self) -> HashMap { + let mut ret = self.db().keys(); + for (k, v) in Self::to_map(self.keys()).into_iter() { + let old = *ret.get(&k).expect("Node in trie is not in database!"); + assert!(old >= v); + match old > v { + true => ret.insert(k, old - v), + _ => ret.remove(&k), + }; + } + ret + } + fn fmt_indent(&self, f: &mut fmt::Formatter, size: usize) -> fmt::Result { for _ in 0..size { try!(write!(f, " ")); @@ -245,6 +292,10 @@ impl TrieDB { Ok(()) } + fn root_node(&self) -> Node { + Node::decoded(self.db.lookup(&self.root).expect("Trie root not found!")) + } + fn get_node<'a>(&'a self, node: &'a [u8]) -> Node { Node::decoded(self.get_raw_or_lookup(node)) } @@ -762,11 +813,14 @@ mod tests { let mut t1 = TrieDB::new_memory(); t1.insert(&[0x01, 0x23], &big_value.to_vec()); t1.insert(&[0x01, 0x34], &big_value.to_vec()); + trace!("keys remaining {:?}", t1.db_items_remaining()); + assert!(t1.db_items_remaining().is_empty()); let mut t2 = TrieDB::new_memory(); t2.insert(&[0x01], &big_value.to_vec()); t2.insert(&[0x01, 0x23], &big_value.to_vec()); t2.insert(&[0x01, 0x34], &big_value.to_vec()); t2.remove(&[0x01]); + assert!(t2.db_items_remaining().is_empty()); /*if t1.root() != t2.root()*/ { trace!("{:?}", t1); trace!("{:?}", t2);