Add try_mutate_exists to StorageDoubleMap. (#6694)

substrate/frame/support/src/lib.rs

@@ -376,7 +376,7 @@ mod tests {
 		DecodeDifferent, StorageEntryMetadata, StorageMetadata, StorageEntryType,
 		StorageEntryModifier, DefaultByteGetter, StorageHasher,
 	};
-	use sp_std::marker::PhantomData;
+	use sp_std::{marker::PhantomData, result};
 	use sp_io::TestExternalities;
 
 	pub trait Trait {
@@ -629,6 +629,36 @@ mod tests {
 		});
 	}
 
+	#[test]
+	fn double_map_try_mutate_exists_should_work() {
+		new_test_ext().execute_with(|| {
+			type DoubleMap = DataDM;
+			type TestResult = result::Result<(), &'static str>;
+
+			let (key1, key2) = (11, 13);
+
+			// mutated if `Ok`
+			assert_ok!(DoubleMap::try_mutate_exists(key1, key2, |v| -> TestResult {
+				*v = Some(1);
+				Ok(())
+			}));
+			assert_eq!(DoubleMap::get(&key1, key2), 1);
+
+			// no-op if `Err`
+			assert_noop!(DoubleMap::try_mutate_exists(key1, key2, |v| -> TestResult {
+				*v = Some(2);
+				Err("nah")
+			}), "nah");
+
+			// removed if mutated to`None`
+			assert_ok!(DoubleMap::try_mutate_exists(key1, key2, |v| -> TestResult {
+				*v = None;
+				Ok(())
+			}));
+			assert!(!DoubleMap::contains_key(&key1, key2));
+		});
+	}
+
 	const EXPECTED_METADATA: StorageMetadata = StorageMetadata {
 		prefix: DecodeDifferent::Encode("Test"),
 		entries: DecodeDifferent::Encode(

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

@@ -246,6 +246,25 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 		ret
 	}
 
+	fn try_mutate_exists<KArg1, KArg2, R, E, F>(k1: KArg1, k2: KArg2, f: F) -> Result<R, E>
+	where
+		KArg1: EncodeLike<K1>,
+		KArg2: EncodeLike<K2>,
+		F: FnOnce(&mut Option<V>) -> Result<R, E>,
+	{
+		let final_key = Self::storage_double_map_final_key(k1, k2);
+		let mut val = unhashed::get(final_key.as_ref());
+
+		let ret = f(&mut val);
+		if ret.is_ok() {
+			match val {
+				Some(ref val) => unhashed::put(final_key.as_ref(), val),
+				None => unhashed::kill(final_key.as_ref()),
+			}
+		}
+		ret
+	}
+
 	fn append<Item, EncodeLikeItem, KArg1, KArg2>(
 		k1: KArg1,
 		k2: KArg2,

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

@@ -366,6 +366,13 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		KArg2: EncodeLike<K2>,
 		F: FnOnce(&mut Self::Query) -> Result<R, E>;
 
+	/// Mutate the item, only if an `Ok` value is returned. Deletes the item if mutated to a `None`.
+	fn try_mutate_exists<KArg1, KArg2, R, E, F>(k1: KArg1, k2: KArg2, f: F) -> Result<R, E>
+	where
+		KArg1: EncodeLike<K1>,
+		KArg2: EncodeLike<K2>,
+		F: FnOnce(&mut Option<V>) -> Result<R, E>;
+
 	/// Append the given item to the value in the storage.
 	///
 	/// `V` is required to implement [`StorageAppend`].