diff --git a/substrate/frame/support/src/lib.rs b/substrate/frame/support/src/lib.rs index 06a8ce856dd6431a70b447d67a056fa818a26cd0..bb8aacd1a4886fb1c92c6066b55474e7553103cb 100644 --- a/substrate/frame/support/src/lib.rs +++ b/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( diff --git a/substrate/frame/support/src/storage/generator/double_map.rs b/substrate/frame/support/src/storage/generator/double_map.rs index ff83aaf8ec8297f6dc77af8bfb1b4cc2870996c4..8fbef16204f4e7710b16e888bb3638dcc05565fc 100644 --- a/substrate/frame/support/src/storage/generator/double_map.rs +++ b/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, diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs index b8b08c5dc028c8436468f1e883441c26ece9dc70..4623f81859b35dfd00e7a55e4232dc192ea4fbd9 100644 --- a/substrate/frame/support/src/storage/mod.rs +++ b/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`].