Add translate function for storage prefixed map. (#4555)

* translace values for prefixed storages

* improve doc

* new impl

* update test

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

@@ -403,6 +403,7 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	/// Storage prefix. Used for generating final key.
 	fn storage_prefix() -> &'static [u8];
 
+	/// Final full prefix that prefixes all keys.
 	fn final_prefix() -> [u8; 32] {
 		let mut final_key = [0u8; 32];
 		final_key[0..16].copy_from_slice(&Twox128::hash(Self::module_prefix()));
@@ -410,10 +411,12 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 		final_key
 	}
 
+	/// Remove all value of the storage.
 	fn remove_all() {
 		sp_io::storage::clear_prefix(&Self::final_prefix())
 	}
 
+	/// Iter over all value of the storage.
 	fn iter() -> PrefixIterator<Value> {
 		let prefix = Self::final_prefix();
 		PrefixIterator {
@@ -422,6 +425,51 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 			phantom_data: Default::default(),
 		}
 	}
+
+	/// Translate the values from some previous `OldValue` to the current type.
+	///
+	/// `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.
+	///
+	/// # Warning
+	///
+	/// This function must be used with care, before being updated the storage still contains the
+	/// old type, thus other calls (such as `get`) will fail at decoding it.
+	///
+	/// # Usage
+	///
+	/// This would typically be called inside the module implementation of on_initialize, while
+	/// ensuring **no usage of this storage are made before the call to `on_initialize`**. (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
+	{
+		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;
+			}
+
+			previous_key = next_key;
+		}
+
+		if errors == 0 {
+			Ok(())
+		} else {
+			Err(errors)
+		}
+	}
 }
 
 #[cfg(test)]
@@ -463,6 +511,7 @@ mod test {
 			let k = [twox_128(b"MyModule"), twox_128(b"MyStorage")].concat();
 			assert_eq!(MyStorage::final_prefix().to_vec(), k);
 
+			// test iteration
 			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![]);
 
 			unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u64);
@@ -472,9 +521,32 @@ mod test {
 
 			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 2, 3, 4]);
 
+			// test removal
 			MyStorage::remove_all();
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![]);
+
+			// test migration
+			unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u32);
+			unhashed::put(&[&k[..], &vec![8][..]].concat(), &2u32);
 
 			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![]);
+			MyStorage::translate_values(|v: u32| v as u64).unwrap();
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 2]);
+			MyStorage::remove_all();
+
+			// test migration 2
+			unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u128);
+			unhashed::put(&[&k[..], &vec![1, 1][..]].concat(), &2u64);
+			unhashed::put(&[&k[..], &vec![8][..]].concat(), &3u128);
+			unhashed::put(&[&k[..], &vec![10][..]].concat(), &4u32);
+
+			// (contains some value that successfully decoded to u64)
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 2, 3]);
+			assert_eq!(MyStorage::translate_values(|v: u128| v as u64), Err(2));
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 3]);
+			MyStorage::remove_all();
+
+			// test that other values are not modified.
 			assert_eq!(unhashed::get(&key_before[..]), Some(32u64));
 			assert_eq!(unhashed::get(&key_after[..]), Some(33u64));
 		});