From c0179cb6e78ba9089428b53a1f35cf62d3519045 Mon Sep 17 00:00:00 2001
From: Shaopeng Wang <spxwang@gmail.com>
Date: Tue, 21 Jul 2020 20:50:37 +1200
Subject: [PATCH] Add try_mutate_exists to StorageDoubleMap. (#6694)
---
substrate/frame/support/src/lib.rs | 32 ++++++++++++++++++-
.../src/storage/generator/double_map.rs | 19 +++++++++++
substrate/frame/support/src/storage/mod.rs | 7 ++++
3 files changed, 57 insertions(+), 1 deletion(-)
diff --git a/substrate/frame/support/src/lib.rs b/substrate/frame/support/src/lib.rs
index 06a8ce856dd..bb8aacd1a48 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 ff83aaf8ec8..8fbef16204f 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 b8b08c5dc02..4623f81859b 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`].
--
GitLab