diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index 3b985ec2b99f6..c0a620120bff2 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/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}, - KeySpacedDBMut, PrefixedMemoryDB, TrieMut, + trie_types::{TrieDB, TrieDBMutV0, TrieDBMutV1}, + 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::>(); + + let trie = test_trie(state_version); + let keys = trie.keys(&[]); + + assert_eq!(expected, keys); + } } diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 0bbea7004848a..11cac92efd2f4 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/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 { - let collect_all = || -> sp_std::result::Result<_, Box>> { - let trie = TrieDB::::new(self, &self.root)?; - let mut v = Vec::new(); - 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() + let mut keys = Vec::new(); + self.for_keys_with_prefix(prefix, |k| keys.push(k.to_vec())); + keys } /// Return the storage root after applying the given `delta`.