Fix linked_map in decl_storage! for option values (#3323)

substrate/Cargo.lock

@@ -4416,8 +4416,8 @@ dependencies = [
 name = "substrate-consensus-babe-primitives"
 version = "2.0.0"
 dependencies = [
- "schnorrkel 0.8.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "parity-scale-codec 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)",
+ "schnorrkel 0.8.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "sr-primitives 2.0.0",
  "sr-std 2.0.0",
  "substrate-client 2.0.0",

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

@@ -200,10 +200,13 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 		} = instance_opts;
 
 		let final_prefix = if let Some(instance) = instance {
-			let const_name = syn::Ident::new(&format!("{}{}", PREFIX_FOR, name.to_string()), proc_macro2::Span::call_site());
-			quote!{ #instance::#const_name.as_bytes() }
+			let const_name = syn::Ident::new(
+				&format!("{}{}", PREFIX_FOR, name.to_string()),
+				proc_macro2::Span::call_site(),
+			);
+			quote! { #instance::#const_name.as_bytes() }
 		} else {
-			quote!{ #prefix.as_bytes() }
+			quote! { #prefix.as_bytes() }
 		};
 
 		let trait_required = ext::type_contains_ident(value_type, traitinstance)
@@ -352,6 +355,12 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 			}
 		};
 
+		let mutate_map = if type_infos.is_option {
+			quote! { .map(|(data, linkage)| (Some(data), Some(linkage))) }
+		} else {
+			quote! { .map(|(data, linkage)| (data, Some(linkage))) }
+		};
+
 		let trait_required = ext::type_contains_ident(value_type, traitinstance)
 			|| ext::type_contains_ident(kty, traitinstance);
 
@@ -442,7 +451,7 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 					fn remove_linkage<S: #scrate::HashedStorage<#scrate::#hasher>>(linkage: Linkage<#kty>, storage: &mut S);
 
 					/// Read the contained data and it's linkage.
-					fn read_with_linkage<S>(storage: &S, key: &[u8]) -> Option<(#value_type, Linkage<#kty>)>
+					fn read_with_linkage<S>(storage: &S, key: &[u8]) -> Option<(#typ, Linkage<#kty>)>
 					where
 						S: #scrate::HashedStorage<#scrate::#hasher>;
 
@@ -504,7 +513,7 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 				fn read_with_linkage<S: #scrate::HashedStorage<#scrate::#hasher>>(
 					storage: &S,
 					key: &[u8],
-				) -> Option<(#value_type, self::#inner_module::Linkage<#kty>)> {
+				) -> Option<(#typ, self::#inner_module::Linkage<#kty>)> {
 					storage.get(key)
 				}
 
@@ -586,14 +595,12 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 				fn take<S: #scrate::HashedStorage<#scrate::#hasher>>(key: &#kty, storage: &mut S) -> Self::Query {
 					use self::#inner_module::Utils;
 
-					let res: Option<(#value_type, self::#inner_module::Linkage<#kty>)> = storage.take(&*#as_map::key_for(key));
-					match res {
-						Some((data, linkage)) => {
-							Self::remove_linkage(linkage, storage);
-							data
-						},
-						None => #fielddefault,
-					}
+					let res: Option<(#typ, self::#inner_module::Linkage<#kty>)> = storage.take(&*#as_map::key_for(key));
+
+					res.map(|(d, l)| {
+						Self::remove_linkage(l, storage);
+						d
+					}).#option_simple_1(|| #fielddefault)
 				}
 
 				/// Remove the value under a key.
@@ -602,7 +609,11 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 				}
 
 				/// Store a value to be associated with the given key from the map.
-				fn insert<S: #scrate::HashedStorage<#scrate::#hasher>>(key: &#kty, val: &#typ, storage: &mut S) {
+				fn insert<S: #scrate::HashedStorage<#scrate::#hasher>>(
+					key: &#kty,
+					val: &#typ,
+					storage: &mut S,
+				) {
 					use self::#inner_module::Utils;
 
 					let key_for = &*#as_map::key_for(key);
@@ -612,6 +623,7 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 						// create new linkage
 						None => Self::new_head_linkage(storage, key),
 					};
+
 					storage.put(key_for, &(val, linkage))
 				}
 
@@ -633,6 +645,7 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 						// create new linkage
 						None => Self::new_head_linkage(storage, key),
 					};
+
 					storage.put(key_for, &(val, linkage))
 				}
 
@@ -646,11 +659,11 @@ impl<'a, I: Iterator<Item=syn::Meta>> Impls<'a, I> {
 
 					let key_for = &*#as_map::key_for(key);
 					let (mut val, linkage) = Self::read_with_linkage(storage, key_for)
-						.map(|(data, linkage)| (data, Some(linkage)))
+						#mutate_map
 						.unwrap_or_else(|| (#fielddefault, None));
 
 					let ret = f(&mut val);
-					#mutate_impl ;
+					#mutate_impl;
 					ret
 				}
 			}

substrate/srml/support/src/lib.rs

@@ -308,6 +308,7 @@ mod tests {
 	decl_storage! {
 		trait Store for Module<T: Trait> as Example {
 			pub Data get(data) build(|_| vec![(15u32, 42u64)]): linked_map hasher(twox_64_concat) u32 => u64;
+			pub OptionLinkedMap: linked_map u32 => Option<u32>;
 			pub GenericData get(generic_data): linked_map hasher(twox_128) T::BlockNumber => T::BlockNumber;
 			pub GenericData2 get(generic_data2): linked_map T::BlockNumber => Option<T::BlockNumber>;
 
@@ -331,6 +332,16 @@ mod tests {
 
 	type Map = Data;
 
+	#[test]
+	fn linked_map_issue_3318() {
+		with_externalities(&mut new_test_ext(), || {
+			OptionLinkedMap::insert(1, 1);
+			assert_eq!(OptionLinkedMap::get(1), Some(1));
+			OptionLinkedMap::insert(1, 2);
+			assert_eq!(OptionLinkedMap::get(1), Some(2));
+		});
+	}
+
 	#[test]
 	fn linked_map_basic_insert_remove_should_work() {
 		with_externalities(&mut new_test_ext(), || {
@@ -472,13 +483,29 @@ mod tests {
 					modifier: StorageEntryModifier::Default,
 					ty: StorageEntryType::Map{
 						hasher: StorageHasher::Twox64Concat,
-						key: DecodeDifferent::Encode("u32"), value: DecodeDifferent::Encode("u64"), is_linked: true
+						key: DecodeDifferent::Encode("u32"),
+						value: DecodeDifferent::Encode("u64"),
+						is_linked: true,
 					},
 					default: DecodeDifferent::Encode(
 						DefaultByteGetter(&__GetByteStructData(PhantomData::<Test>))
 					),
 					documentation: DecodeDifferent::Encode(&[]),
 				},
+				StorageEntryMetadata {
+					name: DecodeDifferent::Encode("OptionLinkedMap"),
+					modifier: StorageEntryModifier::Optional,
+					ty: StorageEntryType::Map {
+						hasher: StorageHasher::Blake2_256,
+						key: DecodeDifferent::Encode("u32"),
+						value: DecodeDifferent::Encode("u32"),
+						is_linked: true,
+					},
+					default: DecodeDifferent::Encode(
+						DefaultByteGetter(&__GetByteStructOptionLinkedMap(PhantomData::<Test>))
+					),
+					documentation: DecodeDifferent::Encode(&[]),
+				},
 				StorageEntryMetadata {
 					name: DecodeDifferent::Encode("GenericData"),
 					modifier: StorageEntryModifier::Default,
@@ -574,6 +601,6 @@ mod tests {
 	#[test]
 	fn store_metadata() {
 		let metadata = Module::<Test>::storage_metadata();
-		assert_eq!(EXPECTED_METADATA, metadata);
+		pretty_assertions::assert_eq!(EXPECTED_METADATA, metadata);
 	}
 }