diff --git a/frame/support/src/storage/generator/double_map.rs b/frame/support/src/storage/generator/double_map.rs index c95dcee9d7e5c..b60bb79566d51 100644 --- a/frame/support/src/storage/generator/double_map.rs +++ b/frame/support/src/storage/generator/double_map.rs @@ -233,6 +233,13 @@ where .into() } + fn contains_prefix(k1: KArg1) -> bool + where + KArg1: EncodeLike, + { + unhashed::contains_prefixed_key(Self::storage_double_map_final_key1(k1).as_ref()) + } + fn iter_prefix_values(k1: KArg1) -> storage::PrefixIterator where KArg1: ?Sized + EncodeLike, diff --git a/frame/support/src/storage/generator/nmap.rs b/frame/support/src/storage/generator/nmap.rs index 79f3d72044e28..8cf1459431171 100755 --- a/frame/support/src/storage/generator/nmap.rs +++ b/frame/support/src/storage/generator/nmap.rs @@ -208,6 +208,13 @@ where ) } + fn contains_prefix(partial_key: KP) -> bool + where + K: HasKeyPrefix, + { + unhashed::contains_prefixed_key(&Self::storage_n_map_partial_key(partial_key)) + } + fn iter_prefix_values(partial_key: KP) -> PrefixIterator where K: HasKeyPrefix, diff --git a/frame/support/src/storage/mod.rs b/frame/support/src/storage/mod.rs index 8c0d6207c3f4d..9cb06eef29f1c 100644 --- a/frame/support/src/storage/mod.rs +++ b/frame/support/src/storage/mod.rs @@ -557,6 +557,12 @@ pub trait StorageDoubleMap { where KArg1: ?Sized + EncodeLike; + /// Does any value under the first key `k1` (explicitly) exist in storage? + /// Might have unexpected behaviour with empty keys, e.g. `[]`. + fn contains_prefix(k1: KArg1) -> bool + where + KArg1: EncodeLike; + /// Iterate over values that share the first key. fn iter_prefix_values(k1: KArg1) -> PrefixIterator where @@ -733,6 +739,12 @@ pub trait StorageNMap { where K: HasKeyPrefix; + /// Does any value under a `partial_key` prefix (explicitly) exist in storage? + /// Might have unexpected behaviour with empty keys, e.g. `[]`. + fn contains_prefix(partial_key: KP) -> bool + where + K: HasKeyPrefix; + /// Iterate over values that share the partial prefix key. fn iter_prefix_values(partial_key: KP) -> PrefixIterator where @@ -1479,7 +1491,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 _; @@ -1774,6 +1786,39 @@ mod test { #[crate::storage_alias] type FooDoubleMap = StorageDoubleMap>>; + #[crate::storage_alias] + type FooTripleMap = StorageNMap< + Prefix, + (NMapKey, NMapKey, NMapKey), + 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() { diff --git a/frame/support/src/storage/unhashed.rs b/frame/support/src/storage/unhashed.rs index 850e93e7d7fe4..8388c5f885c51 100644 --- a/frame/support/src/storage/unhashed.rs +++ b/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> { sp_io::storage::get(key).map(|value| value.to_vec())