From 83895ca109bc9e49cc478dcbe6da5773464dfef7 Mon Sep 17 00:00:00 2001 From: debris Date: Tue, 1 Dec 2015 16:39:44 +0100 Subject: [PATCH] trie.at() impl --- src/memorydb.rs | 2 +- src/trie.rs | 169 ++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 166 insertions(+), 5 deletions(-) diff --git a/src/memorydb.rs b/src/memorydb.rs index 14fae12d0..85fdc7c19 100644 --- a/src/memorydb.rs +++ b/src/memorydb.rs @@ -189,4 +189,4 @@ fn memorydb_denote() { } assert_eq!(m.lookup(&hash).unwrap(), b"Hello world!"); -} \ No newline at end of file +} diff --git a/src/trie.rs b/src/trie.rs index 9f0b209be..30a186e80 100644 --- a/src/trie.rs +++ b/src/trie.rs @@ -16,11 +16,86 @@ pub trait Trie { // TODO: consider returning &[u8]... fn contains(&self, key: &[u8]) -> bool; - fn at(&self, key: &[u8]) -> Option<&[u8]>; + fn at<'a>(&'a self, key: &'a [u8]) -> Option<&'a [u8]>; fn insert(&mut self, key: &[u8], value: &[u8]); fn remove(&mut self, key: &[u8]); } +#[derive(Eq, PartialEq, Debug)] +pub enum Node<'a> { + Leaf(NibbleSlice<'a>, &'a[u8]), + ExtensionRaw(NibbleSlice<'a>, &'a[u8]), + ExtensionSha3(NibbleSlice<'a>, &'a[u8]), + Branch(Option<[&'a[u8]; 16]>, &'a [u8]) +} + +impl <'a>Node<'a> { + pub fn decoded(node_rlp: &'a [u8]) -> Node<'a> { + let r = Rlp::new(node_rlp); + match r.prototype() { + // either leaf or extension - decode first item with NibbleSlice::??? + // and use is_leaf return to figure out which. + // if leaf, second item is a value (is_data()) + // if extension, second item is a node (either SHA3 to be looked up and + // 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!(); } + } + }, + // 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()) + }, + // an empty branch index. + Prototype::Data(0) => Node::Branch(None, r.data()), + // something went wrong. + _ => panic!("Rlp is not valid.") + } + } + + pub fn encoded(&self) -> Bytes { + match *self { + Node::Leaf(ref slice, ref value) => { + let mut stream = RlpStream::new_list(2); + stream.append(&slice.encoded(true)); + stream.append(value); + stream.out() + }, + Node::ExtensionRaw(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); } + stream.append(value); + stream.out() + }, + Node::Branch(_, _) => { + let mut stream = RlpStream::new(); + stream.append_empty_data(); + stream.out() + } + } + } +} + enum Operation { New(H256, Bytes), Delete(H256), @@ -98,6 +173,30 @@ impl TrieDB { } } + fn get<'a>(&'a self, key: &NibbleSlice<'a>) -> Option<&'a [u8]> { + let root_rlp = self.db.lookup(&self.root).expect("Trie root not found!"); + self.get_from_node(&root_rlp, key) + } + + 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("not found!"); + self.get_from_node(rlp, &key.mid(slice.len())) + }, + Node::Branch(Some(ref nodes), ref value) => match key.is_empty() { + true => Some(value), + // we don't really need to do lookup for nodes[key.at(0)] in db? + // if its empty hash return Nonem without lookup + false => self.get_from_node(nodes[key.at(0) as usize], &key.mid(1)) + }, + _ => None + } + } + fn add(&mut self, key: &NibbleSlice, value: &[u8]) { // determine what the new root is, insert new nodes and remove old as necessary. let mut todo: Diff = Diff::new(); @@ -311,8 +410,8 @@ impl Trie for TrieDB { unimplemented!(); } - fn at(&self, _key: &[u8]) -> Option<&[u8]> { - unimplemented!(); + fn at<'a>(&'a self, key: &'a [u8]) -> Option<&'a [u8]> { + self.get(&NibbleSlice::new(key)) } fn insert(&mut self, key: &[u8], value: &[u8]) { @@ -329,8 +428,70 @@ mod tests { use memorydb::*; use triehash::*; use super::*; + use nibbleslice::*; + use rlp; use env_logger; + #[test] + fn test_node_leaf() { + let k = vec![0x20u8, 0x01, 0x23, 0x45]; + let v: Vec = From::from("cat"); + let (slice, is_leaf) = NibbleSlice::from_encoded(&k); + assert_eq!(is_leaf, true); + let leaf = Node::Leaf(slice, &v); + let rlp = leaf.encoded(); + let leaf2 = Node::decoded(&rlp); + assert_eq!(leaf, leaf2); + } + + #[test] + fn test_node_extension() { + let k = vec![0x00u8, 0x01, 0x23, 0x45]; + // in extension, value must be valid rlp + 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 rlp = ex.encoded(); + let ex2 = Node::decoded(&rlp); + assert_eq!(ex, ex2); + } + + #[test] + fn test_node_empty_branch() { + let branch = Node::Branch(None, &b""[..]); + let rlp = branch.encoded(); + let branch2 = Node::decoded(&rlp); + assert_eq!(branch, branch2); + } + + #[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 v: Vec = From::from("dog"); + let branch = Node::Branch(Some(nodes), &v); + let rlp = branch.encoded(); + let branch2 = Node::decoded(&rlp); + assert_eq!(branch, branch2); + } + + #[test] + fn test_at_empty() { + let mut t = TrieDB::new(MemoryDB::new()); + t.init(); + assert_eq!(t.at(&[0x5]), None); + } + + #[test] + fn test_at_one() { + let mut t = TrieDB::new(MemoryDB::new()); + t.init(); + t.insert(&[0x01u8, 0x23], &[0x01u8, 0x23]); + assert_eq!(t.at(&[0x1, 0x23]).unwrap(), &[0x1u8, 0x23]); + } + #[test] fn playpen() { env_logger::init().unwrap(); @@ -364,4 +525,4 @@ mod tests { t.insert(&[0x01u8, 0x23], &[0x23u8, 0x45]); assert_eq!(*t.root(), trie_root(vec![ (vec![0x01u8, 0x23], vec![0x23u8, 0x45]) ])); } -} \ No newline at end of file +}