Iterator for NibbleSlice and TrieDB.

This commit is contained in:
Gav Wood 2016-01-17 22:18:47 +01:00
parent 6ea8eaa3b5
commit 62a0737e59
4 changed files with 171 additions and 4 deletions

View File

@ -34,6 +34,22 @@ pub struct NibbleSlice<'a> {
offset_encode_suffix: usize, offset_encode_suffix: usize,
} }
pub struct NibbleSliceIterator<'a> {
p: &'a NibbleSlice<'a>,
i: usize,
}
impl<'a> Iterator for NibbleSliceIterator<'a> {
type Item = u8;
fn next(&mut self) -> Option<u8> {
self.i += 1;
match self.i <= self.p.len() {
true => Some(self.p.at(self.i - 1)),
false => None,
}
}
}
impl<'a, 'view> NibbleSlice<'a> where 'a: 'view { impl<'a, 'view> NibbleSlice<'a> where 'a: 'view {
/// Create a new nibble slice with the given byte-slice. /// Create a new nibble slice with the given byte-slice.
pub fn new(data: &[u8]) -> NibbleSlice { NibbleSlice::new_offset(data, 0) } pub fn new(data: &[u8]) -> NibbleSlice { NibbleSlice::new_offset(data, 0) }
@ -41,7 +57,7 @@ impl<'a, 'view> NibbleSlice<'a> where 'a: 'view {
/// Create a new nibble slice with the given byte-slice with a nibble offset. /// Create a new nibble slice with the given byte-slice with a nibble offset.
pub fn new_offset(data: &'a [u8], offset: usize) -> NibbleSlice { NibbleSlice{data: data, offset: offset, data_encode_suffix: &b""[..], offset_encode_suffix: 0} } pub fn new_offset(data: &'a [u8], offset: usize) -> NibbleSlice { NibbleSlice{data: data, offset: offset, data_encode_suffix: &b""[..], offset_encode_suffix: 0} }
/// /// Create a composed nibble slice; one followed by the other.
pub fn new_composed(a: &'a NibbleSlice, b: &'a NibbleSlice) -> NibbleSlice<'a> { NibbleSlice{data: a.data, offset: a.offset, data_encode_suffix: b.data, offset_encode_suffix: b.offset} } pub fn new_composed(a: &'a NibbleSlice, b: &'a NibbleSlice) -> NibbleSlice<'a> { NibbleSlice{data: a.data, offset: a.offset, data_encode_suffix: b.data, offset_encode_suffix: b.offset} }
/*pub fn new_composed_bytes_offset(a: &NibbleSlice, b: &NibbleSlice) -> (Bytes, usize) { /*pub fn new_composed_bytes_offset(a: &NibbleSlice, b: &NibbleSlice) -> (Bytes, usize) {
@ -60,6 +76,10 @@ impl<'a, 'view> NibbleSlice<'a> where 'a: 'view {
(r, a.len() + b.len()) (r, a.len() + b.len())
}*/ }*/
pub fn iter(&'a self) -> NibbleSliceIterator<'a> {
NibbleSliceIterator { p: self, i: 0 }
}
/// Create a new nibble slice from the given HPE encoded data (e.g. output of `encoded()`). /// Create a new nibble slice from the given HPE encoded data (e.g. output of `encoded()`).
pub fn from_encoded(data: &'a [u8]) -> (NibbleSlice, bool) { pub fn from_encoded(data: &'a [u8]) -> (NibbleSlice, bool) {
(Self::new_offset(data, if data[0] & 16 == 16 {1} else {2}), data[0] & 32 == 32) (Self::new_offset(data, if data[0] & 16 == 16 {1} else {2}), data[0] & 32 == 32)
@ -189,6 +209,14 @@ mod tests {
} }
} }
#[test]
fn iterator() {
let n = NibbleSlice::new(D);
let mut nibbles: Vec<u8> = vec![];
nibbles.extend(n.iter());
assert_eq!(nibbles, (0u8..6).collect::<Vec<_>>())
}
#[test] #[test]
fn mid() { fn mid() {
let n = NibbleSlice::new(D); let n = NibbleSlice::new(D);

View File

@ -5,7 +5,7 @@ use rlp::*;
use super::journal::*; use super::journal::*;
/// Type of node in the trie and essential information thereof. /// Type of node in the trie and essential information thereof.
#[derive(Eq, PartialEq, Debug)] #[derive(Clone, Eq, PartialEq, Debug)]
pub enum Node<'a> { pub enum Node<'a> {
Empty, Empty,
Leaf(NibbleSlice<'a>, &'a[u8]), Leaf(NibbleSlice<'a>, &'a[u8]),

View File

@ -198,6 +198,126 @@ impl<'db> TrieDB<'db> {
} }
} }
#[derive(Clone, Eq, PartialEq)]
enum Status {
Entering,
At,
AtChild(usize),
Exiting,
}
#[derive(Clone, Eq, PartialEq)]
struct Crumb<'a> {
node: Node<'a>,
// key: &'a[u8],
status: Status,
}
impl<'a> Crumb<'a> {
/// Move on to next status in the node's sequence.
fn increment(&mut self) {
self.status = match (&self.status, &self.node) {
(_, &Node::Empty) => Status::Exiting,
(&Status::Entering, _) => Status::At,
(&Status::At, &Node::Branch(_, _)) => Status::AtChild(0),
(&Status::AtChild(x), &Node::Branch(_, _)) if x < 15 => Status::AtChild(x + 1),
_ => Status::Exiting,
}
}
}
/// Iterator for going through all values in the trie.
#[derive(Clone)]
pub struct TrieDBIterator<'a> {
db: &'a TrieDB<'a>,
trail: Vec<Crumb<'a>>,
key_nibbles: Bytes,
}
impl<'a> TrieDBIterator<'a> {
/// Create a new iterator.
fn new(db: &'a TrieDB) -> TrieDBIterator<'a> {
let mut r = TrieDBIterator {
db: db,
trail: vec![],
key_nibbles: Vec::new(),
};
r.descend(db.db.lookup(&db.root).expect("Trie root not found!"));
r
}
/// Descend into a payload.
fn descend(&mut self, d: &'a [u8]) {
self.trail.push(Crumb {
status: Status::Entering,
node: self.db.get_node(d)
});
match self.trail.last().unwrap().node {
Node::Leaf(n, _) | Node::Extension(n, _) => {
println!("Entering. Extend key {:?}, {:?}", self.key_nibbles, n.iter().collect::<Vec<_>>());
self.key_nibbles.extend(n.iter());
},
_ => {}
}
}
/// Descend into a payload and get the next item.
fn descend_next(&mut self, d: &'a [u8]) -> Option<(Bytes, &'a [u8])> { self.descend(d); self.next() }
/// The present key.
fn key(&self) -> Bytes {
// collapse the key_nibbles down to bytes.
self.key_nibbles.iter().step(2).zip(self.key_nibbles.iter().skip(1).step(2)).map(|(h, l)| h * 16 + l).collect()
}
}
impl<'a> Iterator for TrieDBIterator<'a> {
type Item = (Bytes, &'a [u8]);
fn next(&mut self) -> Option<Self::Item> {
let b = match self.trail.last_mut() {
Some(ref mut b) => { b.increment(); b.clone() },
None => return None
};
match (b.status, b.node) {
(Status::Exiting, n) => {
match n {
Node::Leaf(n, _) | Node::Extension(n, _) => {
println!("Exiting. Truncate key {:?}, {:?}", self.key_nibbles, n.iter().collect::<Vec<_>>());
let l = self.key_nibbles.len();
self.key_nibbles.truncate(l - n.len());
},
Node::Branch(_, _) => { println!("Exit branch. Pop {:?}.", self.key_nibbles); self.key_nibbles.pop(); },
_ => {}
}
self.trail.pop();
self.next()
},
(Status::At, Node::Leaf(_, v)) => Some((self.key(), v)),
(Status::At, Node::Extension(_, d)) => self.descend_next(d),
(Status::At, Node::Branch(_, Some(v))) => Some((self.key(), v)),
(Status::At, Node::Branch(_, _)) => self.next(),
(Status::AtChild(i), Node::Branch(children, _)) if children[i].len() > 0 => {
match i {
0 => { println!("Enter first child of branch. Push {:?}.", self.key_nibbles); self.key_nibbles.push(0); },
i => *self.key_nibbles.last_mut().unwrap() = i as u8,
}
self.descend_next(children[i])
},
(Status::AtChild(i), Node::Branch(_, _)) => {
if i == 0 { println!("Enter first child of branch. Push {:?}.", self.key_nibbles); self.key_nibbles.push(0); }
self.next()
},
_ => panic!() // Should never see Entering or AtChild without a Branch here.
}
}
}
impl<'db> TrieDB<'db> {
/// Get all keys/values stored in the trie.
pub fn iter(&self) -> TrieDBIterator { TrieDBIterator::new(self) }
}
impl<'db> Trie for TrieDB<'db> { impl<'db> Trie for TrieDB<'db> {
fn root(&self) -> &H256 { &self.root } fn root(&self) -> &H256 { &self.root }
@ -218,3 +338,21 @@ impl<'db> fmt::Debug for TrieDB<'db> {
writeln!(f, "]") writeln!(f, "]")
} }
} }
#[test]
fn iterator() {
use memorydb::*;
use super::triedbmut::*;
let d = vec![ &b"A"[..], &b"AA"[..], &b"AB"[..], &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);
}
}
assert_eq!(d.iter().map(|i|i.to_vec()).collect::<Vec<_>>(), TrieDB::new(&memdb, &root).iter().map(|x|x.0).collect::<Vec<_>>());
}

View File

@ -65,8 +65,9 @@ impl<'db> TrieDBMut<'db> {
r r
} }
/// Create a new trie with the backing database `db` and `root` /// Create a new trie with the backing database `db` and `root`.
/// Panics, if `root` does not exist /// Panics, if `root` does not exist.
// TODO: return Result<Self, TrieError>
pub fn from_existing(db: &'db mut HashDB, root: &'db mut H256) -> Self { pub fn from_existing(db: &'db mut HashDB, root: &'db mut H256) -> Self {
if !db.exists(root) { if !db.exists(root) {
flush(format!("Trie root not found {}", root)); flush(format!("Trie root not found {}", root));