From 0cec68b59e99e981c4e72c596fcc6a7ab42be39a Mon Sep 17 00:00:00 2001 From: debris Date: Tue, 1 Dec 2015 23:14:38 +0100 Subject: [PATCH] fixed trie get and trie print --- src/trie.rs | 75 +++++++++++++++++++++-------------------------------- 1 file changed, 30 insertions(+), 45 deletions(-) diff --git a/src/trie.rs b/src/trie.rs index 86c2e84cd..597dbd085 100644 --- a/src/trie.rs +++ b/src/trie.rs @@ -25,8 +25,7 @@ pub trait Trie { #[derive(Eq, PartialEq, Debug)] pub enum Node<'a> { Leaf(NibbleSlice<'a>, &'a[u8]), - ExtensionRaw(NibbleSlice<'a>, &'a[u8]), - ExtensionSha3(NibbleSlice<'a>, &'a[u8]), + Extension(NibbleSlice<'a>, &'a[u8]), Branch(Option<[&'a[u8]; 16]>, &'a [u8]) } @@ -41,13 +40,7 @@ impl <'a>Node<'a> { // fed back into this function or inline RLP which can be fed back into this function). Prototype::List(2) => match NibbleSlice::from_encoded(r.at(0).data()) { (slice, true) => Node::Leaf(slice, r.at(1).data()), - (slice, false) => match r.at(1).raw().len() { - // its just raw extension - 0...31 => Node::ExtensionRaw(slice, r.at(1).raw()), - // its SHA3 of raw extension + the length of sha3 - 33 => Node::ExtensionSha3(slice, r.at(1).data()), - _ => { panic!(); } - } + (slice, false) => Node::Extension(slice, r.at(1).raw()) }, // branch - first 16 are nodes, 17th is a value (or empty). Prototype::List(17) => { @@ -62,6 +55,7 @@ impl <'a>Node<'a> { } } + // todo: should check length before encoding, cause it may just be sha3 of data pub fn encoded(&self) -> Bytes { match *self { Node::Leaf(ref slice, ref value) => { @@ -70,18 +64,12 @@ impl <'a>Node<'a> { stream.append(value); stream.out() }, - Node::ExtensionRaw(ref slice, ref raw_rlp) => { + Node::Extension(ref slice, ref raw_rlp) => { let mut stream = RlpStream::new_list(2); stream.append(&slice.encoded(false)); stream.append_raw(raw_rlp, 1); stream.out() }, - Node::ExtensionSha3(ref slice, ref sha3) => { - let mut stream = RlpStream::new_list(2); - stream.append(&slice.encoded(false)); - stream.append(sha3); - stream.out() - }, Node::Branch(Some(ref nodes), ref value) => { let mut stream = RlpStream::new_list(17); for i in 0..16 { stream.append_raw(nodes[i], 1); } @@ -200,27 +188,21 @@ impl TrieDB { let node = Node::decoded(node); match node { Node::Leaf(slice, value) => try!(writeln!(f, "Leaf {:?}, {:?}", slice, value.pretty())), - Node::ExtensionRaw(ref slice, ref item) => { - try!(write!(f, "Extension (raw): {:?} ", slice)); - try!(self.fmt_all(item, f, deepness + 1)); - }, - Node::ExtensionSha3(ref slice, sha3) => { - try!(write!(f, "Extension (sha3): {:?} ", slice)); - let rlp = self.db.lookup(&H256::from_slice(sha3)).expect("sha3 not found!"); - try!(self.fmt_all(rlp, f, deepness + 1)); + Node::Extension(ref slice, ref item) => { + try!(write!(f, "Extension: {:?} ,", slice)); + try!(self.fmt_all(self.get_raw_or_lookup(item), f, deepness + 1)); }, Node::Branch(Some(ref nodes), ref value) => { try!(writeln!(f, "Branch: ")); try!(self.fmt_indent(f, deepness + 1)); try!(writeln!(f, ": {:?}", value.pretty())); for i in 0..16 { - match Node::decoded(nodes[i]) { - Node::Branch(None, _) => (), - _ => { - try!(self.fmt_indent(f, deepness + 1)); - try!(write!(f, "{:x}: ", i)); - try!(self.fmt_all(nodes[i], f, deepness + 1)); - } + let rlp = Rlp::new(nodes[i]); + if rlp.is_data() && rlp.size() != 0 { + try!(self.fmt_indent(f, deepness + 1)); + try!(write!(f, "{:x}: ", i)); + try!(self.fmt_all(self.get_raw_or_lookup(nodes[i]), f, deepness + 1)); + try!(writeln!(f, "")); } } }, @@ -240,22 +222,25 @@ impl TrieDB { fn get_from_node<'a>(&'a self, node: &'a [u8], key: &NibbleSlice<'a>) -> Option<&'a [u8]> { match Node::decoded(node) { Node::Leaf(ref slice, ref value) if key == slice => Some(value), - Node::ExtensionRaw(ref slice, ref item) if key.starts_with(slice) => { - self.get_from_node(item, &key.mid(slice.len())) - }, - Node::ExtensionSha3(ref slice, ref sha3) => { - // lookup for this item - let rlp = self.db.lookup(&H256::from_slice(sha3)).expect("sha3 not found!"); - self.get_from_node(rlp, &key.mid(slice.len())) + 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(Some(ref nodes), ref value) => match key.is_empty() { true => Some(value), - false => self.get_from_node(nodes[key.at(0) as usize], &key.mid(1)) + false => self.get_from_node(self.get_raw_or_lookup(nodes[key.at(0) as usize]), &key.mid(1)) }, _ => None } } + fn get_raw_or_lookup<'a>(&'a self, node: &'a [u8]) -> &'a [u8] { + // check if its sha3 + len + match node.len() > 32 { + true => self.db.lookup(&H256::from_slice(&node[1..])).expect("Not found!"), + false => node + } + } + fn add(&mut self, key: &NibbleSlice, value: &[u8]) { trace!("ADD: {:?} {:?}", key, value.pretty()); // determine what the new root is, insert new nodes and remove old as necessary. @@ -524,7 +509,7 @@ mod tests { let v = rlp::encode(&"cat"); let (slice, is_leaf) = NibbleSlice::from_encoded(&k); assert_eq!(is_leaf, false); - let ex = Node::ExtensionRaw(slice, &v); + let ex = Node::Extension(slice, &v); let rlp = ex.encoded(); let ex2 = Node::decoded(&rlp); assert_eq!(ex, ex2); @@ -604,15 +589,15 @@ mod tests { t.insert(&key, &val); } -// trace!("{:?}", t); + println!("{:?}", t); - assert_eq!(*t.root(), trie_root(v)); - - /*for i in 0..v.len() { + for i in 0..v.len() { let key: &[u8]= &v[i].0; let val: &[u8] = &v[i].1; assert_eq!(t.at(&key).unwrap(), val); - }*/ + } + + assert_eq!(*t.root(), trie_root(v)); } #[test]