diff --git a/frame/elections-phragmen/src/lib.rs b/frame/elections-phragmen/src/lib.rs index 9a9e3c3143856..d6b2bb8743b90 100644 --- a/frame/elections-phragmen/src/lib.rs +++ b/frame/elections-phragmen/src/lib.rs @@ -88,7 +88,7 @@ use sp_runtime::{ }; use frame_support::{ decl_storage, decl_event, ensure, decl_module, decl_error, - weights::{SimpleDispatchInfo, Weight, WeighData}, storage::{StorageMap, IterableStorageMap}, + weights::{SimpleDispatchInfo, Weight, WeighData}, storage::StorageMap, traits::{ Currency, Get, LockableCurrency, LockIdentifier, ReservableCurrency, WithdrawReasons, ChangeMembers, OnUnbalanced, WithdrawReason, Contains, BalanceStatus diff --git a/frame/staking/src/lib.rs b/frame/staking/src/lib.rs index b19caec258fa5..a5f7efba9b242 100644 --- a/frame/staking/src/lib.rs +++ b/frame/staking/src/lib.rs @@ -274,7 +274,6 @@ use codec::{HasCompact, Encode, Decode}; use frame_support::{ decl_module, decl_event, decl_storage, ensure, decl_error, debug, weights::{SimpleDispatchInfo, Weight}, - storage::IterableStorageMap, dispatch::{IsSubType, DispatchResult}, traits::{ Currency, LockIdentifier, LockableCurrency, WithdrawReasons, OnUnbalanced, Imbalance, Get, diff --git a/frame/staking/src/mock.rs b/frame/staking/src/mock.rs index 0045e96b3e3a0..905efbf643dfb 100644 --- a/frame/staking/src/mock.rs +++ b/frame/staking/src/mock.rs @@ -25,7 +25,7 @@ use sp_staking::{SessionIndex, offence::{OffenceDetails, OnOffenceHandler}}; use sp_core::H256; use frame_support::{ assert_ok, impl_outer_origin, parameter_types, impl_outer_dispatch, impl_outer_event, - StorageValue, StorageMap, StorageDoubleMap, IterableStorageMap, + StorageValue, StorageMap, StorageDoubleMap, traits::{Currency, Get, FindAuthor, OnFinalize, OnInitialize}, weights::Weight, }; @@ -497,7 +497,7 @@ pub fn active_era() -> EraIndex { } pub fn check_exposure_all(era: EraIndex) { - ErasStakers::::iter_prefix(era).for_each(check_exposure) + ErasStakers::::iter_prefix_values(era).for_each(check_exposure) } pub fn check_nominator_all(era: EraIndex) { diff --git a/frame/staking/src/tests.rs b/frame/staking/src/tests.rs index f54fe05de088d..e9ebef23b220d 100644 --- a/frame/staking/src/tests.rs +++ b/frame/staking/src/tests.rs @@ -346,7 +346,7 @@ fn less_than_needed_candidates_works() { // But the exposure is updated in a simple way. No external votes exists. // This is purely self-vote. assert!( - ErasStakers::::iter_prefix(Staking::active_era().unwrap().index) + ErasStakers::::iter_prefix_values(Staking::active_era().unwrap().index) .all(|exposure| exposure.others.is_empty()) ); check_exposure_all(Staking::active_era().unwrap().index); diff --git a/frame/support/procedural/src/lib.rs b/frame/support/procedural/src/lib.rs index a9662f530a598..78c933f9ec782 100644 --- a/frame/support/procedural/src/lib.rs +++ b/frame/support/procedural/src/lib.rs @@ -67,7 +67,6 @@ 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. You will generally want to use one @@ -106,7 +105,6 @@ use proc_macro::TokenStream; /// * Double map: `Foo: double_map hasher($hash1) u32, hasher($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 chosen with care, see diff --git a/frame/support/procedural/src/storage/mod.rs b/frame/support/procedural/src/storage/mod.rs index e8599c52a9071..c94b05de2f20c 100644 --- a/frame/support/procedural/src/storage/mod.rs +++ b/frame/support/procedural/src/storage/mod.rs @@ -414,7 +414,6 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr StorageValue as _, StorageMap 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 cbd477354e820..0c7008c8e06b6 100644 --- a/frame/support/procedural/src/storage/storage_struct.rs +++ b/frame/support/procedural/src/storage/storage_struct.rs @@ -122,18 +122,6 @@ 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 { @@ -162,18 +150,6 @@ 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/hash.rs b/frame/support/src/hash.rs index f8d6409060f19..693e929a309e3 100644 --- a/frame/support/src/hash.rs +++ b/frame/support/src/hash.rs @@ -92,6 +92,10 @@ impl StorageHasher for Twox64Concat { } impl ReversibleStorageHasher for Twox64Concat { fn reverse(x: &[u8]) -> &[u8] { + if x.len() < 8 { + crate::debug::error!("Invalid reverse: hash length too short"); + return &[] + } &x[8..] } } @@ -110,6 +114,10 @@ impl StorageHasher for Blake2_128Concat { } impl ReversibleStorageHasher for Blake2_128Concat { fn reverse(x: &[u8]) -> &[u8] { + if x.len() < 16 { + crate::debug::error!("Invalid reverse: hash length too short"); + return &[] + } &x[16..] } } diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs index f242efecc401e..de64cf70e03bd 100644 --- a/frame/support/src/lib.rs +++ b/frame/support/src/lib.rs @@ -70,10 +70,7 @@ pub use self::hash::{ Twox256, Twox128, Blake2_256, Blake2_128, Identity, Twox64Concat, Blake2_128Concat, Hashable, StorageHasher }; -pub use self::storage::{ - StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap, IterableStorageMap, - IterableStorageDoubleMap, migration -}; +pub use self::storage::{StorageValue, StorageMap, StorageDoubleMap, migration}; pub use self::dispatch::{Parameter, Callable, IsSubType}; pub use sp_runtime::{self, ConsensusEngineId, print, traits::Printable}; diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index 9d05ff0b2d6d9..e35b7783f6765 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -16,7 +16,7 @@ use sp_std::prelude::*; use sp_std::borrow::Borrow; -use codec::{Ref, FullCodec, FullEncode, Decode, Encode, EncodeLike, EncodeAppend}; +use codec::{Ref, FullCodec, Decode, Encode, EncodeLike, EncodeAppend}; use crate::{storage::{self, unhashed}, traits::Len}; use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher}; @@ -40,7 +40,7 @@ use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher}; /// If the key2s are not trusted (e.g. can be set by a user), a cryptographic `hasher` such as /// `blake2_256` must be used for Hasher2. Otherwise, other items in storage with the same first /// key can be compromised. -pub trait StorageDoubleMap { +pub trait StorageDoubleMap { /// The type that get/take returns. type Query; @@ -124,12 +124,14 @@ pub trait StorageDoubleMap { } impl storage::StorageDoubleMap for G where - K1: FullEncode, - K2: FullEncode, + K1: FullCodec, + K2: FullCodec, V: FullCodec, G: StorageDoubleMap, { type Query = G::Query; + type Hasher1 = G::Hasher1; + type Hasher2 = G::Hasher2; fn hashed_key_for(k1: KArg1, k2: KArg2) -> Vec where KArg1: EncodeLike, @@ -208,15 +210,8 @@ impl storage::StorageDoubleMap for G where unhashed::kill_prefix(Self::storage_double_map_final_key1(k1).as_ref()) } - fn iter_prefix(k1: KArg1) -> storage::PrefixIterator where - KArg1: ?Sized + EncodeLike - { - let prefix = Self::storage_double_map_final_key1(k1); - storage::PrefixIterator:: { - prefix: prefix.clone(), - previous_key: prefix, - phantom_data: Default::default(), - } + fn remove_all() { + sp_io::storage::clear_prefix(&Self::prefix_hash()) } fn mutate(k1: KArg1, k2: KArg2, f: F) -> R where @@ -332,84 +327,75 @@ impl storage::StorageDoubleMap for G where value }) } -} - -/// Utility to iterate through items in a storage map. -pub struct MapIterator { - prefix: Vec, - previous_key: Vec, - drain: bool, - _phantom: ::sp_std::marker::PhantomData<(K, V, Hasher)>, -} -impl< - K: Decode + Sized, - V: Decode + Sized, - Hasher: ReversibleStorageHasher -> Iterator for MapIterator { - type Item = (K, V); + fn iter_prefix(k1: impl EncodeLike) -> storage::PrefixIterator<(K2, V)> where + Self::Hasher2: ReversibleStorageHasher + { + let prefix = G::storage_double_map_final_key1(k1); + storage::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: false, + closure: from_hashed_key2_to_key2_value::<_, _, _, G>, + } + } - fn next(&mut self) -> Option<(K, V)> { - loop { - let maybe_next = sp_io::storage::next_key(&self.previous_key) - .filter(|n| n.starts_with(&self.prefix)); - break match maybe_next { - Some(next) => { - self.previous_key = next; - match unhashed::get::(&self.previous_key) { - Some(value) => { - if self.drain { - unhashed::kill(&self.previous_key) - } - let mut key_material = Hasher::reverse(&self.previous_key[self.prefix.len()..]); - match K::decode(&mut key_material) { - Ok(key) => Some((key, value)), - Err(_) => continue, - } - } - None => continue, - } - } - None => None, - } + fn drain_prefix(k1: impl EncodeLike) -> storage::PrefixIterator<(K2, V)> where + Self::Hasher2: ReversibleStorageHasher + { + let prefix = G::storage_double_map_final_key1(k1); + storage::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: true, + closure: from_hashed_key2_to_key2_value::<_, _, _, G>, } } -} -impl< - K1: FullCodec, - K2: FullCodec, - V: FullCodec, - G: StorageDoubleMap, -> storage::IterableStorageDoubleMap for G where - G::Hasher1: ReversibleStorageHasher, - G::Hasher2: ReversibleStorageHasher -{ - type Iterator = MapIterator; + fn iter_prefix_values(k1: KArg1) -> storage::PrefixIterator where + KArg1: ?Sized + EncodeLike + { + let prefix = Self::storage_double_map_final_key1(k1); + storage::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: false, + closure: |_raw_key, mut raw_value| V::decode(&mut raw_value), + } + } - /// Enumerate all elements in the map. - fn iter(k1: impl EncodeLike) -> Self::Iterator { - let prefix = G::storage_double_map_final_key1(k1); - Self::Iterator { + fn iter() -> storage::PrefixIterator<(K1, K2, V)> where + Self::Hasher1: ReversibleStorageHasher, + Self::Hasher2: ReversibleStorageHasher, + { + let prefix = G::prefix_hash(); + storage::PrefixIterator { prefix: prefix.clone(), previous_key: prefix, drain: false, - _phantom: Default::default(), + closure: |mut raw_key_without_prefix, mut raw_value| { + let mut key1_key2_material = G::Hasher1::reverse(&mut raw_key_without_prefix); + let key1 = K1::decode(&mut key1_key2_material)?; + let mut key2_material = G::Hasher2::reverse(&mut key1_key2_material); + let key2 = K2::decode(&mut key2_material)?; + let value = V::decode(&mut raw_value)?; + + Ok((key1, key2, value)) + } } } - /// Enumerate all elements in the map. - fn drain(k1: impl EncodeLike) -> Self::Iterator { - let prefix = G::storage_double_map_final_key1(k1); - Self::Iterator { + fn iter_values() -> storage::PrefixIterator { + let prefix = G::prefix_hash(); + storage::PrefixIterator { prefix: prefix.clone(), previous_key: prefix, - drain: true, - _phantom: Default::default(), + drain: false, + closure: |_raw_key, mut raw_value| V::decode(&mut raw_value), } } - fn translate Option>(f: F) { + fn translate_values Option>(f: F) { let prefix = G::prefix_hash(); let mut previous_key = prefix.clone(); loop { @@ -422,7 +408,13 @@ impl< Some(new) => unhashed::put::(&previous_key, &new), None => unhashed::kill(&previous_key), }, - None => continue, + None => { + crate::debug::error!( + "next_key returned a key which failed to decode at {:?}", + previous_key + ); + continue + } } } None => return, @@ -431,6 +423,23 @@ impl< } } +fn from_hashed_key2_to_key2_value( + hashed_key2: &[u8], + mut raw_value: &[u8] +) -> Result<(K2, V), codec::Error> where + K1: FullCodec, + K2: FullCodec, + V: FullCodec, + G: StorageDoubleMap, + G::Hasher2: ReversibleStorageHasher, +{ + let mut key2_material = G::Hasher2::reverse(hashed_key2); + let key2 = K2::decode(&mut key2_material)?; + let value = V::decode(&mut raw_value)?; + + Ok((key2, value)) +} + #[cfg(test)] mod test { use sp_io::TestExternalities; @@ -456,8 +465,8 @@ mod test { MyStorage::insert(2, 5, 9); MyStorage::insert(2, 6, 10); - assert_eq!(MyStorage::iter_prefix(1).collect::>(), vec![7, 8]); - assert_eq!(MyStorage::iter_prefix(2).collect::>(), vec![10, 9]); + assert_eq!(MyStorage::iter_prefix_values(1).collect::>(), vec![7, 8]); + assert_eq!(MyStorage::iter_prefix_values(2).collect::>(), vec![10, 9]); }); } } diff --git a/frame/support/src/storage/generator/map.rs b/frame/support/src/storage/generator/map.rs index c29a9a223aacf..49b5703e3957e 100644 --- a/frame/support/src/storage/generator/map.rs +++ b/frame/support/src/storage/generator/map.rs @@ -17,7 +17,7 @@ #[cfg(not(feature = "std"))] use sp_std::prelude::*; use sp_std::borrow::Borrow; -use codec::{FullCodec, FullEncode, Decode, Encode, EncodeLike, Ref, EncodeAppend}; +use codec::{FullCodec, Decode, Encode, EncodeLike, Ref, EncodeAppend}; use crate::{storage::{self, unhashed}, traits::Len}; use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher}; @@ -32,7 +32,7 @@ use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher}; /// /// If the keys are not trusted (e.g. can be set by a user), a cryptographic `hasher` such as /// `blake2_256` must be used. Otherwise, other values in storage can be compromised. -pub trait StorageMap { +pub trait StorageMap { /// The type that get/take returns. type Query; @@ -87,109 +87,9 @@ pub trait StorageMap { } } -/// Utility to iterate through items in a storage map. -pub struct StorageMapIterator { - prefix: Vec, - previous_key: Vec, - drain: bool, - _phantom: ::sp_std::marker::PhantomData<(K, V, Hasher)>, -} - -impl< - K: Decode + Sized, - V: Decode + Sized, - Hasher: ReversibleStorageHasher -> Iterator for StorageMapIterator { - type Item = (K, V); - - fn next(&mut self) -> Option<(K, V)> { - loop { - let maybe_next = sp_io::storage::next_key(&self.previous_key) - .filter(|n| n.starts_with(&self.prefix)); - break match maybe_next { - Some(next) => { - self.previous_key = next; - match unhashed::get::(&self.previous_key) { - Some(value) => { - if self.drain { - unhashed::kill(&self.previous_key) - } - let mut key_material = Hasher::reverse(&self.previous_key[self.prefix.len()..]); - match K::decode(&mut key_material) { - Ok(key) => Some((key, value)), - Err(_) => continue, - } - } - None => continue, - } - } - None => None, - } - } - } -} - -impl< - K: FullCodec, - V: FullCodec, - G: StorageMap, -> storage::IterableStorageMap for G where - G::Hasher: ReversibleStorageHasher -{ - type Iterator = StorageMapIterator; - - /// Enumerate all elements in the map. - fn iter() -> Self::Iterator { - let prefix = G::prefix_hash(); - Self::Iterator { - prefix: prefix.clone(), - previous_key: prefix, - drain: false, - _phantom: Default::default(), - } - } - - /// Enumerate all elements in the map. - fn drain() -> Self::Iterator { - let prefix = G::prefix_hash(); - Self::Iterator { - prefix: prefix.clone(), - previous_key: prefix, - drain: true, - _phantom: Default::default(), - } - } - - fn translate Option>(f: F) { - let prefix = G::prefix_hash(); - let mut previous_key = prefix.clone(); - loop { - match sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix)) { - Some(next) => { - previous_key = next; - let maybe_value = unhashed::get::(&previous_key); - match maybe_value { - Some(value) => { - let mut key_material = G::Hasher::reverse(&previous_key[prefix.len()..]); - match K::decode(&mut key_material) { - Ok(key) => match f(key, value) { - Some(new) => unhashed::put::(&previous_key, &new), - None => unhashed::kill(&previous_key), - }, - Err(_) => continue, - } - } - None => continue, - } - } - None => return, - } - } - } -} - -impl> storage::StorageMap for G { +impl> storage::StorageMap for G { type Query = G::Query; + type Hasher = G::Hasher; fn hashed_key_for>(key: KeyArg) -> Vec { Self::storage_map_final_key(key) @@ -367,4 +267,130 @@ impl> storage::StorageMap value }) } + + fn iter() -> storage::PrefixIterator<(K, V)> where + Self::Hasher: ReversibleStorageHasher + { + let prefix = G::prefix_hash(); + storage::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: false, + closure: from_hashed_key_to_key_value::<_, _, G>, + } + } + + fn drain() -> storage::PrefixIterator<(K, V)> where + Self::Hasher: ReversibleStorageHasher + { + let prefix = G::prefix_hash(); + storage::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: true, + closure: from_hashed_key_to_key_value::<_, _, G>, + } + } + + fn iter_values() -> storage::PrefixIterator { + let prefix = Self::prefix_hash(); + storage::PrefixIterator { + prefix: prefix.to_vec(), + previous_key: prefix.to_vec(), + drain: false, + closure: |_raw_key, mut raw_value| V::decode(&mut raw_value), + } + } + + fn translate Option>(f: F) where + Self::Hasher: ReversibleStorageHasher + { + let prefix = G::prefix_hash(); + let mut previous_key = prefix.clone(); + loop { + match sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix)) { + Some(next) => { + previous_key = next; + let maybe_value = unhashed::get::(&previous_key); + let value = match maybe_value { + Some(value) => value, + None => { + crate::debug::error!( + "next_key returned a key with no value at {:?}", previous_key + ); + continue + }, + }; + + // Slice in bound because already check by `starts_with` + let mut hashed_key = &previous_key[prefix.len()..]; + + let mut key_material = G::Hasher::reverse(&mut hashed_key); + let key = match K::decode(&mut key_material) { + Ok(key) => key, + Err(e) => { + crate::debug::error!( + "old key failed to decode at {:?}: {:?}", + previous_key, e + ); + continue + }, + }; + + match f(key, value) { + Some(new) => unhashed::put::(&previous_key, &new), + None => unhashed::kill(&previous_key), + } + } + None => return, + } + } + } + + fn translate_values(translate_val: TranslateV) -> Result<(), u32> + where OldV: Decode, TranslateV: Fn(OldV) -> V + { + let prefix = Self::prefix_hash(); + let mut previous_key = prefix.to_vec(); + let mut errors = 0; + while let Some(next_key) = sp_io::storage::next_key(&previous_key) + .filter(|n| n.starts_with(&prefix[..])) + { + if let Some(value) = unhashed::get(&next_key) { + unhashed::put(&next_key[..], &translate_val(value)); + } else { + // We failed to read the value. Remove the key and increment errors. + unhashed::kill(&next_key[..]); + errors += 1; + } + + previous_key = next_key; + } + + if errors == 0 { + Ok(()) + } else { + Err(errors) + } + } + + fn remove_all() { + sp_io::storage::clear_prefix(&Self::prefix_hash()) + } +} + +fn from_hashed_key_to_key_value( + hashed_key: &[u8], + mut raw_value: &[u8] +) -> Result<(K, V), codec::Error> where + K: FullCodec, + V: FullCodec, + G: StorageMap, + G::Hasher: ReversibleStorageHasher, +{ + let mut key_material = G::Hasher::reverse(hashed_key); + let key = K::decode(&mut key_material)?; + let value = V::decode(&mut raw_value)?; + + Ok((key, value)) } diff --git a/frame/support/src/storage/generator/mod.rs b/frame/support/src/storage/generator/mod.rs index 687d8a3c9361b..df919300ab740 100644 --- a/frame/support/src/storage/generator/mod.rs +++ b/frame/support/src/storage/generator/mod.rs @@ -36,7 +36,7 @@ pub use value::StorageValue; mod tests { use sp_io::TestExternalities; use codec::Encode; - use crate::storage::{unhashed, generator::StorageValue, IterableStorageMap}; + use crate::storage::{unhashed, generator::StorageValue}; struct Runtime {} pub trait Trait { diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index efec36b540abc..205a335a761af 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 sp_std::{prelude::*, marker::PhantomData}; +use sp_std::prelude::*; use codec::{FullCodec, FullEncode, Encode, EncodeAppend, EncodeLike, Decode}; -use crate::{traits::Len, hash::{Twox128, StorageHasher}}; +use crate::{traits::Len, hash::{StorageHasher, ReversibleStorageHasher}}; pub mod unhashed; pub mod hashed; @@ -131,6 +131,9 @@ pub trait StorageMap { /// The type that get/take return. type Query; + /// Hasher used to hash the key. + type Hasher; + /// Get the storage key used to fetch a value corresponding to a specific key. fn hashed_key_for>(key: KeyArg) -> Vec; @@ -149,6 +152,9 @@ pub trait StorageMap { /// Remove the value under a key. fn remove>(key: KeyArg); + /// Remove all values. + fn remove_all(); + /// Mutate the value under a key. fn mutate, R, F: FnOnce(&mut Self::Query) -> R>(key: KeyArg, f: F) -> R; @@ -214,48 +220,44 @@ pub trait StorageMap { fn migrate_key_from_blake>(key: KeyArg) -> Option { Self::migrate_key::(key) } -} - -/// A strongly-typed map in storage whose keys and values can be iterated over. -pub trait IterableStorageMap: StorageMap { - /// The type that iterates over all `(key, value)`. - type Iterator: Iterator; /// Enumerate all elements in the map in no particular order. If you alter the map while doing /// this, you'll get undefined results. - fn iter() -> Self::Iterator; + fn iter() -> PrefixIterator<(K, V)> where Self::Hasher: ReversibleStorageHasher; /// Remove all elements from the map and iterate through them in no particular order. If you /// add elements to the map while doing this, you'll get undefined results. - fn drain() -> Self::Iterator; - - /// Translate the values of all elements by a function `f`, in the map in no particular order. - /// By returning `None` from `f` for an element, you'll remove it from the map. - fn translate Option>(f: F); -} - -/// A strongly-typed double map in storage whose secondary keys and values can be iterated over. -pub trait IterableStorageDoubleMap< - K1: FullCodec, - K2: FullCodec, - V: FullCodec ->: StorageDoubleMap { - /// The type that iterates over all `(key, value)`. - type Iterator: Iterator; - - /// Enumerate all elements in the map with first key `k1` in no particular order. If you add or - /// remove values whose first key is `k1` to the map while doing this, you'll get undefined - /// results. - fn iter(k1: impl EncodeLike) -> Self::Iterator; + fn drain() -> PrefixIterator<(K, V)> where Self::Hasher: ReversibleStorageHasher; - /// Remove all elements from the map with first key `k1` and iterate through them in no - /// particular order. If you add elements with first key `k1` to the map while doing this, - /// you'll get undefined results. - fn drain(k1: impl EncodeLike) -> Self::Iterator; + /// Enumerate all elements in the map in no particular order. If you alter the map while doing + /// this, you'll get undefined results. + fn iter_values() -> PrefixIterator; /// Translate the values of all elements by a function `f`, in the map in no particular order. /// By returning `None` from `f` for an element, you'll remove it from the map. - fn translate Option>(f: F); + fn translate Option>(f: F) where + Self::Hasher: ReversibleStorageHasher; + + /// Translate the values from some previous `OldValue` to the current type. + /// + /// `TV` translates values. + /// + /// Returns `Err` if the map could not be interpreted as the old type, and Ok if it could. + /// The `Err` contains the number of value that couldn't be interpreted, those value are + /// removed from the map. + /// + /// # Warning + /// + /// This function must be used with care, before being updated the storage still contains the + /// old type, thus other calls (such as `get`) will fail at decoding it. + /// + /// # Usage + /// + /// This would typically be called inside the module implementation of on_runtime_upgrade, while + /// ensuring **no usage of this storage are made before the call to `on_runtime_upgrade`**. (More + /// precisely prior initialized modules doesn't make use of this storage). + fn translate_values(translate_val: TranslateV) -> Result<(), u32> + where OldV: Decode, TranslateV: Fn(OldV) -> V; } /// An implementation of a map with a two keys. @@ -269,6 +271,12 @@ pub trait StorageDoubleMap { /// The type that get/take returns. type Query; + /// Hasher used to hash first key. + type Hasher1; + + /// Hasher used to hash second key. + type Hasher2; + fn hashed_key_for(k1: KArg1, k2: KArg2) -> Vec where KArg1: EncodeLike, @@ -310,8 +318,7 @@ pub trait StorageDoubleMap { fn remove_prefix(k1: KArg1) where KArg1: ?Sized + EncodeLike; - fn iter_prefix(k1: KArg1) -> PrefixIterator - where KArg1: ?Sized + EncodeLike; + fn remove_all(); fn mutate(k1: KArg1, k2: KArg2, f: F) -> R where @@ -370,201 +377,85 @@ pub trait StorageDoubleMap { KeyArg1: EncodeLike, KeyArg2: EncodeLike, >(key1: KeyArg1, key2: KeyArg2) -> Option; -} - -/// 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 sp_io::storage::next_key(&self.previous_key) - .filter(|n| n.starts_with(&self.prefix[..])) - { - Some(next_key) => { - 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; + /// Translate the values of all elements by a function `f`, in the map in no particular order. + /// By returning `None` from `f` for an element, you'll remove it from the map. + fn translate_values Option>(f: F); - value - }, - _ => None, - } - } -} + /// Enumerate all elements in the map with first key `k1` in no particular order. If you add or + /// remove values whose first key is `k1` to the map while doing this, you'll get undefined + /// results. + fn iter_prefix(k1: impl EncodeLike) -> PrefixIterator<(K2, V)> where + Self::Hasher2: ReversibleStorageHasher; -/// 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]; - - /// Final full prefix that prefixes all keys. - 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 - } + /// Remove all elements from the map with first key `k1` and iterate through them in no + /// particular order. If you add elements with first key `k1` to the map while doing this, + /// you'll get undefined results. + fn drain_prefix(k1: impl EncodeLike) -> PrefixIterator<(K2, V)> where + Self::Hasher2: ReversibleStorageHasher; - /// Remove all value of the storage. - fn remove_all() { - sp_io::storage::clear_prefix(&Self::final_prefix()) - } + /// Iter over all value associated to the first key. + fn iter_prefix_values(k1: KArg1) -> PrefixIterator where + KArg1: ?Sized + EncodeLike; - /// Iter over all value of the storage. - fn iter_values() -> PrefixIterator { - let prefix = Self::final_prefix(); - PrefixIterator { - prefix: prefix.to_vec(), - previous_key: prefix.to_vec(), - phantom_data: Default::default(), - } - } + /// Iter over all value and decode the first key and the second key. + fn iter() -> PrefixIterator<(K1, K2, V)> where + Self::Hasher1: ReversibleStorageHasher, + Self::Hasher2: ReversibleStorageHasher; - /// Translate the values from some previous `OldValue` to the current type. - /// - /// `TV` translates values. - /// - /// Returns `Err` if the map could not be interpreted as the old type, and Ok if it could. - /// The `Err` contains the number of value that couldn't be interpreted, those value are - /// removed from the map. - /// - /// # Warning - /// - /// This function must be used with care, before being updated the storage still contains the - /// old type, thus other calls (such as `get`) will fail at decoding it. - /// - /// # Usage - /// - /// This would typically be called inside the module implementation of on_runtime_upgrade, while - /// ensuring **no usage of this storage are made before the call to `on_runtime_upgrade`**. (More - /// precisely prior initialized modules doesn't make use of this storage). - fn translate_values(translate_val: TV) -> Result<(), u32> - where OldValue: Decode, TV: Fn(OldValue) -> Value - { - let prefix = Self::final_prefix(); - let mut previous_key = prefix.to_vec(); - let mut errors = 0; - while let Some(next_key) = sp_io::storage::next_key(&previous_key) - .filter(|n| n.starts_with(&prefix[..])) - { - if let Some(value) = unhashed::get(&next_key) { - unhashed::put(&next_key[..], &translate_val(value)); - } else { - // We failed to read the value. Remove the key and increment errors. - unhashed::kill(&next_key[..]); - errors += 1; - } - - previous_key = next_key; - } + /// Iter over all value. + fn iter_values() -> PrefixIterator; +} - if errors == 0 { - Ok(()) - } else { - Err(errors) - } - } +/// Iterate over a prefix and decode raw_key and raw_value into `T`. +pub struct PrefixIterator { + prefix: Vec, + previous_key: Vec, + /// If true then value are removed while iterating + drain: bool, + /// Function that take `(raw_key_without_prefix, raw_value)` and decode `T`. + /// `raw_key_without_prefix` is the raw storage key without the prefix iterated on. + closure: fn(&[u8], &[u8]) -> Result, } -#[cfg(test)] -mod test { - use sp_core::hashing::twox_128; - use sp_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" - } +impl Iterator for PrefixIterator { + type Item = T; - fn storage_prefix() -> &'static [u8] { - b"MyStorage" + fn next(&mut self) -> Option { + loop { + let maybe_next = sp_io::storage::next_key(&self.previous_key) + .filter(|n| n.starts_with(&self.prefix)); + break match maybe_next { + Some(next) => { + self.previous_key = next; + match unhashed::get_raw(&self.previous_key) { + Some(raw_value) => { + if self.drain { + unhashed::kill(&self.previous_key) + } + let raw_key_without_prefix = &self.previous_key[self.prefix.len()..]; + match (self.closure)(raw_key_without_prefix, &raw_value[..]) { + Ok(t) => Some(t), + Err(e) => { + crate::debug::error!( + "(key, value) failed to decode at {:?}: {:?}", + self.previous_key, e + ); + continue + } + } + } + None => { + crate::debug::error!( + "next_key returned a key with no value at {:?}", + self.previous_key + ); + continue + } + } } + None => None, } - - 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); - - // test iteration - assert_eq!(MyStorage::iter_values().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_values().collect::>(), vec![1, 2, 3, 4]); - - // test removal - MyStorage::remove_all(); - assert_eq!(MyStorage::iter_values().collect::>(), vec![]); - - // test migration - unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u32); - unhashed::put(&[&k[..], &vec![8][..]].concat(), &2u32); - - assert_eq!(MyStorage::iter_values().collect::>(), vec![]); - MyStorage::translate_values(|v: u32| v as u64).unwrap(); - assert_eq!(MyStorage::iter_values().collect::>(), vec![1, 2]); - MyStorage::remove_all(); - - // test migration 2 - unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u128); - unhashed::put(&[&k[..], &vec![1, 1][..]].concat(), &2u64); - unhashed::put(&[&k[..], &vec![8][..]].concat(), &3u128); - unhashed::put(&[&k[..], &vec![10][..]].concat(), &4u32); - - // (contains some value that successfully decoded to u64) - assert_eq!(MyStorage::iter_values().collect::>(), vec![1, 2, 3]); - assert_eq!(MyStorage::translate_values(|v: u128| v as u64), Err(2)); - assert_eq!(MyStorage::iter_values().collect::>(), vec![1, 3]); - MyStorage::remove_all(); - - // test that other values are not modified. - assert_eq!(unhashed::get(&key_before[..]), Some(32u64)); - assert_eq!(unhashed::get(&key_after[..]), Some(33u64)); - }); + } } } diff --git a/frame/support/test/tests/decl_storage.rs b/frame/support/test/tests/decl_storage.rs index ea9b09f9d7bc2..4746921fbed8c 100644 --- a/frame/support/test/tests/decl_storage.rs +++ b/frame/support/test/tests/decl_storage.rs @@ -620,3 +620,133 @@ mod test_append_and_len { }); } } + +/// Test iterators for StorageMap and StorageDoubleMap with key information +#[cfg(test)] +#[allow(dead_code)] +mod test_iterators { + use codec::{Encode, Decode}; + use frame_support::storage::{generator::{StorageMap, StorageDoubleMap}, unhashed}; + + pub trait Trait { + type Origin; + type BlockNumber; + } + + frame_support::decl_module! { + pub struct Module for enum Call where origin: T::Origin {} + } + + #[derive(PartialEq, Eq, Clone, Encode, Decode)] + struct NoDef(u32); + + frame_support::decl_storage! { + trait Store for Module as Test { + MapRev: map hasher(blake2_128_concat) u16 => u64; + + DoubleMapRevRev: double_map hasher(blake2_128_concat) u16, hasher(blake2_128_concat) u32 => u64; + } + } + + struct Test {} + + impl Trait for Test { + type Origin = u32; + type BlockNumber = u32; + } + + fn key_before_prefix(mut prefix: Vec) -> Vec { + let last = prefix.iter_mut().last().unwrap(); + *last = last.checked_sub(1).unwrap_or_else(|| unimplemented!()); + prefix + } + + fn key_after_prefix(mut prefix: Vec) -> Vec { + let last = prefix.iter_mut().last().unwrap(); + *last = last.checked_add(1).unwrap_or_else(|| unimplemented!()); + prefix + } + + fn key_in_prefix(mut prefix: Vec) -> Vec { + prefix.push(0); + prefix + } + + #[test] + fn map_info_iteration() { + sp_io::TestExternalities::default().execute_with(|| { + let prefix = MapRev::prefix_hash(); + + unhashed::put(&key_before_prefix(prefix.clone()), &1u64); + unhashed::put(&key_after_prefix(prefix.clone()), &1u64); + // Put an invalid storage value inside the map. + unhashed::put(&key_in_prefix(prefix), &1u8); + + for i in 0..4 { + MapRev::insert(i as u16, i as u64); + } + + assert_eq!( + MapRev::iter().collect::>(), + vec![(3, 3), (0, 0), (2, 2), (1, 1)], + ); + + assert_eq!( + MapRev::iter_values().collect::>(), + vec![3, 0, 2, 1], + ); + }) + } + + #[test] + fn double_map_reversible_reversible_iteration() { + sp_io::TestExternalities::default().execute_with(|| { + // All map iterator + let prefix = DoubleMapRevRev::prefix_hash(); + + unhashed::put(&key_before_prefix(prefix.clone()), &1u64); + unhashed::put(&key_after_prefix(prefix.clone()), &1u64); + // Put an invalid storage value inside the map. + unhashed::put(&key_in_prefix(prefix), &1u8); + + for i in 0..4 { + DoubleMapRevRev::insert(i as u16, i as u32, i as u64); + } + + assert_eq!( + DoubleMapRevRev::iter().collect::>(), + vec![(3, 3, 3), (0, 0, 0), (2, 2, 2), (1, 1, 1)], + ); + + assert_eq!( + DoubleMapRevRev::iter_values().collect::>(), + vec![3, 0, 2, 1], + ); + + DoubleMapRevRev::remove_all(); + + // Prefix iterator + let k1 = 3 << 8; + let prefix = DoubleMapRevRev::storage_double_map_final_key1(k1); + + unhashed::put(&key_before_prefix(prefix.clone()), &1u64); + unhashed::put(&key_after_prefix(prefix.clone()), &1u64); + // Put an invalid storage value inside the map. + unhashed::put(&key_in_prefix(prefix), &1u8); + + for i in 0..4 { + DoubleMapRevRev::insert(k1, i as u32, i as u64); + } + + assert_eq!( + DoubleMapRevRev::iter_prefix(k1).collect::>(), + vec![(0, 0), (2, 2), (1, 1), (3, 3)], + ); + + assert_eq!( + DoubleMapRevRev::iter_prefix_values(k1).collect::>(), + vec![0, 2, 1, 3], + ); + }) + } +} diff --git a/frame/support/test/tests/final_keys.rs b/frame/support/test/tests/final_keys.rs index ae23c5a64c200..82573ddd49bdb 100644 --- a/frame/support/test/tests/final_keys.rs +++ b/frame/support/test/tests/final_keys.rs @@ -16,7 +16,7 @@ use frame_support::storage::unhashed; use codec::Encode; -use frame_support::{StorageDoubleMap, StorageMap, StorageValue, StoragePrefixedMap}; +use frame_support::{StorageDoubleMap, StorageMap, StorageValue}; use sp_io::{TestExternalities, hashing::{twox_64, twox_128, blake2_128}}; mod no_instance { @@ -102,27 +102,23 @@ fn final_keys_no_instance() { let mut k = [twox_128(b"FinalKeysNone"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_128_concat)); 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_64_concat)); 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_128_concat)); k.extend(2u32.using_encoded(blake2_128_concat)); 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_64_concat)); k.extend(2u32.using_encoded(twox_64_concat)); assert_eq!(unhashed::get::(&k), Some(3u32)); - assert_eq!(&k[..32], &::final_prefix()); }); } @@ -137,27 +133,23 @@ fn final_keys_default_instance() { let mut k = [twox_128(b"FinalKeysSome"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_128_concat)); 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_64_concat)); 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_128_concat)); k.extend(2u32.using_encoded(blake2_128_concat)); 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_64_concat)); k.extend(2u32.using_encoded(twox_64_concat)); assert_eq!(unhashed::get::(&k), Some(3u32)); - assert_eq!(&k[..32], &>::final_prefix()); }); } @@ -172,26 +164,22 @@ fn final_keys_instance_2() { let mut k = [twox_128(b"Instance2FinalKeysSome"), twox_128(b"Map")].concat(); k.extend(1u32.using_encoded(blake2_128_concat)); 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_64_concat)); 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_128_concat)); k.extend(2u32.using_encoded(blake2_128_concat)); 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_64_concat)); k.extend(2u32.using_encoded(twox_64_concat)); assert_eq!(unhashed::get::(&k), Some(3u32)); - assert_eq!(&k[..32], &>::final_prefix()); }); } diff --git a/utils/frame/rpc/support/src/lib.rs b/utils/frame/rpc/support/src/lib.rs index 118f5709a6b70..911caedd13aac 100644 --- a/utils/frame/rpc/support/src/lib.rs +++ b/utils/frame/rpc/support/src/lib.rs @@ -22,7 +22,7 @@ use core::marker::PhantomData; use futures::compat::Future01CompatExt; use jsonrpc_client_transports::RpcError; -use codec::{DecodeAll, FullCodec, FullEncode}; +use codec::{DecodeAll, FullCodec}; use serde::{de::DeserializeOwned, Serialize}; use frame_support::storage::generator::{ StorageDoubleMap, StorageMap, StorageValue @@ -104,7 +104,7 @@ impl StorageQuery { } /// Create a storage query for a value in a StorageMap. - pub fn map, K: FullEncode>(key: K) -> Self { + pub fn map, K: FullCodec>(key: K) -> Self { Self { key: StorageKey(St::storage_map_final_key(key)), _spook: PhantomData, @@ -112,7 +112,7 @@ impl StorageQuery { } /// Create a storage query for a value in a StorageDoubleMap. - pub fn double_map, K1: FullEncode, K2: FullEncode>( + pub fn double_map, K1: FullCodec, K2: FullCodec>( key1: K1, key2: K2, ) -> Self {