Add methods to PrefixIterator to support iterating from a specific key

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

@@ -351,10 +351,7 @@ impl<
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			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)?))
-			},
+			closure: <Self as storage::IterableStorageDoubleMap<K1, K2, V>>::partial_key_value_decode_fn,
 		}
 	}
 
@@ -364,10 +361,7 @@ impl<
 			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)
-			}
+			closure: <Self as storage::IterableStorageDoubleMap<K1, K2, V>>::partial_key_decode_fn,
 		}
 	}
 
@@ -383,13 +377,7 @@ impl<
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			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)?))
-			},
+			closure: <Self as storage::IterableStorageDoubleMap<K1, K2, V>>::full_key_value_decode_fn,
 		}
 	}
 
@@ -399,13 +387,7 @@ impl<
 			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))
-			}
+			closure: <Self as storage::IterableStorageDoubleMap<K1, K2, V>>::full_key_decode_fn,
 		}
 	}
 
@@ -415,6 +397,38 @@ impl<
 		iterator
 	}
 
+	fn partial_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		mut raw_value: &[u8],
+	) -> Result<(K2, V), codec::Error> {
+		let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
+		Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?))
+	}
+
+	fn partial_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K2, codec::Error> {
+		let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
+		K2::decode(&mut key_material)
+	}
+
+	fn full_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		mut raw_value: &[u8],
+	) -> Result<(K1, K2, V), codec::Error> {
+		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)?))
+	}
+
+	fn full_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<(K1, K2), codec::Error> {
+		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 translate<O: Decode, F: FnMut(K1, K2, O) -> Option<V>>(mut f: F) {
 		let prefix = G::prefix_hash();
 		let mut previous_key = prefix.clone();

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

@@ -149,10 +149,7 @@ impl<
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			closure: |raw_key_without_prefix, mut raw_value| {
-				let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
-				Ok((K::decode(&mut key_material)?, V::decode(&mut raw_value)?))
-			},
+			closure: <Self as storage::IterableStorageMap<K, V>>::key_value_decode_fn,
 		}
 	}
 
@@ -163,13 +160,25 @@ impl<
 			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)
-			}
+			closure: <Self as storage::IterableStorageMap<K, V>>::key_decode_fn,
 		}
 	}
 
+	/// Decode function used in iter.
+	fn key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		mut raw_value: &[u8],
+	) -> Result<(K, V), codec::Error> {
+		let mut key_material = G::Hasher::reverse(raw_key_without_prefix);
+		Ok((K::decode(&mut key_material)?, V::decode(&mut raw_value)?))
+	}
+
+	/// Decode function used in iter_keys.
+	fn key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K, codec::Error> {
+		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();

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

@@ -322,10 +322,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			closure: |raw_key_without_prefix, mut raw_value| {
-				let partial_key = K::decode_partial_key(raw_key_without_prefix)?;
-				Ok((partial_key, V::decode(&mut raw_value)?))
-			},
+			closure: <Self as storage::IterableStorageNMap<K, V>>::partial_key_value_decode_fn::<KP>,
 		}
 	}
 
@@ -338,7 +335,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			closure: K::decode_partial_key,
+			closure: <Self as storage::IterableStorageNMap<K, V>>::partial_key_decode_fn::<KP>,
 		}
 	}
 
@@ -357,10 +354,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			closure: |raw_key_without_prefix, mut raw_value| {
-				let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
-				Ok((final_key, V::decode(&mut raw_value)?))
-			},
+			closure: <Self as storage::IterableStorageNMap<K, V>>::full_key_value_decode_fn,
 		}
 	}
 
@@ -370,10 +364,7 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 			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)
-			}
+			closure: <Self as storage::IterableStorageNMap<K, V>>::full_key_decode_fn,
 		}
 	}
 
@@ -383,6 +374,39 @@ impl<K: ReversibleKeyGenerator, V: FullCodec, G: StorageNMap<K, V>>
 		iterator
 	}
 
+	fn partial_key_value_decode_fn<KP>(
+		raw_key_without_prefix: &[u8],
+		mut raw_value: &[u8],
+	) -> Result<(<K as HasKeyPrefix<KP>>::Suffix, V), codec::Error>
+	where
+		K: HasReversibleKeyPrefix<KP>,
+	{
+		let partial_key = K::decode_partial_key(raw_key_without_prefix)?;
+		Ok((partial_key, V::decode(&mut raw_value)?))
+	}
+
+	fn partial_key_decode_fn<KP>(
+		raw_key_without_prefix: &[u8],
+	) -> Result<<K as HasKeyPrefix<KP>>::Suffix, codec::Error>
+	where
+		K: HasReversibleKeyPrefix<KP>,
+	{
+		K::decode_partial_key(raw_key_without_prefix)
+	}
+
+	fn full_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		mut raw_value: &[u8],
+	) -> Result<(K::Key, V), codec::Error> {
+		let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
+		Ok((final_key, V::decode(&mut raw_value)?))
+	}
+
+	fn full_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K::Key, codec::Error> {
+		let (final_key, _) = K::decode_final_key(raw_key_without_prefix)?;
+		Ok(final_key)
+	}
+
 	fn translate<O: Decode, F: FnMut(K::Key, O) -> Option<V>>(mut f: F) {
 		let prefix = G::prefix_hash();
 		let mut previous_key = prefix.clone();