use std::fmt; use hashdb::*; use hash::*; use nibbleslice::*; use bytes::*; use rlp::*; use std::collections::HashMap; use super::trietraits::*; use super::node::*; /// A `Trie` implementation using a generic `HashDB` backing database. /// /// Use it as a `Trie` trait object. You can use `db()` to get the backing database object, `keys` /// to get the keys belonging to the trie in the backing database, and `db_items_remaining()` to get /// which items in the backing database do not belong to this trie. If this is the only trie in the /// backing database, then `db_items_remaining()` should be empty. /// /// # Example /// ``` /// extern crate ethcore_util as util; /// use util::trie::*; /// use util::hashdb::*; /// use util::memorydb::*; /// use util::hash::*; /// use util::rlp::*; /// /// fn main() { /// let mut memdb = MemoryDB::new(); /// let mut root = H256::new(); /// TrieDBMut::new(&mut memdb, &mut root).insert(b"foo", b"bar"); /// let t = TrieDB::new(&memdb, &root); /// assert!(t.contains(b"foo")); /// assert_eq!(t.get(b"foo").unwrap(), b"bar"); /// assert!(t.db_items_remaining().is_empty()); /// } /// ``` pub struct TrieDB<'db> { db: &'db HashDB, root: &'db H256, pub hash_count: usize, } impl<'db> TrieDB<'db> { /// Create a new trie with the backing database `db` and `root` /// Panics, if `root` does not exist pub fn new(db: &'db HashDB, root: &'db H256) -> Self { assert!(db.exists(root)); TrieDB { db: db, root: root, hash_count: 0 } } /// Get the backing database. pub fn db(&'db self) -> &'db HashDB { self.db } /// Determine all the keys in the backing database that belong to the trie. 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 } /// Convert a vector of hashes to a hashmap of hash to occurances. pub 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 } /// Determine occurances of items in the backing database which are not related to this /// trie. 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 keycount = *ret.get(&k).unwrap_or(&0); match keycount <= v as i32 { true => ret.remove(&k), _ => ret.insert(k, keycount - v as i32), }; } ret } /// Recursion helper for `keys`. 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(p.as_val()); } 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) }, _ => {}, } } /// Get the root node's RLP. fn root_node(&self) -> Node { Node::decoded(self.db.lookup(&self.root).expect("Trie root not found!")) } /// Get the root node as a `Node`. fn get_node<'a>(&'a self, node: &'a [u8]) -> Node { Node::decoded(self.get_raw_or_lookup(node)) } /// Indentation helper for `formal_all`. fn fmt_indent(&self, f: &mut fmt::Formatter, size: usize) -> fmt::Result { for _ in 0..size { try!(write!(f, " ")); } Ok(()) } /// Recursion helper for implementation of formatting trait. fn fmt_all(&self, node: Node, f: &mut fmt::Formatter, deepness: usize) -> fmt::Result { match node { Node::Leaf(slice, value) => try!(writeln!(f, "'{:?}: {:?}.", slice, value.pretty())), Node::Extension(ref slice, ref item) => { try!(write!(f, "'{:?} ", slice)); try!(self.fmt_all(self.get_node(item), f, deepness)); }, Node::Branch(ref nodes, ref value) => { try!(writeln!(f, "")); match value { &Some(v) => { try!(self.fmt_indent(f, deepness + 1)); try!(writeln!(f, "=: {:?}", v.pretty())) }, &None => {} } for i in 0..16 { match self.get_node(nodes[i]) { Node::Empty => {}, n => { try!(self.fmt_indent(f, deepness + 1)); try!(write!(f, "'{:x} ", i)); try!(self.fmt_all(n, f, deepness + 1)); } } } }, // empty Node::Empty => { try!(writeln!(f, "")); } }; Ok(()) } /// Return optional data for a key given as a `NibbleSlice`. Returns `None` if no data exists. fn do_lookup<'a, 'key>(&'a self, key: &NibbleSlice<'key>) -> Option<&'a [u8]> where 'a: 'key { let root_rlp = self.db.lookup(&self.root).expect("Trie root not found!"); self.get_from_node(&root_rlp, key) } /// Recursible function to retrieve the value given a `node` and a partial `key`. `None` if no /// value exists for the key. /// /// Note: Not a public API; use Trie trait functions. fn get_from_node<'a, 'key>(&'a self, node: &'a [u8], key: &NibbleSlice<'key>) -> Option<&'a [u8]> where 'a: 'key { match Node::decoded(node) { Node::Leaf(ref slice, ref value) if key == slice => Some(value), Node::Extension(ref slice, ref item) if key.starts_with(slice) => { self.get_from_node(self.get_raw_or_lookup(item), &key.mid(slice.len())) }, Node::Branch(ref nodes, value) => match key.is_empty() { true => value, false => self.get_from_node(self.get_raw_or_lookup(nodes[key.at(0) as usize]), &key.mid(1)) }, _ => None } } /// Given some node-describing data `node`, return the actual node RLP. /// This could be a simple identity operation in the case that the node is sufficiently small, but /// may require a database lookup. fn get_raw_or_lookup<'a>(&'a self, node: &'a [u8]) -> &'a [u8] { // check if its sha3 + len let r = Rlp::new(node); match r.is_data() && r.size() == 32 { true => self.db.lookup(&r.as_val::()).expect("Not found!"), false => node } } } impl<'db> Trie for TrieDB<'db> { fn root(&self) -> &H256 { &self.root } fn contains(&self, key: &[u8]) -> bool { self.get(key).is_some() } fn get<'a, 'key>(&'a self, key: &'key [u8]) -> Option<&'a [u8]> where 'a: 'key { self.do_lookup(&NibbleSlice::new(key)) } } impl<'db> fmt::Debug for TrieDB<'db> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { try!(writeln!(f, "c={:?} [", self.hash_count)); let root_rlp = self.db.lookup(&self.root).expect("Trie root not found!"); try!(self.fmt_all(Node::decoded(root_rlp), f, 0)); writeln!(f, "]") } }