Skip to content

Commit 9f5a73f

Browse files
KiChjanggui1117
KiChjang
and
gui1117
authored
Add ChildTriePrefixIterator and methods (paritytech#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>
1 parent 67f882c commit 9f5a73f

File tree

3 files changed

+307
-5
lines changed

3 files changed

+307
-5
lines changed

‎frame/support/src/storage/migration.rs

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,14 @@ pub struct StorageIterator<T> {
3434

3535
impl<T> StorageIterator<T> {
3636
/// Construct iterator to iterate over map items in `module` for the map called `item`.
37+
#[deprecated(note="Please use the storage_iter or storage_iter_with_suffix functions instead")]
3738
pub fn new(module: &[u8], item: &[u8]) -> Self {
39+
#[allow(deprecated)]
3840
Self::with_suffix(module, item, &[][..])
3941
}
4042

4143
/// Construct iterator to iterate over map items in `module` for the map called `item`.
44+
#[deprecated(note="Please use the storage_iter or storage_iter_with_suffix functions instead")]
4245
pub fn with_suffix(module: &[u8], item: &[u8], suffix: &[u8]) -> Self {
4346
let mut prefix = Vec::new();
4447
prefix.extend_from_slice(&Twox128::hash(module));
@@ -92,11 +95,14 @@ pub struct StorageKeyIterator<K, T, H: ReversibleStorageHasher> {
9295

9396
impl<K, T, H: ReversibleStorageHasher> StorageKeyIterator<K, T, H> {
9497
/// Construct iterator to iterate over map items in `module` for the map called `item`.
98+
#[deprecated(note="Please use the storage_key_iter or storage_key_iter_with_suffix functions instead")]
9599
pub fn new(module: &[u8], item: &[u8]) -> Self {
100+
#[allow(deprecated)]
96101
Self::with_suffix(module, item, &[][..])
97102
}
98103

99104
/// Construct iterator to iterate over map items in `module` for the map called `item`.
105+
#[deprecated(note="Please use the storage_key_iter or storage_key_iter_with_suffix functions instead")]
100106
pub fn with_suffix(module: &[u8], item: &[u8], suffix: &[u8]) -> Self {
101107
let mut prefix = Vec::new();
102108
prefix.extend_from_slice(&Twox128::hash(module));
@@ -148,6 +154,58 @@ impl<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher> Iterator
148154
}
149155
}
150156

157+
/// Construct iterator to iterate over map items in `module` for the map called `item`.
158+
pub fn storage_iter<T: Decode + Sized>(module: &[u8], item: &[u8]) -> PrefixIterator<(Vec<u8>, T)> {
159+
storage_iter_with_suffix(module, item, &[][..])
160+
}
161+
162+
/// Construct iterator to iterate over map items in `module` for the map called `item`.
163+
pub fn storage_iter_with_suffix<T: Decode + Sized>(
164+
module: &[u8],
165+
item: &[u8],
166+
suffix: &[u8],
167+
) -> PrefixIterator<(Vec<u8>, T)> {
168+
let mut prefix = Vec::new();
169+
prefix.extend_from_slice(&Twox128::hash(module));
170+
prefix.extend_from_slice(&Twox128::hash(item));
171+
prefix.extend_from_slice(suffix);
172+
let previous_key = prefix.clone();
173+
let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
174+
let value = T::decode(&mut &raw_value[..])?;
175+
Ok((raw_key_without_prefix.to_vec(), value))
176+
};
177+
178+
PrefixIterator { prefix, previous_key, drain: false, closure }
179+
}
180+
181+
/// Construct iterator to iterate over map items in `module` for the map called `item`.
182+
pub fn storage_key_iter<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher>(
183+
module: &[u8],
184+
item: &[u8],
185+
) -> PrefixIterator<(K, T)> {
186+
storage_key_iter_with_suffix::<K, T, H>(module, item, &[][..])
187+
}
188+
189+
/// Construct iterator to iterate over map items in `module` for the map called `item`.
190+
pub fn storage_key_iter_with_suffix<K: Decode + Sized, T: Decode + Sized, H: ReversibleStorageHasher>(
191+
module: &[u8],
192+
item: &[u8],
193+
suffix: &[u8],
194+
) -> PrefixIterator<(K, T)> {
195+
let mut prefix = Vec::new();
196+
prefix.extend_from_slice(&Twox128::hash(module));
197+
prefix.extend_from_slice(&Twox128::hash(item));
198+
prefix.extend_from_slice(suffix);
199+
let previous_key = prefix.clone();
200+
let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
201+
let mut key_material = H::reverse(raw_key_without_prefix);
202+
let key = K::decode(&mut key_material)?;
203+
let value = T::decode(&mut &raw_value[..])?;
204+
Ok((key, value))
205+
};
206+
PrefixIterator { prefix, previous_key, drain: false, closure }
207+
}
208+
151209
/// Get a particular value in storage by the `module`, the map's `item` name and the key `hash`.
152210
pub fn have_storage_value(module: &[u8], item: &[u8], hash: &[u8]) -> bool {
153211
get_storage_value::<()>(module, item, hash).is_some()
@@ -294,7 +352,13 @@ mod tests {
294352
hash::StorageHasher,
295353
};
296354
use sp_io::TestExternalities;
297-
use super::{move_prefix, move_pallet, move_storage_from_pallet};
355+
use super::{
356+
move_prefix,
357+
move_pallet,
358+
move_storage_from_pallet,
359+
storage_iter,
360+
storage_key_iter,
361+
};
298362

299363
struct OldPalletStorageValuePrefix;
300364
impl frame_support::traits::StorageInstance for OldPalletStorageValuePrefix {
@@ -386,4 +450,31 @@ mod tests {
386450
assert_eq!(NewStorageMap::iter().collect::<Vec<_>>(), vec![(1, 2), (3, 4)]);
387451
})
388452
}
453+
454+
#[test]
455+
fn test_storage_iter() {
456+
TestExternalities::new_empty().execute_with(|| {
457+
OldStorageValue::put(3);
458+
OldStorageMap::insert(1, 2);
459+
OldStorageMap::insert(3, 4);
460+
461+
assert_eq!(
462+
storage_key_iter::<i32, i32, Twox64Concat>(b"my_old_pallet", b"foo_map").collect::<Vec<_>>(),
463+
vec![(1, 2), (3, 4)],
464+
);
465+
466+
assert_eq!(
467+
storage_iter(b"my_old_pallet", b"foo_map").drain().map(|t| t.1).collect::<Vec<i32>>(),
468+
vec![2, 4],
469+
);
470+
assert_eq!(OldStorageMap::iter().collect::<Vec<_>>(), vec![]);
471+
472+
// Empty because storage iterator skips over the entry under the first key
473+
assert_eq!(
474+
storage_iter::<i32>(b"my_old_pallet", b"foo_value").drain().next(),
475+
None
476+
);
477+
assert_eq!(OldStorageValue::get(), Some(3));
478+
});
479+
}
389480
}

‎frame/support/src/storage/mod.rs

Lines changed: 212 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,10 @@
1717

1818
//! Stuff to do with the runtime's storage.
1919
20+
use sp_core::storage::ChildInfo;
2021
use sp_std::prelude::*;
2122
use codec::{FullCodec, FullEncode, Encode, EncodeLike, Decode};
22-
use crate::hash::{Twox128, StorageHasher};
23+
use crate::hash::{Twox128, StorageHasher, ReversibleStorageHasher};
2324
use sp_runtime::generic::{Digest, DigestItem};
2425
pub use sp_runtime::TransactionOutcome;
2526

@@ -519,6 +520,14 @@ pub struct PrefixIterator<T> {
519520
closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
520521
}
521522

523+
impl<T> PrefixIterator<T> {
524+
/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
525+
pub fn drain(mut self) -> Self {
526+
self.drain = true;
527+
self
528+
}
529+
}
530+
522531
impl<T> Iterator for PrefixIterator<T> {
523532
type Item = T;
524533

@@ -563,6 +572,133 @@ impl<T> Iterator for PrefixIterator<T> {
563572
}
564573
}
565574

575+
/// Iterate over a prefix of a child trie and decode raw_key and raw_value into `T`.
576+
///
577+
/// If any decoding fails it skips the key and continues to the next one.
578+
pub struct ChildTriePrefixIterator<T> {
579+
/// The prefix iterated on
580+
prefix: Vec<u8>,
581+
/// child info for child trie
582+
child_info: ChildInfo,
583+
/// The last key iterated on
584+
previous_key: Vec<u8>,
585+
/// If true then values are removed while iterating
586+
drain: bool,
587+
/// Whether or not we should fetch the previous key
588+
fetch_previous_key: bool,
589+
/// Function that takes `(raw_key_without_prefix, raw_value)` and decode `T`.
590+
/// `raw_key_without_prefix` is the raw storage key without the prefix iterated on.
591+
closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
592+
}
593+
594+
impl<T> ChildTriePrefixIterator<T> {
595+
/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
596+
pub fn drain(mut self) -> Self {
597+
self.drain = true;
598+
self
599+
}
600+
}
601+
602+
impl<T: Decode + Sized> ChildTriePrefixIterator<(Vec<u8>, T)> {
603+
/// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`.
604+
///
605+
/// NOTE: Iterator with [`Self::drain`] will remove any value who failed to decode
606+
pub fn with_prefix(child_info: &ChildInfo, prefix: &[u8]) -> Self {
607+
let prefix = prefix.to_vec();
608+
let previous_key = prefix.clone();
609+
let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
610+
let value = T::decode(&mut &raw_value[..])?;
611+
Ok((raw_key_without_prefix.to_vec(), value))
612+
};
613+
614+
Self {
615+
prefix,
616+
child_info: child_info.clone(),
617+
previous_key,
618+
drain: false,
619+
fetch_previous_key: true,
620+
closure,
621+
}
622+
}
623+
}
624+
625+
impl<K: Decode + Sized, T: Decode + Sized> ChildTriePrefixIterator<(K, T)> {
626+
/// Construct iterator to iterate over child trie items in `child_info` with the prefix `prefix`.
627+
///
628+
/// NOTE: Iterator with [`Self::drain`] will remove any key or value who failed to decode
629+
pub fn with_prefix_over_key<H: ReversibleStorageHasher>(child_info: &ChildInfo, prefix: &[u8]) -> Self {
630+
let prefix = prefix.to_vec();
631+
let previous_key = prefix.clone();
632+
let closure = |raw_key_without_prefix: &[u8], raw_value: &[u8]| {
633+
let mut key_material = H::reverse(raw_key_without_prefix);
634+
let key = K::decode(&mut key_material)?;
635+
let value = T::decode(&mut &raw_value[..])?;
636+
Ok((key, value))
637+
};
638+
639+
Self {
640+
prefix,
641+
child_info: child_info.clone(),
642+
previous_key,
643+
drain: false,
644+
fetch_previous_key: true,
645+
closure,
646+
}
647+
}
648+
}
649+
650+
impl<T> Iterator for ChildTriePrefixIterator<T> {
651+
type Item = T;
652+
653+
fn next(&mut self) -> Option<Self::Item> {
654+
loop {
655+
let maybe_next = if self.fetch_previous_key {
656+
self.fetch_previous_key = false;
657+
Some(self.previous_key.clone())
658+
} else {
659+
sp_io::default_child_storage::next_key(
660+
&self.child_info.storage_key(),
661+
&self.previous_key,
662+
)
663+
.filter(|n| n.starts_with(&self.prefix))
664+
};
665+
break match maybe_next {
666+
Some(next) => {
667+
self.previous_key = next;
668+
let raw_value = match child::get_raw(&self.child_info, &self.previous_key) {
669+
Some(raw_value) => raw_value,
670+
None => {
671+
log::error!(
672+
"next_key returned a key with no value at {:?}",
673+
self.previous_key,
674+
);
675+
continue
676+
}
677+
};
678+
if self.drain {
679+
child::kill(&self.child_info, &self.previous_key)
680+
}
681+
let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
682+
let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
683+
Ok(item) => item,
684+
Err(e) => {
685+
log::error!(
686+
"(key, value) failed to decode at {:?}: {:?}",
687+
self.previous_key,
688+
e,
689+
);
690+
continue
691+
}
692+
};
693+
694+
Some(item)
695+
}
696+
None => None,
697+
}
698+
}
699+
}
700+
}
701+
566702
/// Trait for maps that store all its value after a unique prefix.
567703
///
568704
/// By default the final prefix is:
@@ -689,6 +825,7 @@ impl<Hash: Encode> StorageAppend<DigestItem<Hash>> for Digest<Hash> {}
689825
mod test {
690826
use super::*;
691827
use sp_core::hashing::twox_128;
828+
use crate::hash::Identity;
692829
use sp_io::TestExternalities;
693830
use generator::StorageValue as _;
694831

@@ -825,4 +962,78 @@ mod test {
825962
});
826963
});
827964
}
965+
966+
#[test]
967+
fn child_trie_prefixed_map_works() {
968+
TestExternalities::default().execute_with(|| {
969+
let child_info_a = child::ChildInfo::new_default(b"a");
970+
child::put(&child_info_a, &[1, 2, 3], &8u16);
971+
child::put(&child_info_a, &[2], &8u16);
972+
child::put(&child_info_a, &[2, 1, 3], &8u8);
973+
child::put(&child_info_a, &[2, 2, 3], &8u16);
974+
child::put(&child_info_a, &[3], &8u16);
975+
976+
assert_eq!(
977+
ChildTriePrefixIterator::with_prefix(&child_info_a, &[2])
978+
.collect::<Vec<(Vec<u8>, u16)>>(),
979+
vec![
980+
(vec![], 8),
981+
(vec![2, 3], 8),
982+
],
983+
);
984+
985+
assert_eq!(
986+
ChildTriePrefixIterator::with_prefix(&child_info_a, &[2])
987+
.drain()
988+
.collect::<Vec<(Vec<u8>, u16)>>(),
989+
vec![
990+
(vec![], 8),
991+
(vec![2, 3], 8),
992+
],
993+
);
994+
995+
// The only remaining is the ones outside prefix
996+
assert_eq!(
997+
ChildTriePrefixIterator::with_prefix(&child_info_a, &[])
998+
.collect::<Vec<(Vec<u8>, u8)>>(),
999+
vec![
1000+
(vec![1, 2, 3], 8),
1001+
(vec![3], 8),
1002+
],
1003+
);
1004+
1005+
child::put(&child_info_a, &[1, 2, 3], &8u16);
1006+
child::put(&child_info_a, &[2], &8u16);
1007+
child::put(&child_info_a, &[2, 1, 3], &8u8);
1008+
child::put(&child_info_a, &[2, 2, 3], &8u16);
1009+
child::put(&child_info_a, &[3], &8u16);
1010+
1011+
assert_eq!(
1012+
ChildTriePrefixIterator::with_prefix_over_key::<Identity>(&child_info_a, &[2])
1013+
.collect::<Vec<(u16, u16)>>(),
1014+
vec![
1015+
(u16::decode(&mut &[2, 3][..]).unwrap(), 8),
1016+
],
1017+
);
1018+
1019+
assert_eq!(
1020+
ChildTriePrefixIterator::with_prefix_over_key::<Identity>(&child_info_a, &[2])
1021+
.drain()
1022+
.collect::<Vec<(u16, u16)>>(),
1023+
vec![
1024+
(u16::decode(&mut &[2, 3][..]).unwrap(), 8),
1025+
],
1026+
);
1027+
1028+
// The only remaining is the ones outside prefix
1029+
assert_eq!(
1030+
ChildTriePrefixIterator::with_prefix(&child_info_a, &[])
1031+
.collect::<Vec<(Vec<u8>, u8)>>(),
1032+
vec![
1033+
(vec![1, 2, 3], 8),
1034+
(vec![3], 8),
1035+
],
1036+
);
1037+
});
1038+
}
8281039
}

0 commit comments

Comments
 (0)