mirror of
https://github.com/pezkuwichain/pezkuwi-subxt.git
synced 2026-06-25 08:21:10 +00:00
Implement iter_keys function for all types of storage maps (#9238)
* Implement `iter_keys` function for all types of storage maps * Remove draining iterator API * Rename associated key iterator types * Simplify iteration code * add test for `iter_keys().drain()` Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
This commit is contained in:
@@ -18,7 +18,7 @@
|
||||
use sp_std::prelude::*;
|
||||
use sp_std::borrow::Borrow;
|
||||
use codec::{FullCodec, FullEncode, Decode, Encode, EncodeLike};
|
||||
use crate::{storage::{self, unhashed, StorageAppend, PrefixIterator}, Never};
|
||||
use crate::{storage::{self, unhashed, KeyPrefixIterator, StorageAppend, PrefixIterator}, Never};
|
||||
use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
|
||||
|
||||
/// Generator for `StorageDoubleMap` used by `decl_storage`.
|
||||
@@ -340,7 +340,9 @@ impl<
|
||||
G::Hasher1: ReversibleStorageHasher,
|
||||
G::Hasher2: ReversibleStorageHasher
|
||||
{
|
||||
type PartialKeyIterator = KeyPrefixIterator<K2>;
|
||||
type PrefixIterator = PrefixIterator<(K2, V)>;
|
||||
type FullKeyIterator = KeyPrefixIterator<(K1, K2)>;
|
||||
type Iterator = PrefixIterator<(K1, K2, V)>;
|
||||
|
||||
fn iter_prefix(k1: impl EncodeLike<K1>) -> Self::PrefixIterator {
|
||||
@@ -356,6 +358,19 @@ impl<
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_key_prefix(k1: impl EncodeLike<K1>) -> Self::PartialKeyIterator {
|
||||
let prefix = G::storage_double_map_final_key1(k1);
|
||||
Self::PartialKeyIterator {
|
||||
prefix: prefix.clone(),
|
||||
previous_key: prefix,
|
||||
drain: false,
|
||||
closure: |raw_key_without_prefix| {
|
||||
let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
|
||||
K2::decode(&mut key_material)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn drain_prefix(k1: impl EncodeLike<K1>) -> Self::PrefixIterator {
|
||||
let mut iterator = Self::iter_prefix(k1);
|
||||
iterator.drain = true;
|
||||
@@ -378,6 +393,22 @@ impl<
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_keys() -> Self::FullKeyIterator {
|
||||
let prefix = G::prefix_hash();
|
||||
Self::FullKeyIterator {
|
||||
prefix: prefix.clone(),
|
||||
previous_key: prefix,
|
||||
drain: false,
|
||||
closure: |raw_key_without_prefix| {
|
||||
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))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn drain() -> Self::Iterator {
|
||||
let mut iterator = Self::iter();
|
||||
iterator.drain = true;
|
||||
@@ -485,6 +516,11 @@ mod test_iterators {
|
||||
vec![(3, 3, 3), (0, 0, 0), (2, 2, 2), (1, 1, 1)],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
DoubleMap::iter_keys().collect::<Vec<_>>(),
|
||||
vec![(3, 3), (0, 0), (2, 2), (1, 1)],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
DoubleMap::iter_values().collect::<Vec<_>>(),
|
||||
vec![3, 0, 2, 1],
|
||||
@@ -515,6 +551,11 @@ mod test_iterators {
|
||||
vec![(1, 1), (2, 2), (0, 0), (3, 3)],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
DoubleMap::iter_key_prefix(k1).collect::<Vec<_>>(),
|
||||
vec![1, 2, 0, 3],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
DoubleMap::iter_prefix_values(k1).collect::<Vec<_>>(),
|
||||
vec![1, 2, 0, 3],
|
||||
|
||||
@@ -20,7 +20,7 @@ use sp_std::prelude::*;
|
||||
use sp_std::borrow::Borrow;
|
||||
use codec::{FullCodec, FullEncode, Decode, Encode, EncodeLike};
|
||||
use crate::{
|
||||
storage::{self, unhashed, StorageAppend, PrefixIterator},
|
||||
storage::{self, unhashed, KeyPrefixIterator, StorageAppend, PrefixIterator},
|
||||
Never, hash::{StorageHasher, Twox128, ReversibleStorageHasher},
|
||||
};
|
||||
|
||||
@@ -140,6 +140,7 @@ impl<
|
||||
G::Hasher: ReversibleStorageHasher
|
||||
{
|
||||
type Iterator = PrefixIterator<(K, V)>;
|
||||
type KeyIterator = KeyPrefixIterator<K>;
|
||||
|
||||
/// Enumerate all elements in the map.
|
||||
fn iter() -> Self::Iterator {
|
||||
@@ -155,6 +156,20 @@ impl<
|
||||
}
|
||||
}
|
||||
|
||||
/// Enumerate all keys in the map.
|
||||
fn iter_keys() -> Self::KeyIterator {
|
||||
let prefix = G::prefix_hash();
|
||||
KeyPrefixIterator {
|
||||
prefix: prefix.clone(),
|
||||
previous_key: prefix,
|
||||
drain: false,
|
||||
closure: |raw_key_without_prefix| {
|
||||
let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
|
||||
K::decode(&mut key_material)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Enumerate all elements in the map.
|
||||
fn drain() -> Self::Iterator {
|
||||
let mut iterator = Self::iter();
|
||||
@@ -378,6 +393,8 @@ mod test_iterators {
|
||||
|
||||
assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 3), (0, 0), (2, 2), (1, 1)]);
|
||||
|
||||
assert_eq!(Map::iter_keys().collect::<Vec<_>>(), vec![3, 0, 2, 1]);
|
||||
|
||||
assert_eq!(Map::iter_values().collect::<Vec<_>>(), vec![3, 0, 2, 1]);
|
||||
|
||||
assert_eq!(Map::drain().collect::<Vec<_>>(), vec![(3, 3), (0, 0), (2, 2), (1, 1)]);
|
||||
|
||||
@@ -37,7 +37,7 @@ use crate::{
|
||||
EncodeLikeTuple, HasKeyPrefix, HasReversibleKeyPrefix, KeyGenerator,
|
||||
ReversibleKeyGenerator, TupleToEncodedIter,
|
||||
},
|
||||
unhashed, PrefixIterator, StorageAppend,
|
||||
unhashed, KeyPrefixIterator, PrefixIterator, StorageAppend,
|
||||
},
|
||||
Never,
|
||||
};
|
||||
@@ -310,6 +310,7 @@ where
|
||||
impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
|
||||
storage::IterableStorageNMap<K, V> for G
|
||||
{
|
||||
type KeyIterator = KeyPrefixIterator<K::Key>;
|
||||
type Iterator = PrefixIterator<(K::Key, V)>;
|
||||
|
||||
fn iter_prefix<KP>(kp: KP) -> PrefixIterator<(<K as HasKeyPrefix<KP>>::Suffix, V)>
|
||||
@@ -328,6 +329,19 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_key_prefix<KP>(kp: KP) -> KeyPrefixIterator<<K as HasKeyPrefix<KP>>::Suffix>
|
||||
where
|
||||
K: HasReversibleKeyPrefix<KP>,
|
||||
{
|
||||
let prefix = G::storage_n_map_partial_key(kp);
|
||||
KeyPrefixIterator {
|
||||
prefix: prefix.clone(),
|
||||
previous_key: prefix,
|
||||
drain: false,
|
||||
closure: K::decode_partial_key,
|
||||
}
|
||||
}
|
||||
|
||||
fn drain_prefix<KP>(kp: KP) -> PrefixIterator<(<K as HasKeyPrefix<KP>>::Suffix, V)>
|
||||
where
|
||||
K: HasReversibleKeyPrefix<KP>,
|
||||
@@ -350,6 +364,19 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
|
||||
}
|
||||
}
|
||||
|
||||
fn iter_keys() -> Self::KeyIterator {
|
||||
let prefix = G::prefix_hash();
|
||||
Self::KeyIterator {
|
||||
prefix: prefix.clone(),
|
||||
previous_key: prefix,
|
||||
drain: false,
|
||||
closure: |raw_key_without_prefix| {
|
||||
let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
|
||||
Ok(final_key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn drain() -> Self::Iterator {
|
||||
let mut iterator = Self::iter();
|
||||
iterator.drain = true;
|
||||
@@ -471,6 +498,11 @@ mod test_iterators {
|
||||
vec![((3, 3), 3), ((0, 0), 0), ((2, 2), 2), ((1, 1), 1)],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
NMap::iter_keys().collect::<Vec<_>>(),
|
||||
vec![(3, 3), (0, 0), (2, 2), (1, 1)],
|
||||
);
|
||||
|
||||
assert_eq!(NMap::iter_values().collect::<Vec<_>>(), vec![3, 0, 2, 1],);
|
||||
|
||||
assert_eq!(
|
||||
@@ -501,6 +533,11 @@ mod test_iterators {
|
||||
vec![(1, 1), (2, 2), (0, 0), (3, 3)],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
NMap::iter_key_prefix((k1,)).collect::<Vec<_>>(),
|
||||
vec![1, 2, 0, 3],
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
NMap::iter_prefix_values((k1,)).collect::<Vec<_>>(),
|
||||
vec![1, 2, 0, 3],
|
||||
|
||||
Reference in New Issue
Block a user