CountedNMap implementation

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

@@ -133,8 +133,12 @@ where
 	}
 
 	/// 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)
+	pub fn set<KeyArg: EncodeLike<Key>>(key: KeyArg, query: QueryKind::Query) {
+		let option = QueryKind::from_query_to_optional_value(query);
+		if option.is_none() {
+			CounterFor::<Prefix>::mutate(|value| value.saturating_dec());
+		}
+		<Self as MapWrapper>::Map::set(key, QueryKind::from_optional_value_to_query(option))
 	}
 
 	/// Swap the values of two keys.
@@ -143,10 +147,7 @@ where
 	}
 
 	/// Store a value to be associated with the given key from the map.
-	pub fn insert<KeyArg: EncodeLike<Key> + Clone, ValArg: EncodeLike<Value>>(
-		key: KeyArg,
-		val: ValArg,
-	) {
+	pub fn insert<KeyArg: EncodeLike<Key>, ValArg: EncodeLike<Value>>(key: KeyArg, val: ValArg) {
 		if !<Self as MapWrapper>::Map::contains_key(Ref::from(&key)) {
 			CounterFor::<Prefix>::mutate(|value| value.saturating_inc());
 		}
@@ -602,8 +603,9 @@ mod test {
 			assert_eq!(A::count(), 2);
 
 			// Insert an existing key, shouldn't increment counted values.
-			A::insert(3, 11);
+			A::insert(3, 12);
 
+			assert_eq!(A::try_get(3), Ok(12));
 			assert_eq!(A::count(), 2);
 
 			// Remove non-existing.
@@ -696,17 +698,17 @@ mod test {
 			// Try succeed mutate existing to existing.
 			A::try_mutate_exists(1, |query| {
 				assert_eq!(*query, Some(43));
-				*query = Some(43);
+				*query = Some(45);
 				Result::<(), ()>::Ok(())
 			})
 			.unwrap();
 
-			assert_eq!(A::try_get(1), Ok(43));
+			assert_eq!(A::try_get(1), Ok(45));
 			assert_eq!(A::count(), 4);
 
 			// Try succeed mutate existing to non-existing.
 			A::try_mutate_exists(1, |query| {
-				assert_eq!(*query, Some(43));
+				assert_eq!(*query, Some(45));
 				*query = None;
 				Result::<(), ()>::Ok(())
 			})