trie: Optimize keys function (#11457)

* trie: Optimize `keys` function

Instead of iterating the entire state and collecting all keys that match the given prefix, we can
directly use the optimized prefix iterator.

* Add a test

substrate/primitives/state-machine/src/trie_backend.rs

@@ -190,8 +190,8 @@ pub mod tests {
 	use sp_core::H256;
 	use sp_runtime::traits::BlakeTwo256;
 	use sp_trie::{
-		trie_types::{TrieDBMutV0, TrieDBMutV1},
+		trie_types::{TrieDB, TrieDBMutV0, TrieDBMutV1},
-		KeySpacedDBMut, PrefixedMemoryDB, TrieMut,
+		KeySpacedDBMut, PrefixedMemoryDB, Trie, TrieMut,
 	};
 	use std::{collections::HashSet, iter};
 
@@ -369,4 +369,24 @@ pub mod tests {
 		expected.insert(b"value2".to_vec());
 		assert_eq!(seen, expected);
 	}
+
+	#[test]
+	fn keys_with_empty_prefix_returns_all_keys() {
+		keys_with_empty_prefix_returns_all_keys_inner(StateVersion::V0);
+		keys_with_empty_prefix_returns_all_keys_inner(StateVersion::V1);
+	}
+	fn keys_with_empty_prefix_returns_all_keys_inner(state_version: StateVersion) {
+		let (test_db, test_root) = test_db(state_version);
+		let expected = TrieDB::new(&test_db, &test_root)
+			.unwrap()
+			.iter()
+			.unwrap()
+			.map(|d| d.unwrap().0.to_vec())
+			.collect::<Vec<_>>();
+
+		let trie = test_trie(state_version);
+		let keys = trie.keys(&[]);
+
+		assert_eq!(expected, keys);
+	}
 }

substrate/primitives/state-machine/src/trie_backend_essence.rs

@@ -453,22 +453,9 @@ where
 
 	/// Returns all keys that start with the given `prefix`.
 	pub fn keys(&self, prefix: &[u8]) -> Vec<StorageKey> {
-		let collect_all = || -> sp_std::result::Result<_, Box<TrieError<H::Out>>> {
+		let mut keys = Vec::new();
-			let trie = TrieDB::<H>::new(self, &self.root)?;
+		self.for_keys_with_prefix(prefix, |k| keys.push(k.to_vec()));
-			let mut v = Vec::new();
+		keys
-			for x in trie.iter()? {
-				let (key, _) = x?;
-				if key.starts_with(prefix) {
-					v.push(key.to_vec());
-				}
-			}
-
-			Ok(v)
-		};
-
-		collect_all()
-			.map_err(|e| debug!(target: "trie", "Error extracting trie keys: {}", e))
-			.unwrap_or_default()
 	}
 
 	/// Return the storage root after applying the given `delta`.