Add ChildTriePrefixIterator and methods (#8478)

* Make use of PrefixIterator underneath Storage[Key]Iterator

* Add ChildTriePrefixIterator and methods

* Add documentation on ChilTriePrefixIterator fields

* Deprecate Storage[Key]Iterator API instead of removing them

* Allow fetching for the prefix as an option for ChildTriePrefixIterator

* Rename prefix_fetch to fetch_previous_key

* fix implementation + test

* make gitdiff better

* Add test for storage_iter and storage_key_iter

Co-authored-by: thiolliere <gui.thiolliere@gmail.com>
This commit is contained in:
Keith Yeung
2021-04-01 07:20:43 -07:00
committed by GitHub
parent 643d2b669f
commit 1fdc8fa1c1
3 changed files with 307 additions and 5 deletions
@@ -34,11 +34,14 @@ pub struct StorageIterator<T> {
impl<T> StorageIterator<T> {
/// 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<K, T, H: ReversibleStorageHasher> {
impl<K, T, H: ReversibleStorageHasher> StorageKeyIterator<K, T, H> {
/// 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<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher> Iterator
}
}
/// Construct iterator to iterate over map items in `module` for the map called `item`.
pub fn storage_iter<T: Decode + Sized>(module: &[u8], item: &[u8]) -> PrefixIterator<(Vec<u8>, 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<T: Decode + Sized>(
module: &[u8],
item: &[u8],
suffix: &[u8],
) -> PrefixIterator<(Vec<u8>, 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<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher>(
module: &[u8],
item: &[u8],
) -> PrefixIterator<(K, T)> {
storage_key_iter_with_suffix::<K, T, H>(module, item, &[][..])
}
/// Construct iterator to iterate over map items in `module` for the map called `item`.
pub fn storage_key_iter_with_suffix<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher>(
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<_>>(), 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::<i32, i32, Twox64Concat>(b"my_old_pallet", b"foo_map").collect::<Vec<_>>(),
vec![(1, 2), (3, 4)],
);
assert_eq!(
storage_iter(b"my_old_pallet", b"foo_map").drain().map(|t| t.1).collect::<Vec<i32>>(),
vec![2, 4],
);
assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
// Empty because storage iterator skips over the entry under the first key
assert_eq!(
storage_iter::<i32>(b"my_old_pallet", b"foo_value").drain().next(),
None
);
assert_eq!(OldStorageValue::get(), Some(3));
});
}
}
+212 -1
View File
@@ -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<T> {
closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
}
impl<T> PrefixIterator<T> {
/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
pub fn drain(mut self) -> Self {
self.drain = true;
self
}
}
impl<T> Iterator for PrefixIterator<T> {
type Item = T;
@@ -563,6 +572,133 @@ impl<T> Iterator for PrefixIterator<T> {
}
}
/// 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<T> {
/// The prefix iterated on
prefix: Vec<u8>,
/// child info for child trie
child_info: ChildInfo,
/// The last key iterated on
previous_key: Vec<u8>,
/// 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<T, codec::Error>,
}
impl<T> ChildTriePrefixIterator<T> {
/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
pub fn drain(mut self) -> Self {
self.drain = true;
self
}
}
impl<T: Decode + Sized> ChildTriePrefixIterator<(Vec<u8>, 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<K: Decode + Sized, T: Decode + Sized> 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<H: ReversibleStorageHasher>(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<T> Iterator for ChildTriePrefixIterator<T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
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<Hash: Encode> StorageAppend<DigestItem<Hash>> for Digest<Hash> {}
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::<Vec<(Vec<u8>, u16)>>(),
vec![
(vec![], 8),
(vec![2, 3], 8),
],
);
assert_eq!(
ChildTriePrefixIterator::with_prefix(&child_info_a, &[2])
.drain()
.collect::<Vec<(Vec<u8>, 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::<Vec<(Vec<u8>, 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::<Identity>(&child_info_a, &[2])
.collect::<Vec<(u16, u16)>>(),
vec![
(u16::decode(&mut &[2, 3][..]).unwrap(), 8),
],
);
assert_eq!(
ChildTriePrefixIterator::with_prefix_over_key::<Identity>(&child_info_a, &[2])
.drain()
.collect::<Vec<(u16, u16)>>(),
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::<Vec<(Vec<u8>, u8)>>(),
vec![
(vec![1, 2, 3], 8),
(vec![3], 8),
],
);
});
}
}