support storage map start with random position (#280)

utilities/Cargo.toml

@@ -13,13 +13,13 @@ codec = { package = "parity-scale-codec", version = "1.3.0", default-features =
 frame-support = { version = "2.0.0-rc6", default-features = false }
 sp-runtime = { version = "2.0.0-rc6", default-features = false }
 sp-std = { version = "2.0.0-rc6", default-features = false }
+sp-io = { version = "2.0.0-rc6", default-features = false }
 
 integer-sqrt = "=0.1.3" # https://github.com/derekdreery/integer-sqrt-rs/issues/12
 
 [dev-dependencies]
 serde_json = "1.0.53"
 frame-system = { version = "2.0.0-rc6" }
-sp-io = { version = "2.0.0-rc6" }
 
 clear_on_drop = { version = "0.2.4", features = ["no_cc"] }	# https://github.com/paritytech/substrate/issues/4179
 
@@ -31,4 +31,5 @@ std = [
 	"frame-support/std",
 	"sp-runtime/std",
 	"sp-std/std",
+	"sp-io/std",
 ]

utilities/src/iterator.rs

+use codec::{Decode, EncodeLike, FullCodec, FullEncode};
+use frame_support::{
+	storage::{
+		generator::{StorageDoubleMap as StorageDoubleMapT, StorageMap as StorageMapT},
+		unhashed, StorageDoubleMap, StorageMap,
+	},
+	ReversibleStorageHasher,
+};
+use sp_std::prelude::*;
+
+/// Utility to iterate through items in a storage map.
+/// Forks from substrate, expose previous_key field
+pub struct StorageMapIterator<K, V, Hasher> {
+	prefix: Vec<u8>,
+	pub previous_key: Vec<u8>,
+	drain: bool,
+	_phantom: sp_std::marker::PhantomData<(K, V, Hasher)>,
+}
+
+/// Forks from substrate
+impl<K: Decode + Sized, V: Decode + Sized, Hasher: ReversibleStorageHasher> Iterator
+	for StorageMapIterator<K, V, Hasher>
+{
+	type Item = (K, V);
+
+	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,
+			};
+		}
+	}
+}
+
+/// Shim for StorageMapIterator, add more features
+pub struct StorageMapIteratorShim<K, V, H> {
+	pub storage_map_iterator: StorageMapIterator<K, V, H>,
+	pub remain_iterator_count: Option<u32>,
+	pub finished: bool,
+}
+
+impl<K: Decode + Sized, V: Decode + Sized, H: ReversibleStorageHasher> Iterator for StorageMapIteratorShim<K, V, H> {
+	type Item = <StorageMapIterator<K, V, H> as Iterator>::Item;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		// check accumulated iteration count
+		if let Some(remain_iterator_count) = self.remain_iterator_count {
+			if remain_iterator_count == 0 {
+				return None;
+			} else {
+				self.remain_iterator_count = Some(remain_iterator_count - 1);
+			}
+		}
+
+		self.storage_map_iterator.next().or_else(|| {
+			// mark this map iterator has already finished
+			self.finished = true;
+			None
+		})
+	}
+}
+
+/// A strongly-typed map in storage whose keys and values can be iterated over.
+pub trait IterableStorageMapExtended<H, K: FullEncode, V: FullCodec>: StorageMap<K, V> {
+	/// The type that iterates over all `(key, value)`.
+	type Iterator: Iterator<Item = (K, V)>;
+
+	/// Enumerate all elements in the map in no particular order. If you alter
+	/// the map while doing this, you'll get undefined results.
+	fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
+
+	/// Remove all elements from the map and iterate through them in no
+	/// particular order. If you add elements to the map while doing this,
+	/// you'll get undefined results.
+	fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
+}
+
+impl<K: FullCodec, V: FullCodec, G: StorageMapT<K, V>> IterableStorageMapExtended<G::Hasher, K, V> for G
+where
+	G::Hasher: ReversibleStorageHasher,
+{
+	type Iterator = StorageMapIteratorShim<K, V, G::Hasher>;
+
+	/// Enumerate all elements in the map.
+	fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
+		let prefix = G::prefix_hash();
+		let previous_key = start_key
+			.filter(|k| k.starts_with(&prefix))
+			.unwrap_or_else(|| prefix.clone());
+		let storage_map_iterator = StorageMapIterator {
+			prefix,
+			previous_key,
+			drain: false,
+			_phantom: Default::default(),
+		};
+
+		StorageMapIteratorShim {
+			storage_map_iterator,
+			remain_iterator_count: max_iterations,
+			finished: false,
+		}
+	}
+
+	/// Enumerate all elements in the map.
+	fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
+		let prefix = G::prefix_hash();
+
+		let previous_key = start_key
+			.filter(|k| k.starts_with(&prefix))
+			.unwrap_or_else(|| prefix.clone());
+		let storage_map_iterator = StorageMapIterator {
+			prefix,
+			previous_key,
+			drain: true,
+			_phantom: Default::default(),
+		};
+
+		StorageMapIteratorShim {
+			storage_map_iterator,
+			remain_iterator_count: max_iterations,
+			finished: false,
+		}
+	}
+}
+
+/// Iterate over a prefix and decode raw_key and raw_value into `T`.
+/// Forks from substrate, expose previous_key field
+pub struct MapIterator<T> {
+	prefix: Vec<u8>,
+	pub previous_key: Vec<u8>,
+	/// If true then value are removed while iterating
+	drain: bool,
+	/// Function that take `(raw_key_without_prefix, raw_value)` and decode `T`.
+	/// `raw_key_without_prefix` is the raw storage key without the prefix
+	/// iterated on.
+	closure: fn(&[u8], &[u8]) -> Result<T, codec::Error>,
+}
+
+/// Forks from substrate
+impl<T> Iterator for MapIterator<T> {
+	type Item = T;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		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;
+					let raw_value = match unhashed::get_raw(&self.previous_key) {
+						Some(raw_value) => raw_value,
+						None => {
+							frame_support::print("ERROR: next_key returned a key with no value in MapIterator");
+							continue;
+						}
+					};
+					if self.drain {
+						unhashed::kill(&self.previous_key)
+					}
+					let raw_key_without_prefix = &self.previous_key[self.prefix.len()..];
+					let item = match (self.closure)(raw_key_without_prefix, &raw_value[..]) {
+						Ok(item) => item,
+						Err(_e) => {
+							frame_support::print("ERROR: (key, value) failed to decode in MapIterator");
+							continue;
+						}
+					};
+
+					Some(item)
+				}
+				None => None,
+			};
+		}
+	}
+}
+
+/// Shim for MapIterator, add more features
+pub struct MapIteratorShim<T> {
+	pub map_iterator: MapIterator<T>,
+	pub remain_iterator_count: Option<u32>,
+	pub finished: bool,
+}
+
+impl<T> Iterator for MapIteratorShim<T> {
+	type Item = <MapIterator<T> as Iterator>::Item;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		// check accumulated iteration count
+		if let Some(remain_iterator_count) = self.remain_iterator_count {
+			if remain_iterator_count == 0 {
+				return None;
+			} else {
+				self.remain_iterator_count = Some(remain_iterator_count - 1);
+			}
+		}
+
+		self.map_iterator.next().or_else(|| {
+			// mark this map iterator has already finished
+			self.finished = true;
+			None
+		})
+	}
+}
+
+/// A strongly-typed map in storage whose keys and values can be iterated over.
+pub trait IterableStorageDoubleMapExtended<K1: FullCodec, K2: FullCodec, V: FullCodec>:
+	StorageDoubleMap<K1, K2, V>
+{
+	/// The type that iterates over all `(key2, value)`.
+	type PrefixIterator: Iterator<Item = (K2, V)>;
+
+	/// The type that iterates over all `(key1, key2, value)`.
+	type Iterator: Iterator<Item = (K1, K2, V)>;
+
+	/// Enumerate all elements in the map with first key `k1` in no particular
+	/// order. If you add or remove values whose first key is `k1` to the map
+	/// while doing this, you'll get undefined results.
+	fn iter_prefix(
+		k1: impl EncodeLike<K1>,
+		max_iterations: Option<u32>,
+		start_key: Option<Vec<u8>>,
+	) -> Self::PrefixIterator;
+
+	/// Remove all elements from the map with first key `k1` and iterate through
+	/// them in no particular order. If you add elements with first key `k1` to
+	/// the map while doing this, you'll get undefined results.
+	fn drain_prefix(
+		k1: impl EncodeLike<K1>,
+		max_iterations: Option<u32>,
+		start_key: Option<Vec<u8>>,
+	) -> Self::PrefixIterator;
+
+	/// Enumerate all elements in the map in no particular order. If you add or
+	/// remove values to the map while doing this, you'll get undefined results.
+	fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
+
+	/// Remove all elements from the map and iterate through them in no
+	/// particular order. If you add elements to the map while doing this,
+	/// you'll get undefined results.
+	fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator;
+}
+
+impl<K1: FullCodec, K2: FullCodec, V: FullCodec, G: StorageDoubleMapT<K1, K2, V>>
+	IterableStorageDoubleMapExtended<K1, K2, V> for G
+where
+	G::Hasher1: ReversibleStorageHasher,
+	G::Hasher2: ReversibleStorageHasher,
+{
+	type PrefixIterator = MapIteratorShim<(K2, V)>;
+	type Iterator = MapIteratorShim<(K1, K2, V)>;
+
+	fn iter_prefix(
+		k1: impl EncodeLike<K1>,
+		max_iterations: Option<u32>,
+		start_key: Option<Vec<u8>>,
+	) -> Self::PrefixIterator {
+		let prefix = G::storage_double_map_final_key1(k1);
+		let previous_key = start_key
+			.filter(|k| k.starts_with(&prefix))
+			.unwrap_or_else(|| prefix.clone());
+
+		let map_iterator = MapIterator {
+			prefix,
+			previous_key,
+			drain: false,
+			closure: |raw_key_without_prefix, mut raw_value| {
+				let mut key_material = G::Hasher2::reverse(raw_key_without_prefix);
+				Ok((K2::decode(&mut key_material)?, V::decode(&mut raw_value)?))
+			},
+		};
+
+		MapIteratorShim {
+			map_iterator,
+			remain_iterator_count: max_iterations,
+			finished: false,
+		}
+	}
+
+	fn drain_prefix(
+		k1: impl EncodeLike<K1>,
+		max_iterations: Option<u32>,
+		start_key: Option<Vec<u8>>,
+	) -> Self::PrefixIterator {
+		let mut shim = Self::iter_prefix(k1, max_iterations, start_key);
+		shim.map_iterator.drain = true;
+		shim
+	}
+
+	fn iter(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
+		let prefix = G::prefix_hash();
+		let previous_key = start_key
+			.filter(|k| k.starts_with(&prefix))
+			.unwrap_or_else(|| prefix.clone());
+
+		let map_iterator = MapIterator {
+			prefix,
+			previous_key,
+			drain: false,
+			closure: |raw_key_without_prefix, mut raw_value| {
+				let mut k1_k2_material = G::Hasher1::reverse(raw_key_without_prefix);
+				let k1 = K1::decode(&mut k1_k2_material)?;
+				let mut k2_material = G::Hasher2::reverse(k1_k2_material);
+				let k2 = K2::decode(&mut k2_material)?;
+				Ok((k1, k2, V::decode(&mut raw_value)?))
+			},
+		};
+
+		MapIteratorShim {
+			map_iterator,
+			remain_iterator_count: max_iterations,
+			finished: false,
+		}
+	}
+
+	fn drain(max_iterations: Option<u32>, start_key: Option<Vec<u8>>) -> Self::Iterator {
+		let mut shim = Self::iter(max_iterations, start_key);
+		shim.map_iterator.drain = true;
+		shim
+	}
+}

utilities/src/lib.rs

@@ -4,8 +4,12 @@ use frame_support::storage::{with_transaction, TransactionOutcome};
 use sp_runtime::DispatchError;
 use sp_std::result::Result;
 
+pub mod iterator;
+pub mod offchain_worker;
 pub mod ordered_set;
 
+pub use iterator::{IterableStorageDoubleMapExtended, IterableStorageMapExtended};
+pub use offchain_worker::OffchainErr;
 pub use ordered_set::OrderedSet;
 
 /// Execute the supplied function in a new storage transaction.

utilities/src/offchain_worker/mod.rs

+/// Error which may occur while executing the off-chain code.
+#[derive(PartialEq)]
+pub enum OffchainErr {
+	OffchainStore,
+	SubmitTransaction,
+	NotValidator,
+	OffchainLock,
+}
+
+impl sp_std::fmt::Debug for OffchainErr {
+	fn fmt(&self, fmt: &mut sp_std::fmt::Formatter) -> sp_std::fmt::Result {
+		match *self {
+			OffchainErr::OffchainStore => write!(fmt, "Failed to manipulate offchain store"),
+			OffchainErr::SubmitTransaction => write!(fmt, "Failed to submit transaction"),
+			OffchainErr::NotValidator => write!(fmt, "Is not validator"),
+			OffchainErr::OffchainLock => write!(fmt, "Failed to manipulate offchain lock"),
+		}
+	}
+}