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/backend.rs

@@ -16,7 +16,7 @@
 
 //! State machine backends. These manage the code and storage of contracts.
 
-use std::{error, fmt, cmp::Ord, collections::HashMap, marker::PhantomData};
+use std::{error, fmt, cmp::Ord, collections::{HashMap, BTreeMap}, marker::PhantomData, ops};
 use log::warn;
 use hash_db::Hasher;
 use crate::trie_backend::TrieBackend;
@@ -67,6 +67,16 @@ pub trait Backend<H: Hasher>: std::fmt::Debug {
 		Ok(self.child_storage(storage_key, key)?.is_some())
 	}
 
+	/// Return the next key in storage in lexicographic order or `None` if there is no value.
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error>;
+
+	/// Return the next key in child storage in lexicographic order or `None` if there is no value.
+	fn next_child_storage_key(
+		&self,
+		storage_key: &[u8],
+		key: &[u8]
+	) -> Result<Option<Vec<u8>>, Self::Error>;
+
 	/// Retrieve all entries keys of child storage and call `f` for each of those keys.
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F);
 
@@ -171,6 +181,14 @@ impl<'a, T: Backend<H>, H: Hasher> Backend<H> for &'a T {
 		(*self).child_storage(storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		(*self).next_storage_key(key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		(*self).next_child_storage_key(storage_key, key)
+	}
+
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {
 		(*self).for_keys_in_child_storage(storage_key, f)
 	}
@@ -250,7 +268,7 @@ impl error::Error for Void {
 /// In-memory backend. Fully recomputes tries each time `as_trie_backend` is called but useful for
 /// tests and proof checking.
 pub struct InMemory<H: Hasher> {
-	inner: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>,
+	inner: HashMap<Option<Vec<u8>>, BTreeMap<Vec<u8>, Vec<u8>>>,
 	// This field is only needed for returning reference in `as_trie_backend`.
 	trie: Option<TrieBackend<MemoryDB<H>, H>>,
 	_hasher: PhantomData<H>,
@@ -291,7 +309,7 @@ impl<H: Hasher> PartialEq for InMemory<H> {
 impl<H: Hasher> InMemory<H> where H::Out: Codec {
 	/// Copy the state, with applied updates
 	pub fn update(&self, changes: <Self as Backend<H>>::Transaction) -> Self {
-		let mut inner: HashMap<_, _> = self.inner.clone();
+		let mut inner = self.inner.clone();
 		for (storage_key, key, val) in changes {
 			match val {
 				Some(v) => { inner.entry(storage_key).or_default().insert(key, v); },
@@ -303,8 +321,8 @@ impl<H: Hasher> InMemory<H> where H::Out: Codec {
 	}
 }
 
-impl<H: Hasher> From<HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>> for InMemory<H> {
-	fn from(inner: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>) -> Self {
+impl<H: Hasher> From<HashMap<Option<Vec<u8>>, BTreeMap<Vec<u8>, Vec<u8>>>> for InMemory<H> {
+	fn from(inner: HashMap<Option<Vec<u8>>, BTreeMap<Vec<u8>, Vec<u8>>>) -> Self {
 		InMemory {
 			inner: inner,
 			trie: None,
@@ -314,14 +332,14 @@ impl<H: Hasher> From<HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>> for In
 }
 
 impl<H: Hasher> From<(
-		HashMap<Vec<u8>, Vec<u8>>,
-		HashMap<Vec<u8>, HashMap<Vec<u8>, Vec<u8>>>,
+		BTreeMap<Vec<u8>, Vec<u8>>,
+		HashMap<Vec<u8>, BTreeMap<Vec<u8>, Vec<u8>>>,
 )> for InMemory<H> {
 	fn from(inners: (
-		HashMap<Vec<u8>, Vec<u8>>,
-		HashMap<Vec<u8>, HashMap<Vec<u8>, Vec<u8>>>,
+		BTreeMap<Vec<u8>, Vec<u8>>,
+		HashMap<Vec<u8>, BTreeMap<Vec<u8>, Vec<u8>>>,
 	)) -> Self {
-		let mut inner: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>>
+		let mut inner: HashMap<Option<Vec<u8>>, BTreeMap<Vec<u8>, Vec<u8>>>
 			= inners.1.into_iter().map(|(k, v)| (Some(k), v)).collect();
 		inner.insert(None, inners.0);
 		InMemory {
@@ -332,8 +350,8 @@ impl<H: Hasher> From<(
 	}
 }
 
-impl<H: Hasher> From<HashMap<Vec<u8>, Vec<u8>>> for InMemory<H> {
-	fn from(inner: HashMap<Vec<u8>, Vec<u8>>) -> Self {
+impl<H: Hasher> From<BTreeMap<Vec<u8>, Vec<u8>>> for InMemory<H> {
+	fn from(inner: BTreeMap<Vec<u8>, Vec<u8>>) -> Self {
 		let mut expanded = HashMap::new();
 		expanded.insert(None, inner);
 		InMemory {
@@ -346,7 +364,7 @@ impl<H: Hasher> From<HashMap<Vec<u8>, Vec<u8>>> for InMemory<H> {
 
 impl<H: Hasher> From<Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>> for InMemory<H> {
 	fn from(inner: Vec<(Option<Vec<u8>>, Vec<u8>, Option<Vec<u8>>)>) -> Self {
-		let mut expanded: HashMap<Option<Vec<u8>>, HashMap<Vec<u8>, Vec<u8>>> = HashMap::new();
+		let mut expanded: HashMap<Option<Vec<u8>>, BTreeMap<Vec<u8>, Vec<u8>>> = HashMap::new();
 		for (child_key, key, value) in inner {
 			if let Some(value) = value {
 				expanded.entry(child_key).or_default().insert(key, value);
@@ -380,6 +398,22 @@ impl<H: Hasher> Backend<H> for InMemory<H> where H::Out: Codec {
 		Ok(self.inner.get(&None).map(|map| map.get(key).is_some()).unwrap_or(false))
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
+		let next_key = self.inner.get(&None)
+			.and_then(|map| map.range::<[u8], _>(range).next().map(|(k, _)| k).cloned());
+
+		Ok(next_key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		let range = (ops::Bound::Excluded(key), ops::Bound::Unbounded);
+		let next_key = self.inner.get(&Some(storage_key.to_vec()))
+			.and_then(|map| map.range::<[u8], _>(range).next().map(|(k, _)| k).cloned());
+
+		Ok(next_key)
+	}
+
 	fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
 		self.inner.get(&None).map(|map| map.keys().filter(|key| key.starts_with(prefix)).map(|k| &**k).for_each(f));
 	}

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

@@ -16,7 +16,9 @@
 
 //! Basic implementation for Externalities.
 
-use std::{collections::HashMap, any::{TypeId, Any}, iter::FromIterator};
+use std::{
+	collections::{HashMap, BTreeMap}, any::{TypeId, Any}, iter::FromIterator, mem, ops::Bound
+};
 use crate::backend::{Backend, InMemory};
 use hash_db::Hasher;
 use trie::{TrieConfiguration, default_child_trie_root};
@@ -54,8 +56,8 @@ impl BasicExternalities {
 
 	/// Consume self and returns inner storages
 	pub fn into_storages(self) -> (
-		HashMap<Vec<u8>, Vec<u8>>,
-		HashMap<Vec<u8>, HashMap<Vec<u8>, Vec<u8>>>,
+		BTreeMap<Vec<u8>, Vec<u8>>,
+		HashMap<Vec<u8>, BTreeMap<Vec<u8>, Vec<u8>>>,
 	) {
 		(self.top, self.children)
 	}
@@ -68,8 +70,8 @@ impl BasicExternalities {
 		f: impl FnOnce() -> R,
 	) -> R {
 		let mut ext = Self {
-			top: storage.0.drain().collect(),
-			children: storage.1.drain().collect(),
+			top: mem::replace(&mut storage.0, BTreeMap::default()),
+			children: mem::replace(&mut storage.1, HashMap::default()),
 		};
 
 		let r = ext.execute_with(f);
@@ -105,8 +107,8 @@ impl Default for BasicExternalities {
 	fn default() -> Self { Self::new(Default::default(), Default::default()) }
 }
 
-impl From<HashMap<Vec<u8>, Vec<u8>>> for BasicExternalities {
-	fn from(hashmap: HashMap<Vec<u8>, Vec<u8>>) -> Self {
+impl From<BTreeMap<Vec<u8>, Vec<u8>>> for BasicExternalities {
+	fn from(hashmap: BTreeMap<Vec<u8>, Vec<u8>>) -> Self {
 		BasicExternalities {
 			top: hashmap,
 			children: Default::default(),
@@ -151,6 +153,17 @@ impl Externalities for BasicExternalities {
 		Externalities::child_storage(self, storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Option<Vec<u8>> {
+		let range = (Bound::Excluded(key), Bound::Unbounded);
+		self.top.range::<[u8], _>(range).next().map(|(k, _)| k).cloned()
+	}
+
+	fn next_child_storage_key(&self, storage_key: ChildStorageKey, key: &[u8]) -> Option<Vec<u8>> {
+		let range = (Bound::Excluded(key), Bound::Unbounded);
+		self.children.get(storage_key.as_ref())
+			.and_then(|child| child.range::<[u8], _>(range).next().map(|(k, _)| k).cloned())
+	}
+
 	fn place_storage(&mut self, key: Vec<u8>, maybe_value: Option<Vec<u8>>) {
 		if is_child_storage_key(&key) {
 			warn!(target: "trie", "Refuse to set child storage key via main storage");
@@ -190,12 +203,28 @@ impl Externalities for BasicExternalities {
 			return;
 		}
 
-		self.top.retain(|key, _| !key.starts_with(prefix));
+		let to_remove = self.top.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
+			.map(|(k, _)| k)
+			.take_while(|k| k.starts_with(prefix))
+			.cloned()
+			.collect::<Vec<_>>();
+
+		for key in to_remove {
+			self.top.remove(&key);
+		}
 	}
 
 	fn clear_child_prefix(&mut self, storage_key: ChildStorageKey, prefix: &[u8]) {
 		if let Some(child) = self.children.get_mut(storage_key.as_ref()) {
-			child.retain(|key, _| !key.starts_with(prefix));
+			let to_remove = child.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
+				.map(|(k, _)| k)
+				.take_while(|k| k.starts_with(prefix))
+				.cloned()
+				.collect::<Vec<_>>();
+
+			for key in to_remove {
+				child.remove(&key);
+			}
 		}
 	}
 

substrate/primitives/state-machine/src/changes_trie/build.rs

@@ -352,7 +352,7 @@ mod test {
 			(vec![103], vec![255]),
 			(vec![104], vec![255]),
 			(vec![105], vec![255]),
-		].into_iter().collect::<::std::collections::HashMap<_, _>>().into();
+		].into_iter().collect::<::std::collections::BTreeMap<_, _>>().into();
 		let child_trie_key1 = b"1".to_vec();
 		let child_trie_key2 = b"2".to_vec();
 		let storage = InMemoryStorage::with_inputs(vec![

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

@@ -336,6 +336,40 @@ where
 		result
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Option<Vec<u8>> {
+		let next_backend_key = self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
+		let next_overlay_key_change = self.overlay.next_storage_key_change(key);
+
+		match (next_backend_key, next_overlay_key_change) {
+			(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
+			(backend_key, None) => backend_key,
+			(_, Some(overlay_key)) => if overlay_key.1.value.is_some() {
+				Some(overlay_key.0.to_vec())
+			} else {
+				self.next_storage_key(&overlay_key.0[..])
+			},
+		}
+	}
+
+	fn next_child_storage_key(&self, storage_key: ChildStorageKey, key: &[u8]) -> Option<Vec<u8>> {
+		let next_backend_key = self.backend.next_child_storage_key(storage_key.as_ref(), key)
+			.expect(EXT_NOT_ALLOWED_TO_FAIL);
+		let next_overlay_key_change = self.overlay.next_child_storage_key_change(
+			storage_key.as_ref(),
+			key
+		);
+
+		match (next_backend_key, next_overlay_key_change) {
+			(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
+			(backend_key, None) => backend_key,
+			(_, Some(overlay_key)) => if overlay_key.1.value.is_some() {
+				Some(overlay_key.0.to_vec())
+			} else {
+				self.next_child_storage_key(storage_key, &overlay_key.0[..])
+			},
+		}
+	}
+
 	fn place_storage(&mut self, key: Vec<u8>, value: Option<Vec<u8>>) {
 		trace!(target: "state-trace", "{:04x}: Put {}={:?}",
 			self.id,
@@ -619,4 +653,71 @@ mod tests {
 			Some(hex!("96f5aae4690e7302737b6f9b7f8567d5bbb9eac1c315f80101235a92d9ec27f4").to_vec()),
 		);
 	}
+
+	#[test]
+	fn next_storage_key_works() {
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_storage(vec![20], None);
+		overlay.set_storage(vec![30], Some(vec![31]));
+		let backend = vec![
+			(None, vec![10], Some(vec![10])),
+			(None, vec![20], Some(vec![20])),
+			(None, vec![40], Some(vec![40])),
+		].into();
+
+		let ext = TestExt::new(&mut overlay, &backend, None, None);
+
+		// next_backend < next_overlay
+		assert_eq!(ext.next_storage_key(&[5]), Some(vec![10]));
+
+		// next_backend == next_overlay but next_overlay is a delete
+		assert_eq!(ext.next_storage_key(&[10]), Some(vec![30]));
+
+		// next_overlay < next_backend
+		assert_eq!(ext.next_storage_key(&[20]), Some(vec![30]));
+
+		// next_backend exist but next_overlay doesn't exist
+		assert_eq!(ext.next_storage_key(&[30]), Some(vec![40]));
+
+		drop(ext);
+		overlay.set_storage(vec![50], Some(vec![50]));
+		let ext = TestExt::new(&mut overlay, &backend, None, None);
+
+		// next_overlay exist but next_backend doesn't exist
+		assert_eq!(ext.next_storage_key(&[40]), Some(vec![50]));
+	}
+
+	#[test]
+	fn next_child_storage_key_works() {
+		let child = || ChildStorageKey::from_slice(b":child_storage:default:Child1").unwrap();
+		let mut overlay = OverlayedChanges::default();
+		overlay.set_child_storage(child().as_ref().to_vec(), vec![20], None);
+		overlay.set_child_storage(child().as_ref().to_vec(), vec![30], Some(vec![31]));
+		let backend = vec![
+			(Some(child().as_ref().to_vec()), vec![10], Some(vec![10])),
+			(Some(child().as_ref().to_vec()), vec![20], Some(vec![20])),
+			(Some(child().as_ref().to_vec()), vec![40], Some(vec![40])),
+		].into();
+
+		let ext = TestExt::new(&mut overlay, &backend, None, None);
+
+		// next_backend < next_overlay
+		assert_eq!(ext.next_child_storage_key(child(), &[5]), Some(vec![10]));
+
+		// next_backend == next_overlay but next_overlay is a delete
+		assert_eq!(ext.next_child_storage_key(child(), &[10]), Some(vec![30]));
+
+		// next_overlay < next_backend
+		assert_eq!(ext.next_child_storage_key(child(), &[20]), Some(vec![30]));
+
+		// next_backend exist but next_overlay doesn't exist
+		assert_eq!(ext.next_child_storage_key(child(), &[30]), Some(vec![40]));
+
+		drop(ext);
+		overlay.set_child_storage(child().as_ref().to_vec(), vec![50], Some(vec![50]));
+		let ext = TestExt::new(&mut overlay, &backend, None, None);
+
+		// next_overlay exist but next_backend doesn't exist
+		assert_eq!(ext.next_child_storage_key(child(), &[40]), Some(vec![50]));
+	}
 }

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

@@ -731,7 +731,7 @@ fn try_read_overlay_value<H, B>(
 
 #[cfg(test)]
 mod tests {
-	use std::collections::HashMap;
+	use std::collections::BTreeMap;
 	use codec::Encode;
 	use overlayed_changes::OverlayedValue;
 	use super::*;
@@ -921,7 +921,7 @@ mod tests {
 
 	#[test]
 	fn clear_prefix_in_ext_works() {
-		let initial: HashMap<_, _> = map![
+		let initial: BTreeMap<_, _> = map![
 			b"aaa".to_vec() => b"0".to_vec(),
 			b"abb".to_vec() => b"1".to_vec(),
 			b"abc".to_vec() => b"2".to_vec(),

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]));
+	}
 }

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

@@ -272,6 +272,14 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		self.0.child_storage(storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.0.next_storage_key(key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.0.next_child_storage_key(storage_key, key)
+	}
+
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {
 		self.0.for_keys_in_child_storage(storage_key, f)
 	}

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

@@ -16,7 +16,7 @@
 
 //! Test implementation for Externalities.
 
-use std::{collections::HashMap, any::{Any, TypeId}};
+use std::{collections::{HashMap, BTreeMap}, any::{Any, TypeId}};
 use hash_db::Hasher;
 use crate::{
 	backend::{InMemory, Backend}, OverlayedChanges,
@@ -35,7 +35,7 @@ use primitives::{
 use codec::Encode;
 use externalities::{Extensions, Extension};
 
-type StorageTuple = (HashMap<Vec<u8>, Vec<u8>>, HashMap<Vec<u8>, HashMap<Vec<u8>, Vec<u8>>>);
+type StorageTuple = (BTreeMap<Vec<u8>, Vec<u8>>, HashMap<Vec<u8>, BTreeMap<Vec<u8>, Vec<u8>>>);
 
 /// Simple HashMap-based Externalities impl.
 pub struct TestExternalities<H: Hasher<Out=H256>=Blake2Hasher, N: ChangesTrieBlockNumber=u64> {

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

@@ -79,6 +79,14 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 		self.essence.child_storage(storage_key, key)
 	}
 
+	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.essence.next_storage_key(key)
+	}
+
+	fn next_child_storage_key(&self, storage_key: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
+		self.essence.next_child_storage_key(storage_key, key)
+	}
+
 	fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], f: F) {
 		self.essence.for_keys_with_prefix(prefix, f)
 	}

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

@@ -64,6 +64,76 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 		self.storage
 	}
 
+	/// Return the next key in the trie i.e. the minimum key that is strictly superior to `key` in
+	/// lexicographic order.
+	pub fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
+		self.next_storage_key_from_root(&self.root, key)
+	}
+
+	/// Return the next key in the child trie i.e. the minimum key that is strictly superior to
+	/// `key` in lexicographic order.
+	pub fn next_child_storage_key(
+		&self,
+		storage_key: &[u8],
+		key: &[u8],
+	) -> Result<Option<Vec<u8>>, String> {
+		let child_root = match self.storage(storage_key)? {
+			Some(child_root) => child_root,
+			None => return Ok(None),
+		};
+
+		let mut hash = H::Out::default();
+
+		if child_root.len() != hash.as_ref().len() {
+			return Err(format!("Invalid child storage hash at {:?}", storage_key));
+		}
+		// note: child_root and hash must be same size, panics otherwise.
+		hash.as_mut().copy_from_slice(&child_root[..]);
+
+		self.next_storage_key_from_root(&hash, key)
+	}
+
+	/// Return next key from main trie or child trie by providing corresponding root.
+	fn next_storage_key_from_root(
+		&self,
+		root: &H::Out,
+		key: &[u8],
+	) -> Result<Option<Vec<u8>>, String> {
+		let mut read_overlay = S::Overlay::default();
+		let eph = Ephemeral {
+			storage: &self.storage,
+			overlay: &mut read_overlay,
+		};
+
+		let trie = TrieDB::<H>::new(&eph, root)
+			.map_err(|e| format!("TrieDB creation error: {}", e))?;
+		let mut iter = trie.iter()
+			.map_err(|e| format!("TrieDB iteration error: {}", e))?;
+
+		// The key just after the one given in input, basically `key++0`.
+		// Note: We are sure this is the next key if:
+		// * size of key has no limit (i.e. we can always add 0 to the path),
+		// * and no keys can be inserted between `key` and `key++0` (this is ensured by sr-io).
+		let mut potential_next_key = Vec::with_capacity(key.len() + 1);
+		potential_next_key.extend_from_slice(key);
+		potential_next_key.push(0);
+
+		iter.seek(&potential_next_key)
+			.map_err(|e| format!("TrieDB iterator seek error: {}", e))?;
+
+		let next_element = iter.next();
+
+		let next_key = if let Some(next_element) = next_element {
+			let (next_key, _) = next_element
+				.map_err(|e| format!("TrieDB iterator next error: {}", e))?;
+			Some(next_key)
+		} else {
+			None
+		};
+
+		Ok(next_key)
+	}
+
 	/// Get the value of storage at given key.
 	pub fn storage(&self, key: &[u8]) -> Result<Option<Vec<u8>>, String> {
 		let mut read_overlay = S::Overlay::default();
@@ -345,3 +415,47 @@ impl<H: Hasher> TrieBackendStorage<H> for MemoryDB<H> {
 		Ok(hash_db::HashDB::get(self, key, prefix))
 	}
 }
+
+#[cfg(test)]
+mod test {
+	use primitives::{Blake2Hasher, H256};
+	use trie::{TrieMut, PrefixedMemoryDB, trie_types::TrieDBMut};
+	use super::*;
+
+	#[test]
+	fn next_storage_key_and_next_child_storage_key_work() {
+		// Contains values
+		let mut root_1 = H256::default();
+		// Contains child trie
+		let mut root_2 = H256::default();
+
+		let mut mdb = PrefixedMemoryDB::<Blake2Hasher>::default();
+		{
+			let mut trie = TrieDBMut::new(&mut mdb, &mut root_1);
+			trie.insert(b"3", &[1]).expect("insert failed");
+			trie.insert(b"4", &[1]).expect("insert failed");
+			trie.insert(b"6", &[1]).expect("insert failed");
+		}
+		{
+			let mut trie = TrieDBMut::new(&mut mdb, &mut root_2);
+			trie.insert(b"MyChild", root_1.as_ref()).expect("insert failed");
+		};
+
+		let essence_1 = TrieBackendEssence::new(mdb, root_1);
+
+		assert_eq!(essence_1.next_storage_key(b"2"), Ok(Some(b"3".to_vec())));
+		assert_eq!(essence_1.next_storage_key(b"3"), Ok(Some(b"4".to_vec())));
+		assert_eq!(essence_1.next_storage_key(b"4"), Ok(Some(b"6".to_vec())));
+		assert_eq!(essence_1.next_storage_key(b"5"), Ok(Some(b"6".to_vec())));
+		assert_eq!(essence_1.next_storage_key(b"6"), Ok(None));
+
+		let mdb = essence_1.into_storage();
+		let essence_2 = TrieBackendEssence::new(mdb, root_2);
+
+		assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"2"), Ok(Some(b"3".to_vec())));
+		assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"3"), Ok(Some(b"4".to_vec())));
+		assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"4"), Ok(Some(b"6".to_vec())));
+		assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"5"), Ok(Some(b"6".to_vec())));
+		assert_eq!(essence_2.next_child_storage_key(b"MyChild", b"6"), Ok(None));
+	}
+}