Introduce prefixed storage with enumeration (#4185)

* Introduce storage_next allowing iteration.  (without childtries)

* Implement prefixed storage

* impl cache in client_storage_cache (needs test)

* switch overlay change to btreemap

* Revert "impl cache in client_storage_cache"

This reverts commit c91a484.

the storage cache cannot be used this way

* Revert "Implement prefixed storage"

This reverts commit 4931088.

* Impl StoragePrefixedMap for all map storages

* remove comment

* Move all overlays to BTreeMap

* btreemap iteration improvment

* impl for child tries

* impl tests for childs

* fix

* remove cache comment

* Fix grumble

client/db/src/lib.rs

@@ -151,6 +151,14 @@ impl<B: BlockT> StateBackend<Blake2Hasher> for RefTrackingState<B> {
 		self.state.exists_child_storage(storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.next_storage_key(key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.next_child_storage_key(storage_key, key)
+	}
+
 	fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
 		self.state.for_keys_with_prefix(prefix, f)
 	}

client/db/src/storage_cache.rs

@@ -544,6 +544,14 @@ impl<H: Hasher, S: StateBackend<H>, B: BlockT> StateBackend<H> for CachingState<
 		self.state.exists_child_storage(storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.next_storage_key(key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.state.next_child_storage_key(storage_key, key)
+	}
+
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {
 		self.state.for_keys_in_child_storage(storage_key, f)
 	}

client/src/light/backend.rs

@@ -365,6 +365,22 @@ impl<H: Hasher> StateBackend<H> for GenesisOrUnavailableState<H>
 		}
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		match *self {
+			GenesisOrUnavailableState::Genesis(ref state) =>
+				Ok(state.next_storage_key(key).expect(IN_MEMORY_EXPECT_PROOF)),
+			GenesisOrUnavailableState::Unavailable => Err(ClientError::NotAvailableOnLightClient),
+		}
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		match *self {
+			GenesisOrUnavailableState::Genesis(ref state) =>
+				Ok(state.next_child_storage_key(storage_key, key).expect(IN_MEMORY_EXPECT_PROOF)),
+			GenesisOrUnavailableState::Unavailable => Err(ClientError::NotAvailableOnLightClient),
+		}
+	}
+
 	fn for_keys_with_prefix<A: FnMut(&[u8])>(&self, prefix: &[u8], action: A) {
 		match *self {
 			GenesisOrUnavailableState::Genesis(ref state) => state.for_keys_with_prefix(prefix, action),

frame/support/procedural/src/lib.rs

@@ -69,6 +69,7 @@ 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.
@@ -89,6 +90,7 @@ use proc_macro::TokenStream;
 /// * Linked map: `Foo: linked_map hasher($hash) type => type`: Implements the
 ///   [`StorageLinkedMap`](../frame_support/storage/trait.StorageLinkedMap.html) trait using the
 ///   [`StorageLinkedMap generator`](../frame_support/storage/generator/trait.StorageLinkedMap.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.
@@ -118,6 +120,7 @@ use proc_macro::TokenStream;
 /// * Double map: `Foo: double_map hasher($hash1) u32, $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 choosen with care, see

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

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

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

@@ -122,6 +122,18 @@ 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
 					{
@@ -155,6 +167,18 @@ pub fn decl_and_impl(scrate: &TokenStream, def: &DeclStorageDefExt) -> TokenStre
 				);
 
 				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
 					{
@@ -191,6 +215,18 @@ 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/lib.rs

@@ -67,7 +67,9 @@ pub mod traits;
 pub mod weights;
 
 pub use self::hash::{Twox256, Twox128, Blake2_256, Blake2_128, Twox64Concat, Hashable};
-pub use self::storage::{StorageValue, StorageMap, StorageLinkedMap, StorageDoubleMap};
+pub use self::storage::{
+	StorageValue, StorageMap, StorageLinkedMap, StorageDoubleMap, StoragePrefixedMap
+};
 pub use self::dispatch::{Parameter, Callable, IsSubType};
 pub use sp_runtime::{self, ConsensusEngineId, print, traits::Printable};
 

frame/support/src/storage/mod.rs

@@ -16,9 +16,9 @@
 
 //! Stuff to do with the runtime's storage.
 
-use rstd::prelude::*;
+use rstd::{prelude::*, marker::PhantomData};
 use codec::{FullCodec, FullEncode, Encode, EncodeAppend, EncodeLike, Decode};
-use crate::traits::Len;
+use crate::{traits::Len, hash::{Twox128, StorageHasher}};
 
 pub mod unhashed;
 pub mod hashed;
@@ -352,3 +352,126 @@ pub trait StorageDoubleMap<K1: FullEncode, K2: FullEncode, V: FullCodec> {
 			KArg2: EncodeLike<K2>,
 			V: codec::DecodeLength + Len;
 }
+
+/// Iterator for prefixed map.
+pub struct PrefixIterator<Value> {
+	prefix: Vec<u8>,
+	previous_key: Vec<u8>,
+	phantom_data: PhantomData<Value>,
+}
+
+impl<Value: Decode> Iterator for PrefixIterator<Value> {
+	type Item = Value;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		match runtime_io::storage::next_key(&self.previous_key) {
+			Some(next_key) if next_key.starts_with(&self.prefix[..]) => {
+				let value = unhashed::get(&next_key);
+
+				if value.is_none() {
+					runtime_print!(
+						"ERROR: returned next_key has no value:\nkey is {:?}\nnext_key is {:?}",
+						&self.previous_key, &next_key,
+					);
+				}
+
+				self.previous_key = next_key;
+
+				value
+			},
+			_ => None,
+		}
+	}
+}
+
+/// Trait for maps that store all its value after a unique prefix.
+///
+/// By default the final prefix is:
+/// ```nocompile
+/// Twox128(module_prefix) ++ Twox128(storage_prefix)
+/// ```
+pub trait StoragePrefixedMap<Value: FullCodec> {
+
+	/// Module prefix. Used for generating final key.
+	fn module_prefix() -> &'static [u8];
+
+	/// Storage prefix. Used for generating final key.
+	fn storage_prefix() -> &'static [u8];
+
+	fn final_prefix() -> [u8; 32] {
+		let mut final_key = [0u8; 32];
+		final_key[0..16].copy_from_slice(&Twox128::hash(Self::module_prefix()));
+		final_key[16..32].copy_from_slice(&Twox128::hash(Self::storage_prefix()));
+		final_key
+	}
+
+	fn remove_all() {
+		runtime_io::storage::clear_prefix(&Self::final_prefix())
+	}
+
+	fn iter() -> PrefixIterator<Value> {
+		let prefix = Self::final_prefix();
+		PrefixIterator {
+			prefix: prefix.to_vec(),
+			previous_key: prefix.to_vec(),
+			phantom_data: Default::default(),
+		}
+	}
+}
+
+#[cfg(test)]
+mod test {
+	use primitives::hashing::twox_128;
+	use runtime_io::TestExternalities;
+	use crate::storage::{unhashed, StoragePrefixedMap};
+
+	#[test]
+	fn prefixed_map_works() {
+		TestExternalities::default().execute_with(|| {
+			struct MyStorage;
+			impl StoragePrefixedMap<u64> for MyStorage {
+				fn module_prefix() -> &'static [u8] {
+					b"MyModule"
+				}
+
+				fn storage_prefix() -> &'static [u8] {
+					b"MyStorage"
+				}
+			}
+
+			let key_before = {
+				let mut k = MyStorage::final_prefix();
+				let last = k.iter_mut().last().unwrap();
+				*last = last.checked_sub(1).unwrap();
+				k
+			};
+			let key_after = {
+				let mut k = MyStorage::final_prefix();
+				let last = k.iter_mut().last().unwrap();
+				*last = last.checked_add(1).unwrap();
+				k
+			};
+
+			unhashed::put(&key_before[..], &32u64);
+			unhashed::put(&key_after[..], &33u64);
+
+			let k = [twox_128(b"MyModule"), twox_128(b"MyStorage")].concat();
+			assert_eq!(MyStorage::final_prefix().to_vec(), k);
+
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![]);
+
+			unhashed::put(&[&k[..], &vec![1][..]].concat(), &1u64);
+			unhashed::put(&[&k[..], &vec![1, 1][..]].concat(), &2u64);
+			unhashed::put(&[&k[..], &vec![8][..]].concat(), &3u64);
+			unhashed::put(&[&k[..], &vec![10][..]].concat(), &4u64);
+
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![1, 2, 3, 4]);
+
+			MyStorage::remove_all();
+
+			assert_eq!(MyStorage::iter().collect::<Vec<_>>(), vec![]);
+			assert_eq!(unhashed::get(&key_before[..]), Some(32u64));
+			assert_eq!(unhashed::get(&key_after[..]), Some(33u64));
+		});
+	}
+}