use codec::{Decode, EncodeLike, FullCodec, FullEncode};
use frame_support::{
storage::{
generator::{StorageDoubleMap as StorageDoubleMapT, StorageMap as StorageMapT},
unhashed, StorageDoubleMap, StorageMap,
},
ReversibleStorageHasher,
};
use sp_std::prelude::*;
pub struct StorageMapIterator<K, V, Hasher> {
prefix: Vec<u8>,
pub previous_key: Vec<u8>,
drain: bool,
_phantom: sp_std::marker::PhantomData<(K, V, Hasher)>,
}
impl<K: Decode + Sized, V: Decode + Sized, Hasher: ReversibleStorageHasher> Iterator
for StorageMapIterator<K, V, Hasher>
{
type Item = (K, V);
fn next(&mut self) -> Option<(K, V)> {
loop {
let maybe_next = sp_io::storage::next_key(&self.previous_key).filter(|n| n.starts_with(&self.prefix));
break match maybe_next {
Some(next) => {
self.previous_key = next;
match unhashed::get::<V>(&self.previous_key) {
Some(value) => {
if self.drain {
unhashed::kill(&self.previous_key)
}
let mut key_material = Hasher::reverse(&self.previous_key[self.prefix.len()..]);
match K::decode(&mut key_material) {
Ok(key) => Some((key, value)),
Err(_) => continue,
}
}
None => continue,
}
}
None => None,
};
}
}
}
pub struct StorageMapIteratorShim<K, V, H> {
pub storage_map_iterator: StorageMapIterator<K, V, H>,
pub remain_iterator_count: Option<u32>,
pub finished: bool,
}
impl<K: Decode + Sized, V: Decode + Sized, H: ReversibleStorageHasher> Iterator for StorageMapIteratorShim<K, V, H> {
type Item = <StorageMapIterator<K, V, H> as Iterator>::Item;
fn next(&mut self) -> Option<Self::Item> {
if let Some(remain_iterator_count) = self.remain_iterator_count {
if remain_iterator_count == 0 {
return None;
} else {
self.remain_iterator_count = Some(remain_iterator_count - 1);
}
}
self.storage_map_iterator.next().or_else(|| {
self.finished = true;
None
})
}
}
pub trait IterableStorageMapExtended<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
type Iterator: Iterator<Item = (K, V)>;
fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
}
impl<K: FullCodec, V: FullCodec, G: StorageMapT<K, V>> IterableStorageMapExtended<K, V> for G
where
G::Hasher: ReversibleStorageHasher,
{
type Iterator = StorageMapIteratorShim<K, V, G::Hasher>;
fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
let prefix = G::prefix_hash();
let previous_key = start_key
.filter(|k| k.starts_with(&prefix))
.unwrap_or_else(|| prefix.clone());
let storage_map_iterator = StorageMapIterator {
prefix,
previous_key,
drain: false,
_phantom: Default::default(),
};
StorageMapIteratorShim {
storage_map_iterator,
remain_iterator_count: max_iterations,
finished: false,
}
}
fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
let prefix = G::prefix_hash();
let previous_key = start_key
.filter(|k| k.starts_with(&prefix))
.unwrap_or_else(|| prefix.clone());
let storage_map_iterator = StorageMapIterator {
prefix,
previous_key,
drain: true,
_phantom: Default::default(),
};
StorageMapIteratorShim {
storage_map_iterator,
remain_iterator_count: max_iterations,
finished: false,
}
}
}
pub struct MapIterator<T> {
prefix: Vec<u8>,
pub previous_key: Vec<u8>,
drain: bool,
closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
}
impl<T> Iterator for MapIterator<T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
loop {
let maybe_next = sp_io::storage::next_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 unhashed::get_raw(&self.previous_key) {
Some(raw_value) => raw_value,
None => {
frame_support::print("ERROR: next_key returned a key with no value in MapIterator");
continue;
}
};
if self.drain {
unhashed::kill(&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) => {
frame_support::print("ERROR: (key, value) failed to decode in MapIterator");
continue;
}
};
Some(item)
}
None => None,
};
}
}
}
pub struct MapIteratorShim<T> {
pub map_iterator: MapIterator<T>,
pub remain_iterator_count: Option<u32>,
pub finished: bool,
}
impl<T> Iterator for MapIteratorShim<T> {
type Item = <MapIterator<T> as Iterator>::Item;
fn next(&mut self) -> Option<Self::Item> {
if let Some(remain_iterator_count) = self.remain_iterator_count {
if remain_iterator_count == 0 {
return None;
} else {
self.remain_iterator_count = Some(remain_iterator_count - 1);
}
}
self.map_iterator.next().or_else(|| {
self.finished = true;
None
})
}
}
pub trait IterableStorageDoubleMapExtended<K1: FullCodec, K2: FullCodec, V: FullCodec>:
StorageDoubleMap<K1, K2, V>
{
type PrefixIterator: Iterator<Item = (K2, V)>;
type Iterator: Iterator<Item = (K1, K2, V)>;
fn iter_prefix(
k1: impl EncodeLike<K1>,
max_iterations: Option<u32>,
start_key: Option<Vec<u8>>,
) -> Self::PrefixIterator;
fn drain_prefix(
k1: impl EncodeLike<K1>,
max_iterations: Option<u32>,
start_key: Option<Vec<u8>>,
) -> Self::PrefixIterator;
fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
}
impl<K1: FullCodec, K2: FullCodec, V: FullCodec, G: StorageDoubleMapT<K1, K2, V>>
IterableStorageDoubleMapExtended<K1, K2, V> for G
where
G::Hasher1: ReversibleStorageHasher,
G::Hasher2: ReversibleStorageHasher,
{
type PrefixIterator = MapIteratorShim<(K2, V)>;
type Iterator = MapIteratorShim<(K1, K2, V)>;
fn iter_prefix(
k1: impl EncodeLike<K1>,
max_iterations: Option<u32>,
start_key: Option<Vec<u8>>,
) -> Self::PrefixIterator {
let prefix = G::storage_double_map_final_key1(k1);
let previous_key = start_key
.filter(|k| k.starts_with(&prefix))
.unwrap_or_else(|| prefix.clone());
let map_iterator = MapIterator {
prefix,
previous_key,
drain: false,
closure: |raw_key_without_prefix, mut raw_value| {
let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?))
},
};
MapIteratorShim {
map_iterator,
remain_iterator_count: max_iterations,
finished: false,
}
}
fn drain_prefix(
k1: impl EncodeLike<K1>,
max_iterations: Option<u32>,
start_key: Option<Vec<u8>>,
) -> Self::PrefixIterator {
let mut shim = Self::iter_prefix(k1, max_iterations, start_key);
shim.map_iterator.drain = true;
shim
}
fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
let prefix = G::prefix_hash();
let previous_key = start_key
.filter(|k| k.starts_with(&prefix))
.unwrap_or_else(|| prefix.clone());
let map_iterator = MapIterator {
prefix,
previous_key,
drain: false,
closure: |raw_key_without_prefix, mut raw_value| {
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, V::decode(&mut raw_value)?))
},
};
MapIteratorShim {
map_iterator,
remain_iterator_count: max_iterations,
finished: false,
}
}
fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
let mut shim = Self::iter(max_iterations, start_key);
shim.map_iterator.drain = true;
shim
}
}