implemented contains_prefix for StorageDoubleMap and StorageNMap (p…

…aritytech#13232)

* implemented `contains_prefix` for StorageDoubleMap and StorageNMap

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

* match prefix to next_key

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

* warning unexpected behaviour with empty keys

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

* clarifications for unhashed::contains_prefixed_key

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

* added tests for StorageNMap

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

---------

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

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

@@ -233,6 +233,13 @@ where
 		.into()
 	}
 
+	fn contains_prefix<KArg1>(k1: KArg1) -> bool
+	where
+		KArg1: EncodeLike<K1>,
+	{
+		unhashed::contains_prefixed_key(Self::storage_double_map_final_key1(k1).as_ref())
+	}
+
 	fn iter_prefix_values<KArg1>(k1: KArg1) -> storage::PrefixIterator<V>
 	where
 		KArg1: ?Sized + EncodeLike<K1>,

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

@@ -208,6 +208,13 @@ where
 		)
 	}
 
+	fn contains_prefix<KP>(partial_key: KP) -> bool
+	where
+		K: HasKeyPrefix<KP>,
+	{
+		unhashed::contains_prefixed_key(&Self::storage_n_map_partial_key(partial_key))
+	}
+
 	fn iter_prefix_values<KP>(partial_key: KP) -> PrefixIterator<V>
 	where
 		K: HasKeyPrefix<KP>,

frame/support/src/storage/mod.rs

@@ -563,6 +563,12 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 	where
 		KArg1: ?Sized + EncodeLike<K1>;
 
+	/// Does any value under the first key `k1` (explicitly) exist in storage?
+	/// Might have unexpected behaviour with empty keys, e.g. `[]`.
+	fn contains_prefix<KArg1>(k1: KArg1) -> bool
+	where
+		KArg1: EncodeLike<K1>;
+
 	/// Iterate over values that share the first key.
 	fn iter_prefix_values<KArg1>(k1: KArg1) -> PrefixIterator<V>
 	where
@@ -739,6 +745,12 @@ pub trait StorageNMap<K: KeyGenerator, V: FullCodec> {
 	where
 		K: HasKeyPrefix<KP>;
 
+	/// Does any value under a `partial_key` prefix (explicitly) exist in storage?
+	/// Might have unexpected behaviour with empty keys, e.g. `[]`.
+	fn contains_prefix<KP>(partial_key: KP) -> bool
+	where
+		K: HasKeyPrefix<KP>;
+
 	/// Iterate over values that share the partial prefix key.
 	fn iter_prefix_values<KP>(partial_key: KP) -> PrefixIterator<V>
 	where
@@ -1485,7 +1497,7 @@ pub fn storage_prefix(pallet_name: &[u8], storage_name: &[u8]) -> [u8; 32] {
 #[cfg(test)]
 mod test {
 	use super::*;
-	use crate::{assert_ok, hash::Identity, Twox128};
+	use crate::{assert_ok, hash::Identity, pallet_prelude::NMapKey, Twox128};
 	use bounded_vec::BoundedVec;
 	use frame_support::traits::ConstU32;
 	use generator::StorageValue as _;
@@ -1780,6 +1792,39 @@ mod test {
 	#[crate::storage_alias]
 	type FooDoubleMap =
 		StorageDoubleMap<Prefix, Twox128, u32, Twox128, u32, BoundedVec<u32, ConstU32<7>>>;
+	#[crate::storage_alias]
+	type FooTripleMap = StorageNMap<
+		Prefix,
+		(NMapKey<Twox128, u32>, NMapKey<Twox128, u32>, NMapKey<Twox128, u32>),
+		u64,
+	>;
+
+	#[test]
+	fn contains_prefix_works() {
+		TestExternalities::default().execute_with(|| {
+			// Test double maps
+			assert!(FooDoubleMap::iter_prefix_values(1).next().is_none());
+			assert_eq!(FooDoubleMap::contains_prefix(1), false);
+
+			assert_ok!(FooDoubleMap::try_append(1, 1, 4));
+			assert_ok!(FooDoubleMap::try_append(2, 1, 4));
+			assert!(FooDoubleMap::iter_prefix_values(1).next().is_some());
+			assert!(FooDoubleMap::contains_prefix(1));
+			FooDoubleMap::remove(1, 1);
+			assert_eq!(FooDoubleMap::contains_prefix(1), false);
+
+			// Test N Maps
+			assert!(FooTripleMap::iter_prefix_values((1,)).next().is_none());
+			assert_eq!(FooTripleMap::contains_prefix((1,)), false);
+
+			FooTripleMap::insert((1, 1, 1), 4);
+			FooTripleMap::insert((2, 1, 1), 4);
+			assert!(FooTripleMap::iter_prefix_values((1,)).next().is_some());
+			assert!(FooTripleMap::contains_prefix((1,)));
+			FooTripleMap::remove((1, 1, 1));
+			assert_eq!(FooTripleMap::contains_prefix((1,)), false);
+		});
+	}
 
 	#[test]
 	fn try_append_works() {

frame/support/src/storage/unhashed.rs

@@ -154,6 +154,16 @@ pub fn clear_prefix(
 	MultiRemovalResults { maybe_cursor, backend: i, unique: i, loops: i }
 }
 
+/// Returns `true` if the storage contains any key, which starts with a certain prefix,
+/// and is longer than said prefix.
+/// This means that a key which equals the prefix will not be counted.
+pub fn contains_prefixed_key(prefix: &[u8]) -> bool {
+	match sp_io::storage::next_key(prefix) {
+		Some(key) => key.starts_with(prefix),
+		None => false,
+	}
+}
+
 /// Get a Vec of bytes from storage.
 pub fn get_raw(key: &[u8]) -> Option<Vec<u8>> {
 	sp_io::storage::get(key).map(|value| value.to_vec())