trie node ref counter. good for testing.

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<H256, u32>;
+
 	/// Look up a given hash into the bytes that hash to it, returning None if the 
 	/// hash is not known.
 	///

src/memorydb.rs

@@ -116,6 +116,10 @@ impl MemoryDB {
 		}
 		self.data.get(key).unwrap()
 	}
+
+	pub fn raw_keys(&self) -> HashMap<H256, i32> {
+		self.data.iter().filter_map(|(k, v)| if v.1 != 0 {Some((k.clone(), v.1))} else {None}).collect::<HashMap<H256, i32>>()
+	}
 }
 
 impl HashDB for MemoryDB {
@@ -126,6 +130,10 @@ impl HashDB for MemoryDB {
 		}
 	}
 
+	fn keys(&self) -> HashMap<H256, u32> {
+		self.data.iter().filter_map(|(k, v)| if v.1 > 0 {Some((k.clone(), v.1 as u32))} else {None} ).collect::<HashMap<H256, u32>>()
+	}
+
 	fn exists(&self, key: &H256) -> bool {
 		match self.data.get(key) {
 			Some(&(_, x)) if x > 0 => true,

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<H256, u32> {
+		let mut ret: HashMap<H256, u32> = 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.

src/rlp.rs

@@ -125,6 +125,7 @@ impl<'a> From<Rlp<'a>> for UntrustedRlp<'a> {
 	}
 }
 
+#[derive(Debug)]
 pub enum Prototype {
 	Null,
 	Data(usize),

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<H256> {
+		let mut ret: Vec<H256> = 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<H256>) {
+		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<H256>) -> HashMap<H256, u32> {
+		let mut r: HashMap<H256, u32> = 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<H256, u32> {
+		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);