mirror of
https://github.com/pezkuwichain/pezkuwi-subxt.git
synced 2026-04-27 11:38:01 +00:00
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:
@@ -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));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
],
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user