Implement iter_keys function for all types of storage maps

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

@@ -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],

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

@@ -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)]);

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

@@ -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],