This commit is contained in:
debris 2015-12-02 00:22:47 +01:00
commit 60dc3decaa
2 changed files with 63 additions and 44 deletions

View File

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

View File

@ -62,7 +62,8 @@ pub fn ordered_trie_root(input: Vec<Vec<u8>>) -> H256 {
/// assert_eq!(trie_root(v), H256::from_str(root).unwrap()); /// assert_eq!(trie_root(v), H256::from_str(root).unwrap());
/// } /// }
/// ``` /// ```
pub fn trie_root(input: Vec<(Vec<u8>, Vec<u8>)>) -> H256 { pub fn trie_root(mut input: Vec<(Vec<u8>, Vec<u8>)>) -> H256 {
input.sort();
let gen_input = input let gen_input = input
.into_iter() .into_iter()
.map(|(k, v)| (as_nibbles(&k), v)) .map(|(k, v)| (as_nibbles(&k), v))