diff --git a/src/trie.rs b/src/trie.rs index 30f8c152e..d6afb3bc8 100644 --- a/src/trie.rs +++ b/src/trie.rs @@ -24,9 +24,10 @@ pub trait Trie { #[derive(Eq, PartialEq, Debug)] pub enum Node<'a> { + NullRoot, Leaf(NibbleSlice<'a>, &'a[u8]), Extension(NibbleSlice<'a>, &'a[u8]), - Branch(Option<[&'a[u8]; 16]>, &'a [u8]) + Branch([Option<&'a[u8]>; 16], Option<&'a [u8]>) } impl <'a>Node<'a> { @@ -40,16 +41,18 @@ 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) => Node::Extension(slice, r.at(1).raw()) + (slice, false) => Node::Extension(slice, r.at(1).raw()), }, // branch - first 16 are nodes, 17th is a value (or empty). Prototype::List(17) => { - let mut nodes: [&'a [u8]; 16] = unsafe { ::std::mem::uninitialized() }; - for i in 0..16 { nodes[i] = r.at(i).raw(); } - Node::Branch(Some(nodes), r.at(16).data()) + let mut nodes: [Option<&'a [u8]>; 16] = unsafe { ::std::mem::uninitialized() }; + for i in 0..16 { + nodes[i] = if r.at(i).is_empty() { None } else { Some(r.at(i).raw()) } + } + Node::Branch(nodes, if r.at(16).is_empty() { None } else { Some(r.at(16).data()) }) }, // an empty branch index. - Prototype::Data(0) => Node::Branch(None, r.data()), + Prototype::Data(0) => Node::NullRoot, // something went wrong. _ => panic!("Rlp is not valid.") } @@ -70,13 +73,21 @@ impl <'a>Node<'a> { stream.append_raw(raw_rlp, 1); stream.out() }, - Node::Branch(Some(ref nodes), ref value) => { + Node::Branch(ref nodes, ref value) => { let mut stream = RlpStream::new_list(17); - for i in 0..16 { stream.append_raw(nodes[i], 1); } - stream.append(value); + for i in 0..16 { + match nodes[i] { + Some(n) => { stream.append_raw(n, 1); }, + None => { stream.append_empty_data(); }, + } + } + match *value { + Some(n) => { stream.append(&n); }, + None => { stream.append_empty_data(); }, + } stream.out() }, - Node::Branch(_, _) => { + Node::NullRoot => { let mut stream = RlpStream::new(); stream.append_empty_data(); stream.out() @@ -187,28 +198,34 @@ impl TrieDB { fn fmt_all(&self, node: &[u8], f: &mut fmt::Formatter, deepness: usize) -> fmt::Result { let node = Node::decoded(node); match node { - Node::Leaf(slice, value) => try!(writeln!(f, "Leaf {:?}, {:?}", slice, value.pretty())), + Node::Leaf(slice, value) => try!(writeln!(f, "-{:?}: {:?}.", slice, value.pretty())), Node::Extension(ref slice, ref item) => { - try!(write!(f, "Extension: {:?} ,", slice)); - try!(self.fmt_all(self.get_raw_or_lookup(item), f, deepness + 1)); + try!(write!(f, "-{:?}- ", slice)); + try!(self.fmt_all(self.get_raw_or_lookup(item), f, deepness)); }, - 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 { - let rlp = Rlp::new(nodes[i]); - if rlp.is_data() && rlp.size() != 0 { + Node::Branch(ref nodes, ref value) => { + try!(writeln!(f, "")); + match value { + &Some(v) => { 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, "")); + try!(writeln!(f, "=: {:?}", v.pretty())) + }, + &None => {} + } + for i in 0..16 { + match nodes[i] { + Some(n) => { + try!(self.fmt_indent(f, deepness + 1)); + try!(write!(f, "{:x}: ", i)); + try!(self.fmt_all(self.get_raw_or_lookup(n), f, deepness + 1)); + }, + None => {}, } } }, // empty - Node::Branch(_, _) => { - // do nothing + Node::NullRoot => { + try!(writeln!(f, "")); } }; Ok(()) @@ -225,9 +242,14 @@ impl TrieDB { 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(self.get_raw_or_lookup(nodes[key.at(0) as usize]), &key.mid(1)) + Node::Branch(ref nodes, value) => match key.is_empty() { + true => value, + false => match nodes[key.at(0) as usize] { + Some(payload) => { + self.get_from_node(self.get_raw_or_lookup(payload), &key.mid(1)) + }, + None => None + } }, _ => None } @@ -235,8 +257,9 @@ impl TrieDB { 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!"), + let r = Rlp::new(node); + match r.is_data() && r.size() == 32 { + true => self.db.lookup(&H256::decode(&r)).expect("Not found!"), false => node } } @@ -378,7 +401,7 @@ impl TrieDB { (true, i) if orig.at(i).is_empty() => // easy - original had empty slot. diff.new_node(Self::compose_leaf(&partial.mid(1), value), &mut s), (true, i) => { // harder - original has something there already - let new = self.augmented(orig.at(i).raw(), &partial.mid(1), value, diff); + let new = self.augmented(self.take_node(&orig.at(i), diff), &partial.mid(1), value, diff); diff.replace_node(&orig.at(i), new, &mut s); } (false, i) => { s.append_raw(orig.at(i).raw(), 1); }, @@ -457,7 +480,7 @@ impl TrieDB { trace!("empty: COMPOSE"); Self::compose_leaf(partial, value) }, - _ => panic!("Invalid RLP for node."), + _ => panic!("Invalid RLP for node: {:?}", old.pretty()), } } } @@ -518,7 +541,7 @@ mod tests { #[test] fn test_node_empty_branch() { - let branch = Node::Branch(None, &b""[..]); + let branch = Node::Branch([None; 16], None); let rlp = branch.encoded(); let branch2 = Node::decoded(&rlp); assert_eq!(branch, branch2); @@ -527,10 +550,10 @@ mod tests { #[test] fn test_node_branch() { let k = rlp::encode(&"cat"); - let mut nodes: [&[u8]; 16] = unsafe { ::std::mem::uninitialized() }; - for i in 0..16 { nodes[i] = &k; } + let mut nodes: [Option<&[u8]>; 16] = unsafe { ::std::mem::uninitialized() }; + for i in 0..16 { nodes[i] = Some(&k); } let v: Vec = From::from("dog"); - let branch = Node::Branch(Some(nodes), &v); + let branch = Node::Branch(nodes, Some(&v)); let rlp = branch.encoded(); let branch2 = Node::decoded(&rlp); assert_eq!(branch, branch2); @@ -575,21 +598,16 @@ mod tests { fn test_all(v: Vec<(Vec, Vec)>) { let mut t = TrieDB::new_memory(); - + for i in 0..v.len() { let key: &[u8]= &v[i].0; let val: &[u8] = &v[i].1; t.insert(&key, &val); } + trace!("{:?}", t); println!("{:?}", t); - 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); - } - // check lifetime let _q = t.at(&[b'd', b'o']).unwrap(); diff --git a/src/triehash.rs b/src/triehash.rs index 8be23e210..afc07a7bf 100644 --- a/src/triehash.rs +++ b/src/triehash.rs @@ -62,7 +62,8 @@ pub fn ordered_trie_root(input: Vec>) -> H256 { /// assert_eq!(trie_root(v), H256::from_str(root).unwrap()); /// } /// ``` -pub fn trie_root(input: Vec<(Vec, Vec)>) -> H256 { +pub fn trie_root(mut input: Vec<(Vec, Vec)>) -> H256 { + input.sort(); let gen_input = input .into_iter() .map(|(k, v)| (as_nibbles(&k), v))