From e0766bb97dd33e2b877625737f78b0f41dce222a Mon Sep 17 00:00:00 2001
From: CJ13th <48095175+CJ13th@users.noreply.github.com>
Date: Mon, 23 Sep 2024 22:42:17 +0200
Subject: [PATCH] Implement try_append for StorageNMap (#5745)

# Description

Closes #5722

Added an implementation of the `try_append` functionality which is
present on the other storage map types but currently missing from
StorageNMap.

---------

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>
---
 prdoc/pr_5745.prdoc                           | 14 ++++
 substrate/frame/support/src/storage/mod.rs    | 76 +++++++++++++++++++
 .../frame/support/src/storage/types/nmap.rs   | 14 ++++
 3 files changed, 104 insertions(+)
 create mode 100644 prdoc/pr_5745.prdoc

diff --git a/prdoc/pr_5745.prdoc b/prdoc/pr_5745.prdoc
new file mode 100644
index 000000000000..7463589378a0
--- /dev/null
+++ b/prdoc/pr_5745.prdoc
@@ -0,0 +1,14 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Implement `try_append` for `StorageNMap`
+
+doc:
+  - audience: Runtime Dev
+    description: |
+      This PR introduces the `try_append` api which is available on other storage map types,
+      but missing on `StorageNMap`.
+
+crates:
+  - name: frame-support
+    bump: minor
diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs
index 7fb991d37792..619392563035 100644
--- a/substrate/frame/support/src/storage/mod.rs
+++ b/substrate/frame/support/src/storage/mod.rs
@@ -1693,6 +1693,46 @@ where
 	}
 }
 
+/// Storage N map that is capable of [`StorageTryAppend`].
+pub trait TryAppendNMap<K: KeyGenerator, T: StorageTryAppend<I>, I: Encode> {
+	/// Try and append the `item` into the storage N map at the given `key`.
+	///
+	/// This might fail if bounds are not respected.
+	fn try_append<
+		LikeK: EncodeLikeTuple<K::KArg> + TupleToEncodedIter + Clone,
+		LikeI: EncodeLike<I>,
+	>(
+		key: LikeK,
+		item: LikeI,
+	) -> Result<(), ()>;
+}
+
+impl<K, T, I, StorageNMapT> TryAppendNMap<K, T, I> for StorageNMapT
+where
+	K: KeyGenerator,
+	T: FullCodec + StorageTryAppend<I>,
+	I: Encode,
+	StorageNMapT: generator::StorageNMap<K, T>,
+{
+	fn try_append<
+		LikeK: EncodeLikeTuple<K::KArg> + TupleToEncodedIter + Clone,
+		LikeI: EncodeLike<I>,
+	>(
+		key: LikeK,
+		item: LikeI,
+	) -> Result<(), ()> {
+		let bound = T::bound();
+		let current = Self::decode_len(key.clone()).unwrap_or_default();
+		if current < bound {
+			let key = Self::storage_n_map_final_key::<K, _>(key);
+			sp_io::storage::append(&key, item.encode());
+			Ok(())
+		} else {
+			Err(())
+		}
+	}
+}
+
 /// Returns the storage prefix for a specific pallet name and storage name.
 ///
 /// The storage prefix is `concat(twox_128(pallet_name), twox_128(storage_name))`.
@@ -2019,6 +2059,17 @@ mod test {
 		(NMapKey<Twox128, u32>, NMapKey<Twox128, u32>, NMapKey<Twox128, u32>),
 		u64,
 	>;
+	#[crate::storage_alias]
+	type FooQuadMap = StorageNMap<
+		Prefix,
+		(
+			NMapKey<Twox128, u32>,
+			NMapKey<Twox128, u32>,
+			NMapKey<Twox128, u32>,
+			NMapKey<Twox128, u32>,
+		),
+		BoundedVec<u32, ConstU32<7>>,
+	>;
 
 	#[test]
 	fn contains_prefix_works() {
@@ -2109,6 +2160,31 @@ mod test {
 				BoundedVec::<u32, ConstU32<7>>::try_from(vec![4, 5]).unwrap(),
 			);
 		});
+
+		TestExternalities::default().execute_with(|| {
+			let bounded: BoundedVec<u32, ConstU32<7>> = vec![1, 2, 3].try_into().unwrap();
+			FooQuadMap::insert((1, 1, 1, 1), bounded);
+
+			assert_ok!(FooQuadMap::try_append((1, 1, 1, 1), 4));
+			assert_ok!(FooQuadMap::try_append((1, 1, 1, 1), 5));
+			assert_ok!(FooQuadMap::try_append((1, 1, 1, 1), 6));
+			assert_ok!(FooQuadMap::try_append((1, 1, 1, 1), 7));
+			assert_eq!(FooQuadMap::decode_len((1, 1, 1, 1)).unwrap(), 7);
+			assert!(FooQuadMap::try_append((1, 1, 1, 1), 8).is_err());
+
+			// append to a non-existing
+			assert!(FooQuadMap::get((2, 1, 1, 1)).is_none());
+			assert_ok!(FooQuadMap::try_append((2, 1, 1, 1), 4));
+			assert_eq!(
+				FooQuadMap::get((2, 1, 1, 1)).unwrap(),
+				BoundedVec::<u32, ConstU32<7>>::try_from(vec![4]).unwrap(),
+			);
+			assert_ok!(FooQuadMap::try_append((2, 1, 1, 1), 5));
+			assert_eq!(
+				FooQuadMap::get((2, 1, 1, 1)).unwrap(),
+				BoundedVec::<u32, ConstU32<7>>::try_from(vec![4, 5]).unwrap(),
+			);
+		});
 	}
 
 	#[crate::storage_alias]
diff --git a/substrate/frame/support/src/storage/types/nmap.rs b/substrate/frame/support/src/storage/types/nmap.rs
index 9ee012f86286..0fc22b35352d 100755
--- a/substrate/frame/support/src/storage/types/nmap.rs
+++ b/substrate/frame/support/src/storage/types/nmap.rs
@@ -25,6 +25,7 @@ use crate::{
 			StorageEntryMetadataBuilder, TupleToEncodedIter,
 		},
 		KeyGenerator, PrefixIterator, StorageAppend, StorageDecodeLength, StoragePrefixedMap,
+		StorageTryAppend,
 	},
 	traits::{Get, GetDefault, StorageInfo, StorageInstance},
 };
@@ -338,6 +339,19 @@ where
 		<Self as crate::storage::StorageNMap<Key, Value>>::append(key, item)
 	}
 
+	/// Try and append the given item to the value in the storage.
+	///
+	/// Is only available if `Value` of the storage implements [`StorageTryAppend`].
+	pub fn try_append<KArg, Item, EncodeLikeItem>(key: KArg, item: EncodeLikeItem) -> Result<(), ()>
+	where
+		KArg: EncodeLikeTuple<Key::KArg> + TupleToEncodedIter + Clone,
+		Item: Encode,
+		EncodeLikeItem: EncodeLike<Item>,
+		Value: StorageTryAppend<Item>,
+	{
+		<Self as crate::storage::TryAppendNMap<Key, Value, Item>>::try_append(key, item)
+	}
+
 	/// Read the length of the storage value without decoding the entire value under the
 	/// given `key1` and `key2`.
 	///