Some iterator utils for double maps.

frame/support/procedural/src/storage/parse.rs

@@ -28,14 +28,17 @@ mod keyword {
 	syn::custom_keyword!(get);
 	syn::custom_keyword!(map);
 	syn::custom_keyword!(double_map);
-	syn::custom_keyword!(blake2_256);
-	syn::custom_keyword!(blake2_128);
+	syn::custom_keyword!(opaque_blake2_256);
+	syn::custom_keyword!(opaque_blake2_128);
 	syn::custom_keyword!(blake2_128_concat);
-	syn::custom_keyword!(twox_256);
-	syn::custom_keyword!(twox_128);
+	syn::custom_keyword!(opaque_twox_256);
+	syn::custom_keyword!(opaque_twox_128);
 	syn::custom_keyword!(twox_64_concat);
 	syn::custom_keyword!(identity);
 	syn::custom_keyword!(hasher);
+	syn::custom_keyword!(tainted);
+	syn::custom_keyword!(natural);
+	syn::custom_keyword!(prehashed);
 }
 
 /// Specific `Opt` to implement structure with optional parsing
@@ -233,32 +236,38 @@ struct DeclStorageDoubleMap {
 
 #[derive(ToTokens, Debug)]
 enum Hasher {
-	Blake2_256(keyword::blake2_256),
-	Blake2_128(keyword::blake2_128),
+	Blake2_256(keyword::opaque_blake2_256),
+	Blake2_128(keyword::opaque_blake2_128),
 	Blake2_128Concat(keyword::blake2_128_concat),
-	Twox256(keyword::twox_256),
-	Twox128(keyword::twox_128),
+	Twox256(keyword::opaque_twox_256),
+	Twox128(keyword::opaque_twox_128),
 	Twox64Concat(keyword::twox_64_concat),
 	Identity(keyword::identity),
 }
 
 impl syn::parse::Parse for Hasher {
 	fn parse(input: syn::parse::ParseStream) -> syn::parse::Result<Self> {
 		let lookahead = input.lookahead1();
-		if lookahead.peek(keyword::blake2_256) {
+		if lookahead.peek(keyword::opaque_blake2_256) {
 			Ok(Self::Blake2_256(input.parse()?))
-		} else if lookahead.peek(keyword::blake2_128) {
+		} else if lookahead.peek(keyword::opaque_blake2_128) {
 			Ok(Self::Blake2_128(input.parse()?))
 		} else if lookahead.peek(keyword::blake2_128_concat) {
 			Ok(Self::Blake2_128Concat(input.parse()?))
-		} else if lookahead.peek(keyword::twox_256) {
+		} else if lookahead.peek(keyword::opaque_twox_256) {
 			Ok(Self::Twox256(input.parse()?))
-		} else if lookahead.peek(keyword::twox_128) {
+		} else if lookahead.peek(keyword::opaque_twox_128) {
 			Ok(Self::Twox128(input.parse()?))
 		} else if lookahead.peek(keyword::twox_64_concat) {
 			Ok(Self::Twox64Concat(input.parse()?))
 		} else if lookahead.peek(keyword::identity) {
 			Ok(Self::Identity(input.parse()?))
+		} else if lookahead.peek(keyword::tainted) {
+			Ok(Self::Blake2_128Concat(input.parse()?))
+		} else if lookahead.peek(keyword::natural) {
+			Ok(Self::Twox64Concat(input.parse()?))
+		} else if lookahead.peek(keyword::prehashed) {
+			Ok(Self::Identity(input.parse()?))
 		} else {
 			Err(lookahead.error())
 		}

frame/support/src/lib.rs

@@ -71,8 +71,8 @@ pub use self::hash::{
 	StorageHasher
 };
 pub use self::storage::{
-	StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap,
-	IterableStorageMap
+	StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap, IterableStorageMap,
+	IterableStorageDoubleMap,
 };
 pub use self::dispatch::{Parameter, Callable, IsSubType};
 pub use sp_runtime::{self, ConsensusEngineId, print, traits::Printable};

frame/support/src/storage/generator/double_map.rs

@@ -16,8 +16,9 @@
 
 use sp_std::prelude::*;
 use sp_std::borrow::Borrow;
-use codec::{Ref, FullCodec, FullEncode, Encode, EncodeLike, EncodeAppend};
-use crate::{storage::{self, unhashed}, hash::{StorageHasher, Twox128}, traits::Len};
+use codec::{Ref, FullCodec, FullEncode, Decode, Encode, EncodeLike, EncodeAppend};
+use crate::{storage::{self, unhashed}, traits::Len};
+use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
 /// Generator for `StorageDoubleMap` used by `decl_storage`.
 ///
@@ -55,6 +56,22 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 	/// Storage prefix. Used for generating final key.
 	fn storage_prefix() -> &'static [u8];
 
+	/// The full prefix; just the hash of `module_prefix` concatenated to the hash of
+	/// `storage_prefix`.
+	fn prefix_hash() -> Vec<u8> {
+		let module_prefix_hashed = Twox128::hash(Self::module_prefix());
+		let storage_prefix_hashed = Twox128::hash(Self::storage_prefix());
+
+		let mut result = Vec::with_capacity(
+			module_prefix_hashed.len() + storage_prefix_hashed.len()
+		);
+
+		result.extend_from_slice(&module_prefix_hashed[..]);
+		result.extend_from_slice(&storage_prefix_hashed[..]);
+
+		result
+	}
+
 	/// Convert an optional value retrieved from storage to the type queried.
 	fn from_optional_value_to_query(v: Option<V>) -> Self::Query;
 
@@ -314,7 +331,103 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 			unhashed::put(Self::storage_double_map_final_key(key1, key2).as_ref(), &value);
 			value
 		})
+	}
+}
+
+/// Utility to iterate through items in a storage map.
+pub struct MapIterator<K, V, Hasher> {
+	prefix: Vec<u8>,
+	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 MapIterator<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,
+			}
+		}
+	}
+}
 
+impl<
+	K1: FullCodec,
+	K2: FullCodec,
+	V: FullCodec,
+	G: StorageDoubleMap<K1, K2, V>,
+> storage::IterableStorageDoubleMap<K1, K2, V> for G where
+	G::Hasher1: ReversibleStorageHasher,
+	G::Hasher2: ReversibleStorageHasher
+{
+	type Iterator = MapIterator<K2, V, G::Hasher2>;
+
+	/// Enumerate all elements in the map.
+	fn iter(k1: impl EncodeLike<K1>) -> Self::Iterator {
+		let prefix = G::storage_double_map_final_key1(k1);
+		Self::Iterator {
+			prefix: prefix.clone(),
+			previous_key: prefix,
+			drain: false,
+			_phantom: Default::default(),
+		}
+	}
+
+	/// Enumerate all elements in the map.
+	fn drain(k1: impl EncodeLike<K1>) -> Self::Iterator {
+		let prefix = G::storage_double_map_final_key1(k1);
+		Self::Iterator {
+			prefix: prefix.clone(),
+			previous_key: prefix,
+			drain: true,
+			_phantom: Default::default(),
+		}
+	}
+
+	fn translate<O: Decode, F: Fn(O) -> Option<V>>(f: F) {
+		let prefix = G::prefix_hash();
+		let mut previous_key = prefix.clone();
+		loop {
+			match sp_io::storage::next_key(&previous_key).filter(|n| n.starts_with(&prefix)) {
+				Some(next) => {
+					previous_key = next;
+					let maybe_value = unhashed::get::<O>(&previous_key);
+					match maybe_value {
+						Some(value) => match f(value) {
+							Some(new) => unhashed::put::<V>(&previous_key, &new),
+							None => unhashed::kill(&previous_key),
+						},
+						None => continue,
+					}
+				}
+				None => return,
+			}
+		}
 	}
 }
 

frame/support/src/storage/generator/map.rs

@@ -18,8 +18,8 @@
 use sp_std::prelude::*;
 use sp_std::borrow::Borrow;
 use codec::{FullCodec, FullEncode, Decode, Encode, EncodeLike, Ref, EncodeAppend};
-use crate::{storage::{self, unhashed}, hash::{StorageHasher, Twox128}, traits::Len};
-use crate::hash::ReversibleStorageHasher;
+use crate::{storage::{self, unhashed}, traits::Len};
+use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
 /// Generator for `StorageMap` used by `decl_storage`.
 ///
@@ -185,7 +185,6 @@ impl<
 				None => return,
 			}
 		}
-
 	}
 }
 

frame/support/src/storage/mod.rs

@@ -234,6 +234,30 @@ pub trait IterableStorageMap<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
 	fn translate<O: Decode, F: Fn(K, O) -> Option<V>>(f: F);
 }
 
+/// A strongly-typed double map in storage whose secondary keys and values can be iterated over.
+pub trait IterableStorageDoubleMap<
+	K1: FullCodec,
+	K2: FullCodec,
+	V: FullCodec
+>: StorageDoubleMap<K1, K2, V> {
+	/// The type that iterates over all `(key, value)`.
+	type Iterator: Iterator<Item = (K2, V)>;
+
+	/// Enumerate all elements in the map with first key `k1` in no particular order. If you add or
+	/// remove values whose first key is `k1` to the map while doing this, you'll get undefined
+	/// results.
+	fn iter(k1: impl EncodeLike<K1>) -> Self::Iterator;
+
+	/// Remove all elements from the map with first key `k1` and iterate through them in no
+	/// particular order. If you add elements with first key `k1` to the map while doing this,
+	/// you'll get undefined results.
+	fn drain(k1: impl EncodeLike<K1>) -> Self::Iterator;
+
+	/// Translate the values of all elements by a function `f`, in the map in no particular order.
+	/// By returning `None` from `f` for an element, you'll remove it from the map.
+	fn translate<O: Decode, F: Fn(O) -> Option<V>>(f: F);
+}
+
 /// An implementation of a map with a two keys.
 ///
 /// It provides an important ability to efficiently remove all entries