diff --git a/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.compact.wasm b/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.compact.wasm index f466a1023391e..4f2c3f852e2af 100644 Binary files a/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.compact.wasm and b/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.compact.wasm differ diff --git a/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.wasm b/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.wasm index 9685d89fa8496..a17a73795b990 100755 Binary files a/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.wasm and b/demo/runtime/wasm/target/wasm32-unknown-unknown/release/demo_runtime.wasm differ diff --git a/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.compact.wasm b/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.compact.wasm index f114ea539bc69..7082ac45b4fc4 100644 Binary files a/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.compact.wasm and b/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.compact.wasm differ diff --git a/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.wasm b/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.wasm index 0794dc20ffe32..4505fd75188b8 100755 Binary files a/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.wasm and b/polkadot/runtime/wasm/target/wasm32-unknown-unknown/release/polkadot_runtime.wasm differ diff --git a/substrate/client/src/light/backend.rs b/substrate/client/src/light/backend.rs index dce924f6ded48..1311618e994ec 100644 --- a/substrate/client/src/light/backend.rs +++ b/substrate/client/src/light/backend.rs @@ -145,6 +145,10 @@ impl StateBackend for OnDemandState where Block: BlockT, F: Err(ClientErrorKind::NotAvailableOnLightClient.into()) // TODO: fetch from remote node } + fn for_keys_with_prefix(&self, _prefix: &[u8], _action: A) { + // whole state is not available on light node + } + fn storage_root(&self, _delta: I) -> ([u8; 32], Self::Transaction) where I: IntoIterator, Option>)> { ([0; 32], ()) diff --git a/substrate/executor/src/wasm_executor.rs b/substrate/executor/src/wasm_executor.rs index d230cc19865d1..f66b81bedf918 100644 --- a/substrate/executor/src/wasm_executor.rs +++ b/substrate/executor/src/wasm_executor.rs @@ -225,6 +225,11 @@ impl_function_executor!(this: FunctionExecutor<'e, E>, this.ext.clear_storage(&key); Ok(()) }, + ext_clear_prefix(prefix_data: *const u8, prefix_len: u32) => { + let prefix = this.memory.get(prefix_data, prefix_len as usize).map_err(|_| DummyUserError)?; + this.ext.clear_prefix(&prefix); + Ok(()) + }, // return 0 and place u32::max_value() into written_out if no value exists for the key. ext_get_allocated_storage(key_data: *const u8, key_len: u32, written_out: *mut u32) -> *mut u8 => { let key = this.memory.get(key_data, key_len as usize).map_err(|_| DummyUserError)?; @@ -577,6 +582,29 @@ mod tests { assert_eq!(expected, ext); } + #[test] + fn clear_prefix_should_work() { + let mut ext = TestExternalities::default(); + ext.set_storage(b"aaa".to_vec(), b"1".to_vec()); + ext.set_storage(b"aab".to_vec(), b"2".to_vec()); + ext.set_storage(b"aba".to_vec(), b"3".to_vec()); + ext.set_storage(b"abb".to_vec(), b"4".to_vec()); + ext.set_storage(b"bbb".to_vec(), b"5".to_vec()); + let test_code = include_bytes!("../wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm"); + + // This will clear all entries which prefix is "ab". + let output = WasmExecutor.call(&mut ext, &test_code[..], "test_clear_prefix", b"ab").unwrap(); + + assert_eq!(output, b"all ok!".to_vec()); + + let expected: HashMap<_, _> = map![ + b"aaa".to_vec() => b"1".to_vec(), + b"aab".to_vec() => b"2".to_vec(), + b"bbb".to_vec() => b"5".to_vec() + ]; + assert_eq!(expected, ext); + } + #[test] fn blake2_256_should_work() { let mut ext = TestExternalities::default(); diff --git a/substrate/executor/wasm/src/lib.rs b/substrate/executor/wasm/src/lib.rs index 804dd3b89cb77..126d1bbf86c61 100644 --- a/substrate/executor/wasm/src/lib.rs +++ b/substrate/executor/wasm/src/lib.rs @@ -11,7 +11,7 @@ extern crate substrate_runtime_io as runtime_io; extern crate substrate_runtime_sandbox as sandbox; use runtime_io::{ - set_storage, storage, print, blake2_256, + set_storage, storage, clear_prefix, print, blake2_256, twox_128, twox_256, ed25519_verify, enumerated_trie_root }; @@ -29,6 +29,10 @@ impl_stubs!( print("finished!"); b"all ok!".to_vec() }, + test_clear_prefix NO_DECODE => |input| { + clear_prefix(input); + b"all ok!".to_vec() + }, test_empty_return NO_DECODE => |_| Vec::new(), test_panic NO_DECODE => |_| panic!("test panic"), test_conditional_panic NO_DECODE => |input: &[u8]| { diff --git a/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm b/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm index 470961909f0b5..b8dfdedeabbfb 100644 Binary files a/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm and b/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.compact.wasm differ diff --git a/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.wasm b/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.wasm index a7b0bdd9fc55b..7c7e9361eeea7 100755 Binary files a/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.wasm and b/substrate/executor/wasm/target/wasm32-unknown-unknown/release/runtime_test.wasm differ diff --git a/substrate/runtime-io/with_std.rs b/substrate/runtime-io/with_std.rs index 69a7a3665949f..93ac7cffa45ab 100644 --- a/substrate/runtime-io/with_std.rs +++ b/substrate/runtime-io/with_std.rs @@ -68,6 +68,13 @@ pub fn clear_storage(key: &[u8]) { ); } +/// Clear the storage entries key of which starts with the given prefix. +pub fn clear_prefix(prefix: &[u8]) { + ext::with(|ext| + ext.clear_prefix(prefix) + ); +} + /// The current relay chain identifier. pub fn chain_id() -> u64 { ext::with(|ext| @@ -211,4 +218,23 @@ mod std_tests { assert_eq!(&w, b"Hello world"); }); } + + #[test] + fn clear_prefix_works() { + let mut t: TestExternalities = map![ + b":a".to_vec() => b"\x0b\0\0\0Hello world".to_vec(), + b":abcd".to_vec() => b"\x0b\0\0\0Hello world".to_vec(), + b":abc".to_vec() => b"\x0b\0\0\0Hello world".to_vec(), + b":abdd".to_vec() => b"\x0b\0\0\0Hello world".to_vec() + ]; + + with_externalities(&mut t, || { + clear_prefix(b":abc"); + + assert!(storage(b":a").is_some()); + assert!(storage(b":abdd").is_some()); + assert!(storage(b":abcd").is_none()); + assert!(storage(b":abc").is_none()); + }); + } } diff --git a/substrate/runtime-io/without_std.rs b/substrate/runtime-io/without_std.rs index 3f8e1f31072a2..272f2fc04f5e1 100644 --- a/substrate/runtime-io/without_std.rs +++ b/substrate/runtime-io/without_std.rs @@ -56,6 +56,7 @@ extern "C" { fn ext_print_num(value: u64); fn ext_set_storage(key_data: *const u8, key_len: u32, value_data: *const u8, value_len: u32); fn ext_clear_storage(key_data: *const u8, key_len: u32); + fn ext_clear_prefix(prefix_data: *const u8, prefix_len: u32); fn ext_get_allocated_storage(key_data: *const u8, key_len: u32, written_out: *mut u32) -> *mut u8; fn ext_get_storage_into(key_data: *const u8, key_len: u32, value_data: *mut u8, value_len: u32, value_offset: u32) -> u32; fn ext_storage_root(result: *mut u8); @@ -80,7 +81,7 @@ pub fn storage(key: &[u8]) -> Option> { } } -/// Set the storage to some particular key. +/// Set the storage of some particular key to Some value. pub fn set_storage(key: &[u8], value: &[u8]) { unsafe { ext_set_storage( @@ -90,7 +91,7 @@ pub fn set_storage(key: &[u8], value: &[u8]) { } } -/// Set the storage to some particular key. +/// Clear the storage of some particular key. pub fn clear_storage(key: &[u8]) { unsafe { ext_clear_storage( @@ -99,6 +100,16 @@ pub fn clear_storage(key: &[u8]) { } } +/// Clear the storage entries key of which starts with the given prefix. +pub fn clear_prefix(prefix: &[u8]) { + unsafe { + ext_clear_prefix( + prefix.as_ptr(), + prefix.len() as u32 + ); + } +} + /// Get `key` from storage, placing the value into `value_out` (as much as possible) and return /// the number of bytes that the key in storage was beyond the offset. pub fn read_storage(key: &[u8], value_out: &mut [u8], value_offset: usize) -> Option { diff --git a/substrate/runtime-support/src/storage/mod.rs b/substrate/runtime-support/src/storage/mod.rs index 635a30848f5a2..2c9e2a5887fce 100644 --- a/substrate/runtime-support/src/storage/mod.rs +++ b/substrate/runtime-support/src/storage/mod.rs @@ -459,6 +459,11 @@ pub mod unhashed { runtime_io::clear_storage(key); } + /// Ensure keys with the given `prefix` have no entries in storage. + pub fn kill_prefix(prefix: &[u8]) { + runtime_io::clear_prefix(prefix); + } + /// Get a Vec of bytes from storage. pub fn get_raw(key: &[u8]) -> Option> { runtime_io::storage(key) diff --git a/substrate/runtime/staking/src/account_db.rs b/substrate/runtime/staking/src/account_db.rs index e7d883a656b8a..b41ce55458e65 100644 --- a/substrate/runtime/staking/src/account_db.rs +++ b/substrate/runtime/staking/src/account_db.rs @@ -20,6 +20,7 @@ use rstd::prelude::*; use rstd::cell::RefCell; use rstd::collections::btree_map::{BTreeMap, Entry}; use runtime_support::StorageMap; +use double_map::StorageDoubleMap; use super::*; pub struct ChangeEntry { @@ -61,7 +62,7 @@ pub trait AccountDb { pub struct DirectAccountDb; impl AccountDb for DirectAccountDb { fn get_storage(&self, account: &T::AccountId, location: &[u8]) -> Option> { - >::get(&(account.clone(), location.to_vec())) + >::get(account.clone(), location.to_vec()) } fn get_code(&self, account: &T::AccountId) -> Vec { >::get(account) @@ -105,9 +106,9 @@ impl AccountDb for DirectAccountDb { } for (k, v) in changed.storage.into_iter() { if let Some(value) = v { - >::insert((address.clone(), k), &value); + >::insert(address.clone(), k, value); } else { - >::remove((address.clone(), k)); + >::remove(address.clone(), k); } } } diff --git a/substrate/runtime/staking/src/double_map.rs b/substrate/runtime/staking/src/double_map.rs new file mode 100644 index 0000000000000..36c4be98d559c --- /dev/null +++ b/substrate/runtime/staking/src/double_map.rs @@ -0,0 +1,90 @@ +// Copyright 2017 Parity Technologies (UK) Ltd. +// This file is part of Substrate Demo. + +// Substrate Demo is free software: you can redistribute it and/or modify +// it under the terms of the GNU General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. + +// Substrate Demo is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU General Public License for more details. + +// You should have received a copy of the GNU General Public License +// along with Substrate Demo. If not, see . + +//! An implementation of double map backed by storage. +//! +//! This implementation is somewhat specialized to the tracking of the storage of accounts. + +use rstd::prelude::*; +use codec::Slicable; +use runtime_support::storage::unhashed; +use runtime_io::{blake2_256, twox_128}; + +/// Returns only a first part of the storage key. +/// +/// Hashed by XX. +fn first_part_of_key(k1: M::Key1) -> [u8; 16] { + let mut raw_prefix = Vec::new(); + raw_prefix.extend(M::PREFIX); + raw_prefix.extend(Slicable::encode(&k1)); + twox_128(&raw_prefix) +} + +/// Returns a compound key that consist of the two parts: (prefix, `k1`) and `k2`. +/// +/// The first part is hased by XX and then concatenated with a blake2 hash of `k2`. +fn full_key(k1: M::Key1, k2: M::Key2) -> Vec { + let first_part = first_part_of_key::(k1); + let second_part = blake2_256(&Slicable::encode(&k2)); + + let mut k = Vec::new(); + k.extend(&first_part); + k.extend(&second_part); + k +} + +/// An implementation of a map with a two keys. +/// +/// It provides an important ability to efficiently remove all entries +/// that have a common first key. +/// +/// # Mapping of keys to a storage path +/// +/// The storage key (i.e. the key under which the `Value` will be stored) is created from two parts. +/// The first part is a XX hash of a concatenation of the `PREFIX` and `Key1`. And the second part +/// is a blake2 hash of a `Key2`. +/// +/// Blake2 is used for `Key2` is because it will be used as a for a key for contract's storage and +/// thus will be susceptible for a untrusted input. +pub trait StorageDoubleMap { + type Key1: Slicable; + type Key2: Slicable; + type Value: Slicable + Default; + + const PREFIX: &'static [u8]; + + /// Insert an entry into this map. + fn insert(k1: Self::Key1, k2: Self::Key2, val: Self::Value) { + unhashed::put(&full_key::(k1, k2)[..], &val); + } + + /// Remove an entry from this map. + fn remove(k1: Self::Key1, k2: Self::Key2) { + unhashed::kill(&full_key::(k1, k2)[..]); + } + + /// Get an entry from this map. + /// + /// If there is entry stored under the given keys, returns `None`. + fn get(k1: Self::Key1, k2: Self::Key2) -> Option { + unhashed::get(&full_key::(k1, k2)[..]) + } + + /// Removes all entries that shares the `k1` as the first key. + fn remove_prefix(k1: Self::Key1) { + unhashed::kill_prefix(&first_part_of_key::(k1)) + } +} diff --git a/substrate/runtime/staking/src/lib.rs b/substrate/runtime/staking/src/lib.rs index c3c1800d4b47b..cf64fca9d7426 100644 --- a/substrate/runtime/staking/src/lib.rs +++ b/substrate/runtime/staking/src/lib.rs @@ -57,12 +57,14 @@ use session::OnSessionChange; use primitives::traits::{Zero, One, Bounded, RefInto, SimpleArithmetic, Executable, MakePayment, As, AuxLookup, Hashing as HashingT, Member}; use address::Address as RawAddress; +use double_map::StorageDoubleMap; pub mod address; mod mock; mod tests; mod genesis_config; mod account_db; +mod double_map; #[cfg(feature = "std")] pub use genesis_config::GenesisConfig; @@ -232,16 +234,17 @@ decl_storage! { // The code associated with an account. pub CodeOf: b"sta:cod:" => default map [ T::AccountId => Vec ]; // TODO Vec values should be optimised to not do a length prefix. +} - // The storage items associated with an account/key. - // TODO: keys should also be able to take AsRef to ensure Vecs can be passed as &[u8] - // TODO: This will need to be stored as a double-map, with `T::AccountId` using the usual XX hash - // function, and then the output of this concatenated onto a separate blake2 hash of the `Vec` - // key. We will then need a `remove_prefix` in addition to `set_storage` which removes all - // storage items with a particular prefix otherwise we'll suffer leakage with the removal - // of smart contracts. -// pub StorageOf: b"sta:sto:" => map [ T::AccountId => map(blake2) Vec => Vec ]; - pub StorageOf: b"sta:sto:" => map [ (T::AccountId, Vec) => Vec ]; +/// The storage items associated with an account/key. +/// +/// TODO: keys should also be able to take AsRef to ensure Vecs can be passed as &[u8] +pub(crate) struct StorageOf(::rstd::marker::PhantomData); +impl double_map::StorageDoubleMap for StorageOf { + type Key1 = T::AccountId; + type Key2 = Vec; + type Value = Vec; + const PREFIX: &'static [u8] = b"sta:sto:"; } enum NewAccountOutcome { @@ -623,7 +626,7 @@ impl Module { .map(|v| (Self::voting_balance(&v) + Self::nomination_balance(&v), v)) .collect::>(); intentions.sort_unstable_by(|&(ref b1, _), &(ref b2, _)| b2.cmp(&b1)); - + >::put( if intentions.len() > 0 { let i = (>::get() as usize).min(intentions.len() - 1); @@ -736,7 +739,7 @@ impl Module { >::remove(who); >::remove(who); >::remove(who); - // TODO: >::remove_prefix(address.clone()); + >::remove_prefix(who.clone()); if Self::reserved_balance(who).is_zero() { >::remove(who); diff --git a/substrate/runtime/staking/src/tests.rs b/substrate/runtime/staking/src/tests.rs index f19e84079bd63..0034c261986de 100644 --- a/substrate/runtime/staking/src/tests.rs +++ b/substrate/runtime/staking/src/tests.rs @@ -585,3 +585,35 @@ fn transferring_incomplete_reserved_balance_should_work() { assert_eq!(Staking::free_balance(&2), 42); }); } + +#[test] +fn account_removal_removes_storage() { + with_externalities(&mut new_test_ext(100, 1, 3, 1, false, 0), || { + // Setup two accounts with free balance above than exsistential threshold. + { + >::insert(1, 110); + >::insert(1, b"foo".to_vec(), b"1".to_vec()); + >::insert(1, b"bar".to_vec(), b"2".to_vec()); + + >::insert(2, 110); + >::insert(2, b"hello".to_vec(), b"3".to_vec()); + >::insert(2, b"world".to_vec(), b"4".to_vec()); + } + + // Transfer funds from account 1 of such amount that after this transfer + // the balance of account 1 is will be below than exsistential threshold. + // + // This should lead to the removal of all storage associated with this account. + assert_ok!(Staking::transfer(&1, 2.into(), 20)); + + // Verify that all entries from account 1 is removed, while + // entries from account 2 is in place. + { + assert_eq!(>::get(1, b"foo".to_vec()), None); + assert_eq!(>::get(1, b"bar".to_vec()), None); + + assert_eq!(>::get(2, b"hello".to_vec()), Some(b"3".to_vec())); + assert_eq!(>::get(2, b"world".to_vec()), Some(b"4".to_vec())); + } + }); +} diff --git a/substrate/state-machine/src/backend.rs b/substrate/state-machine/src/backend.rs index cbd2fc6ccf192..457a40dcc5382 100644 --- a/substrate/state-machine/src/backend.rs +++ b/substrate/state-machine/src/backend.rs @@ -35,6 +35,10 @@ pub trait Backend: TryIntoTrieBackend { /// Get keyed storage associated with specific address, or None if there is nothing associated. fn storage(&self, key: &[u8]) -> Result>, Self::Error>; + /// Retrieve all entries keys of which start with the given prefix and + /// call `f` for each of those keys. + fn for_keys_with_prefix(&self, prefix: &[u8], f: F); + /// Calculate the storage root, with given delta over what is already stored in /// the backend, and produce a "transaction" that can be used to commit. fn storage_root(&self, delta: I) -> ([u8; 32], Self::Transaction) @@ -105,6 +109,10 @@ impl Backend for InMemory { Ok(self.inner.get(key).map(Clone::clone)) } + fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { + self.inner.keys().filter(|key| key.starts_with(prefix)).map(|k| &**k).for_each(f); + } + fn storage_root(&self, delta: I) -> ([u8; 32], Self::Transaction) where I: IntoIterator, Option>)> { diff --git a/substrate/state-machine/src/ext.rs b/substrate/state-machine/src/ext.rs index 6ac27a71ef9e3..5c5b9b573627d 100644 --- a/substrate/state-machine/src/ext.rs +++ b/substrate/state-machine/src/ext.rs @@ -74,6 +74,13 @@ impl<'a, B: 'a + Backend> Ext<'a, B> { let _ = self.storage_root(); self.transaction.expect("transaction always set after calling storage root; qed").0 } + + /// Invalidates the currently cached storage root and the db transaction. + /// + /// Called when there are changes that likely will invalidate the storage root. + fn mark_dirty(&mut self) { + self.transaction = None; + } } #[cfg(test)] @@ -101,10 +108,17 @@ impl<'a, B: 'a> Externalities for Ext<'a, B> } fn place_storage(&mut self, key: Vec, value: Option>) { - self.transaction = None; // wipe out the transaction since root will no longer be the same. + self.mark_dirty(); self.overlay.set_storage(key, value); } + fn clear_prefix(&mut self, prefix: &[u8]) { + self.mark_dirty(); + self.backend.for_keys_with_prefix(prefix, |key| { + self.overlay.set_storage(key.to_vec(), None); + }); + } + fn chain_id(&self) -> u64 { 42 } diff --git a/substrate/state-machine/src/lib.rs b/substrate/state-machine/src/lib.rs index c7b4bdb391e8f..e0119edeaaa40 100644 --- a/substrate/state-machine/src/lib.rs +++ b/substrate/state-machine/src/lib.rs @@ -132,6 +132,9 @@ pub trait Externalities { self.place_storage(key.to_vec(), None); } + /// Clear storage entries which keys are start with the given prefix. + fn clear_prefix(&mut self, prefix: &[u8]); + /// Set or clear a storage entry (`key`) of current contract being called (effective immediately). fn place_storage(&mut self, key: Vec, value: Option>); diff --git a/substrate/state-machine/src/proving_backend.rs b/substrate/state-machine/src/proving_backend.rs index 688dba6ef23bb..70961adc96902 100644 --- a/substrate/state-machine/src/proving_backend.rs +++ b/substrate/state-machine/src/proving_backend.rs @@ -69,6 +69,10 @@ impl Backend for ProvingBackend { .get_with(key, &mut *proof_recorder).map(|x| x.map(|val| val.to_vec())).map_err(map_e) } + fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { + self.backend.for_keys_with_prefix(prefix, f) + } + fn pairs(&self) -> Vec<(Vec, Vec)> { self.backend.pairs() } @@ -134,7 +138,7 @@ mod tests { let proving_backend = test_proving(); assert_eq!(trie_backend.storage(b"key").unwrap(), proving_backend.storage(b"key").unwrap()); assert_eq!(trie_backend.pairs(), proving_backend.pairs()); - + let (trie_root, mut trie_mdb) = trie_backend.storage_root(::std::iter::empty()); let (proving_root, mut proving_mdb) = proving_backend.storage_root(::std::iter::empty()); assert_eq!(trie_root, proving_root); diff --git a/substrate/state-machine/src/testing.rs b/substrate/state-machine/src/testing.rs index e9b4a76db61b5..7b85b523cd47a 100644 --- a/substrate/state-machine/src/testing.rs +++ b/substrate/state-machine/src/testing.rs @@ -35,6 +35,12 @@ impl Externalities for TestExternalities { } } + fn clear_prefix(&mut self, prefix: &[u8]) { + self.retain(|key, _| + !key.starts_with(prefix) + ) + } + fn chain_id(&self) -> u64 { 42 } fn storage_root(&mut self) -> [u8; 32] { diff --git a/substrate/state-machine/src/trie_backend.rs b/substrate/state-machine/src/trie_backend.rs index 3419bb8f6753b..628bb04fbf941 100644 --- a/substrate/state-machine/src/trie_backend.rs +++ b/substrate/state-machine/src/trie_backend.rs @@ -99,6 +99,37 @@ impl Backend for TrieBackend { .get(key).map(|x| x.map(|val| val.to_vec())).map_err(map_e) } + fn for_keys_with_prefix(&self, prefix: &[u8], mut f: F) { + let mut read_overlay = MemoryDB::default(); + let eph = Ephemeral { + storage: &self.storage, + overlay: &mut read_overlay, + }; + + let mut iter = move || -> Result<(), Box> { + let trie = TrieDB::new(&eph, &self.root)?; + let mut iter = trie.iter()?; + + iter.seek(prefix)?; + + for x in iter { + let (key, _) = x?; + + if !key.starts_with(prefix) { + break; + } + + f(&key); + } + + Ok(()) + }; + + if let Err(e) = iter() { + debug!(target: "trie", "Error while iterating by prefix: {}", e); + } + } + fn pairs(&self) -> Vec<(Vec, Vec)> { let mut read_overlay = MemoryDB::default(); let eph = Ephemeral { @@ -238,6 +269,7 @@ impl TrieBackendStorage { #[cfg(test)] pub mod tests { use super::*; + use std::collections::HashSet; fn test_db() -> (MemoryDB, TrieH256) { let mut root = TrieH256::default(); @@ -293,4 +325,20 @@ pub mod tests { assert!(!tx.drain().is_empty()); assert!(new_root != test_trie().storage_root(::std::iter::empty()).0); } + + #[test] + fn prefix_walking_works() { + let trie = test_trie(); + + let mut seen = HashSet::new(); + trie.for_keys_with_prefix(b"value", |key| { + let for_first_time = seen.insert(key.to_vec()); + assert!(for_first_time, "Seen key '{:?}' more than once", key); + }); + + let mut expected = HashSet::new(); + expected.insert(b"value1".to_vec()); + expected.insert(b"value2".to_vec()); + assert_eq!(seen, expected); + } } diff --git a/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.compact.wasm b/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.compact.wasm index 80d8c4e431b3b..b7d5c940eb80c 100644 Binary files a/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.compact.wasm and b/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.compact.wasm differ diff --git a/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.wasm b/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.wasm index 36c7255551a0b..732797401ecd7 100755 Binary files a/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.wasm and b/substrate/test-runtime/wasm/target/wasm32-unknown-unknown/release/substrate_test_runtime.wasm differ