diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs index bb8aacd1a4886..f0ffdc90a74e0 100644 --- a/frame/support/src/lib.rs +++ b/frame/support/src/lib.rs @@ -629,6 +629,23 @@ mod tests { }); } + #[test] + fn double_map_mutate_exists_should_work() { + new_test_ext().execute_with(|| { + type DoubleMap = DataDM; + + let (key1, key2) = (11, 13); + + // mutated + DoubleMap::mutate_exists(key1, key2, |v| *v = Some(1)); + assert_eq!(DoubleMap::get(&key1, key2), 1); + + // removed if mutated to `None` + DoubleMap::mutate_exists(key1, key2, |v| *v = None); + assert!(!DoubleMap::contains_key(&key1, key2)); + }); + } + #[test] fn double_map_try_mutate_exists_should_work() { new_test_ext().execute_with(|| { diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index 8fbef16204f4e..3c82f4156a271 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -228,6 +228,15 @@ impl storage::StorageDoubleMap for G where Self::try_mutate(k1, k2, |v| Ok::(f(v))).expect("`Never` can not be constructed; qed") } + fn mutate_exists(k1: KArg1, k2: KArg2, f: F) -> R + where + KArg1: EncodeLike, + KArg2: EncodeLike, + F: FnOnce(&mut Option) -> R, + { + Self::try_mutate_exists(k1, k2, |v| Ok::(f(v))).expect("`Never` can not be constructed; qed") + } + fn try_mutate(k1: KArg1, k2: KArg2, f: F) -> Result where KArg1: EncodeLike, KArg2: EncodeLike, diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 4623f81859b35..347fd814136d7 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -366,6 +366,13 @@ pub trait StorageDoubleMap { KArg2: EncodeLike, F: FnOnce(&mut Self::Query) -> Result; + /// Mutate the value under the given keys. Deletes the item if mutated to a `None`. + fn mutate_exists(k1: KArg1, k2: KArg2, f: F) -> R + where + KArg1: EncodeLike, + KArg2: EncodeLike, + F: FnOnce(&mut Option) -> R; + /// Mutate the item, only if an `Ok` value is returned. Deletes the item if mutated to a `None`. fn try_mutate_exists(k1: KArg1, k2: KArg2, f: F) -> Result where