From 8b3cdc513bc37c9d3c7aee8207ba9a0c06c5db60 Mon Sep 17 00:00:00 2001 From: arkpar Date: Thu, 24 Nov 2016 11:44:24 +0100 Subject: [PATCH] Seekable iterator --- Cargo.lock | 2 +- util/src/trie/fatdb.rs | 10 +++- util/src/trie/mod.rs | 10 +++- util/src/trie/sectriedb.rs | 4 +- util/src/trie/triedb.rs | 97 +++++++++++++++++++++++++++++++++++++- 5 files changed, 114 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 62ef5ea11..232b8cbb2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -224,7 +224,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] name = "elastic-array" version = "0.6.0" -source = "git+https://github.com/ethcore/elastic-array#70e4012e691b732c7c4cb04e9232799e6aa268bc" +source = "git+https://github.com/ethcore/elastic-array#346f1ba5982576dab9d0b8fa178b50e1db0a21cd" dependencies = [ "heapsize 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)", ] diff --git a/util/src/trie/fatdb.rs b/util/src/trie/fatdb.rs index 700156429..3a4ad8fc6 100644 --- a/util/src/trie/fatdb.rs +++ b/util/src/trie/fatdb.rs @@ -17,7 +17,7 @@ use hash::H256; use sha3::Hashable; use hashdb::{HashDB, DBValue}; -use super::{TrieDB, Trie, TrieDBIterator, TrieItem, Recorder}; +use super::{TrieDB, Trie, TrieDBIterator, TrieItem, Recorder, TrieIterator}; /// A `Trie` implementation which hashes keys and uses a generic `HashDB` backing database. /// Additionaly it stores inserted hash-key mappings for later retrieval. @@ -46,7 +46,7 @@ impl<'db> FatDB<'db> { } impl<'db> Trie for FatDB<'db> { - fn iter<'a>(&'a self) -> super::Result + 'a>> { + fn iter<'a>(&'a self) -> super::Result + 'a>> { FatDBIterator::new(&self.raw).map(|iter| Box::new(iter) as Box<_>) } @@ -81,6 +81,12 @@ impl<'db> FatDBIterator<'db> { } } +impl<'db> TrieIterator for FatDBIterator<'db> { + fn seek(&mut self, key: &[u8]) -> super::Result<()> { + self.trie_iterator.seek(&key.sha3()) + } +} + impl<'db> Iterator for FatDBIterator<'db> { type Item = TrieItem<'db>; diff --git a/util/src/trie/mod.rs b/util/src/trie/mod.rs index d4cc04962..71a3e71ec 100644 --- a/util/src/trie/mod.rs +++ b/util/src/trie/mod.rs @@ -102,7 +102,7 @@ pub trait Trie { where 'a: 'b, R: Recorder; /// Returns an iterator over elements of trie. - fn iter<'a>(&'a self) -> Result + 'a>>; + fn iter<'a>(&'a self) -> Result + 'a>>; } /// A key-value datastore implemented as a database-backed modified Merkle tree. @@ -130,6 +130,12 @@ pub trait TrieMut { fn remove(&mut self, key: &[u8]) -> Result<()>; } +/// A trie iterator that also supports random access. +pub trait TrieIterator : Iterator { + /// Position the iterator on the first element with key >= `key` + fn seek(&mut self, key: &[u8]) -> Result<()>; +} + /// Trie types #[derive(Debug, PartialEq, Clone)] pub enum TrieSpec { @@ -193,7 +199,7 @@ impl<'db> Trie for TrieKinds<'db> { wrapper!(self, get_recorded, key, r) } - fn iter<'a>(&'a self) -> Result + 'a>> { + fn iter<'a>(&'a self) -> Result + 'a>> { wrapper!(self, iter,) } } diff --git a/util/src/trie/sectriedb.rs b/util/src/trie/sectriedb.rs index b1d7bbc0c..0861f53f3 100644 --- a/util/src/trie/sectriedb.rs +++ b/util/src/trie/sectriedb.rs @@ -18,7 +18,7 @@ use hash::H256; use sha3::Hashable; use hashdb::{HashDB, DBValue}; use super::triedb::TrieDB; -use super::{Trie, TrieItem, Recorder}; +use super::{Trie, TrieItem, Recorder, TrieIterator}; /// A `Trie` implementation which hashes keys and uses a generic `HashDB` backing database. /// @@ -49,7 +49,7 @@ impl<'db> SecTrieDB<'db> { } impl<'db> Trie for SecTrieDB<'db> { - fn iter<'a>(&'a self) -> super::Result + 'a>> { + fn iter<'a>(&'a self) -> super::Result + 'a>> { TrieDB::iter(&self.raw) } diff --git a/util/src/trie/triedb.rs b/util/src/trie/triedb.rs index d929c9d68..f546fa9d7 100644 --- a/util/src/trie/triedb.rs +++ b/util/src/trie/triedb.rs @@ -20,7 +20,7 @@ use nibbleslice::*; use rlp::*; use super::node::Node; use super::recorder::{Recorder, NoOp}; -use super::{Trie, TrieItem, TrieError}; +use super::{Trie, TrieItem, TrieError, TrieIterator}; /// A `Trie` implementation using a generic `HashDB` backing database. /// @@ -295,6 +295,61 @@ impl<'a> TrieDBIterator<'a> { Ok(r) } + /// 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 seek_recurse<'key> ( + &mut self, + node: &[u8], + key: &NibbleSlice<'key>, + d: u32 + ) -> super::Result<()> { + let n = Node::decoded(node); + match Node::decoded(node) { + Node::Leaf(_, ref item) => { + self.descend(node); + Ok(()) + }, + Node::Extension(ref slice, ref item) => { + let slice = &NibbleSlice::from_encoded(slice).0; + if key.starts_with(slice) { + let mut r = NoOp; + self.trail.push(Crumb { + status: Status::At, + node: n, + }); + self.key_nibbles.extend(slice.iter()); + let data = try!(self.db.get_raw_or_lookup(&*item, &mut r, d)); + self.seek_recurse(&data, &key.mid(slice.len()), d + 1) + } else { + self.descend(node); + Ok(()) + } + }, + Node::Branch(ref nodes, _) => match key.is_empty() { + true => { + self.trail.push(Crumb { + status: Status::Entering, + node: n, + }); + Ok(()) + }, + false => { + let mut r = NoOp; + let node = try!(self.db.get_raw_or_lookup(&*nodes[key.at(0) as usize], &mut r, d)); + self.trail.push(Crumb { + status: Status::AtChild(key.at(0) as usize), + node: n, + }); + self.key_nibbles.push(key.at(0)); + self.seek_recurse(&node, &key.mid(1), d + 1) + } + }, + _ => Ok(()) + } + } + /// Descend into a payload. fn descend(&mut self, d: &[u8]) -> super::Result<()> { self.trail.push(Crumb { @@ -316,6 +371,17 @@ impl<'a> TrieDBIterator<'a> { } } +impl<'a> TrieIterator for TrieDBIterator<'a> { + /// Position the iterator on the first element with key >= `key` + fn seek(&mut self, key: &[u8]) -> super::Result<()> { + self.trail.clear(); + self.key_nibbles.clear(); + let mut r = NoOp; + let root_rlp = try!(self.db.root_data(&mut r)); + self.seek_recurse(&root_rlp, &NibbleSlice::new(key), 1) + } +} + impl<'a> Iterator for TrieDBIterator<'a> { type Item = TrieItem<'a>; @@ -372,7 +438,7 @@ impl<'a> Iterator for TrieDBIterator<'a> { } impl<'db> Trie for TrieDB<'db> { - fn iter<'a>(&'a self) -> super::Result + 'a>> { + fn iter<'a>(&'a self) -> super::Result + 'a>> { TrieDBIterator::new(self).map(|iter| Box::new(iter) as Box<_>) } @@ -415,3 +481,30 @@ fn iterator() { assert_eq!(d.iter().map(|i| i.clone().to_vec()).collect::>(), t.iter().unwrap().map(|x| x.unwrap().0).collect::>()); assert_eq!(d, t.iter().unwrap().map(|x| x.unwrap().1).collect::>()); } + +#[test] +fn iterator_seek() { + use memorydb::*; + use super::TrieMut; + use super::triedbmut::*; + + let d = vec![ DBValue::from_slice(b"A"), DBValue::from_slice(b"AA"), DBValue::from_slice(b"AB"), DBValue::from_slice(b"B") ]; + + let mut memdb = MemoryDB::new(); + let mut root = H256::new(); + { + let mut t = TrieDBMut::new(&mut memdb, &mut root); + for x in &d { + t.insert(x, x).unwrap(); + } + } + + let t = TrieDB::new(&memdb, &root).unwrap(); + let mut iter = t.iter().unwrap(); + //assert_eq!(iter.next(), Some(Ok((b"A".to_vec(), DBValue::from_slice(b"A"))))); + //iter.seek(b"!").unwrap(); + //assert_eq!(d, iter.map(|x| x.unwrap().1).collect::>()); + let mut iter = t.iter().unwrap(); + iter.seek(b"A").unwrap(); + assert_eq!(d, iter.map(|x| x.unwrap().1).collect::>()); +}