diff --git a/client/db/src/lib.rs b/client/db/src/lib.rs index d9365ce4bfba7..9f1316253b87e 100644 --- a/client/db/src/lib.rs +++ b/client/db/src/lib.rs @@ -151,6 +151,14 @@ impl StateBackend for RefTrackingState { self.state.exists_child_storage(storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + self.state.next_storage_key(key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + self.state.next_child_storage_key(storage_key, key) + } + fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { self.state.for_keys_with_prefix(prefix, f) } diff --git a/client/db/src/storage_cache.rs b/client/db/src/storage_cache.rs index 99266c7b61882..2b8e356d9cc82 100644 --- a/client/db/src/storage_cache.rs +++ b/client/db/src/storage_cache.rs @@ -544,6 +544,14 @@ impl, B: BlockT> StateBackend for CachingState< self.state.exists_child_storage(storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + self.state.next_storage_key(key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + self.state.next_child_storage_key(storage_key, key) + } + fn for_keys_in_child_storage(&self, storage_key: &[u8], f: F) { self.state.for_keys_in_child_storage(storage_key, f) } diff --git a/client/src/light/backend.rs b/client/src/light/backend.rs index b8dc0c34d7e35..0d974411a9d38 100644 --- a/client/src/light/backend.rs +++ b/client/src/light/backend.rs @@ -365,6 +365,22 @@ impl StateBackend for GenesisOrUnavailableState } } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + match *self { + GenesisOrUnavailableState::Genesis(ref state) => + Ok(state.next_storage_key(key).expect(IN_MEMORY_EXPECT_PROOF)), + GenesisOrUnavailableState::Unavailable => Err(ClientError::NotAvailableOnLightClient), + } + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + match *self { + GenesisOrUnavailableState::Genesis(ref state) => + Ok(state.next_child_storage_key(storage_key, key).expect(IN_MEMORY_EXPECT_PROOF)), + GenesisOrUnavailableState::Unavailable => Err(ClientError::NotAvailableOnLightClient), + } + } + fn for_keys_with_prefix(&self, prefix: &[u8], action: A) { match *self { GenesisOrUnavailableState::Genesis(ref state) => state.for_keys_with_prefix(prefix, action), diff --git a/frame/support/procedural/src/lib.rs b/frame/support/procedural/src/lib.rs index 506e55fa17597..d9c2fe03c9bf5 100644 --- a/frame/support/procedural/src/lib.rs +++ b/frame/support/procedural/src/lib.rs @@ -69,6 +69,7 @@ use proc_macro::TokenStream; /// * Map: `Foo: map hasher($hash) type => type`: Implements the /// [`StorageMap`](../frame_support/storage/trait.StorageMap.html) trait using the /// [`StorageMap generator`](../frame_support/storage/generator/trait.StorageMap.html). +/// And [`StoragePrefixedMap`](../frame_support/storage/trait.StoragePrefixedMap.html). /// /// `$hash` representing a choice of hashing algorithms available in the /// [`Hashable`](../frame_support/trait.Hashable.html) trait. @@ -89,6 +90,7 @@ use proc_macro::TokenStream; /// * Linked map: `Foo: linked_map hasher($hash) type => type`: Implements the /// [`StorageLinkedMap`](../frame_support/storage/trait.StorageLinkedMap.html) trait using the /// [`StorageLinkedMap generator`](../frame_support/storage/generator/trait.StorageLinkedMap.html). +/// And [`StoragePrefixedMap`](../frame_support/storage/trait.StoragePrefixedMap.html). /// /// `$hash` representing a choice of hashing algorithms available in the /// [`Hashable`](../frame_support/trait.Hashable.html) trait. @@ -118,6 +120,7 @@ use proc_macro::TokenStream; /// * Double map: `Foo: double_map hasher($hash1) u32, $hash2(u32) => u32`: Implements the /// [`StorageDoubleMap`](../frame_support/storage/trait.StorageDoubleMap.html) trait using the /// [`StorageDoubleMap generator`](../frame_support/storage/generator/trait.StorageDoubleMap.html). +/// And [`StoragePrefixedMap`](../frame_support/storage/trait.StoragePrefixedMap.html). /// /// `$hash1` and `$hash2` representing choices of hashing algorithms available in the /// [`Hashable`](../frame_support/trait.Hashable.html) trait. They must be choosen with care, see diff --git a/frame/support/procedural/src/storage/mod.rs b/frame/support/procedural/src/storage/mod.rs index 0ec266a8a0c54..8fbb97d916ba9 100644 --- a/frame/support/procedural/src/storage/mod.rs +++ b/frame/support/procedural/src/storage/mod.rs @@ -418,7 +418,8 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr StorageValue as _, StorageMap as _, StorageLinkedMap as _, - StorageDoubleMap as _ + StorageDoubleMap as _, + StoragePrefixedMap as _, }; #scrate_decl diff --git a/frame/support/procedural/src/storage/storage_struct.rs b/frame/support/procedural/src/storage/storage_struct.rs index 017e6cf2ff159..97aea3567d270 100644 --- a/frame/support/procedural/src/storage/storage_struct.rs +++ b/frame/support/procedural/src/storage/storage_struct.rs @@ -122,6 +122,18 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre StorageLineTypeDef::Map(map) => { let hasher = map.hasher.to_storage_hasher_struct(); quote!( + impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type> + for #storage_struct #optional_storage_where_clause + { + fn module_prefix() -> &'static [u8] { + #instance_or_inherent::PREFIX.as_bytes() + } + + fn storage_prefix() -> &'static [u8] { + #storage_name_str.as_bytes() + } + } + impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct #optional_storage_where_clause { @@ -155,6 +167,18 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre ); quote!( + impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type> + for #storage_struct #optional_storage_where_clause + { + fn module_prefix() -> &'static [u8] { + #instance_or_inherent::PREFIX.as_bytes() + } + + fn storage_prefix() -> &'static [u8] { + #storage_name_str.as_bytes() + } + } + impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct #optional_storage_where_clause { @@ -191,6 +215,18 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre let hasher1 = map.hasher1.to_storage_hasher_struct(); let hasher2 = map.hasher2.to_storage_hasher_struct(); quote!( + impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type> + for #storage_struct #optional_storage_where_clause + { + fn module_prefix() -> &'static [u8] { + #instance_or_inherent::PREFIX.as_bytes() + } + + fn storage_prefix() -> &'static [u8] { + #storage_name_str.as_bytes() + } + } + impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct #optional_storage_where_clause { diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs index 0a7a967db85f3..35962eba88d5c 100644 --- a/frame/support/src/lib.rs +++ b/frame/support/src/lib.rs @@ -67,7 +67,9 @@ pub mod traits; pub mod weights; pub use self::hash::{Twox256, Twox128, Blake2_256, Blake2_128, Twox64Concat, Hashable}; -pub use self::storage::{StorageValue, StorageMap, StorageLinkedMap, StorageDoubleMap}; +pub use self::storage::{ + StorageValue, StorageMap, StorageLinkedMap, StorageDoubleMap, StoragePrefixedMap +}; pub use self::dispatch::{Parameter, Callable, IsSubType}; pub use sp_runtime::{self, ConsensusEngineId, print, traits::Printable}; diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 7c8ce9c24f868..08d45571d1553 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -16,9 +16,9 @@ //! Stuff to do with the runtime's storage. -use rstd::prelude::*; +use rstd::{prelude::*, marker::PhantomData}; use codec::{FullCodec, FullEncode, Encode, EncodeAppend, EncodeLike, Decode}; -use crate::traits::Len; +use crate::{traits::Len, hash::{Twox128, StorageHasher}}; pub mod unhashed; pub mod hashed; @@ -352,3 +352,126 @@ pub trait StorageDoubleMap { KArg2: EncodeLike, V: codec::DecodeLength + Len; } + +/// Iterator for prefixed map. +pub struct PrefixIterator { + prefix: Vec, + previous_key: Vec, + phantom_data: PhantomData, +} + +impl Iterator for PrefixIterator { + type Item = Value; + + fn next(&mut self) -> Option { + match runtime_io::storage::next_key(&self.previous_key) { + Some(next_key) if next_key.starts_with(&self.prefix[..]) => { + let value = unhashed::get(&next_key); + + if value.is_none() { + runtime_print!( + "ERROR: returned next_key has no value:\nkey is {:?}\nnext_key is {:?}", + &self.previous_key, &next_key, + ); + } + + self.previous_key = next_key; + + value + }, + _ => None, + } + } +} + +/// Trait for maps that store all its value after a unique prefix. +/// +/// By default the final prefix is: +/// ```nocompile +/// Twox128(module_prefix) ++ Twox128(storage_prefix) +/// ``` +pub trait StoragePrefixedMap { + + /// Module prefix. Used for generating final key. + fn module_prefix() -> &'static [u8]; + + /// Storage prefix. Used for generating final key. + fn storage_prefix() -> &'static [u8]; + + fn final_prefix() -> [u8; 32] { + let mut final_key = [0u8; 32]; + final_key[0..16].copy_from_slice(&Twox128::hash(Self::module_prefix())); + final_key[16..32].copy_from_slice(&Twox128::hash(Self::storage_prefix())); + final_key + } + + fn remove_all() { + runtime_io::storage::clear_prefix(&Self::final_prefix()) + } + + fn iter() -> PrefixIterator { + let prefix = Self::final_prefix(); + PrefixIterator { + prefix: prefix.to_vec(), + previous_key: prefix.to_vec(), + phantom_data: Default::default(), + } + } +} + +#[cfg(test)] +mod test { + use primitives::hashing::twox_128; + use runtime_io::TestExternalities; + use crate::storage::{unhashed, StoragePrefixedMap}; + + #[test] + fn prefixed_map_works() { + TestExternalities::default().execute_with(|| { + struct MyStorage; + impl StoragePrefixedMap for MyStorage { + fn module_prefix() -> &'static [u8] { + b"MyModule" + } + + fn storage_prefix() -> &'static [u8] { + b"MyStorage" + } + } + + let key_before = { + let mut k = MyStorage::final_prefix(); + let last = k.iter_mut().last().unwrap(); + *last = last.checked_sub(1).unwrap(); + k + }; + let key_after = { + let mut k = MyStorage::final_prefix(); + let last = k.iter_mut().last().unwrap(); + *last = last.checked_add(1).unwrap(); + k + }; + + unhashed::put(&key_before[..], &32u64); + unhashed::put(&key_after[..], &33u64); + + let k = [twox_128(b"MyModule"), twox_128(b"MyStorage")].concat(); + assert_eq!(MyStorage::final_prefix().to_vec(), k); + + assert_eq!(MyStorage::iter().collect::>(), vec![]); + + unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u64); + unhashed::put(&[&k[..], &vec![1, 1][..]].concat(), &2u64); + unhashed::put(&[&k[..], &vec![8][..]].concat(), &3u64); + unhashed::put(&[&k[..], &vec![10][..]].concat(), &4u64); + + assert_eq!(MyStorage::iter().collect::>(), vec![1, 2, 3, 4]); + + MyStorage::remove_all(); + + assert_eq!(MyStorage::iter().collect::>(), vec![]); + assert_eq!(unhashed::get(&key_before[..]), Some(32u64)); + assert_eq!(unhashed::get(&key_after[..]), Some(33u64)); + }); + } +} diff --git a/frame/support/test/tests/final_keys.rs b/frame/support/test/tests/final_keys.rs index 365aa5779d11c..3c9e1058a8cfa 100644 --- a/frame/support/test/tests/final_keys.rs +++ b/frame/support/test/tests/final_keys.rs @@ -16,7 +16,7 @@ use support::storage::unhashed; use codec::Encode; -use support::{StorageDoubleMap, StorageLinkedMap, StorageMap, StorageValue}; +use support::{StorageDoubleMap, StorageLinkedMap, StorageMap, StorageValue, StoragePrefixedMap}; use runtime_io::{TestExternalities, hashing::{twox_128, blake2_128, blake2_256}}; mod no_instance { @@ -96,11 +96,13 @@ fn final_keys_no_instance() { let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &::final_prefix()); no_instance::Map2::insert(1, 2); let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"Map2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &::final_prefix()); let head = [twox_128(b"FinalKeysNone"), twox_128(b"HeadOfLinkedMap")].concat(); assert_eq!(unhashed::get::(&head), None); @@ -110,23 +112,27 @@ fn final_keys_no_instance() { k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); assert_eq!(unhashed::get::(&head), Some(1u32)); + assert_eq!(&k[..32], &::final_prefix()); no_instance::LinkedMap2::insert(1, 2); let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"LinkedMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &::final_prefix()); no_instance::DoubleMap::insert(&1, &2, &3); let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"DoubleMap")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); k.extend(2u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &::final_prefix()); no_instance::DoubleMap2::insert(&1, &2, &3); let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"DoubleMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); k.extend(2u32.using_encoded(blake2_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &::final_prefix()); }); } @@ -141,11 +147,13 @@ fn final_keys_default_instance() { let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(1, 2); let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"Map2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); let head = [twox_128(b"FinalKeysSome"), twox_128(b"HeadOfLinkedMap")].concat(); assert_eq!(unhashed::get::(&head), None); @@ -155,23 +163,27 @@ fn final_keys_default_instance() { k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); assert_eq!(unhashed::get::(&head), Some(1u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(1, 2); let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"LinkedMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(&1, &2, &3); let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"DoubleMap")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); k.extend(2u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(&1, &2, &3); let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"DoubleMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); k.extend(2u32.using_encoded(blake2_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &>::final_prefix()); }); } @@ -186,11 +198,13 @@ fn final_keys_instance_2() { let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(1, 2); let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"Map2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); let head = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"HeadOfLinkedMap")].concat(); assert_eq!(unhashed::get::(&head), None); @@ -200,22 +214,26 @@ fn final_keys_instance_2() { k.extend(1u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); assert_eq!(unhashed::get::(&head), Some(1u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(1, 2); let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"LinkedMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(2u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(&1, &2, &3); let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"DoubleMap")].concat(); k.extend(1u32.using_encoded(blake2_256).to_vec()); k.extend(2u32.using_encoded(blake2_256).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &>::final_prefix()); >::insert(&1, &2, &3); let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"DoubleMap2")].concat(); k.extend(1u32.using_encoded(twox_128).to_vec()); k.extend(2u32.using_encoded(blake2_128).to_vec()); assert_eq!(unhashed::get::(&k), Some(3u32)); + assert_eq!(&k[..32], &>::final_prefix()); }); } diff --git a/frame/support/test/tests/instance.rs b/frame/support/test/tests/instance.rs index a3ef6ba86881d..fb8575f2eafba 100644 --- a/frame/support/test/tests/instance.rs +++ b/frame/support/test/tests/instance.rs @@ -300,7 +300,7 @@ fn new_test_ext() -> runtime_io::TestExternalities { #[test] fn storage_instance_independance() { - let mut storage = (std::collections::HashMap::new(), std::collections::HashMap::new()); + let mut storage = Default::default(); state_machine::BasicExternalities::execute_with_storage(&mut storage, || { module2::Value::::put(0); module2::Value::::put(0); diff --git a/primitives/core/storage/src/lib.rs b/primitives/core/storage/src/lib.rs index ebb23023b96b2..3fd883788c968 100644 --- a/primitives/core/storage/src/lib.rs +++ b/primitives/core/storage/src/lib.rs @@ -42,7 +42,7 @@ pub struct StorageData( /// A set of key value pairs for storage. #[cfg(feature = "std")] -pub type StorageOverlay = std::collections::HashMap, Vec>; +pub type StorageOverlay = std::collections::BTreeMap, Vec>; /// A set of key value pairs for children storage; #[cfg(feature = "std")] diff --git a/primitives/externalities/src/lib.rs b/primitives/externalities/src/lib.rs index da3fe16d77945..05121f34d3dec 100644 --- a/primitives/externalities/src/lib.rs +++ b/primitives/externalities/src/lib.rs @@ -106,6 +106,12 @@ pub trait Externalities: ExtensionStore { self.child_storage(storage_key, key).is_some() } + /// Returns the key immediately following the given key, if it exists. + fn next_storage_key(&self, key: &[u8]) -> Option>; + + /// Returns the key immediately following the given key, if it exists, in child storage. + fn next_child_storage_key(&self, storage_key: ChildStorageKey, key: &[u8]) -> Option>; + /// Clear an entire child storage. fn kill_child_storage(&mut self, storage_key: ChildStorageKey); diff --git a/primitives/sr-io/src/lib.rs b/primitives/sr-io/src/lib.rs index d4b654ab2931e..7fdab1ab3b885 100644 --- a/primitives/sr-io/src/lib.rs +++ b/primitives/sr-io/src/lib.rs @@ -213,6 +213,17 @@ pub trait Storage { fn changes_root(&mut self, parent_hash: &[u8]) -> Option> { self.storage_changes_root(parent_hash).ok().and_then(|h| h) } + + /// Get the next key in storage after the given one in lexicographic order. + fn next_key(&mut self, key: &[u8]) -> Option> { + self.next_storage_key(&key) + } + + /// Get the next key in storage after the given one in lexicographic order in child storage. + fn child_next_key(&mut self, child_storage_key: &[u8], key: &[u8]) -> Option> { + let storage_key = child_storage_key_or_panic(child_storage_key); + self.next_child_storage_key(storage_key, key) + } } /// Interface that provides trie related functionality. diff --git a/primitives/state-machine/src/backend.rs b/primitives/state-machine/src/backend.rs index 366634ae1a9d7..78e08df6b9f74 100644 --- a/primitives/state-machine/src/backend.rs +++ b/primitives/state-machine/src/backend.rs @@ -16,7 +16,7 @@ //! State machine backends. These manage the code and storage of contracts. -use std::{error, fmt, cmp::Ord, collections::HashMap, marker::PhantomData}; +use std::{error, fmt, cmp::Ord, collections::{HashMap, BTreeMap}, marker::PhantomData, ops}; use log::warn; use hash_db::Hasher; use crate::trie_backend::TrieBackend; @@ -67,6 +67,16 @@ pub trait Backend: std::fmt::Debug { Ok(self.child_storage(storage_key, key)?.is_some()) } + /// Return the next key in storage in lexicographic order or `None` if there is no value. + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error>; + + /// Return the next key in child storage in lexicographic order or `None` if there is no value. + fn next_child_storage_key( + &self, + storage_key: &[u8], + key: &[u8] + ) -> Result>, Self::Error>; + /// Retrieve all entries keys of child storage and call `f` for each of those keys. fn for_keys_in_child_storage(&self, storage_key: &[u8], f: F); @@ -171,6 +181,14 @@ impl<'a, T: Backend, H: Hasher> Backend for &'a T { (*self).child_storage(storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + (*self).next_storage_key(key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + (*self).next_child_storage_key(storage_key, key) + } + fn for_keys_in_child_storage(&self, storage_key: &[u8], f: F) { (*self).for_keys_in_child_storage(storage_key, f) } @@ -250,7 +268,7 @@ impl error::Error for Void { /// In-memory backend. Fully recomputes tries each time `as_trie_backend` is called but useful for /// tests and proof checking. pub struct InMemory { - inner: HashMap>, HashMap, Vec>>, + inner: HashMap>, BTreeMap, Vec>>, // This field is only needed for returning reference in `as_trie_backend`. trie: Option, H>>, _hasher: PhantomData, @@ -291,7 +309,7 @@ impl PartialEq for InMemory { impl InMemory where H::Out: Codec { /// Copy the state, with applied updates pub fn update(&self, changes: >::Transaction) -> Self { - let mut inner: HashMap<_, _> = self.inner.clone(); + let mut inner = self.inner.clone(); for (storage_key, key, val) in changes { match val { Some(v) => { inner.entry(storage_key).or_default().insert(key, v); }, @@ -303,8 +321,8 @@ impl InMemory where H::Out: Codec { } } -impl From>, HashMap, Vec>>> for InMemory { - fn from(inner: HashMap>, HashMap, Vec>>) -> Self { +impl From>, BTreeMap, Vec>>> for InMemory { + fn from(inner: HashMap>, BTreeMap, Vec>>) -> Self { InMemory { inner: inner, trie: None, @@ -314,14 +332,14 @@ impl From>, HashMap, Vec>>> for In } impl From<( - HashMap, Vec>, - HashMap, HashMap, Vec>>, + BTreeMap, Vec>, + HashMap, BTreeMap, Vec>>, )> for InMemory { fn from(inners: ( - HashMap, Vec>, - HashMap, HashMap, Vec>>, + BTreeMap, Vec>, + HashMap, BTreeMap, Vec>>, )) -> Self { - let mut inner: HashMap>, HashMap, Vec>> + let mut inner: HashMap>, BTreeMap, Vec>> = inners.1.into_iter().map(|(k, v)| (Some(k), v)).collect(); inner.insert(None, inners.0); InMemory { @@ -332,8 +350,8 @@ impl From<( } } -impl From, Vec>> for InMemory { - fn from(inner: HashMap, Vec>) -> Self { +impl From, Vec>> for InMemory { + fn from(inner: BTreeMap, Vec>) -> Self { let mut expanded = HashMap::new(); expanded.insert(None, inner); InMemory { @@ -346,7 +364,7 @@ impl From, Vec>> for InMemory { impl From>, Vec, Option>)>> for InMemory { fn from(inner: Vec<(Option>, Vec, Option>)>) -> Self { - let mut expanded: HashMap>, HashMap, Vec>> = HashMap::new(); + let mut expanded: HashMap>, BTreeMap, Vec>> = HashMap::new(); for (child_key, key, value) in inner { if let Some(value) = value { expanded.entry(child_key).or_default().insert(key, value); @@ -380,6 +398,22 @@ impl Backend for InMemory where H::Out: Codec { Ok(self.inner.get(&None).map(|map| map.get(key).is_some()).unwrap_or(false)) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded); + let next_key = self.inner.get(&None) + .and_then(|map| map.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()); + + Ok(next_key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded); + let next_key = self.inner.get(&Some(storage_key.to_vec())) + .and_then(|map| map.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()); + + Ok(next_key) + } + fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { self.inner.get(&None).map(|map| map.keys().filter(|key| key.starts_with(prefix)).map(|k| &**k).for_each(f)); } diff --git a/primitives/state-machine/src/basic.rs b/primitives/state-machine/src/basic.rs index 82a0c21f20969..deae7f2852592 100644 --- a/primitives/state-machine/src/basic.rs +++ b/primitives/state-machine/src/basic.rs @@ -16,7 +16,9 @@ //! Basic implementation for Externalities. -use std::{collections::HashMap, any::{TypeId, Any}, iter::FromIterator}; +use std::{ + collections::{HashMap, BTreeMap}, any::{TypeId, Any}, iter::FromIterator, mem, ops::Bound +}; use crate::backend::{Backend, InMemory}; use hash_db::Hasher; use trie::{TrieConfiguration, default_child_trie_root}; @@ -54,8 +56,8 @@ impl BasicExternalities { /// Consume self and returns inner storages pub fn into_storages(self) -> ( - HashMap, Vec>, - HashMap, HashMap, Vec>>, + BTreeMap, Vec>, + HashMap, BTreeMap, Vec>>, ) { (self.top, self.children) } @@ -68,8 +70,8 @@ impl BasicExternalities { f: impl FnOnce() -> R, ) -> R { let mut ext = Self { - top: storage.0.drain().collect(), - children: storage.1.drain().collect(), + top: mem::replace(&mut storage.0, BTreeMap::default()), + children: mem::replace(&mut storage.1, HashMap::default()), }; let r = ext.execute_with(f); @@ -105,8 +107,8 @@ impl Default for BasicExternalities { fn default() -> Self { Self::new(Default::default(), Default::default()) } } -impl From, Vec>> for BasicExternalities { - fn from(hashmap: HashMap, Vec>) -> Self { +impl From, Vec>> for BasicExternalities { + fn from(hashmap: BTreeMap, Vec>) -> Self { BasicExternalities { top: hashmap, children: Default::default(), @@ -151,6 +153,17 @@ impl Externalities for BasicExternalities { Externalities::child_storage(self, storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Option> { + let range = (Bound::Excluded(key), Bound::Unbounded); + self.top.range::<[u8], _>(range).next().map(|(k, _)| k).cloned() + } + + fn next_child_storage_key(&self, storage_key: ChildStorageKey, key: &[u8]) -> Option> { + let range = (Bound::Excluded(key), Bound::Unbounded); + self.children.get(storage_key.as_ref()) + .and_then(|child| child.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()) + } + fn place_storage(&mut self, key: Vec, maybe_value: Option>) { if is_child_storage_key(&key) { warn!(target: "trie", "Refuse to set child storage key via main storage"); @@ -190,12 +203,28 @@ impl Externalities for BasicExternalities { return; } - self.top.retain(|key, _| !key.starts_with(prefix)); + let to_remove = self.top.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded)) + .map(|(k, _)| k) + .take_while(|k| k.starts_with(prefix)) + .cloned() + .collect::>(); + + for key in to_remove { + self.top.remove(&key); + } } fn clear_child_prefix(&mut self, storage_key: ChildStorageKey, prefix: &[u8]) { if let Some(child) = self.children.get_mut(storage_key.as_ref()) { - child.retain(|key, _| !key.starts_with(prefix)); + let to_remove = child.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded)) + .map(|(k, _)| k) + .take_while(|k| k.starts_with(prefix)) + .cloned() + .collect::>(); + + for key in to_remove { + child.remove(&key); + } } } diff --git a/primitives/state-machine/src/changes_trie/build.rs b/primitives/state-machine/src/changes_trie/build.rs index 6c50c028ca66e..7e082ad83276b 100644 --- a/primitives/state-machine/src/changes_trie/build.rs +++ b/primitives/state-machine/src/changes_trie/build.rs @@ -352,7 +352,7 @@ mod test { (vec![103], vec![255]), (vec![104], vec![255]), (vec![105], vec![255]), - ].into_iter().collect::<::std::collections::HashMap<_, _>>().into(); + ].into_iter().collect::<::std::collections::BTreeMap<_, _>>().into(); let child_trie_key1 = b"1".to_vec(); let child_trie_key2 = b"2".to_vec(); let storage = InMemoryStorage::with_inputs(vec![ diff --git a/primitives/state-machine/src/ext.rs b/primitives/state-machine/src/ext.rs index 57197a4ae3c45..3ac8c19048451 100644 --- a/primitives/state-machine/src/ext.rs +++ b/primitives/state-machine/src/ext.rs @@ -336,6 +336,40 @@ where result } + fn next_storage_key(&self, key: &[u8]) -> Option> { + let next_backend_key = self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL); + let next_overlay_key_change = self.overlay.next_storage_key_change(key); + + match (next_backend_key, next_overlay_key_change) { + (Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key), + (backend_key, None) => backend_key, + (_, Some(overlay_key)) => if overlay_key.1.value.is_some() { + Some(overlay_key.0.to_vec()) + } else { + self.next_storage_key(&overlay_key.0[..]) + }, + } + } + + fn next_child_storage_key(&self, storage_key: ChildStorageKey, key: &[u8]) -> Option> { + let next_backend_key = self.backend.next_child_storage_key(storage_key.as_ref(), key) + .expect(EXT_NOT_ALLOWED_TO_FAIL); + let next_overlay_key_change = self.overlay.next_child_storage_key_change( + storage_key.as_ref(), + key + ); + + match (next_backend_key, next_overlay_key_change) { + (Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key), + (backend_key, None) => backend_key, + (_, Some(overlay_key)) => if overlay_key.1.value.is_some() { + Some(overlay_key.0.to_vec()) + } else { + self.next_child_storage_key(storage_key, &overlay_key.0[..]) + }, + } + } + fn place_storage(&mut self, key: Vec, value: Option>) { trace!(target: "state-trace", "{:04x}: Put {}={:?}", self.id, @@ -619,4 +653,71 @@ mod tests { Some(hex!("96f5aae4690e7302737b6f9b7f8567d5bbb9eac1c315f80101235a92d9ec27f4").to_vec()), ); } + + #[test] + fn next_storage_key_works() { + let mut overlay = OverlayedChanges::default(); + overlay.set_storage(vec![20], None); + overlay.set_storage(vec![30], Some(vec![31])); + let backend = vec![ + (None, vec![10], Some(vec![10])), + (None, vec![20], Some(vec![20])), + (None, vec![40], Some(vec![40])), + ].into(); + + let ext = TestExt::new(&mut overlay, &backend, None, None); + + // next_backend < next_overlay + assert_eq!(ext.next_storage_key(&[5]), Some(vec![10])); + + // next_backend == next_overlay but next_overlay is a delete + assert_eq!(ext.next_storage_key(&[10]), Some(vec![30])); + + // next_overlay < next_backend + assert_eq!(ext.next_storage_key(&[20]), Some(vec![30])); + + // next_backend exist but next_overlay doesn't exist + assert_eq!(ext.next_storage_key(&[30]), Some(vec![40])); + + drop(ext); + overlay.set_storage(vec![50], Some(vec![50])); + let ext = TestExt::new(&mut overlay, &backend, None, None); + + // next_overlay exist but next_backend doesn't exist + assert_eq!(ext.next_storage_key(&[40]), Some(vec![50])); + } + + #[test] + fn next_child_storage_key_works() { + let child = || ChildStorageKey::from_slice(b":child_storage:default:Child1").unwrap(); + let mut overlay = OverlayedChanges::default(); + overlay.set_child_storage(child().as_ref().to_vec(), vec![20], None); + overlay.set_child_storage(child().as_ref().to_vec(), vec![30], Some(vec![31])); + let backend = vec![ + (Some(child().as_ref().to_vec()), vec![10], Some(vec![10])), + (Some(child().as_ref().to_vec()), vec![20], Some(vec![20])), + (Some(child().as_ref().to_vec()), vec![40], Some(vec![40])), + ].into(); + + let ext = TestExt::new(&mut overlay, &backend, None, None); + + // next_backend < next_overlay + assert_eq!(ext.next_child_storage_key(child(), &[5]), Some(vec![10])); + + // next_backend == next_overlay but next_overlay is a delete + assert_eq!(ext.next_child_storage_key(child(), &[10]), Some(vec![30])); + + // next_overlay < next_backend + assert_eq!(ext.next_child_storage_key(child(), &[20]), Some(vec![30])); + + // next_backend exist but next_overlay doesn't exist + assert_eq!(ext.next_child_storage_key(child(), &[30]), Some(vec![40])); + + drop(ext); + overlay.set_child_storage(child().as_ref().to_vec(), vec![50], Some(vec![50])); + let ext = TestExt::new(&mut overlay, &backend, None, None); + + // next_overlay exist but next_backend doesn't exist + assert_eq!(ext.next_child_storage_key(child(), &[40]), Some(vec![50])); + } } diff --git a/primitives/state-machine/src/lib.rs b/primitives/state-machine/src/lib.rs index d4798ad4fb185..594e539b25fe9 100644 --- a/primitives/state-machine/src/lib.rs +++ b/primitives/state-machine/src/lib.rs @@ -731,7 +731,7 @@ fn try_read_overlay_value( #[cfg(test)] mod tests { - use std::collections::HashMap; + use std::collections::BTreeMap; use codec::Encode; use overlayed_changes::OverlayedValue; use super::*; @@ -921,7 +921,7 @@ mod tests { #[test] fn clear_prefix_in_ext_works() { - let initial: HashMap<_, _> = map![ + let initial: BTreeMap<_, _> = map![ b"aaa".to_vec() => b"0".to_vec(), b"abb".to_vec() => b"1".to_vec(), b"abc".to_vec() => b"2".to_vec(), diff --git a/primitives/state-machine/src/overlayed_changes.rs b/primitives/state-machine/src/overlayed_changes.rs index 9bbfd68f67893..d61d14961da36 100644 --- a/primitives/state-machine/src/overlayed_changes.rs +++ b/primitives/state-machine/src/overlayed_changes.rs @@ -18,10 +18,11 @@ #[cfg(test)] use std::iter::FromIterator; -use std::collections::{HashMap, BTreeSet}; +use std::collections::{HashMap, BTreeMap, BTreeSet}; use codec::Decode; use crate::changes_trie::{NO_EXTRINSIC_INDEX, Configuration as ChangesTrieConfig}; use primitives::storage::well_known_keys::EXTRINSIC_INDEX; +use std::{mem, ops}; /// The overlayed changes to state to be queried on top of the backend. /// @@ -54,9 +55,9 @@ pub struct OverlayedValue { #[cfg_attr(test, derive(PartialEq))] pub struct OverlayedChangeSet { /// Top level storage changes. - pub top: HashMap, OverlayedValue>, + pub top: BTreeMap, OverlayedValue>, /// Child storage changes. - pub children: HashMap, HashMap, OverlayedValue>>, + pub children: HashMap, BTreeMap, OverlayedValue>>, } #[cfg(test)] @@ -274,9 +275,10 @@ impl OverlayedChanges { /// Commit prospective changes to state. pub fn commit_prospective(&mut self) { if self.committed.is_empty() { - ::std::mem::swap(&mut self.prospective, &mut self.committed); + mem::swap(&mut self.prospective, &mut self.committed); } else { - for (key, val) in self.prospective.top.drain() { + let top_to_commit = mem::replace(&mut self.prospective.top, BTreeMap::new()); + for (key, val) in top_to_commit.into_iter() { let entry = self.committed.top.entry(key).or_default(); entry.value = val.value; @@ -285,9 +287,9 @@ impl OverlayedChanges { .extend(prospective_extrinsics); } } - for (storage_key, mut map) in self.prospective.children.drain() { + for (storage_key, map) in self.prospective.children.drain() { let map_dest = self.committed.children.entry(storage_key).or_default(); - for (key, val) in map.drain() { + for (key, val) in map.into_iter() { let entry = map_dest.entry(key).or_default(); entry.value = val.value; @@ -339,6 +341,56 @@ impl OverlayedChanges { false => None, } } + + /// Returns the next (in lexicographic order) storage key in the overlayed alongside its value. + /// If no value is next then `None` is returned. + pub fn next_storage_key_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> { + let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded); + + let next_prospective_key = self.prospective.top + .range::<[u8], _>(range) + .next() + .map(|(k, v)| (&k[..], v)); + + let next_committed_key = self.committed.top + .range::<[u8], _>(range) + .next() + .map(|(k, v)| (&k[..], v)); + + match (next_committed_key, next_prospective_key) { + // Committed is strictly less than prospective + (Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 => + Some(committed_key), + (committed_key, None) => committed_key, + // Prospective key is less or equal to committed or committed doesn't exist + (_, prospective_key) => prospective_key, + } + } + + /// Returns the next (in lexicographic order) child storage key in the overlayed alongside its + /// value. If no value is next then `None` is returned. + pub fn next_child_storage_key_change( + &self, + storage_key: &[u8], + key: &[u8] + ) -> Option<(&[u8], &OverlayedValue)> { + let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded); + + let next_prospective_key = self.prospective.children.get(storage_key) + .and_then(|map| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v))); + + let next_committed_key = self.committed.children.get(storage_key) + .and_then(|map| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v))); + + match (next_committed_key, next_prospective_key) { + // Committed is strictly less than prospective + (Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 => + Some(committed_key), + (committed_key, None) => committed_key, + // Prospective key is less or equal to committed or committed doesn't exist + (_, prospective_key) => prospective_key, + } + } } #[cfg(test)] @@ -359,8 +411,8 @@ mod tests { use crate::ext::Ext; use super::*; - fn strip_extrinsic_index(map: &HashMap, OverlayedValue>) - -> HashMap, OverlayedValue> + fn strip_extrinsic_index(map: &BTreeMap, OverlayedValue>) + -> BTreeMap, OverlayedValue> { let mut clone = map.clone(); clone.remove(&EXTRINSIC_INDEX.to_vec()); @@ -397,7 +449,7 @@ mod tests { #[test] fn overlayed_storage_root_works() { - let initial: HashMap<_, _> = vec![ + let initial: BTreeMap<_, _> = vec![ (b"doe".to_vec(), b"reindeer".to_vec()), (b"dog".to_vec(), b"puppyXXX".to_vec()), (b"dogglesworth".to_vec(), b"catXXX".to_vec()), @@ -543,4 +595,79 @@ mod tests { assert_eq!(overlay.prospective, Default::default()); } + + #[test] + fn next_storage_key_change_works() { + let mut overlay = OverlayedChanges::default(); + overlay.set_storage(vec![20], Some(vec![20])); + overlay.set_storage(vec![30], Some(vec![30])); + overlay.set_storage(vec![40], Some(vec![40])); + overlay.commit_prospective(); + overlay.set_storage(vec![10], Some(vec![10])); + overlay.set_storage(vec![30], None); + + // next_prospective < next_committed + let next_to_5 = overlay.next_storage_key_change(&[5]).unwrap(); + assert_eq!(next_to_5.0.to_vec(), vec![10]); + assert_eq!(next_to_5.1.value, Some(vec![10])); + + // next_committed < next_prospective + let next_to_10 = overlay.next_storage_key_change(&[10]).unwrap(); + assert_eq!(next_to_10.0.to_vec(), vec![20]); + assert_eq!(next_to_10.1.value, Some(vec![20])); + + // next_committed == next_prospective + let next_to_20 = overlay.next_storage_key_change(&[20]).unwrap(); + assert_eq!(next_to_20.0.to_vec(), vec![30]); + assert_eq!(next_to_20.1.value, None); + + // next_committed, no next_prospective + let next_to_30 = overlay.next_storage_key_change(&[30]).unwrap(); + assert_eq!(next_to_30.0.to_vec(), vec![40]); + assert_eq!(next_to_30.1.value, Some(vec![40])); + + overlay.set_storage(vec![50], Some(vec![50])); + // next_prospective, no next_committed + let next_to_40 = overlay.next_storage_key_change(&[40]).unwrap(); + assert_eq!(next_to_40.0.to_vec(), vec![50]); + assert_eq!(next_to_40.1.value, Some(vec![50])); + } + + #[test] + fn next_child_storage_key_change_works() { + let child = b"Child1".to_vec(); + let mut overlay = OverlayedChanges::default(); + overlay.set_child_storage(child.clone(), vec![20], Some(vec![20])); + overlay.set_child_storage(child.clone(), vec![30], Some(vec![30])); + overlay.set_child_storage(child.clone(), vec![40], Some(vec![40])); + overlay.commit_prospective(); + overlay.set_child_storage(child.clone(), vec![10], Some(vec![10])); + overlay.set_child_storage(child.clone(), vec![30], None); + + // next_prospective < next_committed + let next_to_5 = overlay.next_child_storage_key_change(&child, &[5]).unwrap(); + assert_eq!(next_to_5.0.to_vec(), vec![10]); + assert_eq!(next_to_5.1.value, Some(vec![10])); + + // next_committed < next_prospective + let next_to_10 = overlay.next_child_storage_key_change(&child, &[10]).unwrap(); + assert_eq!(next_to_10.0.to_vec(), vec![20]); + assert_eq!(next_to_10.1.value, Some(vec![20])); + + // next_committed == next_prospective + let next_to_20 = overlay.next_child_storage_key_change(&child, &[20]).unwrap(); + assert_eq!(next_to_20.0.to_vec(), vec![30]); + assert_eq!(next_to_20.1.value, None); + + // next_committed, no next_prospective + let next_to_30 = overlay.next_child_storage_key_change(&child, &[30]).unwrap(); + assert_eq!(next_to_30.0.to_vec(), vec![40]); + assert_eq!(next_to_30.1.value, Some(vec![40])); + + overlay.set_child_storage(child.clone(), vec![50], Some(vec![50])); + // next_prospective, no next_committed + let next_to_40 = overlay.next_child_storage_key_change(&child, &[40]).unwrap(); + assert_eq!(next_to_40.0.to_vec(), vec![50]); + assert_eq!(next_to_40.1.value, Some(vec![50])); + } } diff --git a/primitives/state-machine/src/proving_backend.rs b/primitives/state-machine/src/proving_backend.rs index 446dc635e1c6c..2a9146dfced4a 100644 --- a/primitives/state-machine/src/proving_backend.rs +++ b/primitives/state-machine/src/proving_backend.rs @@ -272,6 +272,14 @@ impl<'a, S, H> Backend for ProvingBackend<'a, S, H> self.0.child_storage(storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + self.0.next_storage_key(key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + self.0.next_child_storage_key(storage_key, key) + } + fn for_keys_in_child_storage(&self, storage_key: &[u8], f: F) { self.0.for_keys_in_child_storage(storage_key, f) } diff --git a/primitives/state-machine/src/testing.rs b/primitives/state-machine/src/testing.rs index dd4f7c557e987..85c9260fba3db 100644 --- a/primitives/state-machine/src/testing.rs +++ b/primitives/state-machine/src/testing.rs @@ -16,7 +16,7 @@ //! Test implementation for Externalities. -use std::{collections::HashMap, any::{Any, TypeId}}; +use std::{collections::{HashMap, BTreeMap}, any::{Any, TypeId}}; use hash_db::Hasher; use crate::{ backend::{InMemory, Backend}, OverlayedChanges, @@ -35,7 +35,7 @@ use primitives::{ use codec::Encode; use externalities::{Extensions, Extension}; -type StorageTuple = (HashMap, Vec>, HashMap, HashMap, Vec>>); +type StorageTuple = (BTreeMap, Vec>, HashMap, BTreeMap, Vec>>); /// Simple HashMap-based Externalities impl. pub struct TestExternalities=Blake2Hasher, N: ChangesTrieBlockNumber=u64> { diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index a2ca5c5d4af93..f24c47c891bf9 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -79,6 +79,14 @@ impl, H: Hasher> Backend for TrieBackend where self.essence.child_storage(storage_key, key) } + fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { + self.essence.next_storage_key(key) + } + + fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result>, Self::Error> { + self.essence.next_child_storage_key(storage_key, key) + } + fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { self.essence.for_keys_with_prefix(prefix, f) } diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index a8572aff6019a..6a38f56e0be1b 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -64,6 +64,76 @@ impl, H: Hasher> TrieBackendEssence where H::Out: self.storage } + /// Return the next key in the trie i.e. the minimum key that is strictly superior to `key` in + /// lexicographic order. + pub fn next_storage_key(&self, key: &[u8]) -> Result>, String> { + self.next_storage_key_from_root(&self.root, key) + } + + /// Return the next key in the child trie i.e. the minimum key that is strictly superior to + /// `key` in lexicographic order. + pub fn next_child_storage_key( + &self, + storage_key: &[u8], + key: &[u8], + ) -> Result>, String> { + let child_root = match self.storage(storage_key)? { + Some(child_root) => child_root, + None => return Ok(None), + }; + + let mut hash = H::Out::default(); + + if child_root.len() != hash.as_ref().len() { + return Err(format!("Invalid child storage hash at {:?}", storage_key)); + } + // note: child_root and hash must be same size, panics otherwise. + hash.as_mut().copy_from_slice(&child_root[..]); + + self.next_storage_key_from_root(&hash, key) + } + + /// Return next key from main trie or child trie by providing corresponding root. + fn next_storage_key_from_root( + &self, + root: &H::Out, + key: &[u8], + ) -> Result>, String> { + let mut read_overlay = S::Overlay::default(); + let eph = Ephemeral { + storage: &self.storage, + overlay: &mut read_overlay, + }; + + let trie = TrieDB::::new(&eph, root) + .map_err(|e| format!("TrieDB creation error: {}", e))?; + let mut iter = trie.iter() + .map_err(|e| format!("TrieDB iteration error: {}", e))?; + + // The key just after the one given in input, basically `key++0`. + // Note: We are sure this is the next key if: + // * size of key has no limit (i.e. we can always add 0 to the path), + // * and no keys can be inserted between `key` and `key++0` (this is ensured by sr-io). + let mut potential_next_key = Vec::with_capacity(key.len() + 1); + potential_next_key.extend_from_slice(key); + potential_next_key.push(0); + + iter.seek(&potential_next_key) + .map_err(|e| format!("TrieDB iterator seek error: {}", e))?; + + let next_element = iter.next(); + + let next_key = if let Some(next_element) = next_element { + let (next_key, _) = next_element + .map_err(|e| format!("TrieDB iterator next error: {}", e))?; + Some(next_key) + } else { + None + }; + + Ok(next_key) + } + /// Get the value of storage at given key. pub fn storage(&self, key: &[u8]) -> Result>, String> { let mut read_overlay = S::Overlay::default(); @@ -345,3 +415,47 @@ impl TrieBackendStorage for MemoryDB { Ok(hash_db::HashDB::get(self, key, prefix)) } } + +#[cfg(test)] +mod test { + use primitives::{Blake2Hasher, H256}; + use trie::{TrieMut, PrefixedMemoryDB, trie_types::TrieDBMut}; + use super::*; + + #[test] + fn next_storage_key_and_next_child_storage_key_work() { + // Contains values + let mut root_1 = H256::default(); + // Contains child trie + let mut root_2 = H256::default(); + + let mut mdb = PrefixedMemoryDB::::default(); + { + let mut trie = TrieDBMut::new(&mut mdb, &mut root_1); + trie.insert(b"3", &[1]).expect("insert failed"); + trie.insert(b"4", &[1]).expect("insert failed"); + trie.insert(b"6", &[1]).expect("insert failed"); + } + { + let mut trie = TrieDBMut::new(&mut mdb, &mut root_2); + trie.insert(b"MyChild", root_1.as_ref()).expect("insert failed"); + }; + + let essence_1 = TrieBackendEssence::new(mdb, root_1); + + assert_eq!(essence_1.next_storage_key(b"2"), Ok(Some(b"3".to_vec()))); + assert_eq!(essence_1.next_storage_key(b"3"), Ok(Some(b"4".to_vec()))); + assert_eq!(essence_1.next_storage_key(b"4"), Ok(Some(b"6".to_vec()))); + assert_eq!(essence_1.next_storage_key(b"5"), Ok(Some(b"6".to_vec()))); + assert_eq!(essence_1.next_storage_key(b"6"), Ok(None)); + + let mdb = essence_1.into_storage(); + let essence_2 = TrieBackendEssence::new(mdb, root_2); + + assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"2"), Ok(Some(b"3".to_vec()))); + assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"3"), Ok(Some(b"4".to_vec()))); + assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"4"), Ok(Some(b"6".to_vec()))); + assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"5"), Ok(Some(b"6".to_vec()))); + assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"6"), Ok(None)); + } +} diff --git a/test/utils/runtime/client/src/lib.rs b/test/utils/runtime/client/src/lib.rs index cd815b7ea40f3..8511750bc2c72 100644 --- a/test/utils/runtime/client/src/lib.rs +++ b/test/utils/runtime/client/src/lib.rs @@ -23,7 +23,7 @@ pub mod trait_tests; mod block_builder_ext; use std::sync::Arc; -use std::collections::HashMap; +use std::collections::{HashMap, BTreeMap}; pub use block_builder_ext::BlockBuilderExt; pub use generic_test_client::*; pub use runtime; @@ -97,8 +97,8 @@ pub type LightExecutor = client::light::call_executor::GenesisCallExecutor< pub struct GenesisParameters { support_changes_trie: bool, heap_pages_override: Option, - extra_storage: HashMap, Vec>, - child_extra_storage: HashMap, HashMap, Vec>>, + extra_storage: BTreeMap, Vec>, + child_extra_storage: HashMap, BTreeMap, Vec>>, } impl GenesisParameters { diff --git a/test/utils/runtime/src/genesismap.rs b/test/utils/runtime/src/genesismap.rs index eb2ee144049a6..fc4acfefe267d 100644 --- a/test/utils/runtime/src/genesismap.rs +++ b/test/utils/runtime/src/genesismap.rs @@ -16,7 +16,7 @@ //! Tool for creating the genesis block. -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use runtime_io::hashing::{blake2_256, twox_128}; use super::{AuthorityId, AccountId, WASM_BINARY, system}; use codec::{Encode, KeyedVec, Joiner}; @@ -30,8 +30,8 @@ pub struct GenesisConfig { balances: Vec<(AccountId, u64)>, heap_pages_override: Option, /// Additional storage key pairs that will be added to the genesis map. - extra_storage: HashMap, Vec>, - child_extra_storage: HashMap, HashMap, Vec>>, + extra_storage: BTreeMap, Vec>, + child_extra_storage: HashMap, BTreeMap, Vec>>, } impl GenesisConfig { @@ -41,8 +41,8 @@ impl GenesisConfig { endowed_accounts: Vec, balance: u64, heap_pages_override: Option, - extra_storage: HashMap, Vec>, - child_extra_storage: HashMap, HashMap, Vec>>, + extra_storage: BTreeMap, Vec>, + child_extra_storage: HashMap, BTreeMap, Vec>>, ) -> Self { GenesisConfig { changes_trie_config: match support_changes_trie { @@ -58,11 +58,11 @@ impl GenesisConfig { } pub fn genesis_map(&self) -> ( - HashMap, Vec>, - HashMap, HashMap, Vec>>, + BTreeMap, Vec>, + HashMap, BTreeMap, Vec>>, ) { let wasm_runtime = WASM_BINARY.to_vec(); - let mut map: HashMap, Vec> = self.balances.iter() + let mut map: BTreeMap, Vec> = self.balances.iter() .map(|&(ref account, balance)| (account.to_keyed_vec(b"balance:"), vec![].and(&balance))) .map(|(k, v)| (blake2_256(&k[..])[..].to_vec(), v.to_vec())) .chain(vec![ @@ -92,8 +92,8 @@ impl GenesisConfig { pub fn insert_genesis_block( storage: &mut ( - HashMap, Vec>, - HashMap, HashMap, Vec>>, + BTreeMap, Vec>, + HashMap, BTreeMap, Vec>>, ) ) -> primitives::hash::H256 { let child_roots = storage.1.iter().map(|(sk, child_map)| { @@ -111,7 +111,7 @@ pub fn insert_genesis_block( genesis_hash } -pub fn additional_storage_with_genesis(genesis_block: &crate::Block) -> HashMap, Vec> { +pub fn additional_storage_with_genesis(genesis_block: &crate::Block) -> BTreeMap, Vec> { map![ twox_128(&b"latest"[..]).to_vec() => genesis_block.hash().as_fixed_bytes().to_vec() ]