Fix backward iterator seek (#215)

Fix backward iterator seek

test-support/reference-trie/src/lib.rs

@@ -1235,3 +1235,135 @@ mod tests {
 		}
 	}
 }
+
+// This is a bit redundant with other iterator fuzzer
+fn test_iterator<L, DB>(entries: Vec<(Vec<u8>, Vec<u8>)>, keys: Vec<Vec<u8>>, prefix: bool)
+where
+	L: TrieLayout,
+	DB: hash_db::HashDB<L::Hash, DBValue> + hash_db::HashDBRef<L::Hash, DBValue> + Default,
+{
+	let mut ref_tree = std::collections::BTreeMap::new();
+
+	// Populate DB with full trie from entries.
+	let (db, root) = {
+		let mut db = DB::default();
+		let mut root = Default::default();
+		{
+			let mut trie = TrieDBMutBuilder::<L>::new(&mut db, &mut root).build();
+			for (key, value) in entries.into_iter() {
+				trie.insert(&key, &value).unwrap();
+				ref_tree.insert(key, value);
+			}
+		}
+		(db, root)
+	};
+	// Lookup items in trie while recording traversed nodes.
+	let trie = TrieDBBuilder::<L>::new(&db, &root).build();
+	// standard iter
+	let mut iter = trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap();
+	let mut iter_ref = ref_tree.iter();
+	let mut iter_ref2 = ref_tree.iter();
+
+	loop {
+		let n = iter.next();
+		let nb = iter.next_back();
+		let n_ref = iter_ref.next();
+		let nb_ref = iter_ref2.next_back();
+		assert_eq!(n.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), n_ref);
+		assert_eq!(nb.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), nb_ref);
+		if n_ref.is_none() && nb_ref.is_none() {
+			break;
+		}
+	}
+	for key in keys {
+		use trie_db::TrieIterator;
+		let mut iter = if prefix {
+			trie_db::triedb::TrieDBDoubleEndedIterator::new_prefixed(&trie, &key).unwrap()
+		} else {
+			trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap()
+		};
+		if !prefix {
+			iter.seek(&key).unwrap();
+		}
+		let mut iter_ref =
+			ref_tree
+				.iter()
+				.filter(|k| if prefix { k.0.starts_with(&key) } else { k.0 >= &key });
+		let mut iter_ref2 =
+			ref_tree
+				.iter()
+				.filter(|k| if prefix { k.0.starts_with(&key) } else { k.0 <= &key });
+		loop {
+			let n = iter.next();
+			let nb = iter.next_back();
+			let n_ref = iter_ref.next();
+			let nb_ref = iter_ref2.next_back();
+			assert_eq!(n.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), n_ref);
+			assert_eq!(nb.as_ref().map(|v| v.as_ref().map(|v| (&v.0, &v.1)).unwrap()), nb_ref);
+			if n_ref.is_none() && nb_ref.is_none() {
+				break;
+			}
+		}
+	}
+}
+
+pub fn fuzz_double_iter<T, DB>(input: &[u8], prefix: bool)
+where
+	T: TrieLayout,
+	DB: hash_db::HashDB<T::Hash, DBValue> + hash_db::HashDBRef<T::Hash, DBValue> + Default,
+{
+	let mut data = fuzz_to_data(input);
+	// - the first 2/3 are added to the trie.
+	// - the last 1/3 is not added to the trie and use for random seek and prefix.
+	let mut keys = data[(data.len() / 3)..].iter().map(|(key, _)| key.clone()).collect::<Vec<_>>();
+	data.truncate(data.len() * 2 / 3);
+
+	let data = data_sorted_unique(data);
+	keys.sort();
+	keys.dedup();
+
+	test_iterator::<T, DB>(data, keys, prefix);
+}
+
+pub fn fuzz_to_data(input: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {
+	let mut result = Vec::new();
+	// enc = (minkeylen, maxkeylen (min max up to 32), datas)
+	// fix data len 2 bytes
+	let mut minkeylen = if let Some(v) = input.get(0) {
+		let mut v = *v & 31u8;
+		v = v + 1;
+		v
+	} else {
+		return result
+	};
+	let mut maxkeylen = if let Some(v) = input.get(1) {
+		let mut v = *v & 31u8;
+		v = v + 1;
+		v
+	} else {
+		return result
+	};
+
+	if maxkeylen < minkeylen {
+		let v = minkeylen;
+		minkeylen = maxkeylen;
+		maxkeylen = v;
+	}
+	let mut ix = 2;
+	loop {
+		let keylen = if let Some(v) = input.get(ix) {
+			let mut v = *v & 31u8;
+			v = v + 1;
+			v = std::cmp::max(minkeylen, v);
+			v = std::cmp::min(maxkeylen, v);
+			v as usize
+		} else {
+			break
+		};
+		let key = if input.len() > ix + keylen { input[ix..ix + keylen].to_vec() } else { break };
+		ix += keylen;
+		let val = if input.len() > ix + 2 { input[ix..ix + 2].to_vec() } else { break };
+		result.push((key, val));
+	}
+	result
+}

trie-db/CHANGELOG.md

@@ -4,6 +4,9 @@ The format is based on [Keep a Changelog].
 
 [Keep a Changelog]: http://keepachangelog.com/en/1.0.0/
 
+## [0.29.1] - 2024-05-24
+- Fix backward iterator seek [#215](https://github.com/paritytech/trie/pull/215)
+
 ## [0.29.0] - 2024-03-04
 - Implements `DoubleEndedIterator` for trie iterator [#208](https://github.com/paritytech/trie/pull/208)
 

trie-db/Cargo.toml

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

trie-db/fuzz/Cargo.toml

@@ -11,7 +11,7 @@ cargo-fuzz = true
 [dependencies]
 hash-db = { path = "../../hash-db", version = "0.16.0" }
 memory-db = { path = "../../memory-db", version = "0.32.0" }
-reference-trie = { path = "../../test-support/reference-trie", version = "0.29.0" }
+reference-trie = { path = "../../test-support/reference-trie", version = "0.29.1" }
 arbitrary = { version = "1.3.0", features = ["derive"] }
 array-bytes = "6.0.0"
 
@@ -68,3 +68,7 @@ path = "fuzz_targets/trie_proof_invalid.rs"
 [[bin]]
 name = "prefix_seek_iter"
 path = "fuzz_targets/prefix_seek_iter.rs"
+
+[[bin]]
+name = "double_iter"
+path = "fuzz_targets/double_iter.rs"

trie-db/fuzz/fuzz_targets/double_iter.rs

@@ -0,0 +1,14 @@
+#![no_main]
+
+use libfuzzer_sys::fuzz_target;
+use memory_db::{MemoryDB, PrefixedKey};
+use trie_db::{DBValue, TrieLayout};
+use trie_db_fuzz::fuzz_double_iter;
+
+type T = reference_trie::NoExtensionLayout;
+type DB = MemoryDB<<T as TrieLayout>::Hash, PrefixedKey<<T as TrieLayout>::Hash>, DBValue>;
+
+fuzz_target!(|data: &[u8]| {
+	fuzz_double_iter::<T, DB>(data, false);
+	fuzz_double_iter::<T, DB>(data, true);
+});

trie-db/fuzz/src/lib.rs

@@ -15,58 +15,17 @@
 use arbitrary::Arbitrary;
 use hash_db::Hasher;
 use memory_db::{HashKey, MemoryDB, PrefixedKey};
+pub use reference_trie::fuzz_double_iter;
 use reference_trie::{
-	calc_root, compare_insert_remove, reference_trie_root_iter_build as reference_trie_root,
+	calc_root, compare_insert_remove, fuzz_to_data,
+	reference_trie_root_iter_build as reference_trie_root,
 };
 use std::convert::TryInto;
 use trie_db::{
 	proof::{generate_proof, verify_proof},
 	DBValue, Trie, TrieDBBuilder, TrieDBIterator, TrieDBMutBuilder, TrieLayout, TrieMut,
 };
 
-fn fuzz_to_data(input: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {
-	let mut result = Vec::new();
-	// enc = (minkeylen, maxkeylen (min max up to 32), datas)
-	// fix data len 2 bytes
-	let mut minkeylen = if let Some(v) = input.get(0) {
-		let mut v = *v & 31u8;
-		v = v + 1;
-		v
-	} else {
-		return result
-	};
-	let mut maxkeylen = if let Some(v) = input.get(1) {
-		let mut v = *v & 31u8;
-		v = v + 1;
-		v
-	} else {
-		return result
-	};
-
-	if maxkeylen < minkeylen {
-		let v = minkeylen;
-		minkeylen = maxkeylen;
-		maxkeylen = v;
-	}
-	let mut ix = 2;
-	loop {
-		let keylen = if let Some(v) = input.get(ix) {
-			let mut v = *v & 31u8;
-			v = v + 1;
-			v = std::cmp::max(minkeylen, v);
-			v = std::cmp::min(maxkeylen, v);
-			v as usize
-		} else {
-			break
-		};
-		let key = if input.len() > ix + keylen { input[ix..ix + keylen].to_vec() } else { break };
-		ix += keylen;
-		let val = if input.len() > ix + 2 { input[ix..ix + 2].to_vec() } else { break };
-		result.push((key, val));
-	}
-	result
-}
-
 fn fuzz_removal(data: Vec<(Vec<u8>, Vec<u8>)>) -> Vec<(bool, Vec<u8>, Vec<u8>)> {
 	let mut res = Vec::new();
 	let mut torem = None;

trie-db/src/iterator.rs

@@ -173,7 +173,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
 					NodePlan::Leaf { partial: partial_plan, .. } => {
 						let slice = partial_plan.build(node_data);
 						if (fwd && slice < partial) || (!fwd && slice > partial) {
-							crumb.status = Status::Exiting;
+							crumb.status = Status::AftExiting;
 							return Ok(false);
 						}
 						return Ok(slice.starts_with(&partial));
@@ -182,8 +182,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
 						let slice = partial_plan.build(node_data);
 						if !partial.starts_with(&slice) {
 							if (fwd && slice < partial) || (!fwd && slice > partial) {
-								crumb.status = Status::Exiting;
-								self.key_nibbles.append_partial(slice.right());
+								crumb.status = Status::AftExiting;
 								return Ok(false);
 							}
 							return Ok(slice.starts_with(&partial));
@@ -230,9 +229,7 @@ impl<L: TrieLayout> TrieDBRawIterator<L> {
 						let slice = partial_plan.build(node_data);
 						if !partial.starts_with(&slice) {
 							if (fwd && slice < partial) || (!fwd && slice > partial) {
-								crumb.status = Status::Exiting;
-								self.key_nibbles.append_partial(slice.right());
-								self.key_nibbles.push((nibble_ops::NIBBLE_LENGTH - 1) as u8);
+								crumb.status = Status::AftExiting;
 								return Ok(false);
 							}
 							return Ok(slice.starts_with(&partial));

trie-db/src/triedb.rs

@@ -496,6 +496,18 @@ impl<'a, 'cache, L: TrieLayout> TrieDBDoubleEndedIterator<'a, 'cache, L> {
 			back_raw_iter: TrieDBRawIterator::new(db)?,
 		})
 	}
+
+	/// Create a new iterator, but limited to a given prefix.
+	pub fn new_prefixed(
+		db: &'a TrieDB<'a, 'cache, L>,
+		prefix: &[u8],
+	) -> Result<Self, TrieHash<L>, CError<L>> {
+		Ok(Self {
+			db,
+			raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
+			back_raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
+		})
+	}
 }
 
 impl<L: TrieLayout> TrieDoubleEndedIterator<L> for TrieDBDoubleEndedIterator<'_, '_, L> {}

trie-db/test/src/double_ended_iterator.rs

@@ -201,10 +201,7 @@ fn seek_back_works_internal<T: TrieLayout>() {
 	let mut iter = TrieDBNodeDoubleEndedIterator::new(&trie).unwrap();
 
 	<dyn TrieDoubleEndedIterator<T, Item = _>>::seek(&mut iter, &hex!("")[..]).unwrap();
-	match iter.next_back() {
-		Some(Ok((prefix, _, _))) => assert_eq!(prefix, nibble_vec(hex!(""), 0)),
-		_ => panic!("unexpected item"),
-	}
+	assert!(iter.next_back().is_none());
 
 	<dyn TrieDoubleEndedIterator<T, Item = _>>::seek(&mut iter, &hex!("03")[..]).unwrap();
 	match iter.next_back() {
@@ -406,3 +403,36 @@ fn prefix_over_empty_works_internal<T: TrieLayout>() {
 	iter.prefix(&hex!("00")[..]).unwrap();
 	assert!(iter.next_back().is_none());
 }
+
+test_layouts!(next_back_weird_behaviour_1, next_back_weird_behaviour_internal_1);
+fn next_back_weird_behaviour_internal_1<T: TrieLayout>() {
+	use trie_db::TrieIterator;
+
+	let pairs = vec![(vec![11], b"bbbb".to_vec())];
+
+	let (memdb, root) = build_trie_db::<T>(&pairs);
+	let trie = TrieDBBuilder::<T>::new(&memdb, &root).build();
+
+	let mut iter = trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap();
+
+	iter.seek(&[10]).unwrap();
+	assert!(iter.next_back().is_none());
+}
+
+test_layouts!(fuzz_set, fuzz_set_internal);
+fn fuzz_set_internal<T: TrieLayout>() {
+	type DB<T> = memory_db::MemoryDB<
+		<T as TrieLayout>::Hash,
+		memory_db::PrefixedKey<<T as TrieLayout>::Hash>,
+		trie_db::DBValue,
+	>;
+
+	let fuzz_inputs = [
+		vec![32, 65, 255, 254, 255, 213, 0, 0, 0, 254, 255, 0, 0, 235, 0, 0, 35],
+		vec![0, 5, 0, 0, 43, 0, 5, 0],
+	];
+	for i in fuzz_inputs {
+		reference_trie::fuzz_double_iter::<T, DB<T>>(&i, false);
+		reference_trie::fuzz_double_iter::<T, DB<T>>(&i, true);
+	}
+}