Safe and sane multi-item storage removal (#11490)

* Fix overlay prefix removal result

* Second part of the overlay prefix removal fix.

* Report only items deleted from storage in clear_prefix

* Fix kill_prefix

* Formatting

* Remove unused code

* Fixes

* Fixes

* Introduce clear_prefix host function v3

* Formatting

* Use v2 for now

* Fixes

* Formatting

* Docs

* Child prefix removal should also hide v3 for now

* Fixes

* Fixes

* Formatting

* Fixes

* apply_to_keys_whle takes start_at

* apply_to_keys_whle takes start_at

* apply_to_keys_whle takes start_at

* Cursor API; force limits

* Use unsafe deprecated functions

* Formatting

* Fixes

* Grumbles

* Fixes

* Docs

* Some nitpicks 🙈

* Update primitives/externalities/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Formatting

* Fixes

* cargo fmt

* Fixes

* Update primitives/io/src/lib.rs

Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

* Formatting

* Fixes

Co-authored-by: Bastian Köcher <info@kchr.de>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Keith Yeung <kungfukeith11@gmail.com>

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

@@ -112,6 +112,7 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	);
 

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

@@ -29,7 +29,7 @@ use sp_core::{
 	traits::Externalities,
 	Blake2Hasher,
 };
-use sp_externalities::{Extension, Extensions};
+use sp_externalities::{Extension, Extensions, MultiRemovalResults};
 use sp_trie::{empty_child_trie_root, HashKey, LayoutV0, LayoutV1, TrieConfiguration};
 use std::{
 	any::{Any, TypeId},
@@ -209,23 +209,34 @@ impl Externalities for BasicExternalities {
 		}
 	}
 
-	fn kill_child_storage(&mut self, child_info: &ChildInfo, _limit: Option<u32>) -> (bool, u32) {
-		let num_removed = self
+	fn kill_child_storage(
+		&mut self,
+		child_info: &ChildInfo,
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
+		let count = self
 			.inner
 			.children_default
 			.remove(child_info.storage_key())
 			.map(|c| c.data.len())
-			.unwrap_or(0);
-		(true, num_removed as u32)
+			.unwrap_or(0) as u32;
+		MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 	}
 
-	fn clear_prefix(&mut self, prefix: &[u8], _limit: Option<u32>) -> (bool, u32) {
+	fn clear_prefix(
+		&mut self,
+		prefix: &[u8],
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		if is_child_storage_key(prefix) {
 			warn!(
 				target: "trie",
 				"Refuse to clear prefix that is part of child storage key via main storage"
 			);
-			return (false, 0)
+			let maybe_cursor = Some(prefix.to_vec());
+			return MultiRemovalResults { maybe_cursor, backend: 0, unique: 0, loops: 0 }
 		}
 
 		let to_remove = self
@@ -237,19 +248,20 @@ impl Externalities for BasicExternalities {
 			.cloned()
 			.collect::<Vec<_>>();
 
-		let num_removed = to_remove.len();
+		let count = to_remove.len() as u32;
 		for key in to_remove {
 			self.inner.top.remove(&key);
 		}
-		(true, num_removed as u32)
+		MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 	}
 
 	fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
-		_limit: Option<u32>,
-	) -> (bool, u32) {
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		if let Some(child) = self.inner.children_default.get_mut(child_info.storage_key()) {
 			let to_remove = child
 				.data
@@ -259,13 +271,13 @@ impl Externalities for BasicExternalities {
 				.cloned()
 				.collect::<Vec<_>>();
 
-			let num_removed = to_remove.len();
+			let count = to_remove.len() as u32;
 			for key in to_remove {
 				child.data.remove(&key);
 			}
-			(true, num_removed as u32)
+			MultiRemovalResults { maybe_cursor: None, backend: count, unique: count, loops: count }
 		} else {
-			(true, 0)
+			MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
 		}
 	}
 
@@ -434,7 +446,7 @@ mod tests {
 		ext.clear_child_storage(child_info, b"dog");
 		assert_eq!(ext.child_storage(child_info, b"dog"), None);
 
-		ext.kill_child_storage(child_info, None);
+		let _ = ext.kill_child_storage(child_info, None, None);
 		assert_eq!(ext.child_storage(child_info, b"doe"), None);
 	}
 
@@ -456,8 +468,8 @@ mod tests {
 			],
 		});
 
-		let res = ext.kill_child_storage(child_info, None);
-		assert_eq!(res, (true, 3));
+		let res = ext.kill_child_storage(child_info, None, None);
+		assert_eq!(res.deconstruct(), (None, 3, 3, 3));
 	}
 
 	#[test]

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

@@ -27,7 +27,7 @@ use sp_core::hexdisplay::HexDisplay;
 use sp_core::storage::{
 	well_known_keys::is_child_storage_key, ChildInfo, StateVersion, TrackedStorageKey,
 };
-use sp_externalities::{Extension, ExtensionStore, Externalities};
+use sp_externalities::{Extension, ExtensionStore, Externalities, MultiRemovalResults};
 use sp_trie::{empty_child_trie_root, LayoutV1};
 
 use crate::{log_error, trace, warn, StorageTransactionCache};
@@ -433,7 +433,12 @@ where
 		self.overlay.set_child_storage(child_info, key, value);
 	}
 
-	fn kill_child_storage(&mut self, child_info: &ChildInfo, limit: Option<u32>) -> (bool, u32) {
+	fn kill_child_storage(
+		&mut self,
+		child_info: &ChildInfo,
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ChildKill",
@@ -442,11 +447,18 @@ where
 		);
 		let _guard = guard();
 		self.mark_dirty();
-		self.overlay.clear_child_storage(child_info);
-		self.limit_remove_from_backend(Some(child_info), None, limit)
+		let overlay = self.overlay.clear_child_storage(child_info);
+		let (maybe_cursor, backend, loops) =
+			self.limit_remove_from_backend(Some(child_info), None, maybe_limit, maybe_cursor);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}
 
-	fn clear_prefix(&mut self, prefix: &[u8], limit: Option<u32>) -> (bool, u32) {
+	fn clear_prefix(
+		&mut self,
+		prefix: &[u8],
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ClearPrefix",
@@ -460,20 +472,23 @@ where
 				target: "trie",
 				"Refuse to directly clear prefix that is part or contains of child storage key",
 			);
-			return (false, 0)
+			return MultiRemovalResults { maybe_cursor: None, backend: 0, unique: 0, loops: 0 }
 		}
 
 		self.mark_dirty();
-		self.overlay.clear_prefix(prefix);
-		self.limit_remove_from_backend(None, Some(prefix), limit)
+		let overlay = self.overlay.clear_prefix(prefix);
+		let (maybe_cursor, backend, loops) =
+			self.limit_remove_from_backend(None, Some(prefix), maybe_limit, maybe_cursor);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}
 
 	fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
-		limit: Option<u32>,
-	) -> (bool, u32) {
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		trace!(
 			target: "state",
 			method = "ChildClearPrefix",
@@ -484,8 +499,14 @@ where
 		let _guard = guard();
 
 		self.mark_dirty();
-		self.overlay.clear_child_prefix(child_info, prefix);
-		self.limit_remove_from_backend(Some(child_info), Some(prefix), limit)
+		let overlay = self.overlay.clear_child_prefix(child_info, prefix);
+		let (maybe_cursor, backend, loops) = self.limit_remove_from_backend(
+			Some(child_info),
+			Some(prefix),
+			maybe_limit,
+			maybe_cursor,
+		);
+		MultiRemovalResults { maybe_cursor, backend, unique: overlay + backend, loops }
 	}
 
 	fn storage_append(&mut self, key: Vec<u8>, value: Vec<u8>) {
@@ -728,45 +749,37 @@ where
 {
 	fn limit_remove_from_backend(
 		&mut self,
-		child_info: Option<&ChildInfo>,
-		prefix: Option<&[u8]>,
-		limit: Option<u32>,
-	) -> (bool, u32) {
-		let mut num_deleted: u32 = 0;
-
-		if let Some(limit) = limit {
-			let mut all_deleted = true;
-			self.backend.apply_to_keys_while(child_info, prefix, |key| {
-				if num_deleted == limit {
-					all_deleted = false;
+		maybe_child: Option<&ChildInfo>,
+		maybe_prefix: Option<&[u8]>,
+		maybe_limit: Option<u32>,
+		maybe_cursor: Option<&[u8]>,
+	) -> (Option<Vec<u8>>, u32, u32) {
+		let mut delete_count: u32 = 0;
+		let mut loop_count: u32 = 0;
+		let mut maybe_next_key = None;
+		self.backend
+			.apply_to_keys_while(maybe_child, maybe_prefix, maybe_cursor, |key| {
+				if maybe_limit.map_or(false, |limit| loop_count == limit) {
+					maybe_next_key = Some(key.to_vec());
 					return false
 				}
-				if let Some(num) = num_deleted.checked_add(1) {
-					num_deleted = num;
-				} else {
-					all_deleted = false;
-					return false
-				}
-				if let Some(child_info) = child_info {
-					self.overlay.set_child_storage(child_info, key.to_vec(), None);
-				} else {
-					self.overlay.set_storage(key.to_vec(), None);
+				let overlay = match maybe_child {
+					Some(child_info) => self.overlay.child_storage(child_info, key),
+					None => self.overlay.storage(key),
+				};
+				if !matches!(overlay, Some(None)) {
+					// not pending deletion from the backend - delete it.
+					if let Some(child_info) = maybe_child {
+						self.overlay.set_child_storage(child_info, key.to_vec(), None);
+					} else {
+						self.overlay.set_storage(key.to_vec(), None);
+					}
+					delete_count = delete_count.saturating_add(1);
 				}
+				loop_count = loop_count.saturating_add(1);
 				true
 			});
-			(all_deleted, num_deleted)
-		} else {
-			self.backend.apply_to_keys_while(child_info, prefix, |key| {
-				num_deleted = num_deleted.saturating_add(1);
-				if let Some(child_info) = child_info {
-					self.overlay.set_child_storage(child_info, key.to_vec(), None);
-				} else {
-					self.overlay.set_storage(key.to_vec(), None);
-				}
-				true
-			});
-			(true, num_deleted)
-		}
+		(maybe_next_key, delete_count, loop_count)
 	}
 }
 
@@ -1085,14 +1098,14 @@ mod tests {
 		not_under_prefix.extend(b"path");
 		ext.set_storage(not_under_prefix.clone(), vec![10]);
 
-		ext.clear_prefix(&[], None);
-		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None);
+		let _ = ext.clear_prefix(&[], None, None);
+		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
 		let mut under_prefix = well_known_keys::CHILD_STORAGE_KEY_PREFIX.to_vec();
 		under_prefix.extend(b"path");
-		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None);
+		let _ = ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], None, None);
 		assert_eq!(ext.child_storage(child_info, &[30]), Some(vec![40]));
 		assert_eq!(ext.storage(not_under_prefix.as_slice()), Some(vec![10]));
-		ext.clear_prefix(&not_under_prefix[..5], None);
+		let _ = ext.clear_prefix(&not_under_prefix[..5], None, None);
 		assert_eq!(ext.storage(not_under_prefix.as_slice()), None);
 	}
 

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

@@ -1348,6 +1348,7 @@ mod execution {
 mod tests {
 	use super::{ext::Ext, *};
 	use crate::{execution::CallResult, in_memory_backend::new_in_mem_hash_key};
+	use assert_matches::assert_matches;
 	use codec::{Decode, Encode};
 	use sp_core::{
 		map,
@@ -1572,7 +1573,7 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
-			ext.clear_prefix(b"ab", None);
+			let _ = ext.clear_prefix(b"ab", None, None);
 		}
 		overlay.commit_transaction().unwrap();
 
@@ -1596,7 +1597,10 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, backend, None);
-			assert_eq!((false, 1), ext.clear_prefix(b"ab", Some(1)));
+			assert_matches!(
+				ext.clear_prefix(b"ab", Some(1), None).deconstruct(),
+				(Some(_), 1, 3, 1)
+			);
 		}
 		overlay.commit_transaction().unwrap();
 
@@ -1638,7 +1642,8 @@ mod tests {
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
-			assert_eq!(ext.kill_child_storage(&child_info, Some(2)), (false, 2));
+			let r = ext.kill_child_storage(&child_info, Some(2), None);
+			assert_matches!(r.deconstruct(), (Some(_), 2, 6, 2));
 		}
 
 		assert_eq!(
@@ -1673,14 +1678,37 @@ mod tests {
 		let mut overlay = OverlayedChanges::default();
 		let mut cache = StorageTransactionCache::default();
 		let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
-		assert_eq!(ext.kill_child_storage(&child_info, Some(0)), (false, 0));
-		assert_eq!(ext.kill_child_storage(&child_info, Some(1)), (false, 1));
-		assert_eq!(ext.kill_child_storage(&child_info, Some(2)), (false, 2));
-		assert_eq!(ext.kill_child_storage(&child_info, Some(3)), (false, 3));
-		assert_eq!(ext.kill_child_storage(&child_info, Some(4)), (true, 4));
-		// Only 4 items to remove
-		assert_eq!(ext.kill_child_storage(&child_info, Some(5)), (true, 4));
-		assert_eq!(ext.kill_child_storage(&child_info, None), (true, 4));
+		let r = ext.kill_child_storage(&child_info, Some(0), None).deconstruct();
+		assert_matches!(r, (Some(_), 0, 0, 0));
+		let r = ext
+			.kill_child_storage(&child_info, Some(1), r.0.as_ref().map(|x| &x[..]))
+			.deconstruct();
+		assert_matches!(r, (Some(_), 1, 1, 1));
+		let r = ext
+			.kill_child_storage(&child_info, Some(4), r.0.as_ref().map(|x| &x[..]))
+			.deconstruct();
+		// Only 3 items remaining to remove
+		assert_matches!(r, (None, 3, 3, 3));
+		let r = ext.kill_child_storage(&child_info, Some(1), None).deconstruct();
+		assert_matches!(r, (Some(_), 0, 0, 1));
+	}
+
+	#[test]
+	fn limited_child_kill_off_by_one_works_without_limit() {
+		let child_info = ChildInfo::new_default(b"sub1");
+		let initial: HashMap<_, BTreeMap<_, _>> = map![
+			Some(child_info.clone()) => map![
+				b"a".to_vec() => b"0".to_vec(),
+				b"b".to_vec() => b"1".to_vec(),
+				b"c".to_vec() => b"2".to_vec(),
+				b"d".to_vec() => b"3".to_vec()
+			],
+		];
+		let backend = InMemoryBackend::<BlakeTwo256>::from((initial, StateVersion::default()));
+		let mut overlay = OverlayedChanges::default();
+		let mut cache = StorageTransactionCache::default();
+		let mut ext = Ext::new(&mut overlay, &mut cache, &backend, None);
+		assert_eq!(ext.kill_child_storage(&child_info, None, None).deconstruct(), (None, 4, 4, 4));
 	}
 
 	#[test]
@@ -1695,7 +1723,7 @@ mod tests {
 
 		ext.set_child_storage(child_info, b"abc".to_vec(), b"def".to_vec());
 		assert_eq!(ext.child_storage(child_info, b"abc"), Some(b"def".to_vec()));
-		ext.kill_child_storage(child_info, None);
+		let _ = ext.kill_child_storage(child_info, None, None);
 		assert_eq!(ext.child_storage(child_info, b"abc"), None);
 	}
 

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

@@ -410,10 +410,15 @@ impl OverlayedChangeSet {
 		&mut self,
 		predicate: impl Fn(&[u8], &OverlayedValue) -> bool,
 		at_extrinsic: Option<u32>,
-	) {
+	) -> u32 {
+		let mut count = 0;
 		for (key, val) in self.changes.iter_mut().filter(|(k, v)| predicate(k, v)) {
+			if val.value_ref().is_some() {
+				count += 1;
+			}
 			val.set(None, insert_dirty(&mut self.dirty_keys, key.clone()), at_extrinsic);
 		}
+		count
 	}
 
 	/// Get the iterator over all changes that follow the supplied `key`.

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

@@ -299,7 +299,7 @@ impl OverlayedChanges {
 	/// Clear child storage of given storage key.
 	///
 	/// Can be rolled back or committed when called inside a transaction.
-	pub(crate) fn clear_child_storage(&mut self, child_info: &ChildInfo) {
+	pub(crate) fn clear_child_storage(&mut self, child_info: &ChildInfo) -> u32 {
 		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
 		let top = &self.top;
@@ -309,20 +309,20 @@ impl OverlayedChanges {
 			.or_insert_with(|| (top.spawn_child(), child_info.clone()));
 		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
-		changeset.clear_where(|_, _| true, extrinsic_index);
+		changeset.clear_where(|_, _| true, extrinsic_index)
 	}
 
 	/// Removes all key-value pairs which keys share the given prefix.
 	///
 	/// Can be rolled back or committed when called inside a transaction.
-	pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) {
-		self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index());
+	pub(crate) fn clear_prefix(&mut self, prefix: &[u8]) -> u32 {
+		self.top.clear_where(|key, _| key.starts_with(prefix), self.extrinsic_index())
 	}
 
 	/// Removes all key-value pairs which keys share the given prefix.
 	///
 	/// Can be rolled back or committed when called inside a transaction
-	pub(crate) fn clear_child_prefix(&mut self, child_info: &ChildInfo, prefix: &[u8]) {
+	pub(crate) fn clear_child_prefix(&mut self, child_info: &ChildInfo, prefix: &[u8]) -> u32 {
 		let extrinsic_index = self.extrinsic_index();
 		let storage_key = child_info.storage_key().to_vec();
 		let top = &self.top;
@@ -332,7 +332,7 @@ impl OverlayedChanges {
 			.or_insert_with(|| (top.spawn_child(), child_info.clone()));
 		let updatable = info.try_update(child_info);
 		debug_assert!(updatable);
-		changeset.clear_where(|key, _| key.starts_with(prefix), extrinsic_index);
+		changeset.clear_where(|key, _| key.starts_with(prefix), extrinsic_index)
 	}
 
 	/// Returns the current nesting depth of the transaction stack.

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

@@ -288,9 +288,10 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	) {
-		self.0.apply_to_keys_while(child_info, prefix, f)
+		self.0.apply_to_keys_while(child_info, prefix, start_at, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {

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

@@ -25,6 +25,7 @@ use sp_core::{
 	traits::Externalities,
 	Blake2Hasher,
 };
+use sp_externalities::MultiRemovalResults;
 use std::{
 	any::{Any, TypeId},
 	marker::PhantomData,
@@ -124,11 +125,21 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		unimplemented!("place_child_storage not supported in ReadOnlyExternalities")
 	}
 
-	fn kill_child_storage(&mut self, _child_info: &ChildInfo, _limit: Option<u32>) -> (bool, u32) {
+	fn kill_child_storage(
+		&mut self,
+		_child_info: &ChildInfo,
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		unimplemented!("kill_child_storage is not supported in ReadOnlyExternalities")
 	}
 
-	fn clear_prefix(&mut self, _prefix: &[u8], _limit: Option<u32>) -> (bool, u32) {
+	fn clear_prefix(
+		&mut self,
+		_prefix: &[u8],
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		unimplemented!("clear_prefix is not supported in ReadOnlyExternalities")
 	}
 
@@ -136,8 +147,9 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		&mut self,
 		_child_info: &ChildInfo,
 		_prefix: &[u8],
-		_limit: Option<u32>,
-	) -> (bool, u32) {
+		_maybe_limit: Option<u32>,
+		_maybe_cursor: Option<&[u8]>,
+	) -> MultiRemovalResults {
 		unimplemented!("clear_child_prefix is not supported in ReadOnlyExternalities")
 	}
 

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

@@ -381,7 +381,10 @@ mod tests {
 		{
 			let mut ext = ext.ext();
 
-			assert!(!ext.kill_child_storage(&child_info, Some(2)).0, "Should not delete all keys");
+			assert!(
+				ext.kill_child_storage(&child_info, Some(2), None).maybe_cursor.is_some(),
+				"Should not delete all keys"
+			);
 
 			assert!(ext.child_storage(&child_info, &b"doe"[..]).is_none());
 			assert!(ext.child_storage(&child_info, &b"dog"[..]).is_none());

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

@@ -123,9 +123,10 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
+		start_at: Option<&[u8]>,
 		f: F,
 	) {
-		self.essence.apply_to_keys_while(child_info, prefix, f)
+		self.essence.apply_to_keys_while(child_info, prefix, start_at, f)
 	}
 
 	fn for_child_keys_with_prefix<F: FnMut(&[u8])>(

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

@@ -267,7 +267,8 @@ where
 		&self,
 		child_info: Option<&ChildInfo>,
 		prefix: Option<&[u8]>,
-		mut f: F,
+		start_at: Option<&[u8]>,
+		f: F,
 	) {
 		let root = if let Some(child_info) = child_info.as_ref() {
 			match self.child_root(child_info) {
@@ -283,7 +284,7 @@ where
 			self.root
 		};
 
-		self.trie_iter_key_inner(&root, prefix, |k| f(k), child_info)
+		self.trie_iter_key_inner(&root, prefix, f, child_info, start_at)
 	}
 
 	/// Execute given closure for all keys starting with prefix.
@@ -311,6 +312,7 @@ where
 				true
 			},
 			Some(child_info),
+			None,
 		)
 	}
 
@@ -324,28 +326,31 @@ where
 				true
 			},
 			None,
+			None,
 		)
 	}
 
 	fn trie_iter_key_inner<F: FnMut(&[u8]) -> bool>(
 		&self,
 		root: &H::Out,
-		prefix: Option<&[u8]>,
+		maybe_prefix: Option<&[u8]>,
 		mut f: F,
 		child_info: Option<&ChildInfo>,
+		maybe_start_at: Option<&[u8]>,
 	) {
 		let mut iter = move |db| -> sp_std::result::Result<(), Box<TrieError<H::Out>>> {
 			let trie = TrieDB::<H>::new(db, root)?;
-			let iter = if let Some(prefix) = prefix.as_ref() {
-				TrieDBKeyIterator::new_prefixed(&trie, prefix)?
-			} else {
-				TrieDBKeyIterator::new(&trie)?
-			};
+			let prefix = maybe_prefix.unwrap_or(&[]);
+			let iter = match maybe_start_at {
+				Some(start_at) =>
+					TrieDBKeyIterator::new_prefixed_then_seek(&trie, prefix, start_at),
+				None => TrieDBKeyIterator::new_prefixed(&trie, prefix),
+			}?;
 
 			for x in iter {
 				let key = x?;
 
-				debug_assert!(prefix
+				debug_assert!(maybe_prefix
 					.as_ref()
 					.map(|prefix| key.starts_with(prefix))
 					.unwrap_or(true));