Add translate function for storage prefixed map.

frame/support/src/storage/mod.rs

@@ -403,17 +403,20 @@ 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()));
 		final_key[16..32].copy_from_slice(&Twox128::hash(Self::storage_prefix()));
 		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,43 @@ 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 first hashed key which could not be migrated.
+	///
+	/// # 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<(), Vec<u8>>
+		where OldValue: Decode, TV: Fn(OldValue) -> Value
+	{
+		let prefix = Self::final_prefix();
+		let mut previous_key = prefix.to_vec();
+		while let Some(next_key) = sp_io::storage::next_key(&previous_key)
+			.filter(|n| n.starts_with(&prefix[..]))
+		{
+			let value: OldValue = unhashed::get(&next_key)
+				// We failed to read the value. Stop the translation and return an error.
+				.ok_or_else(|| next_key.clone())?;
+
+			unhashed::put(&next_key[..], &translate_val(value));
+
+			previous_key = next_key;
+		}
+
+		Ok(())
+	}
 }
 
 #[cfg(test)]
@@ -463,6 +503,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 +513,18 @@ 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![2][..]].concat(), &2u32);
+
+			MyStorage::translate_values(|v: u32| v as u64).unwrap();
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 2]);
+
+			// test that other values are not modified.
 			assert_eq!(unhashed::get(&key_before[..]), Some(32u64));
 			assert_eq!(unhashed::get(&key_after[..]), Some(33u64));
 		});