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>

substrate/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],

substrate/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)]);

substrate/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],

substrate/frame/support/src/storage/mod.rs

@@ -314,11 +314,17 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 pub trait IterableStorageMap<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
 	/// The type that iterates over all `(key, value)`.
 	type Iterator: Iterator<Item = (K, V)>;
+	/// The type that itereates over all `key`s.
+	type KeyIterator: Iterator<Item = K>;
 
 	/// Enumerate all elements in the map in no particular order. If you alter the map while doing
 	/// this, you'll get undefined results.
 	fn iter() -> Self::Iterator;
 
+	/// Enumerate all keys in the map in no particular order, skipping over the elements. If you
+	/// alter the map while doing this, you'll get undefined results.
+	fn iter_keys() -> Self::KeyIterator;
+
 	/// Remove all elements from the map and iterate through them in no particular order. If you
 	/// add elements to the map while doing this, you'll get undefined results.
 	fn drain() -> Self::Iterator;
@@ -336,9 +342,15 @@ pub trait IterableStorageDoubleMap<
 	K2: FullCodec,
 	V: FullCodec
 >: StorageDoubleMap<K1, K2, V> {
+	/// The type that iterates over all `key2`.
+	type PartialKeyIterator: Iterator<Item = K2>;
+
 	/// The type that iterates over all `(key2, value)`.
 	type PrefixIterator: Iterator<Item = (K2, V)>;
 
+	/// The type that iterates over all `(key1, key2)`.
+	type FullKeyIterator: Iterator<Item = (K1, K2)>;
+
 	/// The type that iterates over all `(key1, key2, value)`.
 	type Iterator: Iterator<Item = (K1, K2, V)>;
 
@@ -347,6 +359,11 @@ pub trait IterableStorageDoubleMap<
 	/// results.
 	fn iter_prefix(k1: impl EncodeLike<K1>) -> Self::PrefixIterator;
 
+	/// Enumerate all second keys `k2` in the map with the same 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_key_prefix(k1: impl EncodeLike<K1>) -> Self::PartialKeyIterator;
+
 	/// 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.
@@ -356,6 +373,10 @@ pub trait IterableStorageDoubleMap<
 	/// the map while doing this, you'll get undefined results.
 	fn iter() -> Self::Iterator;
 
+	/// Enumerate all keys `k1` and `k2` in the map in no particular order. If you add or remove
+	/// values to the map while doing this, you'll get undefined results.
+	fn iter_keys() -> Self::FullKeyIterator;
+
 	/// Remove all elements from the map and iterate through them in no particular order. If you
 	/// add elements to the map while doing this, you'll get undefined results.
 	fn drain() -> Self::Iterator;
@@ -370,7 +391,10 @@ pub trait IterableStorageDoubleMap<
 /// A strongly-typed map with arbitrary number of keys in storage whose keys and values can be
 /// iterated over.
 pub trait IterableStorageNMap<K: ReversibleKeyGenerator, V: FullCodec>: StorageNMap<K, V> {
-	/// The type that iterates over all `(key1, (key2, (key3, ... (keyN, ()))), value)` tuples
+	/// The type that iterates over all `(key1, key2, key3, ... keyN)` tuples.
+	type KeyIterator: Iterator<Item = K::Key>;
+
+	/// The type that iterates over all `(key1, key2, key3, ... keyN), value)` tuples.
 	type Iterator: Iterator<Item = (K::Key, V)>;
 
 	/// Enumerate all elements in the map with prefix key `kp` in no particular order. If you add or
@@ -379,6 +403,12 @@ pub trait IterableStorageNMap<K: ReversibleKeyGenerator, V: FullCodec>: StorageN
 	fn iter_prefix<KP>(kp: KP) -> PrefixIterator<(<K as HasKeyPrefix<KP>>::Suffix, V)>
 	where K: HasReversibleKeyPrefix<KP>;
 
+	/// Enumerate all suffix keys in the map with prefix key `kp` in no particular order. If you
+	/// add or remove values whose prefix is `kp` to the map while doing this, you'll get undefined
+	/// results.
+	fn iter_key_prefix<KP>(kp: KP) -> KeyPrefixIterator<<K as HasKeyPrefix<KP>>::Suffix>
+	where K: HasReversibleKeyPrefix<KP>;
+
 	/// Remove all elements from the map with prefix key `kp` and iterate through them in no
 	/// particular order. If you add elements with prefix key `kp` to the map while doing this,
 	/// you'll get undefined results.
@@ -389,6 +419,10 @@ pub trait IterableStorageNMap<K: ReversibleKeyGenerator, V: FullCodec>: StorageN
 	/// the map while doing this, you'll get undefined results.
 	fn iter() -> Self::Iterator;
 
+	/// Enumerate all keys in the map in no particular order. If you add or remove values to the
+	/// map while doing this, you'll get undefined results.
+	fn iter_keys() -> Self::KeyIterator;
+
 	/// Remove all elements from the map and iterate through them in no particular order. If you
 	/// add elements to the map while doing this, you'll get undefined results.
 	fn drain() -> Self::Iterator;
@@ -733,6 +767,56 @@ impl<T> Iterator for PrefixIterator<T> {
 	}
 }
 
+/// Iterate over a prefix and decode raw_key into `T`.
+///
+/// If any decoding fails it skips it and continues to the next key.
+pub struct KeyPrefixIterator<T> {
+	prefix: Vec<u8>,
+	previous_key: Vec<u8>,
+	/// If true then value are removed while iterating
+	drain: bool,
+	/// Function that take `raw_key_without_prefix` and decode `T`.
+	/// `raw_key_without_prefix` is the raw storage key without the prefix iterated on.
+	closure: fn(&[u8]) -> Result<T, codec::Error>,
+}
+
+impl<T> KeyPrefixIterator<T> {
+	/// Mutate this iterator into a draining iterator; items iterated are removed from storage.
+	pub fn drain(mut self) -> Self {
+		self.drain = true;
+		self
+	}
+}
+
+impl<T> Iterator for KeyPrefixIterator<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));
+
+			if let Some(next) = maybe_next {
+				self.previous_key = next;
+				if self.drain {
+					unhashed::kill(&self.previous_key);
+				}
+				let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
+
+				match (self.closure)(raw_key_without_prefix) {
+					Ok(item) => return Some(item),
+					Err(e) => {
+						log::error!("key failed to decode at {:?}: {:?}", self.previous_key, e);
+						continue;
+					}
+				}
+			}
+
+			return None;
+		}
+	}
+}
+
 /// Iterate over a prefix of a child trie and decode raw_key and raw_value into `T`.
 ///
 /// If any decoding fails it skips the key and continues to the next one.
@@ -1276,6 +1360,59 @@ mod test {
 		});
 	}
 
+	#[test]
+	fn key_prefix_iterator_works() {
+		TestExternalities::default().execute_with(|| {
+			use crate::storage::generator::StorageMap;
+			use crate::hash::Twox64Concat;
+			struct MyStorageMap;
+			impl StorageMap<u64, u64> for MyStorageMap {
+				type Query = u64;
+				type Hasher = Twox64Concat;
+
+				fn module_prefix() -> &'static [u8] {
+					b"MyModule"
+				}
+
+				fn storage_prefix() -> &'static [u8] {
+					b"MyStorageMap"
+				}
+
+				fn from_optional_value_to_query(v: Option<u64>) -> Self::Query {
+					v.unwrap_or_default()
+				}
+
+				fn from_query_to_optional_value(v: Self::Query) -> Option<u64> {
+					Some(v)
+				}
+			}
+
+			let k = [twox_128(b"MyModule"), twox_128(b"MyStorageMap")].concat();
+			assert_eq!(MyStorageMap::prefix_hash().to_vec(), k);
+
+			// empty to start
+			assert!(MyStorageMap::iter_keys().collect::<Vec<_>>().is_empty());
+
+			MyStorageMap::insert(1, 10);
+			MyStorageMap::insert(2, 20);
+			MyStorageMap::insert(3, 30);
+			MyStorageMap::insert(4, 40);
+
+			// just looking
+			let mut keys = MyStorageMap::iter_keys().collect::<Vec<_>>();
+			keys.sort();
+			assert_eq!(keys, vec![1, 2, 3, 4]);
+
+			// draining the keys and values
+			let mut drained_keys = MyStorageMap::iter_keys().drain().collect::<Vec<_>>();
+			drained_keys.sort();
+			assert_eq!(drained_keys, vec![1, 2, 3, 4]);
+
+			// empty again
+			assert!(MyStorageMap::iter_keys().collect::<Vec<_>>().is_empty());
+		});
+	}
+
 	#[test]
 	fn child_trie_prefixed_map_works() {
 		TestExternalities::default().execute_with(|| {

substrate/frame/support/src/storage/types/double_map.rs

@@ -387,6 +387,15 @@ where
 		<Self as crate::storage::IterableStorageDoubleMap<Key1, Key2, Value>>::iter_prefix(k1)
 	}
 
+	/// Enumerate all second keys `k2` in the map with the same 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.
+	pub fn iter_key_prefix(k1: impl EncodeLike<Key1>) -> crate::storage::KeyPrefixIterator<Key2> {
+		<Self as crate::storage::IterableStorageDoubleMap<Key1, Key2, Value>>::iter_key_prefix(k1)
+	}
+
 	/// Remove all elements from the map with first key `k1` and iterate through them in no
 	/// particular order.
 	///
@@ -403,6 +412,13 @@ where
 		<Self as crate::storage::IterableStorageDoubleMap<Key1, Key2, Value>>::iter()
 	}
 
+	/// Enumerate all keys `k1` and `k2` in the map in no particular order.
+	///
+	/// If you add or remove values to the map while doing this, you'll get undefined results.
+	pub fn iter_keys() -> crate::storage::KeyPrefixIterator<(Key1, Key2)> {
+		<Self as crate::storage::IterableStorageDoubleMap<Key1, Key2, Value>>::iter_keys()
+	}
+
 	/// Remove all elements from the map and iterate through them in no particular order.
 	///
 	/// If you add elements to the map while doing this, you'll get undefined results.

substrate/frame/support/src/storage/types/map.rs

@@ -297,6 +297,13 @@ where
 		<Self as crate::storage::IterableStorageMap<Key, Value>>::iter()
 	}
 
+	/// Enumerate all keys in the map in no particular order.
+	///
+	/// If you alter the map while doing this, you'll get undefined results.
+	pub fn iter_keys() -> crate::storage::KeyPrefixIterator<Key> {
+		<Self as crate::storage::IterableStorageMap<Key, Value>>::iter_keys()
+	}
+
 	/// Remove all elements from the map and iterate through them in no particular order.
 	///
 	/// If you add elements to the map while doing this, you'll get undefined results.

substrate/frame/support/src/storage/types/nmap.rs

@@ -318,6 +318,19 @@ where
 		<Self as crate::storage::IterableStorageNMap<Key, Value>>::iter_prefix(kp)
 	}
 
+	/// Enumerate all suffix keys in the map with prefix key `kp` in no particular order.
+	///
+	/// If you add or remove values whose prefix key is `kp` to the map while doing this, you'll get
+	/// undefined results.
+	pub fn iter_key_prefix<KP>(
+		kp: KP,
+	) -> crate::storage::KeyPrefixIterator<<Key as HasKeyPrefix<KP>>::Suffix>
+	where
+		Key: HasReversibleKeyPrefix<KP>,
+	{
+		<Self as crate::storage::IterableStorageNMap<Key, Value>>::iter_key_prefix(kp)
+	}
+
 	/// Remove all elements from the map with prefix key `kp` and iterate through them in no
 	/// particular order.
 	///
@@ -339,6 +352,13 @@ where
 		<Self as crate::storage::IterableStorageNMap<Key, Value>>::iter()
 	}
 
+	/// Enumerate all keys in the map in no particular order.
+	///
+	/// If you add or remove values to the map while doing this, you'll get undefined results.
+	pub fn iter_keys() -> crate::storage::KeyPrefixIterator<Key::Key> {
+		<Self as crate::storage::IterableStorageNMap<Key, Value>>::iter_keys()
+	}
+
 	/// Remove all elements from the map and iterate through them in no particular order.
 	///
 	/// If you add elements to the map while doing this, you'll get undefined results.