Implement try_mutate for storage value and storage double map (#5699)

* impl try_mutate for storage value and storage double map

* Docs + Reuse `try_mutate` in `mutate`

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

@@ -17,7 +17,7 @@
 use sp_std::prelude::*;
 use sp_std::borrow::Borrow;
 use codec::{Ref, FullCodec, FullEncode, Decode, Encode, EncodeLike, EncodeAppend};
-use crate::{storage::{self, unhashed}, traits::Len};
+use crate::{storage::{self, unhashed}, traits::Len, Never};
 use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
 /// Generator for `StorageDoubleMap` used by `decl_storage`.
@@ -223,14 +223,24 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>,
 		F: FnOnce(&mut Self::Query) -> R,
+	{
+		Self::try_mutate(k1, k2, |v| Ok::<R, Never>(f(v))).expect("`Never` can not be constructed; qed")
+	}
+
+	fn try_mutate<KArg1, KArg2, R, E, F>(k1: KArg1, k2: KArg2, f: F) -> Result<R, E> where
+		KArg1: EncodeLike<K1>,
+		KArg2: EncodeLike<K2>,
+		F: FnOnce(&mut Self::Query) -> Result<R, E>,
 	{
 		let final_key = Self::storage_double_map_final_key(k1, k2);
 		let mut val = G::from_optional_value_to_query(unhashed::get(final_key.as_ref()));
 
 		let ret = f(&mut val);
-		match G::from_query_to_optional_value(val) {
-			Some(ref val) => unhashed::put(final_key.as_ref(), val),
-			None => unhashed::kill(final_key.as_ref()),
+		if ret.is_ok() {
+			match G::from_query_to_optional_value(val) {
+				Some(ref val) => unhashed::put(final_key.as_ref(), val),
+				None => unhashed::kill(final_key.as_ref()),
+			}
 		}
 		ret
 	}

substrate/frame/support/src/storage/generator/map.rs

@@ -18,7 +18,7 @@
 use sp_std::prelude::*;
 use sp_std::borrow::Borrow;
 use codec::{FullCodec, FullEncode, Decode, Encode, EncodeLike, Ref, EncodeAppend};
-use crate::{storage::{self, unhashed}, traits::Len};
+use crate::{storage::{self, unhashed}, traits::Len, Never};
 use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
 /// Generator for `StorageMap` used by `decl_storage`.
@@ -229,27 +229,11 @@ impl<K: FullEncode, V: FullCodec, G: StorageMap<K, V>> storage::StorageMap<K, V>
 	}
 
 	fn mutate<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Self::Query) -> R>(key: KeyArg, f: F) -> R {
-		let final_key = Self::storage_map_final_key(key);
-		let mut val = G::from_optional_value_to_query(unhashed::get(final_key.as_ref()));
-
-		let ret = f(&mut val);
-		match G::from_query_to_optional_value(val) {
-			Some(ref val) => unhashed::put(final_key.as_ref(), &val),
-			None => unhashed::kill(final_key.as_ref()),
-		}
-		ret
+		Self::try_mutate(key, |v| Ok::<R, Never>(f(v))).expect("`Never` can not be constructed; qed")
 	}
 
 	fn mutate_exists<KeyArg: EncodeLike<K>, R, F: FnOnce(&mut Option<V>) -> R>(key: KeyArg, f: F) -> R {
-		let final_key = Self::storage_map_final_key(key);
-		let mut val = unhashed::get(final_key.as_ref());
-
-		let ret = f(&mut val);
-		match val {
-			Some(ref val) => unhashed::put(final_key.as_ref(), &val),
-			None => unhashed::kill(final_key.as_ref()),
-		}
-		ret
+		Self::try_mutate_exists(key, |v| Ok::<R, Never>(f(v))).expect("`Never` can not be constructed; qed")
 	}
 
 	fn try_mutate<KeyArg: EncodeLike<K>, R, E, F: FnOnce(&mut Self::Query) -> Result<R, E>>(

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

@@ -37,6 +37,7 @@ mod tests {
 	use sp_io::TestExternalities;
 	use codec::Encode;
 	use crate::storage::{unhashed, generator::StorageValue, IterableStorageMap};
+	use crate::{assert_noop, assert_ok};
 
 	struct Runtime {}
 	pub trait Trait {
@@ -57,6 +58,7 @@ mod tests {
 		trait Store for Module<T: Trait> as Runtime {
 			Value get(fn value) config(): (u64, u64);
 			NumberMap: map hasher(identity) u32 => u64;
+			DoubleMap: double_map hasher(identity) u32, hasher(identity) u32 => u64;
 		}
 	}
 
@@ -102,4 +104,54 @@ mod tests {
 			);
 		})
 	}
+
+	#[test]
+	fn try_mutate_works() {
+		let t = GenesisConfig::default().build_storage().unwrap();
+		TestExternalities::new(t).execute_with(|| {
+			assert_eq!(Value::get(), (0, 0));
+			assert_eq!(NumberMap::get(0), 0);
+			assert_eq!(DoubleMap::get(0, 0), 0);
+
+			// `assert_noop` ensures that the state does not change
+			assert_noop!(Value::try_mutate(|value| -> Result<(), &'static str> {
+				*value = (2, 2);
+				Err("don't change value")
+			}), "don't change value");
+
+			assert_noop!(NumberMap::try_mutate(0, |value| -> Result<(), &'static str> {
+				*value = 4;
+				Err("don't change value")
+			}), "don't change value");
+
+			assert_noop!(DoubleMap::try_mutate(0, 0, |value| -> Result<(), &'static str> {
+				*value = 6;
+				Err("don't change value")
+			}), "don't change value");
+
+			// Showing this explicitly for clarity
+			assert_eq!(Value::get(), (0, 0));
+			assert_eq!(NumberMap::get(0), 0);
+			assert_eq!(DoubleMap::get(0, 0), 0);
+
+			assert_ok!(Value::try_mutate(|value| -> Result<(), &'static str> {
+				*value = (2, 2);
+				Ok(())
+			}));
+
+			assert_ok!(NumberMap::try_mutate(0, |value| -> Result<(), &'static str> {
+				*value = 4;
+				Ok(())
+			}));
+
+			assert_ok!(DoubleMap::try_mutate(0, 0, |value| -> Result<(), &'static str> {
+				*value = 6;
+				Ok(())
+			}));
+
+			assert_eq!(Value::get(), (2, 2));
+			assert_eq!(NumberMap::get(0), 4);
+			assert_eq!(DoubleMap::get(0, 0), 6);
+		});
+	}
 }

substrate/frame/support/src/storage/generator/value.rs

@@ -17,7 +17,12 @@
 #[cfg(not(feature = "std"))]
 use sp_std::prelude::*;
 use codec::{FullCodec, Encode, EncodeAppend, EncodeLike, Decode};
-use crate::{storage::{self, unhashed}, hash::{Twox128, StorageHasher}, traits::Len};
+use crate::{
+	Never,
+	storage::{self, unhashed},
+	hash::{Twox128, StorageHasher},
+	traits::Len
+};
 
 /// Generator for `StorageValue` used by `decl_storage`.
 ///
@@ -104,12 +109,18 @@ impl<T: FullCodec, G: StorageValue<T>> storage::StorageValue<T> for G {
 	}
 
 	fn mutate<R, F: FnOnce(&mut G::Query) -> R>(f: F) -> R {
+		Self::try_mutate(|v| Ok::<R, Never>(f(v))).expect("`Never` can not be constructed; qed")
+	}
+
+	fn try_mutate<R, E, F: FnOnce(&mut G::Query) -> Result<R, E>>(f: F) -> Result<R, E> {
 		let mut val = G::get();
 
 		let ret = f(&mut val);
-		match G::from_query_to_optional_value(val) {
-			Some(ref val) => G::put(val),
-			None => G::kill(),
+		if ret.is_ok() {
+			match G::from_query_to_optional_value(val) {
+				Some(ref val) => G::put(val),
+				None => G::kill(),
+			}
 		}
 		ret
 	}

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

@@ -80,6 +80,9 @@ pub trait StorageValue<T: FullCodec> {
 	/// Mutate the value
 	fn mutate<R, F: FnOnce(&mut Self::Query) -> R>(f: F) -> R;
 
+	/// Mutate the value if closure returns `Ok`
+	fn try_mutate<R, E, F: FnOnce(&mut Self::Query) -> Result<R, E>>(f: F) -> Result<R, E>;
+
 	/// Clear the storage value.
 	fn kill();
 
@@ -280,21 +283,25 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 	/// The type that get/take returns.
 	type Query;
 
+	/// Get the storage key used to fetch a value corresponding to a specific key.
 	fn hashed_key_for<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Vec<u8>
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>;
 
+	/// Does the value (explicitly) exist in storage?
 	fn contains_key<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> bool
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>;
 
+	/// Load the value associated with the given key from the double map.
 	fn get<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Self::Query
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>;
 
+	/// Take a value from storage, removing it afterwards.
 	fn take<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Self::Query
 	where
 		KArg1: EncodeLike<K1>,
@@ -308,28 +315,43 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		YKArg1: EncodeLike<K1>,
 		YKArg2: EncodeLike<K2>;
 
+	/// Store a value to be associated with the given keys from the double map.
 	fn insert<KArg1, KArg2, VArg>(k1: KArg1, k2: KArg2, val: VArg)
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>,
 		VArg: EncodeLike<V>;
 
+	/// Remove the value under the given keys.
 	fn remove<KArg1, KArg2>(k1: KArg1, k2: KArg2)
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>;
 
+	/// Remove all values under the first key.
 	fn remove_prefix<KArg1>(k1: KArg1) where KArg1: ?Sized + EncodeLike<K1>;
 
+	/// Iterate over values that share the first key.
 	fn iter_prefix_values<KArg1>(k1: KArg1) -> PrefixIterator<V>
 		where KArg1: ?Sized + EncodeLike<K1>;
 
+	/// Mutate the value under the given keys.
 	fn mutate<KArg1, KArg2, R, F>(k1: KArg1, k2: KArg2, f: F) -> R
 	where
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>,
 		F: FnOnce(&mut Self::Query) -> R;
 
+	/// Mutate the value under the given keys when the closure returns `Ok`.
+	fn try_mutate<KArg1, KArg2, R, E, F>(k1: KArg1, k2: KArg2, f: F) -> Result<R, E>
+	where
+		KArg1: EncodeLike<K1>,
+		KArg2: EncodeLike<K2>,
+		F: FnOnce(&mut Self::Query) -> Result<R, E>;
+
+	/// Append the given item to the value in the storage.
+	///
+	/// `V` is required to implement `codec::EncodeAppend`.
 	fn append<Items, Item, EncodeLikeItem, KArg1, KArg2>(
 		k1: KArg1,
 		k2: KArg2,
@@ -344,6 +366,10 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		Items: IntoIterator<Item=EncodeLikeItem>,
 		Items::IntoIter: ExactSizeIterator;
 
+	/// Safely append the given items to the value in the storage. If a codec error occurs, then the
+	/// old (presumably corrupt) value is replaced with the given `items`.
+	///
+	/// `V` is required to implement `codec::EncodeAppend`.
 	fn append_or_insert<Items, Item, EncodeLikeItem, KArg1, KArg2>(
 		k1: KArg1,
 		k2: KArg2,