mutate_exists for StorageValue with ValueQuery

muraca · muraca · commit c19838924ed9 · 2023-01-26T17:33:51.000-03:00
Signed-off-by: muraca &lt;mmuraca247@gmail.com&gt;
diff --git 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,55 @@ pub mod tests {
 		});
 	}
 
+	#[test]
+	fn storage_value_mutate_exists_should_work() {
+		new_test_ext().execute_with(|| {
+			assert!(!Value::exists());
+
+			Value::mutate_exists(|v| *v = Some(1));
+			assert!(Value::exists());
+			assert_eq!(Value::get(), 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(|| {
+			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(), 1);
+
+			// no-op if `Err`
+			assert_noop!(
+				Value::try_mutate_exists(|v| -> TestResult {
+					*v = Some(2);
+					Err("nah")
+				}),
+				"nah"
+			);
+			assert_eq!(Value::get(), 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(|| {
diff --git 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
@@ -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
@@ -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()
