From 4833b0b969a03669bf082156396dc3f01e1efe41 Mon Sep 17 00:00:00 2001
From: Matteo Muraca <56828990+muraca@users.noreply.github.com>
Date: Sat, 28 Jan 2023 00:47:17 +0100
Subject: [PATCH] mutate_exists for StorageValue with ValueQuery (#13245)

* mutate_exists for StorageValue with ValueQuery

Signed-off-by: muraca <mmuraca247@gmail.com>

* added `#[crate::storage_alias]` to tests

Signed-off-by: muraca <mmuraca247@gmail.com>

* added StorageEntryMetadata

Signed-off-by: muraca <mmuraca247@gmail.com>

* Update frame/support/src/lib.rs

---------

Signed-off-by: muraca <mmuraca247@gmail.com>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
---
 frame/support/src/lib.rs                     | 63 ++++++++++++++++++++
 frame/support/src/storage/generator/value.rs | 24 ++++++++
 frame/support/src/storage/mod.rs             |  6 ++
 frame/support/src/storage/types/value.rs     | 12 ++++
 4 files changed, 105 insertions(+)

diff --git a/frame/support/src/lib.rs b/frame/support/src/lib.rs
index 902893972f0b1..77ca89dc6f800 100644
--- a/frame/support/src/lib.rs
+++ b/frame/support/src/lib.rs
@@ -874,6 +874,7 @@ pub mod tests {
 
 	decl_storage! {
 		trait Store for Module<T: Config> as Test {
+			pub Value get(fn value): u64;
 			pub Data get(fn data) build(|_| vec![(15u32, 42u64)]):
 				map hasher(twox_64_concat) u32 => u64;
 			pub OptionLinkedMap: map hasher(blake2_128_concat) u32 => Option<u32>;
@@ -946,6 +947,61 @@ pub mod tests {
 		});
 	}
 
+	#[test]
+	fn storage_value_mutate_exists_should_work() {
+		new_test_ext().execute_with(|| {
+			#[crate::storage_alias]
+			pub type Value = StorageValue<Test, u32>;
+
+			assert!(!Value::exists());
+
+			Value::mutate_exists(|v| *v = Some(1));
+			assert!(Value::exists());
+			assert_eq!(Value::get(), Some(1));
+
+			// removed if mutated to `None`
+			Value::mutate_exists(|v| *v = None);
+			assert!(!Value::exists());
+		});
+	}
+
+	#[test]
+	fn storage_value_try_mutate_exists_should_work() {
+		new_test_ext().execute_with(|| {
+			#[crate::storage_alias]
+			pub type Value = StorageValue<Test, u32>;
+
+			type TestResult = result::Result<(), &'static str>;
+
+			assert!(!Value::exists());
+
+			// mutated if `Ok`
+			assert_ok!(Value::try_mutate_exists(|v| -> TestResult {
+				*v = Some(1);
+				Ok(())
+			}));
+			assert!(Value::exists());
+			assert_eq!(Value::get(), Some(1));
+
+			// no-op if `Err`
+			assert_noop!(
+				Value::try_mutate_exists(|v| -> TestResult {
+					*v = Some(2);
+					Err("nah")
+				}),
+				"nah"
+			);
+			assert_eq!(Value::get(), Some(1));
+
+			// removed if mutated to`None`
+			assert_ok!(Value::try_mutate_exists(|v| -> TestResult {
+				*v = None;
+				Ok(())
+			}));
+			assert!(!Value::exists());
+		});
+	}
+
 	#[test]
 	fn map_issue_3318() {
 		new_test_ext().execute_with(|| {
@@ -1257,6 +1313,13 @@ pub mod tests {
 		PalletStorageMetadata {
 			prefix: "Test",
 			entries: vec![
+				StorageEntryMetadata {
+					name: "Value",
+					modifier: StorageEntryModifier::Default,
+					ty: StorageEntryType::Plain(scale_info::meta_type::<u64>()),
+					default: vec![0, 0, 0, 0, 0, 0, 0, 0],
+					docs: vec![],
+				},
 				StorageEntryMetadata {
 					name: "Data",
 					modifier: StorageEntryModifier::Default,
diff --git a/frame/support/src/storage/generator/value.rs b/frame/support/src/storage/generator/value.rs
index 55b3487b1324c..4a1fd5c551d3f 100644
--- a/frame/support/src/storage/generator/value.rs
+++ b/frame/support/src/storage/generator/value.rs
@@ -118,6 +118,30 @@ impl<T: FullCodec, G: StorageValue<T>> storage::StorageValue<T> for G {
 		ret
 	}
 
+	fn mutate_exists<R, F>(f: F) -> R
+	where
+		F: FnOnce(&mut Option<T>) -> R,
+	{
+		Self::try_mutate_exists(|v| Ok::<R, Never>(f(v)))
+			.expect("`Never` can not be constructed; qed")
+	}
+
+	fn try_mutate_exists<R, E, F>(f: F) -> Result<R, E>
+	where
+		F: FnOnce(&mut Option<T>) -> Result<R, E>,
+	{
+		let mut val = G::from_query_to_optional_value(Self::get());
+
+		let ret = f(&mut val);
+		if ret.is_ok() {
+			match val {
+				Some(ref val) => Self::put(val),
+				None => Self::kill(),
+			}
+		}
+		ret
+	}
+
 	fn take() -> G::Query {
 		let key = Self::storage_value_final_key();
 		let value = unhashed::get(&key);
diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs
index 8c0d6207c3f4d..28f2dee992281 100644
--- a/frame/support/src/storage/mod.rs
+++ b/frame/support/src/storage/mod.rs
@@ -114,6 +114,12 @@ pub trait StorageValue<T: FullCodec> {
 	/// 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>;
 
+	/// Mutate the value. Deletes the item if mutated to a `None`.
+	fn mutate_exists<R, F: FnOnce(&mut Option<T>) -> R>(f: F) -> R;
+
+	/// Mutate the value if closure returns `Ok`. Deletes the item if mutated to a `None`.
+	fn try_mutate_exists<R, E, F: FnOnce(&mut Option<T>) -> Result<R, E>>(f: F) -> Result<R, E>;
+
 	/// Clear the storage value.
 	fn kill();
 
diff --git a/frame/support/src/storage/types/value.rs b/frame/support/src/storage/types/value.rs
index f145e9fb30414..15290f1b1e085 100644
--- a/frame/support/src/storage/types/value.rs
+++ b/frame/support/src/storage/types/value.rs
@@ -142,6 +142,18 @@ where
 		<Self as crate::storage::StorageValue<Value>>::try_mutate(f)
 	}
 
+	/// Mutate the value. Deletes the item if mutated to a `None`.
+	pub fn mutate_exists<R, F: FnOnce(&mut Option<Value>) -> R>(f: F) -> R {
+		<Self as crate::storage::StorageValue<Value>>::mutate_exists(f)
+	}
+
+	/// Mutate the value if closure returns `Ok`. Deletes the item if mutated to a `None`.
+	pub fn try_mutate_exists<R, E, F: FnOnce(&mut Option<Value>) -> Result<R, E>>(
+		f: F,
+	) -> Result<R, E> {
+		<Self as crate::storage::StorageValue<Value>>::try_mutate_exists(f)
+	}
+
 	/// Clear the storage value.
 	pub fn kill() {
 		<Self as crate::storage::StorageValue<Value>>::kill()