diff --git a/util/src/nibbleslice.rs b/util/src/nibbleslice.rs index 6f4232945..b9028dff3 100644 --- a/util/src/nibbleslice.rs +++ b/util/src/nibbleslice.rs @@ -34,6 +34,22 @@ pub struct NibbleSlice<'a> { 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 { + 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 { /// Create a new nibble slice with the given byte-slice. 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. 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_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()) }*/ + 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()`). 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) @@ -189,6 +209,14 @@ mod tests { } } + #[test] + fn iterator() { + let n = NibbleSlice::new(D); + let mut nibbles: Vec = vec![]; + nibbles.extend(n.iter()); + assert_eq!(nibbles, (0u8..6).collect::>()) + } + #[test] fn mid() { let n = NibbleSlice::new(D); diff --git a/util/src/trie/node.rs b/util/src/trie/node.rs index b5745b66f..b10b0e05e 100644 --- a/util/src/trie/node.rs +++ b/util/src/trie/node.rs @@ -5,7 +5,7 @@ use rlp::*; use super::journal::*; /// Type of node in the trie and essential information thereof. -#[derive(Eq, PartialEq, Debug)] +#[derive(Clone, Eq, PartialEq, Debug)] pub enum Node<'a> { Empty, Leaf(NibbleSlice<'a>, &'a[u8]), diff --git a/util/src/trie/triedb.rs b/util/src/trie/triedb.rs index bd34e710d..0ee95ea4e 100644 --- a/util/src/trie/triedb.rs +++ b/util/src/trie/triedb.rs @@ -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>, + 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::>()); + 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 { + 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::>()); + 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> { fn root(&self) -> &H256 { &self.root } @@ -218,3 +338,21 @@ impl<'db> fmt::Debug for TrieDB<'db> { 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::>(), TrieDB::new(&memdb, &root).iter().map(|x|x.0).collect::>()); +} diff --git a/util/src/trie/triedbmut.rs b/util/src/trie/triedbmut.rs index 832b532f8..0f3dde4fb 100644 --- a/util/src/trie/triedbmut.rs +++ b/util/src/trie/triedbmut.rs @@ -65,8 +65,9 @@ impl<'db> TrieDBMut<'db> { r } - /// Create a new trie with the backing database `db` and `root` - /// Panics, if `root` does not exist + /// Create a new trie with the backing database `db` and `root`. + /// Panics, if `root` does not exist. + // TODO: return Result pub fn from_existing(db: &'db mut HashDB, root: &'db mut H256) -> Self { if !db.exists(root) { flush(format!("Trie root not found {}", root));