Frame remove_all with size limit. (#9106)

* remove prefixed content with limit.

* test match

* factor comment and factor ext limit removal.

* fix benchmark

Co-authored-by: Shawn Tabrizi <shawntabrizi@gmail.com>

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

@@ -93,11 +93,12 @@ pub trait Backend<H: Hasher>: sp_std::fmt::Debug {
 		key: &[u8]
 	) -> Result<Option<StorageKey>, Self::Error>;
 
-	/// Retrieve all entries keys of child storage and call `f` for each of those keys.
+	/// Retrieve all entries keys of storage and call `f` for each of those keys.
 	/// Aborts as soon as `f` returns false.
-	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
+	fn apply_to_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
-		child_info: &ChildInfo,
+		child_info: Option<&ChildInfo>,
+		prefix: Option<&[u8]>,
 		f: F,
 	);
 

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

@@ -216,13 +216,13 @@ impl Externalities for BasicExternalities {
 		(true, num_removed as u32)
 	}
 
-	fn clear_prefix(&mut self, prefix: &[u8]) {
+	fn clear_prefix(&mut self, prefix: &[u8], _limit: Option<u32>) -> (bool, u32) {
 		if is_child_storage_key(prefix) {
 			warn!(
 				target: "trie",
 				"Refuse to clear prefix that is part of child storage key via main storage"
 			);
-			return;
+			return (false, 0);
 		}
 
 		let to_remove = self.inner.top.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
@@ -231,16 +231,19 @@ impl Externalities for BasicExternalities {
 			.cloned()
 			.collect::<Vec<_>>();
 
+		let num_removed = to_remove.len();
 		for key in to_remove {
 			self.inner.top.remove(&key);
 		}
+		(true, num_removed as u32)
 	}
 
 	fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
-	) {
+		_limit: Option<u32>,
+	) -> (bool, u32) {
 		if let Some(child) = self.inner.children_default.get_mut(child_info.storage_key()) {
 			let to_remove = child.data.range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
 				.map(|(k, _)| k)
@@ -248,9 +251,13 @@ impl Externalities for BasicExternalities {
 				.cloned()
 				.collect::<Vec<_>>();
 
+			let num_removed = to_remove.len();
 			for key in to_remove {
 				child.data.remove(&key);
 			}
+			(true, num_removed as u32)
+		} else {
+			(true, 0)
 		}
 	}
 

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

@@ -460,36 +460,10 @@ where
 		let _guard = guard();
 		self.mark_dirty();
 		self.overlay.clear_child_storage(child_info);
-		let mut num_deleted: u32 = 0;
-
-		if let Some(limit) = limit {
-			let mut all_deleted = true;
-			self.backend.apply_to_child_keys_while(child_info, |key| {
-				if num_deleted == limit {
-					all_deleted = false;
-					return false;
-				}
-				if let Some(num) = num_deleted.checked_add(1) {
-					num_deleted = num;
-				} else {
-					all_deleted = false;
-					return false;
-				}
-				self.overlay.set_child_storage(child_info, key.to_vec(), None);
-				true
-			});
-			(all_deleted, num_deleted)
-		} else {
-			self.backend.apply_to_child_keys_while(child_info, |key| {
-				num_deleted = num_deleted.saturating_add(1);
-				self.overlay.set_child_storage(child_info, key.to_vec(), None);
-				true
-			});
-			(true, num_deleted)
-		}
+		self.limit_remove_from_backend(Some(child_info), None, limit)
 	}
 
-	fn clear_prefix(&mut self, prefix: &[u8]) {
+	fn clear_prefix(&mut self, prefix: &[u8], limit: Option<u32>) -> (bool, u32) {
 		trace!(target: "state", "{:04x}: ClearPrefix {}",
 			self.id,
 			HexDisplay::from(&prefix),
@@ -498,21 +472,20 @@ where
 
 		if sp_core::storage::well_known_keys::starts_with_child_storage_key(prefix) {
 			warn!(target: "trie", "Refuse to directly clear prefix that is part or contains of child storage key");
-			return;
+			return (false, 0);
 		}
 
 		self.mark_dirty();
 		self.overlay.clear_prefix(prefix);
-		self.backend.for_keys_with_prefix(prefix, |key| {
-			self.overlay.set_storage(key.to_vec(), None);
-		});
+		self.limit_remove_from_backend(None, Some(prefix), limit)
 	}
 
 	fn clear_child_prefix(
 		&mut self,
 		child_info: &ChildInfo,
 		prefix: &[u8],
-	) {
+		limit: Option<u32>,
+	) -> (bool, u32) {
 		trace!(target: "state", "{:04x}: ClearChildPrefix({}) {}",
 			self.id,
 			HexDisplay::from(&child_info.storage_key()),
@@ -522,9 +495,7 @@ where
 
 		self.mark_dirty();
 		self.overlay.clear_child_prefix(child_info, prefix);
-		self.backend.for_child_keys_with_prefix(child_info, prefix, |key| {
-			self.overlay.set_child_storage(child_info, key.to_vec(), None);
-		});
+		self.limit_remove_from_backend(Some(child_info), Some(prefix), limit)
 	}
 
 	fn storage_append(
@@ -780,6 +751,57 @@ where
 	}
 }
 
+impl<'a, H, N, B> Ext<'a, H, N, B>
+where
+	H: Hasher,
+	H::Out: Ord + 'static + codec::Codec,
+	B: Backend<H>,
+	N: crate::changes_trie::BlockNumber,
+{
+	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;
+					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);
+				}
+				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)
+		}
+	}
+}
+
 /// Implement `Encode` by forwarding the stored raw vec.
 struct EncodeOpaqueValue(Vec<u8>);
 
@@ -1155,14 +1177,14 @@ mod tests {
 		not_under_prefix.extend(b"path");
 		ext.set_storage(not_under_prefix.clone(), vec![10]);
 
-		ext.clear_prefix(&[]);
-		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4]);
+		ext.clear_prefix(&[], None);
+		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], 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]);
+		ext.clear_prefix(&well_known_keys::CHILD_STORAGE_KEY_PREFIX[..4], 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]);
+		ext.clear_prefix(&not_under_prefix[..5], None);
 		assert_eq!(ext.storage(not_under_prefix.as_slice()), None);
 	}
 

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

@@ -1102,6 +1102,7 @@ mod tests {
 		overlay.set_storage(b"abd".to_vec(), Some(b"69".to_vec()));
 		overlay.set_storage(b"bbd".to_vec(), Some(b"42".to_vec()));
 
+		let overlay_limit = overlay.clone();
 		{
 			let mut cache = StorageTransactionCache::default();
 			let mut ext = Ext::new(
@@ -1111,7 +1112,7 @@ mod tests {
 				changes_trie::disabled_state::<_, u64>(),
 				None,
 			);
-			ext.clear_prefix(b"ab");
+			ext.clear_prefix(b"ab", None);
 		}
 		overlay.commit_transaction().unwrap();
 
@@ -1128,6 +1129,33 @@ mod tests {
 				b"bbd".to_vec() => Some(b"42".to_vec()).into()
 			],
 		);
+
+		let mut overlay = overlay_limit;
+		{
+			let mut cache = StorageTransactionCache::default();
+			let mut ext = Ext::new(
+				&mut overlay,
+				&mut cache,
+				backend,
+				changes_trie::disabled_state::<_, u64>(),
+				None,
+			);
+			assert_eq!((false, 1), ext.clear_prefix(b"ab", Some(1)));
+		}
+		overlay.commit_transaction().unwrap();
+
+		assert_eq!(
+			overlay.changes().map(|(k, v)| (k.clone(), v.value().cloned()))
+				.collect::<HashMap<_, _>>(),
+			map![
+				b"abb".to_vec() => None.into(),
+				b"aba".to_vec() => None.into(),
+				b"abd".to_vec() => None.into(),
+
+				b"bab".to_vec() => Some(b"228".to_vec()).into(),
+				b"bbd".to_vec() => Some(b"42".to_vec()).into()
+			],
+		);
 	}
 
 	#[test]

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

@@ -260,12 +260,13 @@ impl<'a, S, H> Backend<H> for ProvingBackend<'a, S, H>
 		self.0.child_storage(child_info, key)
 	}
 
-	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
+	fn apply_to_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
-		child_info: &ChildInfo,
+		child_info: Option<&ChildInfo>,
+		prefix: Option<&[u8]>,
 		f: F,
 	) {
-		self.0.apply_to_child_keys_while(child_info, f)
+		self.0.apply_to_keys_while(child_info, prefix, f)
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {

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

@@ -136,7 +136,7 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		unimplemented!("kill_child_storage is not supported in ReadOnlyExternalities")
 	}
 
-	fn clear_prefix(&mut self, _prefix: &[u8]) {
+	fn clear_prefix(&mut self, _prefix: &[u8], _limit: Option<u32>) -> (bool, u32) {
 		unimplemented!("clear_prefix is not supported in ReadOnlyExternalities")
 	}
 
@@ -144,7 +144,8 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
 		&mut self,
 		_child_info: &ChildInfo,
 		_prefix: &[u8],
-	) {
+		_limit: Option<u32>,
+	) -> (bool, u32) {
 		unimplemented!("clear_child_prefix is not supported in ReadOnlyExternalities")
 	}
 

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

@@ -113,12 +113,13 @@ impl<S: TrieBackendStorage<H>, H: Hasher> Backend<H> for TrieBackend<S, H> where
 		self.essence.for_key_values_with_prefix(prefix, f)
 	}
 
-	fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
+	fn apply_to_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
-		child_info: &ChildInfo,
+		child_info: Option<&ChildInfo>,
+		prefix: Option<&[u8]>,
 		f: F,
 	) {
-		self.essence.apply_to_child_keys_while(child_info, f)
+		self.essence.apply_to_keys_while(child_info, prefix, f)
 	}
 
 	fn for_child_keys_with_prefix<F: FnMut(&[u8])>(

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

@@ -25,7 +25,7 @@ use crate::{warn, debug};
 use hash_db::{self, Hasher, Prefix};
 use sp_trie::{Trie, MemoryDB, PrefixedMemoryDB, DBValue,
 	empty_child_trie_root, read_trie_value, read_child_trie_value,
-	for_keys_in_child_trie, KeySpacedDB, TrieDBIterator};
+	KeySpacedDB, TrieDBIterator};
 use sp_trie::trie_types::{TrieDB, TrieError, Layout};
 use crate::{backend::Consolidate, StorageKey, StorageValue};
 use sp_core::storage::ChildInfo;
@@ -189,29 +189,30 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 			.map_err(map_e)
 	}
 
-	/// Retrieve all entries keys of child storage and call `f` for each of those keys.
+	/// Retrieve all entries keys of a storage and call `f` for each of those keys.
 	/// Aborts as soon as `f` returns false.
-	pub fn apply_to_child_keys_while<F: FnMut(&[u8]) -> bool>(
+	pub fn apply_to_keys_while<F: FnMut(&[u8]) -> bool>(
 		&self,
-		child_info: &ChildInfo,
-		f: F,
+		child_info: Option<&ChildInfo>,
+		prefix: Option<&[u8]>,
+		mut f: F,
 	) {
-		let root = match self.child_root(child_info) {
-			Ok(v) => v.unwrap_or_else(|| empty_child_trie_root::<Layout<H>>().encode()),
-			Err(e) => {
-				debug!(target: "trie", "Error while iterating child storage: {}", e);
-				return;
-			}
+		let mut child_root = H::Out::default();
+		let root = if let Some(child_info) = child_info.as_ref() {
+			let root_vec = match self.child_root(child_info) {
+				Ok(v) => v.unwrap_or_else(|| empty_child_trie_root::<Layout<H>>().encode()),
+				Err(e) => {
+					debug!(target: "trie", "Error while iterating child storage: {}", e);
+					return;
+				}
+			};
+			child_root.as_mut().copy_from_slice(&root_vec);
+			&child_root
+		} else {
+			&self.root
 		};
 
-		if let Err(e) = for_keys_in_child_trie::<Layout<H>, _, _>(
-			child_info.keyspace(),
-			self,
-			&root,
-			f,
-		) {
-			debug!(target: "trie", "Error while iterating child storage: {}", e);
-		}
+		self.trie_iter_inner(root, prefix, |k, _v| f(k), child_info)
 	}
 
 	/// Execute given closure for all keys starting with prefix.
@@ -230,30 +231,38 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 		};
 		let mut root = H::Out::default();
 		root.as_mut().copy_from_slice(&root_vec);
-		self.keys_values_with_prefix_inner(&root, prefix, |k, _v| f(k), Some(child_info))
+		self.trie_iter_inner(&root, Some(prefix), |k, _v| { f(k); true }, Some(child_info))
 	}
 
 	/// Execute given closure for all keys starting with prefix.
 	pub fn for_keys_with_prefix<F: FnMut(&[u8])>(&self, prefix: &[u8], mut f: F) {
-		self.keys_values_with_prefix_inner(&self.root, prefix, |k, _v| f(k), None)
+		self.trie_iter_inner(&self.root, Some(prefix), |k, _v| { f(k); true }, None)
 	}
 
-	fn keys_values_with_prefix_inner<F: FnMut(&[u8], &[u8])>(
+	fn trie_iter_inner<F: FnMut(&[u8], &[u8]) -> bool>(
 		&self,
 		root: &H::Out,
-		prefix: &[u8],
+		prefix: Option<&[u8]>,
 		mut f: F,
 		child_info: Option<&ChildInfo>,
 	) {
 		let mut iter = move |db| -> sp_std::result::Result<(), Box<TrieError<H::Out>>> {
 			let trie = TrieDB::<H>::new(db, root)?;
 
-			for x in TrieDBIterator::new_prefixed(&trie, prefix)? {
+			let iter = if let Some(prefix) = prefix.as_ref() {
+				TrieDBIterator::new_prefixed(&trie, prefix)?
+			} else {
+				TrieDBIterator::new(&trie)?
+			};
+
+			for x in iter {
 				let (key, value) = x?;
 
-				debug_assert!(key.starts_with(prefix));
+				debug_assert!(prefix.as_ref().map(|prefix| key.starts_with(prefix)).unwrap_or(true));
 
-				f(&key, &value);
+				if !f(&key, &value) {
+					break;
+				}
 			}
 
 			Ok(())
@@ -271,8 +280,8 @@ impl<S: TrieBackendStorage<H>, H: Hasher> TrieBackendEssence<S, H> where H::Out:
 	}
 
 	/// Execute given closure for all key and values starting with prefix.
-	pub fn for_key_values_with_prefix<F: FnMut(&[u8], &[u8])>(&self, prefix: &[u8], f: F) {
-		self.keys_values_with_prefix_inner(&self.root, prefix, f, None)
+	pub fn for_key_values_with_prefix<F: FnMut(&[u8], &[u8])>(&self, prefix: &[u8], mut f: F) {
+		self.trie_iter_inner(&self.root, Some(prefix), |k, v| { f(k, v); true }, None)
 	}
 }