make iterations consistent

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

@@ -212,10 +212,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),
 		}
 	}
 
@@ -408,27 +409,6 @@ impl<
 			_phantom: Default::default(),
 		}
 	}
-
-	fn translate<O: Decode, F: Fn(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,
-			}
-		}
-	}
 }
 
 #[cfg(test)]

frame/support/src/storage/mod.rs

@@ -16,7 +16,7 @@
 
 //! Stuff to do with the runtime's storage.
 
-use sp_std::{prelude::*, marker::PhantomData};
+use sp_std::prelude::*;
 use codec::{FullCodec, FullEncode, Encode, EncodeAppend, EncodeLike, Decode};
 use crate::{traits::Len, hash::{Twox128, StorageHasher}};
 
@@ -231,6 +231,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);
 }
 
@@ -252,10 +254,6 @@ pub trait IterableStorageDoubleMap<
 	/// particular order. If you add elements with first key `k1` to the map while doing this,
 	/// you'll get undefined results.
 	fn drain(k1: impl EncodeLike<K1>) -> Self::Iterator;
-
-	/// 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);
 }
 
 /// An implementation of a map with a two keys.
@@ -372,35 +370,56 @@ 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`.
+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);
-
-				if value.is_none() {
-					runtime_print!(
-						"ERROR: returned next_key has no value:\nkey is {:?}\nnext_key is {:?}",
-						&self.previous_key, &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
+						}
+					};
+
+					Some(item)
 				}
-
-				self.previous_key = next_key;
-
-				value
-			},
-			_ => None,
+				None => None,
+			}
 		}
 	}
 }
@@ -433,22 +452,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.
-	///
-	/// `TV` translates values.
+	/// 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.
 	///
-	/// 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
 	///
@@ -457,33 +476,31 @@ 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[..]))
-		{
-			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;
+		let mut previous_key = prefix.clone().to_vec();
+		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::<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
+						},
+					}
+				}
+				None => return,
 			}
-
-			previous_key = next_key;
-		}
-
-		if errors == 0 {
-			Ok(())
-		} else {
-			Err(errors)
 		}
 	}
 }
@@ -546,7 +563,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();
 
@@ -558,8 +575,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.