also fuzz prefix

paritytech · May 13, 2024 · 93ec517 · 93ec517
1 parent b97b355
commit 93ec517
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 11 deletions.
diff --git a/test-support/reference-trie/src/lib.rs b/test-support/reference-trie/src/lib.rs
@@ -1237,7 +1237,7 @@ 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>>)
+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,
@@ -1277,10 +1277,22 @@ where
 	}
 	for key in keys {
 		use trie_db::TrieIterator;
-		let mut iter = trie_db::triedb::TrieDBDoubleEndedIterator::new(&trie).unwrap();
-		iter.seek(&key).unwrap();
-		let mut iter_ref = ref_tree.iter().filter(|k| k.0 >= &key);
-		let mut iter_ref2 = ref_tree.iter().filter(|k| k.0 <= &key);
+		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();
@@ -1291,13 +1303,11 @@ where
 			if n_ref.is_none() && nb_ref.is_none() {
 				break;
 			}
-
-			// TODO use key as prefix
 		}
 	}
 }
 
-pub fn fuzz_double_iter<T, DB>(input: &[u8])
+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,
@@ -1312,7 +1322,7 @@ where
 	keys.sort();
 	keys.dedup();
 
-	test_iterator::<T, DB>(data, keys);
+	test_iterator::<T, DB>(data, keys, prefix);
 }
 
 pub fn fuzz_to_data(input: &[u8]) -> Vec<(Vec<u8>, Vec<u8>)> {

diff --git a/trie-db/fuzz/fuzz_targets/double_iter.rs b/trie-db/fuzz/fuzz_targets/double_iter.rs
@@ -9,5 +9,6 @@ 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);
+	fuzz_double_iter::<T, DB>(data, false);
+	fuzz_double_iter::<T, DB>(data, true);
 });
diff --git a/trie-db/src/triedb.rs b/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> {}

diff --git a/trie-db/test/src/double_ended_iterator.rs b/trie-db/test/src/double_ended_iterator.rs
@@ -432,6 +432,7 @@ fn fuzz_set_internal<T: TrieLayout>() {
 		vec![0, 5, 0, 0, 43, 0, 5, 0],
 	];
 	for i in fuzz_inputs {
-		reference_trie::fuzz_double_iter::<T, DB<T>>(&i);
+		reference_trie::fuzz_double_iter::<T, DB<T>>(&i, false);
+		reference_trie::fuzz_double_iter::<T, DB<T>>(&i, true);
 	}
 }