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 c91a4848916eba87184b3dc4722cea81aec9339d.

the storage cache cannot be used this way

* Revert "Implement prefixed storage"

This reverts commit 4931088126a427082d7310ed7e83b8eea966bc20.

* 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

substrate/primitives/state-machine/src/overlayed_changes.rs

@@ -18,10 +18,11 @@
 
 #[cfg(test)]
 use std::iter::FromIterator;
-use std::collections::{HashMap, BTreeSet};
+use std::collections::{HashMap, BTreeMap, BTreeSet};
 use codec::Decode;
 use crate::changes_trie::{NO_EXTRINSIC_INDEX, Configuration as ChangesTrieConfig};
 use primitives::storage::well_known_keys::EXTRINSIC_INDEX;
+use std::{mem, ops};
 
 /// The overlayed changes to state to be queried on top of the backend.
 ///
@@ -54,9 +55,9 @@ pub struct OverlayedValue {
 #[cfg_attr(test, derive(PartialEq))]
 pub struct OverlayedChangeSet {
 	/// Top level storage changes.
-	pub top: HashMap<Vec<u8>, OverlayedValue>,
+	pub top: BTreeMap<Vec<u8>, OverlayedValue>,
 	/// Child storage changes.
-	pub children: HashMap<Vec<u8>, HashMap<Vec<u8>, OverlayedValue>>,
+	pub children: HashMap<Vec<u8>, BTreeMap<Vec<u8>, OverlayedValue>>,
 }
 
 #[cfg(test)]
@@ -274,9 +275,10 @@ impl OverlayedChanges {
 	/// Commit prospective changes to state.
 	pub fn commit_prospective(&mut self) {
 		if self.committed.is_empty() {
-			::std::mem::swap(&mut self.prospective, &mut self.committed);
+			mem::swap(&mut self.prospective, &mut self.committed);
 		} else {
-			for (key, val) in self.prospective.top.drain() {
+			let top_to_commit = mem::replace(&mut self.prospective.top, BTreeMap::new());
+			for (key, val) in top_to_commit.into_iter() {
 				let entry = self.committed.top.entry(key).or_default();
 				entry.value = val.value;
 
@@ -285,9 +287,9 @@ impl OverlayedChanges {
 						.extend(prospective_extrinsics);
 				}
 			}
-			for (storage_key, mut map) in self.prospective.children.drain() {
+			for (storage_key, map) in self.prospective.children.drain() {
 				let map_dest = self.committed.children.entry(storage_key).or_default();
-				for (key, val) in map.drain() {
+				for (key, val) in map.into_iter() {
 					let entry = map_dest.entry(key).or_default();
 					entry.value = val.value;
 
@@ -339,6 +341,56 @@ impl OverlayedChanges {
 			false => None,
 		}
 	}
+
+	/// Returns the next (in lexicographic order) storage key in the overlayed alongside its value.
+	/// If no value is next then `None` is returned.
+	pub fn next_storage_key_change(&self, key: &[u8]) -> Option<(&[u8], &OverlayedValue)> {
+		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
+
+		let next_prospective_key = self.prospective.top
+			.range::<[u8], _>(range)
+			.next()
+			.map(|(k, v)| (&k[..], v));
+
+		let next_committed_key = self.committed.top
+			.range::<[u8], _>(range)
+			.next()
+			.map(|(k, v)| (&k[..], v));
+
+		match (next_committed_key, next_prospective_key) {
+			// Committed is strictly less than prospective
+			(Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 =>
+				Some(committed_key),
+			(committed_key, None) => committed_key,
+			// Prospective key is less or equal to committed or committed doesn't exist
+			(_, prospective_key) => prospective_key,
+		}
+	}
+
+	/// Returns the next (in lexicographic order) child storage key in the overlayed alongside its
+	/// value.  If no value is next then `None` is returned.
+	pub fn next_child_storage_key_change(
+		&self,
+		storage_key: &[u8],
+		key: &[u8]
+	) -> Option<(&[u8], &OverlayedValue)> {
+		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
+
+		let next_prospective_key = self.prospective.children.get(storage_key)
+			.and_then(|map| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v)));
+
+		let next_committed_key = self.committed.children.get(storage_key)
+			.and_then(|map| map.range::<[u8], _>(range).next().map(|(k, v)| (&k[..], v)));
+
+		match (next_committed_key, next_prospective_key) {
+			// Committed is strictly less than prospective
+			(Some(committed_key), Some(prospective_key)) if committed_key.0 < prospective_key.0 =>
+				Some(committed_key),
+			(committed_key, None) => committed_key,
+			// Prospective key is less or equal to committed or committed doesn't exist
+			(_, prospective_key) => prospective_key,
+		}
+	}
 }
 
 #[cfg(test)]
@@ -359,8 +411,8 @@ mod tests {
 	use crate::ext::Ext;
 	use super::*;
 
-	fn strip_extrinsic_index(map: &HashMap<Vec<u8>, OverlayedValue>)
-		-> HashMap<Vec<u8>, OverlayedValue>
+	fn strip_extrinsic_index(map: &BTreeMap<Vec<u8>, OverlayedValue>)
+		-> BTreeMap<Vec<u8>, OverlayedValue>
 	{
 		let mut clone = map.clone();
 		clone.remove(&EXTRINSIC_INDEX.to_vec());
@@ -397,7 +449,7 @@ mod tests {
 
 	#[test]
 	fn overlayed_storage_root_works() {
-		let initial: HashMap<_, _> = vec![
+		let initial: BTreeMap<_, _> = vec![
 			(b"doe".to_vec(), b"reindeer".to_vec()),
 			(b"dog".to_vec(), b"puppyXXX".to_vec()),
 			(b"dogglesworth".to_vec(), b"catXXX".to_vec()),
@@ -543,4 +595,79 @@ mod tests {
 		assert_eq!(overlay.prospective,
 			Default::default());
 	}
+
+	#[test]
+	fn next_storage_key_change_works() {
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_storage(vec![20], Some(vec![20]));
+		overlay.set_storage(vec![30], Some(vec![30]));
+		overlay.set_storage(vec![40], Some(vec![40]));
+		overlay.commit_prospective();
+		overlay.set_storage(vec![10], Some(vec![10]));
+		overlay.set_storage(vec![30], None);
+
+		// next_prospective < next_committed
+		let next_to_5 = overlay.next_storage_key_change(&[5]).unwrap();
+		assert_eq!(next_to_5.0.to_vec(), vec![10]);
+		assert_eq!(next_to_5.1.value, Some(vec![10]));
+
+		// next_committed < next_prospective
+		let next_to_10 = overlay.next_storage_key_change(&[10]).unwrap();
+		assert_eq!(next_to_10.0.to_vec(), vec![20]);
+		assert_eq!(next_to_10.1.value, Some(vec![20]));
+
+		// next_committed == next_prospective
+		let next_to_20 = overlay.next_storage_key_change(&[20]).unwrap();
+		assert_eq!(next_to_20.0.to_vec(), vec![30]);
+		assert_eq!(next_to_20.1.value, None);
+
+		// next_committed, no next_prospective
+		let next_to_30 = overlay.next_storage_key_change(&[30]).unwrap();
+		assert_eq!(next_to_30.0.to_vec(), vec![40]);
+		assert_eq!(next_to_30.1.value, Some(vec![40]));
+
+		overlay.set_storage(vec![50], Some(vec![50]));
+		// next_prospective, no next_committed
+		let next_to_40 = overlay.next_storage_key_change(&[40]).unwrap();
+		assert_eq!(next_to_40.0.to_vec(), vec![50]);
+		assert_eq!(next_to_40.1.value, Some(vec![50]));
+	}
+
+	#[test]
+	fn next_child_storage_key_change_works() {
+		let child = b"Child1".to_vec();
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_child_storage(child.clone(), vec![20], Some(vec![20]));
+		overlay.set_child_storage(child.clone(), vec![30], Some(vec![30]));
+		overlay.set_child_storage(child.clone(), vec![40], Some(vec![40]));
+		overlay.commit_prospective();
+		overlay.set_child_storage(child.clone(), vec![10], Some(vec![10]));
+		overlay.set_child_storage(child.clone(), vec![30], None);
+
+		// next_prospective < next_committed
+		let next_to_5 = overlay.next_child_storage_key_change(&child, &[5]).unwrap();
+		assert_eq!(next_to_5.0.to_vec(), vec![10]);
+		assert_eq!(next_to_5.1.value, Some(vec![10]));
+
+		// next_committed < next_prospective
+		let next_to_10 = overlay.next_child_storage_key_change(&child, &[10]).unwrap();
+		assert_eq!(next_to_10.0.to_vec(), vec![20]);
+		assert_eq!(next_to_10.1.value, Some(vec![20]));
+
+		// next_committed == next_prospective
+		let next_to_20 = overlay.next_child_storage_key_change(&child, &[20]).unwrap();
+		assert_eq!(next_to_20.0.to_vec(), vec![30]);
+		assert_eq!(next_to_20.1.value, None);
+
+		// next_committed, no next_prospective
+		let next_to_30 = overlay.next_child_storage_key_change(&child, &[30]).unwrap();
+		assert_eq!(next_to_30.0.to_vec(), vec![40]);
+		assert_eq!(next_to_30.1.value, Some(vec![40]));
+
+		overlay.set_child_storage(child.clone(), vec![50], Some(vec![50]));
+		// next_prospective, no next_committed
+		let next_to_40 = overlay.next_child_storage_key_change(&child, &[40]).unwrap();
+		assert_eq!(next_to_40.0.to_vec(), vec![50]);
+		assert_eq!(next_to_40.1.value, Some(vec![50]));
+	}
 }