Expose the decode functions used in iterators for storage maps

paritytech · Jul 9, 2021 · 34f57d9 · 34f57d9
1 parent cfbfad9
commit 34f57d9
Show file tree

Hide file tree

Showing 4 changed files with 153 additions and 43 deletions.
diff --git a/frame/support/src/storage/generator/double_map.rs b/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();

diff --git a/frame/support/src/storage/generator/map.rs b/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();

diff --git a/frame/support/src/storage/generator/nmap.rs b/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();

diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs
@@ -325,6 +325,17 @@ pub trait IterableStorageMap<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
 	/// alter the map while doing this, you'll get undefined results.
 	fn iter_keys() -> Self::KeyIterator;
 
+	/// The decode function used while iterating through the key-value pairs in
+	/// [`IterableStorageMap::iter`].
+	fn key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		raw_value: &[u8],
+	) -> Result<(K, V), codec::Error>;
+
+	/// The decode function used while iterating through the keys in
+	/// [`IterableStorageMap::iter_keys`].
+	fn key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K, codec::Error>;
+
 	/// 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;
@@ -381,6 +392,28 @@ pub trait IterableStorageDoubleMap<
 	/// add elements to the map while doing this, you'll get undefined results.
 	fn drain() -> Self::Iterator;
 
+	/// The decode function used while iterating through the partial key-value pairs in
+	/// [`IterableStorageDoubleMap::iter_prefix`].
+	fn partial_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		raw_value: &[u8],
+	) -> Result<(K2, V), codec::Error>;
+
+	/// The decode function used while iterating through the partial keys in
+	/// [`IterableStorageDoubleMap::iter_key_prefix`].
+	fn partial_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K2, codec::Error>;
+
+	/// The decode function used while iterating through the full key-value pairs in
+	/// [`IterableStorageDoubleMap::iter`].
+	fn full_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		raw_value: &[u8],
+	) -> Result<(K1, K2, V), codec::Error>;
+
+	/// The decode function used while iterating through the full key pairs in
+	/// [`IterableStorageDoubleMap::iter_keys`].
+	fn full_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<(K1, K2), codec::Error>;
+
 	/// 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.
 	///
@@ -427,6 +460,36 @@ pub trait IterableStorageNMap<K: ReversibleKeyGenerator, V: FullCodec>: StorageN
 	/// add elements to the map while doing this, you'll get undefined results.
 	fn drain() -> Self::Iterator;
 
+	/// The decode function used while iterating through the prefix key-value pairs in
+	/// [`IterableStorageNMap::iter_prefix`].
+	///
+	/// A type parameter `KP` denoting the type of the key prefix must be provided. See 
+	fn partial_key_value_decode_fn<KP>(
+		raw_key_without_prefix: &[u8],
+		raw_value: &[u8],
+	) -> Result<(<K as HasKeyPrefix<KP>>::Suffix, V), codec::Error>
+	where K: HasReversibleKeyPrefix<KP>;
+
+	/// The decode function used while iterating through the prefix keys in
+	/// [`IterableStorageNMap::iter_key_prefix`].
+	///
+	/// A type parameter denoting the type of the key prefix must be provided. See
+	fn partial_key_decode_fn<KP>(
+		raw_key_without_prefix: &[u8],
+	) -> Result<<K as HasKeyPrefix<KP>>::Suffix, codec::Error>
+	where K: HasReversibleKeyPrefix<KP>;
+
+	/// The decode function used while iterating through the key-value pairs in
+	/// [`IterableStorageNMap::iter`].
+	fn full_key_value_decode_fn(
+		raw_key_without_prefix: &[u8],
+		raw_value: &[u8],
+	) -> Result<(K::Key, V), codec::Error>;
+
+	/// The decode function used while iterating through the keys in
+	/// [`IterableStorageNMap::iter_keys`].
+	fn full_key_decode_fn(raw_key_without_prefix: &[u8]) -> Result<K::Key, codec::Error>;
+
 	/// 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.
 	///