Frame-support storage: make iterations and translate consistent (#5470)

* implementation and factorisation

* factorize test

* doc

* fix bug and improve test

* address suggestions

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}, Never};
+use crate::{storage::{self, unhashed, StorageAppend, PrefixIterator}, Never};
 use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
 /// Generator for `StorageDoubleMap` used by `decl_storage`.
@@ -213,10 +213,11 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 		KArg1: ?Sized + EncodeLike<K1>
 	{
 		let prefix = Self::storage_double_map_final_key1(k1);
-		storage::PrefixIterator::<V> {
+		storage::PrefixIterator {
 			prefix: prefix.clone(),
 			previous_key: prefix,
-			phantom_data: Default::default(),
+			drain: false,
+			closure: |_raw_key, mut raw_value| V::decode(&mut raw_value),
 		}
 	}
 
@@ -322,54 +323,6 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 	}
 }
 
-/// Iterate over a prefix and decode raw_key and raw_value into `T`.
-pub struct MapIterator<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, raw_value)` and decode `T`.
-	/// `raw_key_without_prefix` is the raw storage key without the prefix iterated on.
-	closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
-}
-
-impl<T> Iterator for MapIterator<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));
-			break match maybe_next {
-				Some(next) => {
-					self.previous_key = next;
-					let raw_value = match unhashed::get_raw(&self.previous_key) {
-						Some(raw_value) => raw_value,
-						None => {
-							frame_support::print("ERROR: next_key returned a key with no value in MapIterator");
-							continue
-						}
-					};
-					if self.drain {
-						unhashed::kill(&self.previous_key)
-					}
-					let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
-					let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
-						Ok(item) => item,
-						Err(_e) => {
-							frame_support::print("ERROR: (key, value) failed to decode in MapIterator");
-							continue
-						}
-					};
-
-					Some(item)
-				}
-				None => None,
-			}
-		}
-	}
-}
-
 impl<
 	K1: FullCodec,
 	K2: FullCodec,
@@ -379,8 +332,8 @@ impl<
 	G::Hasher1: ReversibleStorageHasher,
 	G::Hasher2: ReversibleStorageHasher
 {
-	type PrefixIterator = MapIterator<(K2, V)>;
-	type Iterator = MapIterator<(K1, K2, V)>;
+	type PrefixIterator = PrefixIterator<(K2, V)>;
+	type Iterator = PrefixIterator<(K1, K2, V)>;
 
 	fn iter_prefix(k1: impl EncodeLike<K1>) -> Self::PrefixIterator {
 		let prefix = G::storage_double_map_final_key1(k1);
@@ -423,23 +376,41 @@ impl<
 		iterator
 	}
 
-	fn translate<O: Decode, F: Fn(O) -> Option<V>>(f: F) {
+	fn translate<O: Decode, F: Fn(K1, K2, 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,
+		while let Some(next) = sp_io::storage::next_key(&previous_key)
+			.filter(|n| n.starts_with(&prefix))
+		{
+			previous_key = next;
+			let value = match unhashed::get::<O>(&previous_key) {
+				Some(value) => value,
+				None => {
+					crate::debug::error!("Invalid translate: fail to decode old value");
+					continue
+				},
+			};
+			let mut key_material = G::Hasher1::reverse(&previous_key[prefix.len()..]);
+			let key1 = match K1::decode(&mut key_material) {
+				Ok(key1) => key1,
+				Err(_) => {
+					crate::debug::error!("Invalid translate: fail to decode key1");
+					continue
+				},
+			};
+
+			let mut key2_material = G::Hasher2::reverse(&key_material);
+			let key2 = match K2::decode(&mut key2_material) {
+				Ok(key2) => key2,
+				Err(_) => {
+					crate::debug::error!("Invalid translate: fail to decode key2");
+					continue
+				},
+			};
+
+			match f(key1, key2, value) {
+				Some(new) => unhashed::put::<V>(&previous_key, &new),
+				None => unhashed::kill(&previous_key),
 			}
 		}
 	}
@@ -447,10 +418,12 @@ impl<
 
 /// Test iterators for StorageDoubleMap
 #[cfg(test)]
-#[allow(dead_code)]
 mod test_iterators {
 	use codec::{Encode, Decode};
-	use crate::storage::{generator::StorageDoubleMap, IterableStorageDoubleMap, unhashed};
+	use crate::{
+		hash::StorageHasher,
+		storage::{generator::StorageDoubleMap, IterableStorageDoubleMap, unhashed},
+	};
 
 	pub trait Trait {
 		type Origin;
@@ -466,7 +439,7 @@ mod test_iterators {
 
 	crate::decl_storage! {
 		trait Store for Module<T: Trait> as Test {
-			DoubleMap: double_map hasher(blake2_128_concat) u16, hasher(blake2_128_concat) u32 => u64;
+			DoubleMap: double_map hasher(blake2_128_concat) u16, hasher(twox_64_concat) u32 => u64;
 		}
 	}
 
@@ -484,11 +457,6 @@ mod test_iterators {
 		prefix
 	}
 
-	fn key_in_prefix(mut prefix: Vec<u8>) -> Vec<u8> {
-		prefix.push(0);
-		prefix
-	}
-
 	#[test]
 	fn double_map_reversible_reversible_iteration() {
 		sp_io::TestExternalities::default().execute_with(|| {
@@ -534,22 +502,59 @@ mod test_iterators {
 
 			assert_eq!(
 				DoubleMap::iter_prefix(k1).collect::<Vec<_>>(),
-				vec![(0, 0), (2, 2), (1, 1), (3, 3)],
+				vec![(1, 1), (2, 2), (0, 0), (3, 3)],
 			);
 
 			assert_eq!(
 				DoubleMap::iter_prefix_values(k1).collect::<Vec<_>>(),
-				vec![0, 2, 1, 3],
+				vec![1, 2, 0, 3],
 			);
 
 			assert_eq!(
 				DoubleMap::drain_prefix(k1).collect::<Vec<_>>(),
-				vec![(0, 0), (2, 2), (1, 1), (3, 3)],
+				vec![(1, 1), (2, 2), (0, 0), (3, 3)],
 			);
 
 			assert_eq!(DoubleMap::iter_prefix(k1).collect::<Vec<_>>(), vec![]);
 			assert_eq!(unhashed::get(&key_before_prefix(prefix.clone())), Some(1u64));
 			assert_eq!(unhashed::get(&key_after_prefix(prefix.clone())), Some(1u64));
+
+			// Translate
+			let prefix = DoubleMap::prefix_hash();
+
+			unhashed::put(&key_before_prefix(prefix.clone()), &1u64);
+			unhashed::put(&key_after_prefix(prefix.clone()), &1u64);
+			for i in 0..4 {
+				DoubleMap::insert(i as u16, i as u32, i as u64);
+			}
+
+			// Wrong key1
+			unhashed::put(
+				&[prefix.clone(), vec![1, 2, 3]].concat(),
+				&3u64.encode()
+			);
+
+			// Wrong key2
+			unhashed::put(
+				&[prefix.clone(), crate::Blake2_128Concat::hash(&1u16.encode())].concat(),
+				&3u64.encode()
+			);
+
+			// Wrong value
+			unhashed::put(
+				&[
+					prefix.clone(),
+					crate::Blake2_128Concat::hash(&1u16.encode()),
+					crate::Twox64Concat::hash(&2u32.encode()),
+				].concat(),
+				&vec![1],
+			);
+
+			DoubleMap::translate(|_k1, _k2, v: u64| Some(v*2));
+			assert_eq!(
+				DoubleMap::iter().collect::<Vec<_>>(),
+				vec![(3, 3, 6), (0, 0, 0), (2, 2, 4), (1, 1, 2)],
+			);
 		})
 	}
 }

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},
+	storage::{self, unhashed, StorageAppend, PrefixIterator},
 	Never, hash::{StorageHasher, Twox128, ReversibleStorageHasher},
 };
 
@@ -139,53 +139,56 @@ impl<
 > storage::IterableStorageMap<K, V> for G where
 	G::Hasher: ReversibleStorageHasher
 {
-	type Iterator = StorageMapIterator<K, V, G::Hasher>;
+	type Iterator = PrefixIterator<(K, V)>;
 
 	/// Enumerate all elements in the map.
 	fn iter() -> Self::Iterator {
 		let prefix = G::prefix_hash();
-		Self::Iterator {
+		PrefixIterator {
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			_phantom: Default::default(),
+			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)?))
+			},
 		}
 	}
 
 	/// Enumerate all elements in the map.
 	fn drain() -> Self::Iterator {
-		let prefix = G::prefix_hash();
-		Self::Iterator {
-			prefix: prefix.clone(),
-			previous_key: prefix,
-			drain: true,
-			_phantom: Default::default(),
-		}
+		let mut iterator = Self::iter();
+		iterator.drain = true;
+		iterator
 	}
 
 	fn translate<O: Decode, F: Fn(K, 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) => {
-							let mut key_material = G::Hasher::reverse(&previous_key[prefix.len()..]);
-							match K::decode(&mut key_material) {
-								Ok(key) => match f(key, value) {
-									Some(new) => unhashed::put::<V>(&previous_key, &new),
-									None => unhashed::kill(&previous_key),
-								},
-								Err(_) => continue,
-							}
-						}
-						None => continue,
-					}
-				}
-				None => return,
+		while let Some(next) = sp_io::storage::next_key(&previous_key)
+			.filter(|n| n.starts_with(&prefix))
+		{
+			previous_key = next;
+			let value = match unhashed::get::<O>(&previous_key) {
+				Some(value) => value,
+				None => {
+					crate::debug::error!("Invalid translate: fail to decode old value");
+					continue
+				},
+			};
+
+			let mut key_material = G::Hasher::reverse(&previous_key[prefix.len()..]);
+			let key = match K::decode(&mut key_material) {
+				Ok(key) => key,
+				Err(_) => {
+					crate::debug::error!("Invalid translate: fail to decode key");
+					continue
+				},
+			};
+
+			match f(key, value) {
+				Some(new) => unhashed::put::<V>(&previous_key, &new),
+				None => unhashed::kill(&previous_key),
 			}
 		}
 	}
@@ -312,3 +315,91 @@ impl<K: FullEncode, V: FullCodec, G: StorageMap<K, V>> storage::StorageMap<K, V>
 		})
 	}
 }
+
+/// Test iterators for StorageMap
+#[cfg(test)]
+mod test_iterators {
+	use codec::{Encode, Decode};
+	use crate::{
+		hash::StorageHasher,
+		storage::{generator::StorageMap, IterableStorageMap, unhashed},
+	};
+
+	pub trait Trait {
+		type Origin;
+		type BlockNumber;
+	}
+
+	crate::decl_module! {
+		pub struct Module<T: Trait> for enum Call where origin: T::Origin {}
+	}
+
+	#[derive(PartialEq, Eq, Clone, Encode, Decode)]
+	struct NoDef(u32);
+
+	crate::decl_storage! {
+		trait Store for Module<T: Trait> as Test {
+			Map: map hasher(blake2_128_concat) u16 => u64;
+		}
+	}
+
+	fn key_before_prefix(mut prefix: Vec<u8>) -> Vec<u8> {
+		let last = prefix.iter_mut().last().unwrap();
+		assert!(*last != 0, "mock function not implemented for this prefix");
+		*last -= 1;
+		prefix
+	}
+
+	fn key_after_prefix(mut prefix: Vec<u8>) -> Vec<u8> {
+		let last = prefix.iter_mut().last().unwrap();
+		assert!(*last != 255, "mock function not implemented for this prefix");
+		*last += 1;
+		prefix
+	}
+
+	#[test]
+	fn map_reversible_reversible_iteration() {
+		sp_io::TestExternalities::default().execute_with(|| {
+			// All map iterator
+			let prefix = Map::prefix_hash();
+
+			unhashed::put(&key_before_prefix(prefix.clone()), &1u64);
+			unhashed::put(&key_after_prefix(prefix.clone()), &1u64);
+
+			for i in 0..4 {
+				Map::insert(i as u16, i as u64);
+			}
+
+			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 3), (0, 0), (2, 2), (1, 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)]);
+
+			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![]);
+			assert_eq!(unhashed::get(&key_before_prefix(prefix.clone())), Some(1u64));
+			assert_eq!(unhashed::get(&key_after_prefix(prefix.clone())), Some(1u64));
+
+			// Translate
+			let prefix = Map::prefix_hash();
+
+			unhashed::put(&key_before_prefix(prefix.clone()), &1u64);
+			unhashed::put(&key_after_prefix(prefix.clone()), &1u64);
+			for i in 0..4 {
+				Map::insert(i as u16, i as u64);
+			}
+
+			// Wrong key
+			unhashed::put(&[prefix.clone(), vec![1, 2, 3]].concat(), &3u64.encode());
+
+			// Wrong value
+			unhashed::put(
+				&[prefix.clone(), crate::Blake2_128Concat::hash(&6u16.encode())].concat(),
+				&vec![1],
+			);
+
+			Map::translate(|_k1, v: u64| Some(v*2));
+			assert_eq!(Map::iter().collect::<Vec<_>>(), vec![(3, 6), (0, 0), (2, 4), (1, 2)]);
+		})
+	}
+}

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

@@ -17,7 +17,7 @@
 
 //! Stuff to do with the runtime's storage.
 
-use sp_std::{prelude::*, marker::PhantomData};
+use sp_std::prelude::*;
 use codec::{FullCodec, FullEncode, Encode, EncodeLike, Decode};
 use crate::hash::{Twox128, StorageHasher};
 use sp_runtime::generic::{Digest, DigestItem};
@@ -251,6 +251,8 @@ pub trait IterableStorageMap<K: FullEncode, V: FullCodec>: StorageMap<K, V> {
 
 	/// 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.
+	///
+	/// NOTE: If a value fail to decode because storage is corrupted then it is skipped.
 	fn translate<O: Decode, F: Fn(K, O) -> Option<V>>(f: F);
 }
 
@@ -286,7 +288,9 @@ pub trait IterableStorageDoubleMap<
 
 	/// 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);
+	///
+	/// NOTE: If a value fail to decode because storage is corrupted then it is skipped.
+	fn translate<O: Decode, F: Fn(K1, K2, O) -> Option<V>>(f: F);
 }
 
 /// An implementation of a map with a two keys.
@@ -433,35 +437,58 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 	>(key1: KeyArg1, key2: KeyArg2) -> Option<V>;
 }
 
-/// Iterator for prefixed map.
-pub struct PrefixIterator<Value> {
+/// Iterate over a prefix and decode raw_key and raw_value into `T`.
+///
+/// If any decoding fails it skips it and continues to the next key.
+pub struct PrefixIterator<T> {
 	prefix: Vec<u8>,
 	previous_key: Vec<u8>,
-	phantom_data: PhantomData<Value>,
+	/// If true then value are removed while iterating
+	drain: bool,
+	/// Function that take `(raw_key_without_prefix, raw_value)` and decode `T`.
+	/// `raw_key_without_prefix` is the raw storage key without the prefix iterated on.
+	closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
 }
 
-impl<Value: Decode> Iterator for PrefixIterator<Value> {
-	type Item = Value;
+impl<T> Iterator for PrefixIterator<T> {
+	type Item = T;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		match sp_io::storage::next_key(&self.previous_key)
-			.filter(|n| n.starts_with(&self.prefix[..]))
-		{
-			Some(next_key) => {
-				let value = unhashed::get(&next_key);
+		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;
+					let raw_value = match unhashed::get_raw(&self.previous_key) {
+						Some(raw_value) => raw_value,
+						None => {
+							crate::debug::error!(
+								"next_key returned a key with no value at {:?}",
+								self.previous_key
+							);
+							continue
+						}
+					};
+					if self.drain {
+						unhashed::kill(&self.previous_key)
+					}
+					let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
+					let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
+						Ok(item) => item,
+						Err(e) => {
+							crate::debug::error!(
+								"(key, value) failed to decode at {:?}: {:?}",
+								self.previous_key, e
+							);
+							continue
+						}
+					};
 
-				if value.is_none() {
-					runtime_print!(
-						"ERROR: returned next_key has no value:\nkey is {:?}\nnext_key is {:?}",
-						&self.previous_key, &next_key,
-					);
+					Some(item)
 				}
-
-				self.previous_key = next_key;
-
-				value
-			},
-			_ => None,
+				None => None,
+			}
 		}
 	}
 }
@@ -493,22 +520,22 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	}
 
 	/// Iter over all value of the storage.
+	///
+	/// NOTE: If a value failed to decode becaues storage is corrupted then it is skipped.
 	fn iter_values() -> PrefixIterator<Value> {
 		let prefix = Self::final_prefix();
 		PrefixIterator {
 			prefix: prefix.to_vec(),
 			previous_key: prefix.to_vec(),
-			phantom_data: Default::default(),
+			drain: false,
+			closure: |_raw_key, mut raw_value| Value::decode(&mut raw_value),
 		}
 	}
 
-	/// Translate the values from some previous `OldValue` to the current type.
+	/// 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.
 	///
-	/// `TV` translates values.
-	///
-	/// Returns `Err` if the map could not be interpreted as the old type, and Ok if it could.
-	/// The `Err` contains the number of value that couldn't be interpreted, those value are
-	/// removed from the map.
+	/// NOTE: If a value fail to decode because storage is corrupted then it is skipped.
 	///
 	/// # Warning
 	///
@@ -517,33 +544,28 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	///
 	/// # Usage
 	///
-	/// This would typically be called inside the module implementation of on_runtime_upgrade, while
-	/// ensuring **no usage of this storage are made before the call to `on_runtime_upgrade`**. (More
-	/// precisely prior initialized modules doesn't make use of this storage).
-	fn translate_values<OldValue, TV>(translate_val: TV) -> Result<(), u32>
-		where OldValue: Decode, TV: Fn(OldValue) -> Value
-	{
+	/// This would typically be called inside the module implementation of on_runtime_upgrade.
+	fn translate_values<OldValue: Decode, F: Fn(OldValue) -> Option<Value>>(f: F) {
 		let prefix = Self::final_prefix();
-		let mut previous_key = prefix.to_vec();
-		let mut errors = 0;
-		while let Some(next_key) = sp_io::storage::next_key(&previous_key)
-			.filter(|n| n.starts_with(&prefix[..]))
+		let mut previous_key = prefix.clone().to_vec();
+		while let Some(next) = sp_io::storage::next_key(&previous_key)
+			.filter(|n| n.starts_with(&prefix))
 		{
-			if let Some(value) = unhashed::get(&next_key) {
-				unhashed::put(&next_key[..], &translate_val(value));
-			} else {
-				// We failed to read the value. Remove the key and increment errors.
-				unhashed::kill(&next_key[..]);
-				errors += 1;
+			previous_key = next;
+			let maybe_value = unhashed::get::<OldValue>(&previous_key);
+			match maybe_value {
+				Some(value) => match f(value) {
+					Some(new) => unhashed::put::<Value>(&previous_key, &new),
+					None => unhashed::kill(&previous_key),
+				},
+				None => {
+					crate::debug::error!(
+						"old key failed to decode at {:?}",
+						previous_key
+					);
+					continue
+				},
 			}
-
-			previous_key = next_key;
-		}
-
-		if errors == 0 {
-			Ok(())
-		} else {
-			Err(errors)
 		}
 	}
 }
@@ -652,7 +674,7 @@ mod test {
 			unhashed::put(&[&k[..], &vec![8][..]].concat(), &2u32);
 
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![]);
-			MyStorage::translate_values(|v: u32| v as u64).unwrap();
+			MyStorage::translate_values(|v: u32| Some(v as u64));
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2]);
 			MyStorage::remove_all();
 
@@ -664,8 +686,8 @@ mod test {
 
 			// (contains some value that successfully decoded to u64)
 			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2, 3]);
-			assert_eq!(MyStorage::translate_values(|v: u128| v as u64), Err(2));
-			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 3]);
+			MyStorage::translate_values(|v: u128| Some(v as u64));
+			assert_eq!(MyStorage::iter_values().collect::<Vec<_>>(), vec![1, 2, 3]);
 			MyStorage::remove_all();
 
 			// test that other values are not modified.