Merge IterableStorageMap into StorageMap and IterableStorageDoubleMap into StorageDoubleMap

frame/elections-phragmen/src/lib.rs

@@ -88,7 +88,7 @@ use sp_runtime::{
 };
 use frame_support::{
 	decl_storage, decl_event, ensure, decl_module, decl_error,
-	weights::{SimpleDispatchInfo, Weight, WeighData}, storage::{StorageMap, IterableStorageMap},
+	weights::{SimpleDispatchInfo, Weight, WeighData}, storage::StorageMap,
 	traits::{
 		Currency, Get, LockableCurrency, LockIdentifier, ReservableCurrency, WithdrawReasons,
 		ChangeMembers, OnUnbalanced, WithdrawReason, Contains, BalanceStatus

frame/staking/src/lib.rs

@@ -263,7 +263,8 @@ use sp_std::{prelude::*, result, collections::btree_map::BTreeMap};
 use codec::{HasCompact, Encode, Decode};
 use frame_support::{
 	decl_module, decl_event, decl_storage, ensure, decl_error, weights::SimpleDispatchInfo,
-	dispatch::DispatchResult, storage::IterableStorageMap, traits::{
+	dispatch::DispatchResult,
+	traits::{
 		Currency, LockIdentifier, LockableCurrency, WithdrawReasons, OnUnbalanced, Imbalance, Get,
 		Time
 	}

frame/staking/src/mock.rs

@@ -25,8 +25,7 @@ use sp_staking::{SessionIndex, offence::{OffenceDetails, OnOffenceHandler}};
 use sp_core::{H256, crypto::key_types};
 use sp_io;
 use frame_support::{
-	assert_ok, impl_outer_origin, parameter_types, StorageValue, StorageMap,
-	StorageDoubleMap, IterableStorageMap,
+	assert_ok, impl_outer_origin, parameter_types, StorageValue, StorageMap, StorageDoubleMap,
 	traits::{Currency, Get, FindAuthor, OnFinalize, OnInitialize}, weights::Weight,
 };
 use crate::{
@@ -379,7 +378,7 @@ pub type Timestamp = pallet_timestamp::Module<Test>;
 pub type Staking = Module<Test>;
 
 pub fn check_exposure_all(era: EraIndex) {
-	ErasStakers::<Test>::iter_prefix(era).for_each(check_exposure)
+	ErasStakers::<Test>::iter_prefix_values(era).for_each(check_exposure)
 }
 
 pub fn check_nominator_all(era: EraIndex) {
@@ -400,7 +399,7 @@ pub fn check_exposure(expo: Exposure<AccountId, Balance>) {
 pub fn check_nominator_exposure(era: EraIndex, stash: AccountId) {
 	assert_is_stash(stash);
 	let mut sum = 0;
-	ErasStakers::<Test>::iter_prefix(era)
+	ErasStakers::<Test>::iter_prefix_values(era)
 		.for_each(|exposure| {
 			exposure.others.iter()
 				.filter(|i| i.who == stash)

frame/staking/src/tests.rs

@@ -345,7 +345,7 @@ fn less_than_needed_candidates_works() {
 			// But the exposure is updated in a simple way. No external votes exists.
 			// This is purely self-vote.
 			assert!(
-				ErasStakers::<Test>::iter_prefix(Staking::active_era().unwrap().index)
+				ErasStakers::<Test>::iter_prefix_values(Staking::active_era().unwrap().index)
 					.all(|exposure| exposure.others.is_empty())
 			);
 			check_exposure_all(Staking::active_era().unwrap().index);

frame/support/procedural/src/lib.rs

@@ -67,7 +67,6 @@ use proc_macro::TokenStream;
 /// * Map: `Foo: map hasher($hash) type => type`: Implements the
 ///   [`StorageMap`](../frame_support/storage/trait.StorageMap.html) trait using the
 ///   [`StorageMap generator`](../frame_support/storage/generator/trait.StorageMap.html).
-///   And [`StoragePrefixedMap`](../frame_support/storage/trait.StoragePrefixedMap.html).
 ///
 ///   `$hash` representing a choice of hashing algorithms available in the
 ///   [`Hashable`](../frame_support/trait.Hashable.html) trait. You will generally want to use one
@@ -106,7 +105,6 @@ use proc_macro::TokenStream;
 /// * Double map: `Foo: double_map hasher($hash1) u32, hasher($hash2) u32 => u32`: Implements the
 ///   [`StorageDoubleMap`](../frame_support/storage/trait.StorageDoubleMap.html) trait using the
 ///   [`StorageDoubleMap generator`](../frame_support/storage/generator/trait.StorageDoubleMap.html).
-///   And [`StoragePrefixedMap`](../frame_support/storage/trait.StoragePrefixedMap.html).
 ///
 ///   `$hash1` and `$hash2` representing choices of hashing algorithms available in the
 ///   [`Hashable`](../frame_support/trait.Hashable.html) trait. They must be chosen with care, see

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

@@ -414,7 +414,6 @@ pub fn decl_storage_impl(input: proc_macro::TokenStream) -> proc_macro::TokenStr
 			StorageValue as _,
 			StorageMap as _,
 			StorageDoubleMap as _,
-			StoragePrefixedMap as _,
 		};
 
 		#scrate_decl

frame/support/procedural/src/storage/storage_struct.rs

@@ -122,18 +122,6 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 			StorageLineTypeDef::Map(map) => {
 				let hasher = map.hasher.to_storage_hasher_struct();
 				quote!(
-					impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type>
-						for #storage_struct #optional_storage_where_clause
-					{
-						fn module_prefix() -> &'static [u8] {
-							#instance_or_inherent::PREFIX.as_bytes()
-						}
-
-						fn storage_prefix() -> &'static [u8] {
-							#storage_name_str.as_bytes()
-						}
-					}
-
 					impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct
 					#optional_storage_where_clause
 					{
@@ -162,18 +150,6 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 				let hasher1 = map.hasher1.to_storage_hasher_struct();
 				let hasher2 = map.hasher2.to_storage_hasher_struct();
 				quote!(
-					impl<#impl_trait> #scrate::storage::StoragePrefixedMap<#value_type>
-						for #storage_struct #optional_storage_where_clause
-					{
-						fn module_prefix() -> &'static [u8] {
-							#instance_or_inherent::PREFIX.as_bytes()
-						}
-
-						fn storage_prefix() -> &'static [u8] {
-							#storage_name_str.as_bytes()
-						}
-					}
-
 					impl<#impl_trait> #scrate::#storage_generator_trait for #storage_struct
 					#optional_storage_where_clause
 					{

frame/support/src/hash.rs

@@ -92,6 +92,10 @@ impl StorageHasher for Twox64Concat {
 }
 impl ReversibleStorageHasher for Twox64Concat {
 	fn reverse(x: &[u8]) -> &[u8] {
+		if x.len() < 8 {
+			crate::debug::error!("Invalid reverse: hash length too short");
+			return &[]
+		}
 		&x[8..]
 	}
 }
@@ -110,6 +114,10 @@ impl StorageHasher for Blake2_128Concat {
 }
 impl ReversibleStorageHasher for Blake2_128Concat {
 	fn reverse(x: &[u8]) -> &[u8] {
+		if x.len() < 16 {
+			crate::debug::error!("Invalid reverse: hash length too short");
+			return &[]
+		}
 		&x[16..]
 	}
 }

frame/support/src/lib.rs

@@ -70,10 +70,7 @@ pub use self::hash::{
 	Twox256, Twox128, Blake2_256, Blake2_128, Identity, Twox64Concat, Blake2_128Concat, Hashable,
 	StorageHasher
 };
-pub use self::storage::{
-	StorageValue, StorageMap, StorageDoubleMap, StoragePrefixedMap, IterableStorageMap,
-	IterableStorageDoubleMap, migration
-};
+pub use self::storage::{StorageValue, StorageMap, StorageDoubleMap, migration};
 pub use self::dispatch::{Parameter, Callable, IsSubType};
 pub use sp_runtime::{self, ConsensusEngineId, print, traits::Printable};
 

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

@@ -16,7 +16,7 @@
 
 use sp_std::prelude::*;
 use sp_std::borrow::Borrow;
-use codec::{Ref, FullCodec, FullEncode, Decode, Encode, EncodeLike, EncodeAppend};
+use codec::{Ref, FullCodec, Decode, Encode, EncodeLike, EncodeAppend};
 use crate::{storage::{self, unhashed}, traits::Len};
 use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 
@@ -40,7 +40,7 @@ use crate::hash::{StorageHasher, Twox128, ReversibleStorageHasher};
 /// If the key2s are not trusted (e.g. can be set by a user), a cryptographic `hasher` such as
 /// `blake2_256` must be used for Hasher2. Otherwise, other items in storage with the same first
 /// key can be compromised.
-pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
+pub trait StorageDoubleMap<K1: FullCodec, K2: FullCodec, V: FullCodec> {
 	/// The type that get/take returns.
 	type Query;
 
@@ -124,12 +124,14 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 }
 
 impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
-	K1: FullEncode,
-	K2: FullEncode,
+	K1: FullCodec,
+	K2: FullCodec,
 	V: FullCodec,
 	G: StorageDoubleMap<K1, K2, V>,
 {
 	type Query = G::Query;
+	type Hasher1 = G::Hasher1;
+	type Hasher2 = G::Hasher2;
 
 	fn hashed_key_for<KArg1, KArg2>(k1: KArg1, k2: KArg2) -> Vec<u8> where
 		KArg1: EncodeLike<K1>,
@@ -208,15 +210,8 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 		unhashed::kill_prefix(Self::storage_double_map_final_key1(k1).as_ref())
 	}
 
-	fn iter_prefix<KArg1>(k1: KArg1) -> storage::PrefixIterator<V> where
-		KArg1: ?Sized + EncodeLike<K1>
-	{
-		let prefix = Self::storage_double_map_final_key1(k1);
-		storage::PrefixIterator::<V> {
-			prefix: prefix.clone(),
-			previous_key: prefix,
-			phantom_data: Default::default(),
-		}
+	fn remove_all() {
+		sp_io::storage::clear_prefix(&Self::prefix_hash())
 	}
 
 	fn mutate<KArg1, KArg2, R, F>(k1: KArg1, k2: KArg2, f: F) -> R where
@@ -332,84 +327,75 @@ impl<K1, K2, V, G> storage::StorageDoubleMap<K1, K2, V> for G where
 			value
 		})
 	}
-}
-
-/// Utility to iterate through items in a storage map.
-pub struct MapIterator<K, V, Hasher> {
-	prefix: Vec<u8>,
-	previous_key: Vec<u8>,
-	drain: bool,
-	_phantom: ::sp_std::marker::PhantomData<(K, V, Hasher)>,
-}
 
-impl<
-	K: Decode + Sized,
-	V: Decode + Sized,
-	Hasher: ReversibleStorageHasher
-> Iterator for MapIterator<K, V, Hasher> {
-	type Item = (K, V);
+	fn iter_prefix(k1: impl EncodeLike<K1>) -> storage::PrefixIterator<(K2, V)> where
+		Self::Hasher2: ReversibleStorageHasher
+	{
+		let prefix = G::storage_double_map_final_key1(k1);
+		storage::PrefixIterator {
+			prefix: prefix.clone(),
+			previous_key: prefix,
+			drain: false,
+			closure: from_hashed_key2_to_key2_value::<_, _, _, G>,
+		}
+	}
 
-	fn next(&mut self) -> Option<(K, V)> {
-		loop {
-			let maybe_next = sp_io::storage::next_key(&self.previous_key)
-				.filter(|n| n.starts_with(&self.prefix));
-			break match maybe_next {
-				Some(next) => {
-					self.previous_key = next;
-					match unhashed::get::<V>(&self.previous_key) {
-						Some(value) => {
-							if self.drain {
-								unhashed::kill(&self.previous_key)
-							}
-							let mut key_material = Hasher::reverse(&self.previous_key[self.prefix.len()..]);
-							match K::decode(&mut key_material) {
-								Ok(key) => Some((key, value)),
-								Err(_) => continue,
-							}
-						}
-						None => continue,
-					}
-				}
-				None => None,
-			}
+	fn drain_prefix(k1: impl EncodeLike<K1>) -> storage::PrefixIterator<(K2, V)> where
+		Self::Hasher2: ReversibleStorageHasher
+	{
+		let prefix = G::storage_double_map_final_key1(k1);
+		storage::PrefixIterator {
+			prefix: prefix.clone(),
+			previous_key: prefix,
+			drain: true,
+			closure: from_hashed_key2_to_key2_value::<_, _, _, G>,
 		}
 	}
-}
 
-impl<
-	K1: FullCodec,
-	K2: FullCodec,
-	V: FullCodec,
-	G: StorageDoubleMap<K1, K2, V>,
-> storage::IterableStorageDoubleMap<K1, K2, V> for G where
-	G::Hasher1: ReversibleStorageHasher,
-	G::Hasher2: ReversibleStorageHasher
-{
-	type Iterator = MapIterator<K2, V, G::Hasher2>;
+	fn iter_prefix_values<KArg1>(k1: KArg1) -> storage::PrefixIterator<V> where
+		KArg1: ?Sized + EncodeLike<K1>
+	{
+		let prefix = Self::storage_double_map_final_key1(k1);
+		storage::PrefixIterator {
+			prefix: prefix.clone(),
+			previous_key: prefix,
+			drain: false,
+			closure: |_raw_key, mut raw_value| V::decode(&mut raw_value),
+		}
+	}
 
-	/// Enumerate all elements in the map.
-	fn iter(k1: impl EncodeLike<K1>) -> Self::Iterator {
-		let prefix = G::storage_double_map_final_key1(k1);
-		Self::Iterator {
+	fn iter() -> storage::PrefixIterator<(K1, K2, V)> where
+		Self::Hasher1: ReversibleStorageHasher,
+		Self::Hasher2: ReversibleStorageHasher,
+	{
+		let prefix = G::prefix_hash();
+		storage::PrefixIterator {
 			prefix: prefix.clone(),
 			previous_key: prefix,
 			drain: false,
-			_phantom: Default::default(),
+			closure: |mut raw_key_without_prefix, mut raw_value| {
+				let mut key1_key2_material = G::Hasher1::reverse(&mut raw_key_without_prefix);
+				let key1 = K1::decode(&mut key1_key2_material)?;
+				let mut key2_material = G::Hasher2::reverse(&mut key1_key2_material);
+				let key2 = K2::decode(&mut key2_material)?;
+				let value = V::decode(&mut raw_value)?;
+
+				Ok((key1, key2, value))
+			}
 		}
 	}
 
-	/// Enumerate all elements in the map.
-	fn drain(k1: impl EncodeLike<K1>) -> Self::Iterator {
-		let prefix = G::storage_double_map_final_key1(k1);
-		Self::Iterator {
+	fn iter_values() -> storage::PrefixIterator<V> {
+		let prefix = G::prefix_hash();
+		storage::PrefixIterator {
 			prefix: prefix.clone(),
 			previous_key: prefix,
-			drain: true,
-			_phantom: Default::default(),
+			drain: false,
+			closure: |_raw_key, mut raw_value| V::decode(&mut raw_value),
 		}
 	}
 
-	fn translate<O: Decode, F: Fn(O) -> Option<V>>(f: F) {
+	fn translate_values<O: Decode, F: Fn(O) -> Option<V>>(f: F) {
 		let prefix = G::prefix_hash();
 		let mut previous_key = prefix.clone();
 		loop {
@@ -422,7 +408,13 @@ impl<
 							Some(new) => unhashed::put::<V>(&previous_key, &new),
 							None => unhashed::kill(&previous_key),
 						},
-						None => continue,
+						None => {
+							crate::debug::error!(
+								"next_key returned a key which failed to decode at {:?}",
+								previous_key
+							);
+							continue
+						}
 					}
 				}
 				None => return,
@@ -431,6 +423,23 @@ impl<
 	}
 }
 
+fn from_hashed_key2_to_key2_value<K1, K2, V, G>(
+	hashed_key2: &[u8],
+	mut raw_value: &[u8]
+) -> Result<(K2, V), codec::Error> where
+	K1: FullCodec,
+	K2: FullCodec,
+	V: FullCodec,
+	G: StorageDoubleMap<K1, K2, V>,
+	G::Hasher2: ReversibleStorageHasher,
+{
+	let mut key2_material = G::Hasher2::reverse(hashed_key2);
+	let key2 = K2::decode(&mut key2_material)?;
+	let value = V::decode(&mut raw_value)?;
+
+	Ok((key2, value))
+}
+
 #[cfg(test)]
 mod test {
 	use sp_io::TestExternalities;
@@ -456,8 +465,8 @@ mod test {
 			MyStorage::insert(2, 5, 9);
 			MyStorage::insert(2, 6, 10);
 
-			assert_eq!(MyStorage::iter_prefix(1).collect::<Vec<_>>(), vec![7, 8]);
-			assert_eq!(MyStorage::iter_prefix(2).collect::<Vec<_>>(), vec![10, 9]);
+			assert_eq!(MyStorage::iter_prefix_values(1).collect::<Vec<_>>(), vec![7, 8]);
+			assert_eq!(MyStorage::iter_prefix_values(2).collect::<Vec<_>>(), vec![10, 9]);
 		});
 	}
 }