impl cache in client_storage_cache (needs test)

client/db/src/storage_cache.rs

@@ -266,6 +266,11 @@ struct LocalCache<H: Hasher> {
 	///
 	/// `None` indicates that key is known to be missing.
 	child_storage: HashMap<ChildStorageKey, Option<StorageValue>>,
+	/// Next key cache.
+	///
+	/// It stores the tuple `(Key, OptionNextKey)`
+	/// with `OptionNextKey = backend.next_storage_key(Key)`.
+	next_key: Option<(Vec<u8>, Option<Vec<u8>>)>,
 }
 
 /// Cache changes.
@@ -415,6 +420,7 @@ impl<H: Hasher, S: StateBackend<H>, B: BlockT> CachingState<H, S, B> {
 					storage: Default::default(),
 					hashes: Default::default(),
 					child_storage: Default::default(),
+					next_key: Default::default(),
 				}),
 				parent_hash: parent_hash,
 			},
@@ -545,8 +551,25 @@ impl<H: Hasher, S: StateBackend<H>, B: BlockT> StateBackend<H> for CachingState<
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
-		// TODO TODO: cache
-		self.state.next_storage_key(key)
+		let local_cache = self.cache.local_cache.upgradable_read();
+
+		if let Some((cached_key, cached_next_key)) = local_cache.next_key.as_ref() {
+			let key_after_cache = cached_next_key
+				.as_ref()
+				.map(|cached_next_key| key >= &cached_next_key[..])
+				.unwrap_or(false);
+
+			let key_before_cache = key < &cached_key[..];
+
+			if !key_after_cache && !key_before_cache {
+				return Ok(cached_next_key.clone());
+			}
+		}
+
+		let next_key = self.state.next_storage_key(key)?;
+		let next_key_cache = Some((key.to_vec(), next_key.clone()));
+		RwLockUpgradableReadGuard::upgrade(local_cache).next_key = next_key_cache;
+		Ok(next_key)
 	}
 
 	fn for_keys_in_child_storage<F: FnMut(&[u8])>(&self, storage_key: &[u8], f: F) {

primitives/state-machine/src/ext.rs

@@ -331,17 +331,24 @@ where
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Option<Vec<u8>> {
-		// TODO TODO: cache next_backend_key with a cache: `Option<(Vec<u8>, Vec<u8>)>`
+		// Get next backend key
 		let next_backend_key = self.backend.next_storage_key(key).expect(EXT_NOT_ALLOWED_TO_FAIL);
+
+		// Get next change from overlay, this includes deletetion.
 		let next_overlay_key_change = self.overlay.next_storage_key_change(key);
 
 		match (next_backend_key, next_overlay_key_change) {
+			// Return backend one if strictly less than overlay one
 			(Some(backend_key), Some(overlay_key)) if &backend_key[..] < overlay_key.0 => Some(backend_key),
+			// Return backend one if no overlay one
 			(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[..])
+			// Overlay is less or equal to backend.
+			(_, Some(overlay_key)) => {
+				if overlay_key.1.value.is_some() {
+					Some(overlay_key.0.to_vec())
+				} else {
+					self.next_storage_key(&overlay_key.0[..])
+				}
 			},
 		}
 	}