Introduce set function into storage maps (#11564)

substrate/frame/support/src/storage/generator/double_map.rs

@@ -151,6 +151,13 @@ where
 		unhashed::get(&Self::storage_double_map_final_key(k1, k2)).ok_or(())
 	}
 
+	fn set<KArg1: EncodeLike<K1>, KArg2: EncodeLike<K2>>(k1: KArg1, k2: KArg2, q: Self::Query) {
+		match G::from_query_to_optional_value(q) {
+			Some(v) => Self::insert(k1, k2, v),
+			None => Self::remove(k1, k2),
+		}
+	}
+
 	fn take<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Self::Query
 	where
 		KArg1: EncodeLike<K1>,

substrate/frame/support/src/storage/generator/map.rs

@@ -245,6 +245,13 @@ impl<K: FullEncode, V: FullCodec, G: StorageMap<K, V>> storage::StorageMap<K, V>
 		unhashed::get(Self::storage_map_final_key(key).as_ref()).ok_or(())
 	}
 
+	fn set<KeyArg: EncodeLike<K>>(key: KeyArg, q: Self::Query) {
+		match G::from_query_to_optional_value(q) {
+			Some(v) => Self::insert(key, v),
+			None => Self::remove(key),
+		}
+	}
+
 	fn insert<KeyArg: EncodeLike<K>, ValArg: EncodeLike<V>>(key: KeyArg, val: ValArg) {
 		unhashed::put(Self::storage_map_final_key(key).as_ref(), &val)
 	}

substrate/frame/support/src/storage/generator/nmap.rs

@@ -138,6 +138,13 @@ where
 		unhashed::get(&Self::storage_n_map_final_key::<K, _>(key)).ok_or(())
 	}
 
+	fn set<KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter>(key: KArg, q: Self::Query) {
+		match G::from_query_to_optional_value(q) {
+			Some(v) => Self::insert(key, v),
+			None => Self::remove(key),
+		}
+	}
+
 	fn take<KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter>(key: KArg) -> Self::Query {
 		let final_key = Self::storage_n_map_final_key::<K, _>(key);
 

substrate/frame/support/src/storage/mod.rs

@@ -167,6 +167,9 @@ pub trait StorageMap<K: FullEncode, V: FullCodec> {
 	/// Load the value associated with the given key from the map.
 	fn get<KeyArg: EncodeLike<K>>(key: KeyArg) -> Self::Query;
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	fn set<KeyArg: EncodeLike<K>>(key: KeyArg, query: Self::Query);
+
 	/// Try to get the value for the given key from the map.
 	///
 	/// Returns `Ok` if it exists, `Err` if not.
@@ -481,6 +484,9 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 		KArg1: EncodeLike<K1>,
 		KArg2: EncodeLike<K2>;
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	fn set<KArg1: EncodeLike<K1>, KArg2: EncodeLike<K2>>(k1: KArg1, k2: KArg2, query: Self::Query);
+
 	/// Take a value from storage, removing it afterwards.
 	fn take<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Self::Query
 	where
@@ -657,6 +663,9 @@ pub trait StorageNMap<K: KeyGenerator, V: FullCodec> {
 	/// Returns `Ok` if it exists, `Err` if not.
 	fn try_get<KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter>(key: KArg) -> Result<V, ()>;
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	fn set<KArg: EncodeLikeTuple<K::KArg> + TupleToEncodedIter>(key: KArg, query: Self::Query);
+
 	/// Swap the values of two keys.
 	fn swap<KOther, KArg1, KArg2>(key1: KArg1, key2: KArg2)
 	where

substrate/frame/support/src/storage/types/counted_map.rs

@@ -132,6 +132,11 @@ where
 		<Self as MapWrapper>::Map::try_get(key)
 	}
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	pub fn set<KeyArg: EncodeLike<Key>>(key: KeyArg, q: QueryKind::Query) {
+		<Self as MapWrapper>::Map::set(key, q)
+	}
+
 	/// Swap the values of two keys.
 	pub fn swap<KeyArg1: EncodeLike<Key>, KeyArg2: EncodeLike<Key>>(key1: KeyArg1, key2: KeyArg2) {
 		<Self as MapWrapper>::Map::swap(key1, key2)

substrate/frame/support/src/storage/types/double_map.rs

@@ -203,6 +203,15 @@ where
 		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::try_get(k1, k2)
 	}
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	pub fn set<KArg1: EncodeLike<Key1>, KArg2: EncodeLike<Key2>>(
+		k1: KArg1,
+		k2: KArg2,
+		q: QueryKind::Query,
+	) {
+		<Self as crate::storage::StorageDoubleMap<Key1, Key2, Value>>::set(k1, k2, q)
+	}
+
 	/// Take a value from storage, removing it afterwards.
 	pub fn take<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> QueryKind::Query
 	where

substrate/frame/support/src/storage/types/map.rs

@@ -152,6 +152,11 @@ where
 		<Self as crate::storage::StorageMap<Key, Value>>::swap(key1, key2)
 	}
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	pub fn set<KeyArg: EncodeLike<Key>>(key: KeyArg, q: QueryKind::Query) {
+		<Self as crate::storage::StorageMap<Key, Value>>::set(key, q)
+	}
+
 	/// Store a value to be associated with the given key from the map.
 	pub fn insert<KeyArg: EncodeLike<Key>, ValArg: EncodeLike<Value>>(key: KeyArg, val: ValArg) {
 		<Self as crate::storage::StorageMap<Key, Value>>::insert(key, val)

substrate/frame/support/src/storage/types/nmap.rs

@@ -142,6 +142,14 @@ where
 		<Self as crate::storage::StorageNMap<Key, Value>>::try_get(key)
 	}
 
+	/// Store or remove the value to be associated with `key` so that `get` returns the `query`.
+	pub fn set<KArg: EncodeLikeTuple<Key::KArg> + TupleToEncodedIter>(
+		key: KArg,
+		query: QueryKind::Query,
+	) {
+		<Self as crate::storage::StorageNMap<Key, Value>>::set(key, query)
+	}
+
 	/// Take a value from storage, removing it afterwards.
 	pub fn take<KArg: EncodeLikeTuple<Key::KArg> + TupleToEncodedIter>(
 		key: KArg,