Skip to content
This repository has been archived by the owner on Nov 15, 2023. It is now read-only.

Add methods to PrefixIterator to support iterating from a specific key #9313

Merged
18 commits merged into from
Jul 27, 2021
Merged

Add methods to PrefixIterator to support iterating from a specific key #9313

Show file tree
Hide file tree
Changes from 15 commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
cfbfad9
Add methods to PrefixIterator to support iterating from a specific key
KiChjang Jul 7, 2021
34f57d9
Expose the decode functions used in iterators for storage maps
KiChjang Jul 9, 2021
5b3bbe2
Use associated decode function in tests
KiChjang Jul 9, 2021
981a1ce
Revert "Expose the decode functions used in iterators for storage maps"
KiChjang Jul 12, 2021
6cb9c1b
Fix documentation for next_key
KiChjang Jul 12, 2021
f3b21a0
Add API for iterating from a specified key for all storage map types
KiChjang Jul 13, 2021
2dd16b7
Enhance pagination test
KiChjang Jul 13, 2021
0d7a2eb
Add API methods to storage map types
KiChjang Jul 13, 2021
13527ff
Rename next_key to last_key
KiChjang Jul 13, 2021
e5593c6
Rename last_key to last_raw_key
KiChjang Jul 13, 2021
b14264b
Specify that iteration starts after starting_raw_key
KiChjang Jul 15, 2021
e9a8775
Update documentation on iteration ordering
KiChjang Jul 15, 2021
58e8fcd
Rename next_key to previous_key
KiChjang Jul 15, 2021
d41953a
Enhance pagination unit test
KiChjang Jul 15, 2021
a65a6ed
Create unit tests for all kinds of iter_from methods
KiChjang Jul 16, 2021
fe315f5
Merge remote-tracking branch 'origin/master' into kckyeung/prefix-ite…
KiChjang Jul 26, 2021
3aefbbb
Define iter_from in terms of iter rather than vice versa
KiChjang Jul 26, 2021
447ef77
Cargo fmt
KiChjang Jul 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion frame/support/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ macro_rules! generate_storage_alias {
>;
}
};
($pallet:ident, $name:ident => NMap<$(($key:ty, $hasher:ty),)+ $value:ty>) => {
($pallet:ident, $name:ident => NMap<Key<$(($key:ty, $hasher:ty)),+>, $value:ty>) => {
$crate::paste::paste! {
$crate::generate_storage_alias!(@GENERATE_INSTANCE_STRUCT $pallet, $name);
type $name = $crate::storage::types::StorageNMap<
Expand Down
70 changes: 66 additions & 4 deletions frame/support/src/storage/generator/double_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -358,6 +358,15 @@ impl<
}
}

fn iter_prefix_from(
k1: impl EncodeLike<K1>,
starting_raw_key: Vec<u8>,
) -> Self::PrefixIterator {
let mut iter = Self::iter_prefix(k1);
iter.set_last_raw_key(starting_raw_key);
iter
}
KiChjang marked this conversation as resolved.
Show resolved Hide resolved

fn iter_key_prefix(k1: impl EncodeLike<K1>) -> Self::PartialKeyIterator {
let prefix = G::storage_double_map_final_key1(k1);
Self::PartialKeyIterator {
Expand All @@ -371,17 +380,30 @@ impl<
}
}

fn iter_key_prefix_from(
k1: impl EncodeLike<K1>,
starting_raw_key: Vec<u8>,
) -> Self::PartialKeyIterator {
let mut iter = Self::iter_key_prefix(k1);
iter.set_last_raw_key(starting_raw_key);
iter
}

fn drain_prefix(k1: impl EncodeLike<K1>) -> Self::PrefixIterator {
let mut iterator = Self::iter_prefix(k1);
iterator.drain = true;
iterator
}

fn iter() -> Self::Iterator {
Self::iter_from(G::prefix_hash())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wee bit strange here. iter_from is already reading let prefix = G::prefix_hash();, but we pass in G::prefix_hash() to it again. I'd make the starting_raw_key an Option.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we can just revert this and implement iter_from in terms of iter by calling set_previous_key on the resulting PrefixIterator.

}

fn iter_from(starting_raw_key: Vec<u8>) -> Self::Iterator {
let prefix = G::prefix_hash();
Self::Iterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix, mut raw_value| {
let mut k1_k2_material = G::Hasher1::reverse(raw_key_without_prefix);
Expand All @@ -394,10 +416,14 @@ impl<
}

fn iter_keys() -> Self::FullKeyIterator {
Self::iter_keys_from(G::prefix_hash())
}

fn iter_keys_from(starting_raw_key: Vec<u8>) -> Self::FullKeyIterator {
let prefix = G::prefix_hash();
Self::FullKeyIterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix| {
let mut k1_k2_material = G::Hasher1::reverse(raw_key_without_prefix);
Expand Down Expand Up @@ -498,6 +524,42 @@ mod test_iterators {
prefix
}

#[test]
fn double_map_iter_from() {
sp_io::TestExternalities::default().execute_with(|| {
use crate::hash::Identity;
crate::generate_storage_alias!(
MyModule,
MyDoubleMap => DoubleMap<(u64, Identity), (u64, Identity), u64>
);

MyDoubleMap::insert(1, 10, 100);
MyDoubleMap::insert(1, 21, 201);
MyDoubleMap::insert(1, 31, 301);
MyDoubleMap::insert(1, 41, 401);
MyDoubleMap::insert(2, 20, 200);
MyDoubleMap::insert(3, 30, 300);
MyDoubleMap::insert(4, 40, 400);
MyDoubleMap::insert(5, 50, 500);

let starting_raw_key = MyDoubleMap::storage_double_map_final_key(1, 21);
let iter = MyDoubleMap::iter_key_prefix_from(1, starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![31, 41]);

let starting_raw_key = MyDoubleMap::storage_double_map_final_key(1, 31);
let iter = MyDoubleMap::iter_prefix_from(1, starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(41, 401)]);

let starting_raw_key = MyDoubleMap::storage_double_map_final_key(2, 20);
let iter = MyDoubleMap::iter_keys_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(3, 30), (4, 40), (5, 50)]);

let starting_raw_key = MyDoubleMap::storage_double_map_final_key(3, 30);
let iter = MyDoubleMap::iter_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(4, 40, 400), (5, 50, 500)]);
});
}

#[test]
fn double_map_reversible_reversible_iteration() {
sp_io::TestExternalities::default().execute_with(|| {
Expand Down
40 changes: 36 additions & 4 deletions frame/support/src/storage/generator/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,10 +144,15 @@ impl<

/// Enumerate all elements in the map.
fn iter() -> Self::Iterator {
Self::iter_from(G::prefix_hash())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same grumble.

}

/// Enumerate all elements in the map after a given key.
fn iter_from(starting_raw_key: Vec<u8>) -> Self::Iterator {
let prefix = G::prefix_hash();
PrefixIterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix, mut raw_value| {
let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
Expand All @@ -158,10 +163,15 @@ impl<

/// Enumerate all keys in the map.
fn iter_keys() -> Self::KeyIterator {
Self::iter_keys_from(G::prefix_hash())
}

/// Enumerate all keys in the map after a given key.
fn iter_keys_from(starting_raw_key: Vec<u8>) -> Self::KeyIterator {
let prefix = G::prefix_hash();
KeyPrefixIterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix| {
let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
Expand Down Expand Up @@ -378,6 +388,28 @@ mod test_iterators {
prefix
}

#[test]
fn map_iter_from() {
sp_io::TestExternalities::default().execute_with(|| {
use crate::hash::Identity;
crate::generate_storage_alias!(MyModule, MyMap => Map<(u64, Identity), u64>);

MyMap::insert(1, 10);
MyMap::insert(2, 20);
MyMap::insert(3, 30);
MyMap::insert(4, 40);
MyMap::insert(5, 50);

let starting_raw_key = MyMap::storage_map_final_key(3);
let iter = MyMap::iter_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(4, 40), (5, 50)]);

let starting_raw_key = MyMap::storage_map_final_key(2);
let iter = MyMap::iter_keys_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![3, 4, 5]);
});
}

#[test]
fn map_reversible_reversible_iteration() {
sp_io::TestExternalities::default().execute_with(|| {
Expand Down
2 changes: 1 addition & 1 deletion frame/support/src/storage/generator/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
//!
//! This is internal api and is subject to change.

mod map;
pub(crate) mod map;
mod nmap;
mod double_map;
mod value;
Expand Down
81 changes: 77 additions & 4 deletions frame/support/src/storage/generator/nmap.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,18 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
}
}

fn iter_prefix_from<KP>(
kp: KP,
starting_raw_key: Vec<u8>,
) -> PrefixIterator<(<K as HasKeyPrefix<KP>>::Suffix, V)>
where
K: HasReversibleKeyPrefix<KP>,
{
let mut iter = Self::iter_prefix(kp);
iter.set_last_raw_key(starting_raw_key);
iter
}

fn iter_key_prefix<KP>(kp: KP) -> KeyPrefixIterator<<K as HasKeyPrefix<KP>>::Suffix>
where
K: HasReversibleKeyPrefix<KP>,
Expand All @@ -342,6 +354,18 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
}
}

fn iter_key_prefix_from<KP>(
kp: KP,
starting_raw_key: Vec<u8>,
) -> KeyPrefixIterator<<K as HasKeyPrefix<KP>>::Suffix>
where
K: HasReversibleKeyPrefix<KP>,
{
let mut iter = Self::iter_key_prefix(kp);
iter.set_last_raw_key(starting_raw_key);
iter
}

fn drain_prefix<KP>(kp: KP) -> PrefixIterator<(<K as HasKeyPrefix<KP>>::Suffix, V)>
where
K: HasReversibleKeyPrefix<KP>,
Expand All @@ -352,10 +376,14 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
}

fn iter() -> Self::Iterator {
Self::iter_from(G::prefix_hash())
}

fn iter_from(starting_raw_key: Vec<u8>) -> Self::Iterator {
let prefix = G::prefix_hash();
Self::Iterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix, mut raw_value| {
let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
Expand All @@ -365,10 +393,14 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
}

fn iter_keys() -> Self::KeyIterator {
Self::iter_keys_from(G::prefix_hash())
}

fn iter_keys_from(starting_raw_key: Vec<u8>) -> Self::KeyIterator {
let prefix = G::prefix_hash();
Self::KeyIterator {
prefix: prefix.clone(),
previous_key: prefix,
prefix,
previous_key: starting_raw_key,
drain: false,
closure: |raw_key_without_prefix| {
let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
Expand Down Expand Up @@ -460,6 +492,47 @@ mod test_iterators {
prefix
}

#[test]
fn n_map_iter_from() {
sp_io::TestExternalities::default().execute_with(|| {
use crate::hash::Identity;
use crate::storage::Key as NMapKey;
crate::generate_storage_alias!(
MyModule,
MyNMap => NMap<Key<(u64, Identity), (u64, Identity), (u64, Identity)>, u64>
);

MyNMap::insert((1, 1, 1), 11);
MyNMap::insert((1, 1, 2), 21);
MyNMap::insert((1, 1, 3), 31);
MyNMap::insert((1, 2, 1), 12);
MyNMap::insert((1, 2, 2), 22);
MyNMap::insert((1, 2, 3), 32);
MyNMap::insert((1, 3, 1), 13);
MyNMap::insert((1, 3, 2), 23);
MyNMap::insert((1, 3, 3), 33);
MyNMap::insert((2, 0, 0), 200);

type Key = (NMapKey<Identity, u64>, NMapKey<Identity, u64>, NMapKey<Identity, u64>);

let starting_raw_key = MyNMap::storage_n_map_final_key::<Key, _>((1, 2, 2));
let iter = MyNMap::iter_key_prefix_from((1,), starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(2, 3), (3, 1), (3, 2), (3, 3)]);

let starting_raw_key = MyNMap::storage_n_map_final_key::<Key, _>((1, 3, 1));
let iter = MyNMap::iter_prefix_from((1, 3), starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(2, 23), (3, 33)]);

let starting_raw_key = MyNMap::storage_n_map_final_key::<Key, _>((1, 3, 2));
let iter = MyNMap::iter_keys_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![(1, 3, 3), (2, 0, 0)]);

let starting_raw_key = MyNMap::storage_n_map_final_key::<Key, _>((1, 3, 3));
let iter = MyNMap::iter_from(starting_raw_key);
assert_eq!(iter.collect::<Vec<_>>(), vec![((2, 0, 0), 200)]);
});
}

#[test]
fn n_map_double_map_identical_key() {
sp_io::TestExternalities::default().execute_with(|| {
Expand Down
Loading