Fix overlay (a bit)

primitives/io/src/lib.rs

@@ -1766,6 +1766,11 @@ mod tests {
 			assert!(storage::get(b":abdd").is_some());
 			assert!(storage::get(b":abcd").is_none());
 			assert!(storage::get(b":abc").is_none());
+
+			assert!(matches!(
+				storage::clear_prefix(b":abc", None),
+				KillStorageResult::AllRemoved(0)
+			));
 		});
 	}
 

primitives/state-machine/src/ext.rs

@@ -464,8 +464,10 @@ where
 		}
 
 		self.mark_dirty();
-		self.overlay.clear_prefix(prefix);
-		self.limit_remove_from_backend(None, Some(prefix), limit)
+		// In the future, we should report this back to the caller.
+		let _overlay_count = self.overlay.clear_prefix(prefix);
+		let (all, count) = self.limit_remove_from_backend(None, Some(prefix), limit);
+		(all, count)
 	}
 
 	fn clear_child_prefix(
@@ -484,8 +486,10 @@ 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)
+		// In the future, we should report this back to the caller.
+		let _overlay_count = self.overlay.clear_child_prefix(child_info, prefix);
+		let (all, count) = self.limit_remove_from_backend(Some(child_info), Some(prefix), limit);
+		(all, count)
 	}
 
 	fn storage_append(&mut self, key: Vec<u8>, value: Vec<u8>) {
@@ -741,6 +745,13 @@ where
 					all_deleted = false;
 					return false
 				}
+				if let Some(None) = match child_info {
+					Some(child_info) => self.overlay.child_storage(child_info, key),
+					None => self.overlay.storage(key),
+				} {
+					// already deleted.
+					return true
+				}
 				if let Some(num) = num_deleted.checked_add(1) {
 					num_deleted = num;
 				} else {
@@ -757,6 +768,13 @@ where
 			(all_deleted, num_deleted)
 		} else {
 			self.backend.apply_to_keys_while(child_info, prefix, |key| {
+				if let Some(None) = match child_info {
+					Some(child_info) => self.overlay.child_storage(child_info, key),
+					None => self.overlay.storage(key),
+				} {
+					// already deleted.
+					return true
+				}
 				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);

primitives/state-machine/src/lib.rs

@@ -1675,11 +1675,25 @@ mod tests {
 		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));
+		// Only 3 items remaining to remove
+		assert_eq!(ext.kill_child_storage(&child_info, Some(4)), (true, 3));
+	}
+
+	#[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), (true, 4));
 	}
 

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

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

@@ -315,14 +315,14 @@ impl OverlayedChanges {
 	/// 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.