Seek for prefixed iterator. (#127)

* prefix then seek primitive

* Very relevant warning.

* seek cannot use seek_prefix condition.

* Revert "Very relevant warning."

This reverts commit aa92ade.

* test rename

trie-db/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "trie-db"
-version = "0.22.3"
+version = "0.22.4"
 authors = ["Parity Technologies <[email protected]>"]
 description = "Merkle-Patricia Trie generic over key hasher and node encoding"
 repository = "https://github.com/paritytech/trie"

trie-db/src/iterator.rs

@@ -244,6 +244,41 @@ impl<'a, L: TrieLayout> TrieDBNodeIterator<'a, L> {
 		Ok(())
 	}
 
+	/// Advance the iterator into a prefix, no value out of the prefix will be accessed
+	/// or returned after this operation.
+	pub fn prefix_then_seek(&mut self, prefix: &[u8], seek: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
+		if seek.starts_with(prefix) {
+			self.seek_prefix(seek)?;
+			let prefix_len = prefix.len() * crate::nibble::nibble_ops::NIBBLE_PER_BYTE;
+			let mut len = 0;
+			// look first prefix in trail
+			for i in 0..self.trail.len() {
+				match self.trail[i].node.node_plan() {
+					NodePlan::Empty => {},
+					NodePlan::Branch { .. } => {
+						len += 1;
+					},
+					NodePlan::Leaf { partial, .. } => {
+						len += partial.len();
+					},
+					NodePlan::Extension { partial, .. } => {
+						len += partial.len();
+					},
+					NodePlan::NibbledBranch { partial, .. } => {
+						len += 1;
+						len += partial.len();
+					},
+				}
+				if len > prefix_len {
+					self.trail = self.trail.split_off(i);
+					return Ok(());
+				}
+			}
+		}
+		// default to empty iter
+		self.trail.clear();
+		Ok(())
+	}
 }
 
 impl<'a, L: TrieLayout> TrieIterator<L> for TrieDBNodeIterator<'a, L> {

trie-db/src/triedb.rs

@@ -283,6 +283,21 @@ impl<'a, L: TrieLayout> TrieDBIterator<'a, L> {
 		})
 	}
 
+	/// Create a new iterator, but limited to a given prefix.
+	/// It then do a seek operation from prefixed context (using `seek` lose
+	/// prefix context by default).
+	pub fn new_prefixed_then_seek(
+		db: &'a TrieDB<L>,
+		prefix: &[u8],
+		start_at: &[u8],
+	) -> Result<TrieDBIterator<'a, L>, TrieHash<L>, CError<L>> {
+		let mut inner = TrieDBNodeIterator::new(db)?;
+		inner.prefix_then_seek(prefix, start_at)?;
+
+		Ok(TrieDBIterator {
+			inner,
+		})
+	}
 }
 
 impl<'a, L: TrieLayout> TrieIterator<L> for TrieDBIterator<'a, L> {

trie-db/test/src/triedb.rs

@@ -219,42 +219,58 @@ fn iterator_seek() {
 		b"AB".to_vec(),
 		b"B".to_vec(),
 	];
+	let vals = vec![
+		vec![0; 32],
+		vec![1; 32],
+		vec![2; 32],
+		vec![3; 32],
+	];
 
 	let mut memdb = MemoryDB::<KeccakHasher, PrefixedKey<_>, DBValue>::default();
 	let mut root = Default::default();
 	{
 		let mut t = RefTrieDBMutNoExt::new(&mut memdb, &mut root);
-		for x in &d {
-			t.insert(x, x).unwrap();
+		for (k, val) in d.iter().zip(vals.iter()) {
+			t.insert(k, val.as_slice()).unwrap();
 		}
 	}
 
 	let t = RefTrieDBNoExt::new(&memdb, &root).unwrap();
 	let mut iter = t.iter().unwrap();
-	assert_eq!(iter.next().unwrap().unwrap(), (b"A".to_vec(), b"A".to_vec()));
+	assert_eq!(iter.next().unwrap().unwrap(), (b"A".to_vec(), vals[0].clone()));
 	iter.seek(b"!").unwrap();
-	assert_eq!(d, iter.map(|x| x.unwrap().1).collect::<Vec<_>>());
+	assert_eq!(vals, iter.map(|x| x.unwrap().1).collect::<Vec<_>>());
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"A").unwrap();
-	assert_eq!(d, &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(vals, &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"AA").unwrap();
-	assert_eq!(&d[1..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[1..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	let iter = trie_db::TrieDBIterator::new_prefixed(&t, b"aaaaa").unwrap();
+	assert_eq!(&vals[..0], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	let iter = trie_db::TrieDBIterator::new_prefixed(&t, b"A").unwrap();
+	assert_eq!(&vals[..3], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	let iter = trie_db::TrieDBIterator::new_prefixed_then_seek(&t, b"A", b"AA").unwrap();
+	assert_eq!(&vals[1..3], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	let iter = trie_db::TrieDBIterator::new_prefixed_then_seek(&t, b"A", b"AB").unwrap();
+	assert_eq!(&vals[2..3], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	let iter = trie_db::TrieDBIterator::new_prefixed_then_seek(&t, b"", b"AB").unwrap();
+	assert_eq!(&vals[2..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"A!").unwrap();
-	assert_eq!(&d[1..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[1..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"AB").unwrap();
-	assert_eq!(&d[2..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[2..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"AB!").unwrap();
-	assert_eq!(&d[3..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[3..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"B").unwrap();
-	assert_eq!(&d[3..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[3..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 	let mut iter = t.iter().unwrap();
 	iter.seek(b"C").unwrap();
-	assert_eq!(&d[4..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
+	assert_eq!(&vals[4..], &iter.map(|x| x.unwrap().1).collect::<Vec<_>>()[..]);
 }
 
 #[test]