diff --git a/substrate/frame/support/src/storage/migration.rs b/substrate/frame/support/src/storage/migration.rs index b29a0b8365..b4a1a9225d 100644 --- a/substrate/frame/support/src/storage/migration.rs +++ b/substrate/frame/support/src/storage/migration.rs @@ -34,11 +34,14 @@ pub struct StorageIterator { impl StorageIterator { /// Construct iterator to iterate over map items in `module` for the map called `item`. + #[deprecated(note="Please use the storage_iter or storage_iter_with_suffix functions instead")] pub fn new(module: &[u8], item: &[u8]) -> Self { + #[allow(deprecated)] Self::with_suffix(module, item, &[][..]) } /// Construct iterator to iterate over map items in `module` for the map called `item`. + #[deprecated(note="Please use the storage_iter or storage_iter_with_suffix functions instead")] pub fn with_suffix(module: &[u8], item: &[u8], suffix: &[u8]) -> Self { let mut prefix = Vec::new(); prefix.extend_from_slice(&Twox128::hash(module)); @@ -92,11 +95,14 @@ pub struct StorageKeyIterator { impl StorageKeyIterator { /// Construct iterator to iterate over map items in `module` for the map called `item`. + #[deprecated(note="Please use the storage_key_iter or storage_key_iter_with_suffix functions instead")] pub fn new(module: &[u8], item: &[u8]) -> Self { + #[allow(deprecated)] Self::with_suffix(module, item, &[][..]) } /// Construct iterator to iterate over map items in `module` for the map called `item`. + #[deprecated(note="Please use the storage_key_iter or storage_key_iter_with_suffix functions instead")] pub fn with_suffix(module: &[u8], item: &[u8], suffix: &[u8]) -> Self { let mut prefix = Vec::new(); prefix.extend_from_slice(&Twox128::hash(module)); @@ -148,6 +154,58 @@ impl Iterator } } +/// Construct iterator to iterate over map items in `module` for the map called `item`. +pub fn storage_iter(module: &[u8], item: &[u8]) -> PrefixIterator<(Vec, T)> { + storage_iter_with_suffix(module, item, &[][..]) +} + +/// Construct iterator to iterate over map items in `module` for the map called `item`. +pub fn storage_iter_with_suffix( + module: &[u8], + item: &[u8], + suffix: &[u8], +) -> PrefixIterator<(Vec, T)> { + let mut prefix = Vec::new(); + prefix.extend_from_slice(&Twox128::hash(module)); + prefix.extend_from_slice(&Twox128::hash(item)); + prefix.extend_from_slice(suffix); + let previous_key = prefix.clone(); + let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| { + let value = T::decode(&mut &raw_value[..])?; + Ok((raw_key_without_prefix.to_vec(), value)) + }; + + PrefixIterator { prefix, previous_key, drain: false, closure } +} + +/// Construct iterator to iterate over map items in `module` for the map called `item`. +pub fn storage_key_iter( + module: &[u8], + item: &[u8], +) -> PrefixIterator<(K, T)> { + storage_key_iter_with_suffix::(module, item, &[][..]) +} + +/// Construct iterator to iterate over map items in `module` for the map called `item`. +pub fn storage_key_iter_with_suffix( + module: &[u8], + item: &[u8], + suffix: &[u8], +) -> PrefixIterator<(K, T)> { + let mut prefix = Vec::new(); + prefix.extend_from_slice(&Twox128::hash(module)); + prefix.extend_from_slice(&Twox128::hash(item)); + prefix.extend_from_slice(suffix); + let previous_key = prefix.clone(); + let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| { + let mut key_material = H::reverse(raw_key_without_prefix); + let key = K::decode(&mut key_material)?; + let value = T::decode(&mut &raw_value[..])?; + Ok((key, value)) + }; + PrefixIterator { prefix, previous_key, drain: false, closure } +} + /// Get a particular value in storage by the `module`, the map's `item` name and the key `hash`. pub fn have_storage_value(module: &[u8], item: &[u8], hash: &[u8]) -> bool { get_storage_value::<()>(module, item, hash).is_some() @@ -294,7 +352,13 @@ mod tests { hash::StorageHasher, }; use sp_io::TestExternalities; - use super::{move_prefix, move_pallet, move_storage_from_pallet}; + use super::{ + move_prefix, + move_pallet, + move_storage_from_pallet, + storage_iter, + storage_key_iter, + }; struct OldPalletStorageValuePrefix; impl frame_support::traits::StorageInstance for OldPalletStorageValuePrefix { @@ -386,4 +450,31 @@ mod tests { assert_eq!(NewStorageMap::iter().collect::>(), vec![(1, 2), (3, 4)]); }) } + + #[test] + fn test_storage_iter() { + TestExternalities::new_empty().execute_with(|| { + OldStorageValue::put(3); + OldStorageMap::insert(1, 2); + OldStorageMap::insert(3, 4); + + assert_eq!( + storage_key_iter::(b"my_old_pallet", b"foo_map").collect::>(), + vec![(1, 2), (3, 4)], + ); + + assert_eq!( + storage_iter(b"my_old_pallet", b"foo_map").drain().map(|t| t.1).collect::>(), + vec![2, 4], + ); + assert_eq!(OldStorageMap::iter().collect::>(), vec![]); + + // Empty because storage iterator skips over the entry under the first key + assert_eq!( + storage_iter::(b"my_old_pallet", b"foo_value").drain().next(), + None + ); + assert_eq!(OldStorageValue::get(), Some(3)); + }); + } } diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs index d9820475a7..e00a3fe831 100644 --- a/substrate/frame/support/src/storage/mod.rs +++ b/substrate/frame/support/src/storage/mod.rs @@ -17,9 +17,10 @@ //! Stuff to do with the runtime's storage. +use sp_core::storage::ChildInfo; use sp_std::prelude::*; use codec::{FullCodec, FullEncode, Encode, EncodeLike, Decode}; -use crate::hash::{Twox128, StorageHasher}; +use crate::hash::{Twox128, StorageHasher, ReversibleStorageHasher}; use sp_runtime::generic::{Digest, DigestItem}; pub use sp_runtime::TransactionOutcome; @@ -519,6 +520,14 @@ pub struct PrefixIterator { closure: fn(&[u8], &[u8]) -> Result, } +impl PrefixIterator { + /// Mutate this iterator into a draining iterator; items iterated are removed from storage. + pub fn drain(mut self) -> Self { + self.drain = true; + self + } +} + impl Iterator for PrefixIterator { type Item = T; @@ -563,6 +572,133 @@ impl Iterator for PrefixIterator { } } +/// Iterate over a prefix of a child trie and decode raw_key and raw_value into `T`. +/// +/// If any decoding fails it skips the key and continues to the next one. +pub struct ChildTriePrefixIterator { + /// The prefix iterated on + prefix: Vec, + /// child info for child trie + child_info: ChildInfo, + /// The last key iterated on + previous_key: Vec, + /// If true then values are removed while iterating + drain: bool, + /// Whether or not we should fetch the previous key + fetch_previous_key: bool, + /// Function that takes `(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 ChildTriePrefixIterator { + /// Mutate this iterator into a draining iterator; items iterated are removed from storage. + pub fn drain(mut self) -> Self { + self.drain = true; + self + } +} + +impl ChildTriePrefixIterator<(Vec, T)> { + /// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`. + /// + /// NOTE: Iterator with [`Self::drain`] will remove any value who failed to decode + pub fn with_prefix(child_info: &ChildInfo, prefix: &[u8]) -> Self { + let prefix = prefix.to_vec(); + let previous_key = prefix.clone(); + let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| { + let value = T::decode(&mut &raw_value[..])?; + Ok((raw_key_without_prefix.to_vec(), value)) + }; + + Self { + prefix, + child_info: child_info.clone(), + previous_key, + drain: false, + fetch_previous_key: true, + closure, + } + } +} + +impl ChildTriePrefixIterator<(K, T)> { + /// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`. + /// + /// NOTE: Iterator with [`Self::drain`] will remove any key or value who failed to decode + pub fn with_prefix_over_key(child_info: &ChildInfo, prefix: &[u8]) -> Self { + let prefix = prefix.to_vec(); + let previous_key = prefix.clone(); + let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| { + let mut key_material = H::reverse(raw_key_without_prefix); + let key = K::decode(&mut key_material)?; + let value = T::decode(&mut &raw_value[..])?; + Ok((key, value)) + }; + + Self { + prefix, + child_info: child_info.clone(), + previous_key, + drain: false, + fetch_previous_key: true, + closure, + } + } +} + +impl Iterator for ChildTriePrefixIterator { + type Item = T; + + fn next(&mut self) -> Option { + loop { + let maybe_next = if self.fetch_previous_key { + self.fetch_previous_key = false; + Some(self.previous_key.clone()) + } else { + sp_io::default_child_storage::next_key( + &self.child_info.storage_key(), + &self.previous_key, + ) + .filter(|n| n.starts_with(&self.prefix)) + }; + break match maybe_next { + Some(next) => { + self.previous_key = next; + let raw_value = match child::get_raw(&self.child_info, &self.previous_key) { + Some(raw_value) => raw_value, + None => { + log::error!( + "next_key returned a key with no value at {:?}", + self.previous_key, + ); + continue + } + }; + if self.drain { + child::kill(&self.child_info, &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) => { + log::error!( + "(key, value) failed to decode at {:?}: {:?}", + self.previous_key, + e, + ); + continue + } + }; + + Some(item) + } + None => None, + } + } + } +} + /// Trait for maps that store all its value after a unique prefix. /// /// By default the final prefix is: @@ -689,6 +825,7 @@ impl StorageAppend> for Digest {} mod test { use super::*; use sp_core::hashing::twox_128; + use crate::hash::Identity; use sp_io::TestExternalities; use generator::StorageValue as _; @@ -825,4 +962,78 @@ mod test { }); }); } + + #[test] + fn child_trie_prefixed_map_works() { + TestExternalities::default().execute_with(|| { + let child_info_a = child::ChildInfo::new_default(b"a"); + child::put(&child_info_a, &[1, 2, 3], &8u16); + child::put(&child_info_a, &[2], &8u16); + child::put(&child_info_a, &[2, 1, 3], &8u8); + child::put(&child_info_a, &[2, 2, 3], &8u16); + child::put(&child_info_a, &[3], &8u16); + + assert_eq!( + ChildTriePrefixIterator::with_prefix(&child_info_a, &[2]) + .collect::, u16)>>(), + vec![ + (vec![], 8), + (vec![2, 3], 8), + ], + ); + + assert_eq!( + ChildTriePrefixIterator::with_prefix(&child_info_a, &[2]) + .drain() + .collect::, u16)>>(), + vec![ + (vec![], 8), + (vec![2, 3], 8), + ], + ); + + // The only remaining is the ones outside prefix + assert_eq!( + ChildTriePrefixIterator::with_prefix(&child_info_a, &[]) + .collect::, u8)>>(), + vec![ + (vec![1, 2, 3], 8), + (vec![3], 8), + ], + ); + + child::put(&child_info_a, &[1, 2, 3], &8u16); + child::put(&child_info_a, &[2], &8u16); + child::put(&child_info_a, &[2, 1, 3], &8u8); + child::put(&child_info_a, &[2, 2, 3], &8u16); + child::put(&child_info_a, &[3], &8u16); + + assert_eq!( + ChildTriePrefixIterator::with_prefix_over_key::(&child_info_a, &[2]) + .collect::>(), + vec![ + (u16::decode(&mut &[2, 3][..]).unwrap(), 8), + ], + ); + + assert_eq!( + ChildTriePrefixIterator::with_prefix_over_key::(&child_info_a, &[2]) + .drain() + .collect::>(), + vec![ + (u16::decode(&mut &[2, 3][..]).unwrap(), 8), + ], + ); + + // The only remaining is the ones outside prefix + assert_eq!( + ChildTriePrefixIterator::with_prefix(&child_info_a, &[]) + .collect::, u8)>>(), + vec![ + (vec![1, 2, 3], 8), + (vec![3], 8), + ], + ); + }); + } } diff --git a/substrate/frame/tips/src/lib.rs b/substrate/frame/tips/src/lib.rs index 6d85df33f1..da22bf0b1f 100644 --- a/substrate/frame/tips/src/lib.rs +++ b/substrate/frame/tips/src/lib.rs @@ -550,13 +550,13 @@ impl Module { tips: Vec<(AccountId, Balance)>, } - use frame_support::{Twox64Concat, migration::StorageKeyIterator}; + use frame_support::{Twox64Concat, migration::storage_key_iter}; - for (hash, old_tip) in StorageKeyIterator::< + for (hash, old_tip) in storage_key_iter::< T::Hash, OldOpenTip, T::BlockNumber, T::Hash>, Twox64Concat, - >::new(b"Treasury", b"Tips").drain() + >(b"Treasury", b"Tips").drain() { let (finder, deposit, finders_fee) = match old_tip.finder {