diff --git a/substrate/frame/recovery/src/lib.rs b/substrate/frame/recovery/src/lib.rs index c055f2bd97..7348566d26 100644 --- a/substrate/frame/recovery/src/lib.rs +++ b/substrate/frame/recovery/src/lib.rs @@ -626,7 +626,7 @@ decl_module! { fn remove_recovery(origin) { let who = ensure_signed(origin)?; // Check there are no active recoveries - let mut active_recoveries = >::iter_prefix(&who); + let mut active_recoveries = >::iter_prefix_values(&who); ensure!(active_recoveries.next().is_none(), Error::::StillActive); // Take the recovery configuration for this account. let recovery_config = >::take(&who).ok_or(Error::::NotRecoverable)?; diff --git a/substrate/frame/staking/src/mock.rs b/substrate/frame/staking/src/mock.rs index f649c530c1..ce00193a67 100644 --- a/substrate/frame/staking/src/mock.rs +++ b/substrate/frame/staking/src/mock.rs @@ -528,7 +528,7 @@ fn check_ledgers() { fn check_exposures() { // a check per validator to ensure the exposure struct is always sane. let era = active_era(); - ErasStakers::::iter_prefix(era).for_each(|expo| { + ErasStakers::::iter_prefix_values(era).for_each(|expo| { assert_eq!( expo.total as u128, expo.own as u128 + expo.others.iter().map(|e| e.value as u128).sum::(), diff --git a/substrate/frame/staking/src/tests.rs b/substrate/frame/staking/src/tests.rs index dd00fdcd3c..15afda1e3a 100644 --- a/substrate/frame/staking/src/tests.rs +++ b/substrate/frame/staking/src/tests.rs @@ -340,7 +340,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()) ); }); @@ -461,7 +461,7 @@ fn nominating_and_rewards_should_work() { // ------ check the staked value of all parties. // 30 and 40 are not chosen anymore - assert_eq!(ErasStakers::::iter_prefix(Staking::active_era().unwrap().index).count(), 2); + assert_eq!(ErasStakers::::iter_prefix_values(Staking::active_era().unwrap().index).count(), 2); assert_eq!( Staking::eras_stakers(Staking::active_era().unwrap().index, 11), Exposure { diff --git a/substrate/frame/support/src/hash.rs b/substrate/frame/support/src/hash.rs index 693e929a30..40cb1f612f 100644 --- a/substrate/frame/support/src/hash.rs +++ b/substrate/frame/support/src/hash.rs @@ -60,7 +60,12 @@ pub trait StorageHasher: 'static { } /// Hasher to use to hash keys to insert to storage. +/// +/// Reversible hasher store the encoded key after the hash part. pub trait ReversibleStorageHasher: StorageHasher { + /// Split the hash part out of the input. + /// + /// I.e. for input `&[hash ++ key ++ some]` returns `&[key ++ some]` fn reverse(x: &[u8]) -> &[u8]; } diff --git a/substrate/frame/support/src/lib.rs b/substrate/frame/support/src/lib.rs index d37a438fc6..9a32d9a779 100644 --- a/substrate/frame/support/src/lib.rs +++ b/substrate/frame/support/src/lib.rs @@ -68,7 +68,7 @@ pub mod weights; pub use self::hash::{ Twox256, Twox128, Blake2_256, Blake2_128, Identity, Twox64Concat, Blake2_128Concat, Hashable, - StorageHasher + StorageHasher, ReversibleStorageHasher }; pub use self::storage::{ StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap, IterableStorageMap, diff --git a/substrate/frame/support/src/storage/generator/double_map.rs b/substrate/frame/support/src/storage/generator/double_map.rs index 9d05ff0b2d..e23b332383 100644 --- a/substrate/frame/support/src/storage/generator/double_map.rs +++ b/substrate/frame/support/src/storage/generator/double_map.rs @@ -208,7 +208,7 @@ 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 + fn iter_prefix_values(k1: KArg1) -> storage::PrefixIterator where KArg1: ?Sized + EncodeLike { let prefix = Self::storage_double_map_final_key1(k1); @@ -334,41 +334,47 @@ impl storage::StorageDoubleMap for G where } } -/// Utility to iterate through items in a storage map. -pub struct MapIterator { +/// Iterate over a prefix and decode raw_key and raw_value into `T`. +pub struct MapIterator { prefix: Vec, previous_key: Vec, + /// If true then value are removed while iterating drain: bool, - _phantom: ::sp_std::marker::PhantomData<(K, V, Hasher)>, + /// 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, } -impl< - K: Decode + Sized, - V: Decode + Sized, - Hasher: ReversibleStorageHasher -> Iterator for MapIterator { - type Item = (K, V); +impl Iterator for MapIterator { + type Item = T; - fn next(&mut self) -> Option<(K, V)> { + 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::(&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, - } + let raw_value = match unhashed::get_raw(&self.previous_key) { + Some(raw_value) => raw_value, + None => { + frame_support::print("ERROR: next_key returned a key with no value in MapIterator"); + continue } - None => continue, + }; + if self.drain { + unhashed::kill(&self.previous_key) } + let raw_key_without_prefix = &self.previous_key[self.prefix.len()..]; + let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) { + Ok(item) => item, + Err(_e) => { + frame_support::print("ERROR: (key, value) failed to decode in MapIterator"); + continue + } + }; + + Some(item) } None => None, } @@ -385,28 +391,48 @@ impl< G::Hasher1: ReversibleStorageHasher, G::Hasher2: ReversibleStorageHasher { - type Iterator = MapIterator; + type PrefixIterator = MapIterator<(K2, V)>; + type Iterator = MapIterator<(K1, K2, V)>; - /// Enumerate all elements in the map. - fn iter(k1: impl EncodeLike) -> Self::Iterator { + fn iter_prefix(k1: impl EncodeLike) -> Self::PrefixIterator { let prefix = G::storage_double_map_final_key1(k1); + Self::PrefixIterator { + prefix: prefix.clone(), + previous_key: prefix, + drain: false, + closure: |raw_key_without_prefix, mut raw_value| { + let mut key_material = G::Hasher2::reverse(raw_key_without_prefix); + Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?)) + }, + } + } + + fn drain_prefix(k1: impl EncodeLike) -> Self::PrefixIterator { + let mut iterator = Self::iter_prefix(k1); + iterator.drain = true; + iterator + } + + fn iter() -> Self::Iterator { + let prefix = G::prefix_hash(); Self::Iterator { prefix: prefix.clone(), previous_key: prefix, drain: false, - _phantom: Default::default(), + closure: |raw_key_without_prefix, mut raw_value| { + let mut k1_k2_material = G::Hasher1::reverse(raw_key_without_prefix); + let k1 = K1::decode(&mut k1_k2_material)?; + let mut k2_material = G::Hasher2::reverse(k1_k2_material); + let k2 = K2::decode(&mut k2_material)?; + Ok((k1, k2, V::decode(&mut raw_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 { - prefix: prefix.clone(), - previous_key: prefix, - drain: true, - _phantom: Default::default(), - } + fn drain() -> Self::Iterator { + let mut iterator = Self::iter(); + iterator.drain = true; + iterator } fn translate Option>(f: F) { @@ -431,33 +457,111 @@ impl< } } +/// Test iterators for StorageDoubleMap #[cfg(test)] -mod test { - use sp_io::TestExternalities; - use crate::storage::{self, StorageDoubleMap}; - use crate::hash::Twox128; +#[allow(dead_code)] +mod test_iterators { + use codec::{Encode, Decode}; + use crate::storage::{generator::StorageDoubleMap, IterableStorageDoubleMap, unhashed}; + + pub trait Trait { + type Origin; + type BlockNumber; + } + + crate::decl_module! { + pub struct Module for enum Call where origin: T::Origin {} + } + + #[derive(PartialEq, Eq, Clone, Encode, Decode)] + struct NoDef(u32); + + crate::decl_storage! { + trait Store for Module as Test { + DoubleMap: double_map hasher(blake2_128_concat) u16, hasher(blake2_128_concat) u32 => u64; + } + } + + fn key_before_prefix(mut prefix: Vec) -> Vec { + let last = prefix.iter_mut().last().unwrap(); + assert!(*last != 0, "mock function not implemented for this prefix"); + *last -= 1; + prefix + } + + fn key_after_prefix(mut prefix: Vec) -> Vec { + let last = prefix.iter_mut().last().unwrap(); + assert!(*last != 255, "mock function not implemented for this prefix"); + *last += 1; + prefix + } + + fn key_in_prefix(mut prefix: Vec) -> Vec { + prefix.push(0); + prefix + } #[test] - fn iter_prefix_works() { - TestExternalities::default().execute_with(|| { - struct MyStorage; - impl storage::generator::StorageDoubleMap for MyStorage { - type Query = Option; - fn module_prefix() -> &'static [u8] { b"MyModule" } - fn storage_prefix() -> &'static [u8] { b"MyStorage" } - type Hasher1 = Twox128; - type Hasher2 = Twox128; - fn from_optional_value_to_query(v: Option) -> Self::Query { v } - fn from_query_to_optional_value(v: Self::Query) -> Option { v } + fn double_map_reversible_reversible_iteration() { + sp_io::TestExternalities::default().execute_with(|| { + // All map iterator + let prefix = DoubleMap::prefix_hash(); + + unhashed::put(&key_before_prefix(prefix.clone()), &1u64); + unhashed::put(&key_after_prefix(prefix.clone()), &1u64); + + for i in 0..4 { + DoubleMap::insert(i as u16, i as u32, i as u64); } - MyStorage::insert(1, 3, 7); - MyStorage::insert(1, 4, 8); - MyStorage::insert(2, 5, 9); - MyStorage::insert(2, 6, 10); + assert_eq!( + DoubleMap::iter().collect::>(), + vec![(3, 3, 3), (0, 0, 0), (2, 2, 2), (1, 1, 1)], + ); - assert_eq!(MyStorage::iter_prefix(1).collect::>(), vec![7, 8]); - assert_eq!(MyStorage::iter_prefix(2).collect::>(), vec![10, 9]); - }); + assert_eq!( + DoubleMap::iter_values().collect::>(), + vec![3, 0, 2, 1], + ); + + assert_eq!( + DoubleMap::drain().collect::>(), + vec![(3, 3, 3), (0, 0, 0), (2, 2, 2), (1, 1, 1)], + ); + + assert_eq!(DoubleMap::iter().collect::>(), vec![]); + assert_eq!(unhashed::get(&key_before_prefix(prefix.clone())), Some(1u64)); + assert_eq!(unhashed::get(&key_after_prefix(prefix.clone())), Some(1u64)); + + // Prefix iterator + let k1 = 3 << 8; + let prefix = DoubleMap::storage_double_map_final_key1(k1); + + unhashed::put(&key_before_prefix(prefix.clone()), &1u64); + unhashed::put(&key_after_prefix(prefix.clone()), &1u64); + + for i in 0..4 { + DoubleMap::insert(k1, i as u32, i as u64); + } + + assert_eq!( + DoubleMap::iter_prefix(k1).collect::>(), + vec![(0, 0), (2, 2), (1, 1), (3, 3)], + ); + + assert_eq!( + DoubleMap::iter_prefix_values(k1).collect::>(), + vec![0, 2, 1, 3], + ); + + assert_eq!( + DoubleMap::drain_prefix(k1).collect::>(), + vec![(0, 0), (2, 2), (1, 1), (3, 3)], + ); + + assert_eq!(DoubleMap::iter_prefix(k1).collect::>(), vec![]); + assert_eq!(unhashed::get(&key_before_prefix(prefix.clone())), Some(1u64)); + assert_eq!(unhashed::get(&key_after_prefix(prefix.clone())), Some(1u64)); + }) } } diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs index efec36b540..47201e22e6 100644 --- a/substrate/frame/support/src/storage/mod.rs +++ b/substrate/frame/support/src/storage/mod.rs @@ -240,18 +240,29 @@ pub trait IterableStorageDoubleMap< K2: FullCodec, V: FullCodec >: StorageDoubleMap { - /// The type that iterates over all `(key, value)`. - type Iterator: Iterator; + /// The type that iterates over all `(key2, value)`. + type PrefixIterator: Iterator; + + /// The type that iterates over all `(key1, key2, 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 iter_prefix(k1: impl EncodeLike) -> Self::PrefixIterator; /// 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; + fn drain_prefix(k1: impl EncodeLike) -> Self::PrefixIterator; + + /// Enumerate all elements in the map in no particular order. If you add or remove values to + /// the map while doing this, you'll get undefined results. + fn iter() -> Self::Iterator; + + /// 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. @@ -310,7 +321,7 @@ pub trait StorageDoubleMap { fn remove_prefix(k1: KArg1) where KArg1: ?Sized + EncodeLike; - fn iter_prefix(k1: KArg1) -> PrefixIterator + fn iter_prefix_values(k1: KArg1) -> PrefixIterator where KArg1: ?Sized + EncodeLike; fn mutate(k1: KArg1, k2: KArg2, f: F) -> R