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:
Keith Yeung
2021-07-01 10:20:15 -07:00
committed by GitHub
parent d489bd70b5
commit ae3c3045df
7 changed files with 279 additions and 4 deletions
@@ -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],