Translation for linked map (#4052)

* make linked_map generic functions better

* implement translation of linked map

* proc-macro for linked map updated

* test linked map migration

* account for instances

* address grumbles

* cut map short if migration fails

substrate/srml/support/procedural/src/lib.rs

@@ -86,6 +86,10 @@ use proc_macro::TokenStream;
 ///   * `head_key`: `"head of " ++ $module_prefix ++ " " ++ $storage_name`
 ///   * `Hasher`: $hash
 ///
+///   All key formatting logic can be accessed in a type-agnostic format via the
+///   [`KeyFormat`](../srml_support/storage/generator/trait.KeyFormat.html) trait, which
+///   is implemented for the storage linked map type as well.
+///
 /// * Double map: `Foo: double_map hasher($hash1) u32, $hash2(u32) => u32`: Implements the
 ///   [`StorageDoubleMap`](../srml_support/storage/trait.StorageDoubleMap.html) trait using the
 ///   [`StorageDoubleMap generator`](../srml_support/storage/generator/trait.StorageDoubleMap.html).

substrate/srml/support/procedural/src/storage/storage_struct.rs

@@ -162,15 +162,12 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 					{
 						type Query = #query_type;
 						type Hasher = #scrate::#hasher;
+						type KeyFormat = Self;
 
 						fn prefix() -> &'static [u8] {
 							#final_prefix
 						}
 
-						fn head_key() -> &'static [u8] {
-							#head_key
-						}
-
 						fn from_optional_value_to_query(v: Option<#value_type>) -> Self::Query {
 							#from_optional_value_to_query
 						}
@@ -179,6 +176,14 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 							#from_query_to_optional_value
 						}
 					}
+
+					impl<#impl_trait> #scrate::storage::generator::LinkedMapKeyFormat for #storage_struct {
+						type Hasher = #scrate::#hasher;
+
+						fn head_key() -> &'static [u8] {
+							#head_key
+						}
+					}
 				)
 			},
 			StorageLineTypeDef::DoubleMap(map) => {

substrate/srml/support/src/storage/generator/linked_map.rs

@@ -47,16 +47,21 @@ pub trait StorageLinkedMap<K: FullCodec, V: FullCodec> {
 	/// Hasher used to insert into storage.
 	type Hasher: StorageHasher;
 
+	/// The family of key formats used for this map.
+	type KeyFormat: KeyFormat<Hasher=Self::Hasher>;
+
 	/// Prefix used to prepend each key.
 	fn prefix() -> &'static [u8];
 
-	/// Key used to store linked map head.
-	fn head_key() -> &'static [u8];
+	/// The head key of the linked-map.
+	fn head_key() -> &'static [u8] {
+		<Self::KeyFormat as KeyFormat>::head_key()
+	}
 
-	/// Convert an optionnal value retrieved from storage to the type queried.
+	/// Convert an optional value retrieved from storage to the type queried.
 	fn from_optional_value_to_query(v: Option<V>) -> Self::Query;
 
-	/// Convert a query to an optionnal value into storage.
+	/// Convert a query to an optional value into storage.
 	fn from_query_to_optional_value(v: Self::Query) -> Option<V>;
 
 	/// Generate the full key used in top storage.
@@ -64,14 +69,37 @@ pub trait StorageLinkedMap<K: FullCodec, V: FullCodec> {
 	where
 		KeyArg: EncodeLike<K>,
 	{
-		let mut final_key = Self::prefix().to_vec();
-		key.encode_to(&mut final_key);
-		Self::Hasher::hash(&final_key)
+		<Self::KeyFormat as KeyFormat>::storage_linked_map_final_key::<KeyArg>(Self::prefix(), &key)
 	}
 
 	/// Generate the hashed key for head
 	fn storage_linked_map_final_head_key() -> <Self::Hasher as StorageHasher>::Output {
-		Self::Hasher::hash(Self::head_key())
+		<Self::KeyFormat as KeyFormat>::storage_linked_map_final_head_key()
+	}
+}
+
+/// A type-abstracted key format used for a family of linked-map types.
+pub trait KeyFormat {
+	type Hasher: StorageHasher;
+
+	/// Key used to store linked map head.
+	fn head_key() -> &'static [u8];
+
+	/// Generate the full key used in top storage.
+	fn storage_linked_map_final_key<K>(prefix: &[u8], key: &K)
+		-> <Self::Hasher as StorageHasher>::Output
+	where
+		K: Encode,
+	{
+		let mut final_key = prefix.to_vec();
+		key.encode_to(&mut final_key);
+		<Self::Hasher as StorageHasher>::hash(&final_key)
+	}
+
+	fn storage_linked_map_final_head_key()
+		-> <Self::Hasher as StorageHasher>::Output
+	{
+		<Self::Hasher as StorageHasher>::hash(Self::head_key())
 	}
 }
 
@@ -105,25 +133,39 @@ struct EncodeLikeLinkage<PKey: EncodeLike<Key>, NKey: EncodeLike<Key>, Key: Enco
 }
 
 /// A key-value pair iterator for enumerable map.
-pub struct Enumerator<K: FullCodec, V: FullCodec, G: StorageLinkedMap<K, V>> {
+pub struct Enumerator<K, V, F> {
 	next: Option<K>,
-	_phantom: PhantomData<(G, V)>,
+	prefix: &'static [u8],
+	_phantom: PhantomData<(V, F)>,
 }
 
-impl<K, V, G> Iterator for Enumerator<K, V, G>
+impl<K, V, F> Enumerator<K, V, F> {
+	/// Create an explicit enumerator for testing.
+	#[cfg(test)]
+	pub fn from_head(head: K, prefix: &'static [u8]) -> Self {
+		Enumerator {
+			next: Some(head),
+			prefix,
+			_phantom: Default::default(),
+		}
+	}
+}
+
+impl<K, V, F> Iterator for Enumerator<K, V, F>
 where
 	K: FullCodec,
 	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	F: KeyFormat,
 {
 	type Item = (K, V);
 
 	fn next(&mut self) -> Option<Self::Item> {
 		let next = self.next.take()?;
+
 		let (val, linkage): (V, Linkage<K>) = {
-			let next_full_key = G::storage_linked_map_final_key(&next);
-			unhashed::get(next_full_key.as_ref())
-				.expect("previous/next only contain existing entires;
+			let next_full_key = F::storage_linked_map_final_key(self.prefix, &next);
+			read_with_linkage::<K, V>(next_full_key.as_ref())
+				.expect("previous/next only contains existing entries;
 						we enumerate using next; entry exists; qed")
 		};
 
@@ -136,33 +178,33 @@ where
 ///
 /// Takes care of updating previous and next elements points
 /// as well as updates head if the element is first or last.
-fn remove_linkage<K, V, G>(linkage: Linkage<K>)
+fn remove_linkage<K, V, F>(linkage: Linkage<K>, prefix: &[u8])
 where
 	K: FullCodec,
 	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	F: KeyFormat,
 {
 	let next_key = linkage.next.as_ref()
-		.map(G::storage_linked_map_final_key)
+		.map(|k| F::storage_linked_map_final_key(prefix, k))
 		.map(|x| x.as_ref().to_vec());
 	let prev_key = linkage.previous.as_ref()
-		.map(G::storage_linked_map_final_key)
+		.map(|k| F::storage_linked_map_final_key(prefix, k))
 		.map(|x| x.as_ref().to_vec());
 
 	if let Some(prev_key) = prev_key {
 		// Retrieve previous element and update `next`
-		let mut res = read_with_linkage::<_, _, G>(prev_key.as_ref())
+		let mut res = read_with_linkage::<K, V>(prev_key.as_ref())
 			.expect("Linkage is updated in case entry is removed;
 					it always points to existing keys; qed");
 		res.1.next = linkage.next;
 		unhashed::put(prev_key.as_ref(), &res);
 	} else {
 		// we were first so let's update the head
-		write_head::<_, _, _, G>(linkage.next.as_ref());
+		write_head::<&K, K, F>(linkage.next.as_ref());
 	}
 	if let Some(next_key) = next_key {
 		// Update previous of next element
-		let mut res = read_with_linkage::<_, _, G>(next_key.as_ref())
+		let mut res = read_with_linkage::<K, V>(next_key.as_ref())
 			.expect("Linkage is updated in case entry is removed;
 					it always points to existing keys; qed");
 		res.1.previous = linkage.previous;
@@ -170,12 +212,11 @@ where
 	}
 }
 
-/// Read the contained data and it's linkage.
-fn read_with_linkage<K, V, G>(key: &[u8]) -> Option<(V, Linkage<K>)>
+/// Read the contained data and its linkage.
+pub(super) fn read_with_linkage<K, V>(key: &[u8]) -> Option<(V, Linkage<K>)>
 where
-	K: FullCodec,
-	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	K: Decode,
+	V: Decode,
 {
 	unhashed::get(key)
 }
@@ -183,18 +224,18 @@ where
 /// Generate linkage for newly inserted element.
 ///
 /// Takes care of updating head and previous head's pointer.
-fn new_head_linkage<KeyArg, K, V, G>(key: KeyArg) -> Linkage<K>
+pub(super) fn new_head_linkage<KeyArg, K, V, F>(key: KeyArg, prefix: &[u8]) -> Linkage<K>
 where
 	KeyArg: EncodeLike<K>,
 	K: FullCodec,
 	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	F: KeyFormat,
 {
-	if let Some(head) = read_head::<_, _, G>() {
+	if let Some(head) = read_head::<K, F>() {
 		// update previous head predecessor
 		{
-			let head_key = G::storage_linked_map_final_key(&head);
-			let (data, linkage) = read_with_linkage::<_, _, G>(head_key.as_ref())
+			let head_key = F::storage_linked_map_final_key(prefix, &head);
+			let (data, linkage) = read_with_linkage::<K, V>(head_key.as_ref())
 				.expect("head is set when first element is inserted
 						and unset when last element is removed;
 						if head is Some then it points to existing key; qed");
@@ -206,41 +247,39 @@ where
 			unhashed::put(head_key.as_ref(), &(data, new_linkage));
 		}
 		// update to current head
-		write_head::<_, _, _, G>(Some(key));
+		write_head::<_, _, F>(Some(key));
 		// return linkage with pointer to previous head
 		let mut linkage = Linkage::default();
 		linkage.next = Some(head);
 		linkage
 	} else {
 		// we are first - update the head and produce empty linkage
-		write_head::<_, _, _, G>(Some(key));
+		write_head::<_, _, F>(Some(key));
 		Linkage::default()
 	}
 }
 
 /// Read current head pointer.
-fn read_head<K, V, G>() -> Option<K>
+pub(crate) fn read_head<K, F>() -> Option<K>
 where
-	K: FullCodec,
-	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	K: Decode,
+	F: KeyFormat,
 {
-	unhashed::get(G::storage_linked_map_final_head_key().as_ref())
+	unhashed::get(F::storage_linked_map_final_head_key().as_ref())
 }
 
 /// Overwrite current head pointer.
 ///
 /// If `None` is given head is removed from storage.
-fn write_head<KeyArg, K, V, G>(head: Option<KeyArg>)
+pub(super) fn write_head<KeyArg, K, F>(head: Option<KeyArg>)
 where
 	KeyArg: EncodeLike<K>,
 	K: FullCodec,
-	V: FullCodec,
-	G: StorageLinkedMap<K, V>,
+	F: KeyFormat,
 {
 	match head.as_ref() {
-		Some(head) => unhashed::put(G::storage_linked_map_final_head_key().as_ref(), head),
-		None => unhashed::kill(G::storage_linked_map_final_head_key().as_ref()),
+		Some(head) => unhashed::put(F::storage_linked_map_final_head_key().as_ref(), head),
+		None => unhashed::kill(F::storage_linked_map_final_head_key().as_ref()),
 	}
 }
 
@@ -252,7 +291,7 @@ where
 {
 	type Query = G::Query;
 
-	type Enumerator = Enumerator<K, V, Self>;
+	type Enumerator = Enumerator<K, V, G::KeyFormat>;
 
 	fn exists<KeyArg: EncodeLike<K>>(key: KeyArg) -> bool {
 		unhashed::exists(Self::storage_linked_map_final_key(key).as_ref())
@@ -264,10 +303,11 @@ where
 	}
 
 	fn swap<KeyArg1: EncodeLike<K>, KeyArg2: EncodeLike<K>>(key1: KeyArg1, key2: KeyArg2) {
+		let prefix = Self::prefix();
 		let final_key1 = Self::storage_linked_map_final_key(Ref::from(&key1));
 		let final_key2 = Self::storage_linked_map_final_key(Ref::from(&key2));
-		let full_value_1 = read_with_linkage::<_, _, G>(final_key1.as_ref());
-		let full_value_2 = read_with_linkage::<_, _, G>(final_key2.as_ref());
+		let full_value_1 = read_with_linkage::<K, V>(final_key1.as_ref());
+		let full_value_2 = read_with_linkage::<K, V>(final_key2.as_ref());
 
 		match (full_value_1, full_value_2) {
 			// Just keep linkage in order and only swap values.
@@ -278,13 +318,13 @@ where
 			// Remove key and insert the new one.
 			(Some((value, _linkage)), None) => {
 				Self::remove(key1);
-				let linkage = new_head_linkage::<_, _, _, G>(key2);
+				let linkage = new_head_linkage::<_, _, V, G::KeyFormat>(key2, prefix);
 				unhashed::put(final_key2.as_ref(), &(value, linkage));
 			}
 			// Remove key and insert the new one.
 			(None, Some((value, _linkage))) => {
 				Self::remove(key2);
-				let linkage = new_head_linkage::<_, _, _, G>(key1);
+				let linkage = new_head_linkage::<_, _, V, G::KeyFormat>(key1, prefix);
 				unhashed::put(final_key1.as_ref(), &(value, linkage));
 			}
 			// No-op.
@@ -294,11 +334,11 @@ where
 
 	fn insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(key: KeyArg, val: ValArg) {
 		let final_key = Self::storage_linked_map_final_key(Ref::from(&key));
-		let linkage = match read_with_linkage::<_, _, G>(final_key.as_ref()) {
+		let linkage = match read_with_linkage::<_, V>(final_key.as_ref()) {
 			// overwrite but reuse existing linkage
 			Some((_data, linkage)) => linkage,
 			// create new linkage
-			None => new_head_linkage::<_, _, _, G>(key),
+			None => new_head_linkage::<_, _, V, G::KeyFormat>(key, Self::prefix()),
 		};
 		unhashed::put(final_key.as_ref(), &(val, linkage))
 	}
@@ -310,7 +350,7 @@ where
 	fn mutate<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Self::Query) -> R>(key: KeyArg, f: F) -> R {
 		let final_key = Self::storage_linked_map_final_key(Ref::from(&key));
 
-		let (mut val, _linkage) = read_with_linkage::<_, _, G>(final_key.as_ref())
+		let (mut val, _linkage) = read_with_linkage::<K, V>(final_key.as_ref())
 			.map(|(data, linkage)| (G::from_optional_value_to_query(Some(data)), Some(linkage)))
 			.unwrap_or_else(|| (G::from_optional_value_to_query(None), None));
 
@@ -328,7 +368,7 @@ where
 		let full_value: Option<(V, Linkage<K>)> = unhashed::take(final_key.as_ref());
 
 		let value = full_value.map(|(data, linkage)| {
-			remove_linkage::<_, _, G>(linkage);
+			remove_linkage::<K, V, G::KeyFormat>(linkage, Self::prefix());
 			data
 		});
 
@@ -336,14 +376,15 @@ where
 	}
 
 	fn enumerate() -> Self::Enumerator {
-		Enumerator::<_, _, G> {
-			next: read_head::<_, _, G>(),
+		Enumerator::<_, _, G::KeyFormat> {
+			next: read_head::<_, G::KeyFormat>(),
+			prefix: Self::prefix(),
 			_phantom: Default::default(),
 		}
 	}
 
 	fn head() -> Option<K> {
-		read_head::<_, _, G>()
+		read_head::<_, G::KeyFormat>()
 	}
 
 	fn decode_len<KeyArg: EncodeLike<K>>(key: KeyArg) -> Result<usize, &'static str>
@@ -360,4 +401,67 @@ where
 			Ok(len)
 		}
 	}
+
+	fn translate<K2, V2, TK, TV>(translate_key: TK, translate_val: TV) -> Result<(), Option<K2>>
+		where K2: FullCodec + Clone, V2: Decode, TK: Fn(K2) -> K, TV: Fn(V2) -> V
+	{
+		let head_key = read_head::<K2, G::KeyFormat>().ok_or(None)?;
+		let prefix = G::prefix();
+
+		let mut last_key = None;
+		let mut current_key = head_key.clone();
+
+		write_head::<&K, K, G::KeyFormat>(Some(&translate_key(head_key)));
+
+		let translate_linkage = |old: Linkage<K2>| -> Linkage<K> {
+			Linkage {
+				previous: old.previous.map(&translate_key),
+				next: old.next.map(&translate_key),
+			}
+		};
+
+		loop {
+			let old_raw_key = G::KeyFormat::storage_linked_map_final_key(prefix, &current_key);
+			let x = unhashed::take(old_raw_key.as_ref());
+			let (val, linkage): (V2, Linkage<K2>) = match x {
+				Some(v) => v,
+				None => {
+					// we failed to read value and linkage. Update the last key's linkage
+					// to end the map early, since it's impossible to iterate further.
+					if let Some(last_key) = last_key {
+						let last_raw_key = G::storage_linked_map_final_key(&last_key);
+						if let Some((val, mut linkage))
+							= read_with_linkage::<K, V>(last_raw_key.as_ref())
+						{
+							// defensive: should always happen, since it was just written
+							// in the last iteration of the loop.
+							linkage.next = None;
+							unhashed::put(last_raw_key.as_ref(), &(&val, &linkage));
+						}
+					}
+
+					return Err(Some(current_key));
+				}
+			};
+			let next = linkage.next.clone();
+
+			let val = translate_val(val);
+			let linkage = translate_linkage(linkage);
+
+			// and write in the value and linkage under the new key.
+			let new_key = translate_key(current_key.clone());
+			let new_raw_key = G::storage_linked_map_final_key(&new_key);
+			unhashed::put(new_raw_key.as_ref(), &(&val, &linkage));
+
+			match next {
+				None => break,
+				Some(next) => {
+					last_key = Some(new_key);
+					current_key = next
+				},
+			}
+		}
+
+		Ok(())
+	}
 }

substrate/srml/support/src/storage/generator/mod.rs

@@ -26,7 +26,7 @@ mod map;
 mod double_map;
 mod value;
 
-pub use linked_map::{StorageLinkedMap, Enumerator, Linkage};
+pub use linked_map::{StorageLinkedMap, Enumerator, Linkage, KeyFormat as LinkedMapKeyFormat};
 pub use map::StorageMap;
 pub use double_map::StorageDoubleMap;
 pub use value::StorageValue;
@@ -36,8 +36,8 @@ pub use value::StorageValue;
 #[allow(dead_code)]
 mod tests {
 	use runtime_io::TestExternalities;
-	use codec::Encode;
-	use crate::storage::{unhashed, generator::StorageValue};
+	use codec::{Encode, Decode};
+	use crate::storage::{unhashed, generator::{StorageValue, StorageLinkedMap}};
 
 	struct Runtime {}
 	pub trait Trait {
@@ -54,9 +54,16 @@ mod tests {
 		pub struct Module<T: Trait> for enum Call where origin: T::Origin {}
 	}
 
+	#[derive(Encode, Decode, Clone, Debug, Eq, PartialEq)]
+	struct NumberNumber {
+		a: u32,
+		b: u32,
+	}
+
 	crate::decl_storage! {
 		trait Store for Module<T: Trait> as Runtime {
 			Value get(fn value) config(): (u64, u64);
+			NumberMap: linked_map NumberNumber => u64;
 		}
 	}
 
@@ -78,4 +85,48 @@ mod tests {
 			assert_eq!(Value::get(), (1111, 2222));
 		})
 	}
+
+	#[test]
+	fn linked_map_translate_works() {
+		use super::linked_map::{self, Enumerator, KeyFormat};
+
+		type Format = <NumberMap as StorageLinkedMap<NumberNumber, u64>>::KeyFormat;
+
+		let t = GenesisConfig::default().build_storage().unwrap();
+		TestExternalities::new(t).execute_with(|| {
+			let prefix = NumberMap::prefix();
+
+			// start with a map of u32 -> u32.
+			for i in 0u32..100u32 {
+				let final_key = <Format as KeyFormat>::storage_linked_map_final_key(
+					prefix, &i,
+				);
+
+				let linkage = linked_map::new_head_linkage::<_, u32, u32, Format>(&i, prefix);
+				unhashed::put(final_key.as_ref(), &(&i, linkage));
+			}
+
+			let head = linked_map::read_head::<u32, Format>().unwrap();
+
+			assert_eq!(
+				Enumerator::<u32, u32, Format>::from_head(head, prefix).collect::<Vec<_>>(),
+				(0..100).rev().map(|x| (x, x)).collect::<Vec<_>>(),
+			);
+
+			// do translation.
+			NumberMap::translate(
+				|k: u32| NumberNumber { a: k, b: k },
+				|v: u32| (v as u64) << 32 | v as u64,
+			).unwrap();
+
+			assert!(linked_map::read_head::<NumberNumber, Format>().is_some());
+			assert_eq!(
+				NumberMap::enumerate().collect::<Vec<_>>(),
+				(0..100u32).rev().map(|x| (
+					NumberNumber { a: x, b: x },
+					(x as u64) << 32 | x as u64,
+				)).collect::<Vec<_>>(),
+			);
+		})
+	}
 }

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

@@ -230,6 +230,27 @@ pub trait StorageLinkedMap<K: FullCodec, V: FullCodec> {
 	/// function for this purpose.
 	fn decode_len<KeyArg: EncodeLike<K>>(key: KeyArg) -> Result<usize, &'static str>
 		where V: codec::DecodeLength + Len;
+
+	/// Translate the keys and values from some previous `(K2, V2)` to the current type.
+	///
+	/// `TK` translates keys from the old type, and `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 key which could not be migrated, or `None` if the
+	/// head of the list could not be read.
+	///
+	/// # 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<K2, V2, TK, TV>(translate_key: TK, translate_val: TV) -> Result<(), Option<K2>>
+		where K2: FullCodec + Clone, V2: Decode, TK: Fn(K2) -> K, TV: Fn(V2) -> V;
 }
 
 /// An implementation of a map with a two keys.

substrate/srml/support/test/tests/final_keys.rs

@@ -156,7 +156,7 @@ fn final_keys_default_instance() {
 		assert_eq!(unhashed::get::<u32>(&hashing::blake2_256(&k)), Some(2u32));
 		assert_eq!(unhashed::get::<u32>(&hashing::blake2_256(&head)), Some(1u32));
 
-<		instance::LinkedMap2<instance::DefaultInstance>>::insert(1, 2);
+		<instance::LinkedMap2<instance::DefaultInstance>>::insert(1, 2);
 		let mut k = b"FinalKeysSome LinkedMap2".to_vec();
 		k.extend(1u32.encode());
 		assert_eq!(unhashed::get::<u32>(&hashing::twox_128(&k)), Some(2u32));