From 516a8fa43362906fc8a449ec74f67f6c1e7aa69a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20K=C3=B6cher?= <bkchr@users.noreply.github.com>
Date: Wed, 18 May 2022 21:20:47 +0200
Subject: [PATCH] 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
---
.../state-machine/src/trie_backend.rs | 24 +++++++++++++++++--
.../state-machine/src/trie_backend_essence.rs | 19 +++------------
2 files changed, 25 insertions(+), 18 deletions(-)
diff --git a/substrate/primitives/state-machine/src/trie_backend.rs b/substrate/primitives/state-machine/src/trie_backend.rs
index 3b985ec2b99..c0a620120bf 100644
--- a/substrate/primitives/state-machine/src/trie_backend.rs
+++ b/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},
- 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::<Vec<_>>();
+
+ let trie = test_trie(state_version);
+ let keys = trie.keys(&[]);
+
+ assert_eq!(expected, keys);
+ }
}
diff --git a/substrate/primitives/state-machine/src/trie_backend_essence.rs b/substrate/primitives/state-machine/src/trie_backend_essence.rs
index 0bbea700484..11cac92efd2 100644
--- a/substrate/primitives/state-machine/src/trie_backend_essence.rs
+++ b/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 trie = TrieDB::<H>::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`.
--
GitLab