From 91999155abc31e750456a1eaff98ee9a56110616 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastian=20K=C3=B6cher?= Date: Wed, 18 May 2022 17:39:00 +0200 Subject: [PATCH 1/2] 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. --- .../state-machine/src/trie_backend_essence.rs | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) 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`. From 108901175ff1130115c76cda99a97c65cffb11bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bastian=20K=C3=B6cher?= Date: Wed, 18 May 2022 18:02:30 +0200 Subject: [PATCH 2/2] Add a test --- primitives/state-machine/src/trie_backend.rs | 24 ++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) 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); + } }