[RFC-145] Refactor storage cursor allocation

cumulus/pallets/dmp-queue/src/lib.rs

@@ -246,16 +246,18 @@ pub mod pallet {
 					let hashed_prefix =
 						twox_128(<Pallet<T> as PalletInfoAccess>::name().as_bytes());
 
-					let result = frame_support::storage::unhashed::clear_prefix(
+					let mut removal_cursor = sp_io::MultiRemovalCursor::new();
+					removal_cursor.set_input(cursor.as_ref().map(|c| c.as_ref()));
+					let counters = frame_support::storage::unhashed::clear_prefix(
 						&hashed_prefix,
 						Some(2), // Somehow it does nothing when set to 1, so we set it to 2.
-						cursor.as_ref().map(|c| c.as_ref()),
+						Some(&mut removal_cursor),
 					);
-					Self::deposit_event(Event::CleanedSome { keys_removed: result.backend });
+					Self::deposit_event(Event::CleanedSome { keys_removed: counters.backend });
 
 					// GOTCHA! We deleted *all* pallet storage; hence we also our own
 					// `MigrationState`. BUT we insert it back:
-					if let Some(unbound_cursor) = result.maybe_cursor {
+					if let Some(unbound_cursor) = removal_cursor.into_inner() {
 						if let Ok(cursor) = unbound_cursor.try_into() {
 							log::debug!(target: LOG, "Next cursor: {:?}", &cursor);
 							MigrationStatus::<T>::put(MigrationState::StartedCleanup {

polkadot/runtime/common/src/crowdloan/mod.rs

@@ -705,7 +705,13 @@ impl<T: Config> Pallet<T> {
 	}
 
 	pub fn crowdloan_kill(index: FundIndex) -> child::MultiRemovalResults {
-		child::clear_storage(&Self::id_from_index(index), Some(T::RemoveKeysLimit::get()), None)
+		let mut cursor = child::MultiRemovalCursor::new();
+		child::clear_storage(
+			&Self::id_from_index(index),
+			Some(T::RemoveKeysLimit::get()),
+			Some(&mut cursor),
+		)
+		.into_results(cursor)
 	}
 
 	pub fn contribution_iterator(

substrate/frame/contracts/src/storage.rs

@@ -327,14 +327,13 @@ impl<T: Config> ContractInfo<T> {
 		while remaining_key_budget > 0 {
 			let Some(entry) = queue.next() else { break };
 
-			#[allow(deprecated)]
 			let outcome = child::clear_storage(
 				&ChildInfo::new_default(&entry.trie_id),
 				Some(remaining_key_budget),
 				None,
 			);
 
-			if outcome.maybe_cursor.is_some() {
+			if outcome.more {
 				remaining_key_budget.saturating_reduce(outcome.backend);
 				break;
 			} else {

substrate/frame/migrations/src/benchmarking.rs

@@ -234,7 +234,7 @@ mod benches {
 		// However, the benchmarking PoV results are correctly dependent on the amount of
 		// keys removed.
 
-		if result.maybe_cursor.is_none() {
+		if !result.more {
 			// All the keys removed
 			#[cfg(not(test))]
 			ensure!(result.backend == n, "Not all keys are removed");

substrate/frame/migrations/src/migrations.rs

@@ -100,7 +100,7 @@ where
 
 		meter.consume(T::WeightInfo::reset_pallet_migration(outcome.backend));
 
-		Ok(Some(outcome.maybe_cursor.is_none()))
+		Ok(Some(!outcome.more))
 	}
 
 	#[cfg(feature = "try-runtime")]

substrate/frame/revive/src/storage.rs

@@ -460,7 +460,7 @@ impl<T: Config> ContractInfo<T> {
 				None,
 			);
 
-			if outcome.maybe_cursor.is_some() {
+			if outcome.more {
 				remaining = remaining
 					.saturating_sub(weight_per_trie_key.saturating_mul(outcome.backend.into()));
 				break;

substrate/frame/support/src/migrations.rs

@@ -321,15 +321,14 @@ impl<P: Get<&'static str>, DbWeight: Get<RuntimeDbWeight>> frame_support::traits
 	fn on_runtime_upgrade() -> frame_support::weights::Weight {
 		let hashed_prefix = twox_128(P::get().as_bytes());
 		let r = clear_prefix(&hashed_prefix, None, None);
-		let keys_removed = match r.maybe_cursor {
-			Some(_) => {
-				log::error!(
-					"`clear_prefix` failed to remove all keys for {}. THIS SHOULD NEVER HAPPEN! 🚨",
-					P::get()
-				);
-				r.backend
-			},
-			None => r.backend,
+		let keys_removed = if r.more {
+			log::error!(
+				"`clear_prefix` failed to remove all keys for {}. THIS SHOULD NEVER HAPPEN! 🚨",
+				P::get()
+			);
+			r.backend
+		} else {
+			r.backend
 		} as u64;
 
 		log::info!("Removed {} {} keys 🧹", keys_removed, P::get());
@@ -429,15 +428,14 @@ impl<P: Get<&'static str>, S: Get<&'static str>, DbWeight: Get<RuntimeDbWeight>>
 	fn on_runtime_upgrade() -> frame_support::weights::Weight {
 		let hashed_prefix = storage_prefix(P::get().as_bytes(), S::get().as_bytes());
 		let r = clear_prefix(&hashed_prefix, None, None);
-		let keys_removed = match r.maybe_cursor {
-			Some(_) => {
-				log::error!(
-					"`clear_prefix` failed to remove all keys for storage `{}` from pallet `{}`. THIS SHOULD NEVER HAPPEN! 🚨",
-					S::get(), P::get()
-				);
-				r.backend
-			},
-			None => r.backend,
+		let keys_removed = if r.more {
+			log::error!(
+				"`clear_prefix` failed to remove all keys for storage `{}` from pallet `{}`. THIS SHOULD NEVER HAPPEN! 🚨",
+				S::get(), P::get()
+			);
+			r.backend
+		} else {
+			r.backend
 		} as u64;
 
 		log::info!("Removed `{}` `{}` `{}` keys 🧹", keys_removed, P::get(), S::get());

substrate/frame/support/src/storage/child.rs

@@ -24,7 +24,7 @@
 use alloc::vec::Vec;
 use codec::{Codec, Decode, Encode};
 pub use sp_core::storage::{ChildInfo, ChildType, StateVersion};
-pub use sp_io::{KillStorageResult, MultiRemovalResults};
+pub use sp_io::{KillStorageResult, MultiRemovalCounters, MultiRemovalCursor, MultiRemovalResults};
 
 /// Return the value of the item in storage under `key`, or `None` if there is no explicit entry.
 pub fn get<T: Decode + Sized>(child_info: &ChildInfo, key: &[u8]) -> Option<T> {
@@ -132,19 +132,20 @@ pub fn exists(child_info: &ChildInfo, key: &[u8]) -> bool {
 ///
 /// # Cursor
 ///
-/// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be
-/// passed once (in the initial call) for any attempt to clear storage. In general, subsequent calls
-/// operating on the same prefix should pass `Some` and this value should be equal to the
-/// previous call result's `maybe_cursor` field. The only exception to this is when you can
-/// guarantee that the subsequent call is in a new block; in this case the previous call's result
-/// cursor need not be passed in and a `None` may be passed instead. This exception may be useful
-/// then making this call solely from a block-hook such as `on_initialize`.
-
-/// Returns [`MultiRemovalResults`] to inform about the result. Once the resultant `maybe_cursor`
-/// field is `None`, then no further items remain to be deleted.
+/// To continue the operation across multiple calls, the caller owns a [`MultiRemovalCursor`] and
+/// passes `Some(&mut cursor)`: it should be freshly [created](MultiRemovalCursor::new) for the
+/// initial call and then passed back in unchanged on subsequent calls until
+/// [`MultiRemovalCounters::more`] is `false`. Reusing the same cursor object across iterations
+/// avoids re-allocating the cursor buffer on every call.
 ///
-/// NOTE: After the initial call for any given child storage, it is important that no keys further
-/// keys are inserted. If so, then they may or may not be deleted by subsequent calls.
+/// Pass `None` when you don't need to continue: no cursor is materialized (no allocation), and
+/// [`MultiRemovalCounters::more`] still reports whether keys remain.
+///
+/// Returns [`MultiRemovalCounters`] to inform about the result. Once `more` is `false`, no further
+/// items remain to be deleted.
+///
+/// NOTE: After the initial call for any given child storage, it is important that no further keys
+/// are inserted. If so, then they may or may not be deleted by subsequent calls.
 ///
 /// # Note
 ///
@@ -153,13 +154,13 @@ pub fn exists(child_info: &ChildInfo, key: &[u8]) -> bool {
 pub fn clear_storage(
 	child_info: &ChildInfo,
 	maybe_limit: Option<u32>,
-	maybe_cursor: Option<&[u8]>,
-) -> MultiRemovalResults {
+	cursor: Option<&mut MultiRemovalCursor>,
+) -> MultiRemovalCounters {
 	match child_info.child_type() {
 		ChildType::ParentKeyId => sp_io::default_child_storage::storage_kill(
 			child_info.storage_key(),
 			maybe_limit,
-			maybe_cursor,
+			cursor,
 		),
 	}
 }

substrate/frame/support/src/storage/generator/double_map.rs

@@ -210,8 +210,16 @@ where
 	where
 		KArg1: EncodeLike<K1>,
 	{
-		unhashed::clear_prefix(Self::storage_double_map_final_key1(k1).as_ref(), maybe_limit, None)
-			.into()
+		let counters = unhashed::clear_prefix(
+			Self::storage_double_map_final_key1(k1).as_ref(),
+			maybe_limit,
+			None,
+		);
+		if counters.more {
+			sp_io::KillStorageResult::SomeRemaining(counters.loops)
+		} else {
+			sp_io::KillStorageResult::AllRemoved(counters.loops)
+		}
 	}
 
 	fn clear_prefix<KArg1>(
@@ -222,12 +230,14 @@ where
 	where
 		KArg1: EncodeLike<K1>,
 	{
+		let mut cursor = sp_io::MultiRemovalCursor::new();
+		cursor.set_input(maybe_cursor);
 		unhashed::clear_prefix(
 			Self::storage_double_map_final_key1(k1).as_ref(),
 			Some(limit),
-			maybe_cursor,
+			Some(&mut cursor),
 		)
-		.into()
+		.into_results(cursor)
 	}
 
 	fn contains_prefix<KArg1>(k1: KArg1) -> bool

substrate/frame/support/src/storage/generator/nmap.rs

@@ -186,7 +186,13 @@ where
 	where
 		K: HasKeyPrefix<KP>,
 	{
-		unhashed::clear_prefix(&Self::storage_n_map_partial_key(partial_key), limit, None).into()
+		let counters =
+			unhashed::clear_prefix(&Self::storage_n_map_partial_key(partial_key), limit, None);
+		if counters.more {
+			sp_io::KillStorageResult::SomeRemaining(counters.loops)
+		} else {
+			sp_io::KillStorageResult::AllRemoved(counters.loops)
+		}
 	}
 
 	fn clear_prefix<KP>(
@@ -197,11 +203,14 @@ where
 	where
 		K: HasKeyPrefix<KP>,
 	{
+		let mut cursor = sp_io::MultiRemovalCursor::new();
+		cursor.set_input(maybe_cursor);
 		unhashed::clear_prefix(
 			&Self::storage_n_map_partial_key(partial_key),
 			Some(limit),
-			maybe_cursor,
+			Some(&mut cursor),
 		)
+		.into_results(cursor)
 	}
 
 	fn contains_prefix<KP>(partial_key: KP) -> bool

substrate/frame/support/src/storage/migration.rs

@@ -313,7 +313,10 @@ pub fn clear_storage_prefix(
 	let storage_prefix = storage_prefix(module, item);
 	key[0..32].copy_from_slice(&storage_prefix);
 	key[32..].copy_from_slice(hash);
-	frame_support::storage::unhashed::clear_prefix(&key, maybe_limit, maybe_cursor)
+	let mut cursor = sp_io::MultiRemovalCursor::new();
+	cursor.set_input(maybe_cursor);
+	frame_support::storage::unhashed::clear_prefix(&key, maybe_limit, Some(&mut cursor))
+		.into_results(cursor)
 }
 
 /// Take a particular item in storage by the `module`, the map's `item` name and the key `hash`.

substrate/frame/support/src/storage/mod.rs

@@ -1389,7 +1389,12 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	/// overlay are not taken into account when deleting keys in the backend.
 	#[deprecated = "Use `clear` instead"]
 	fn remove_all(limit: Option<u32>) -> sp_io::KillStorageResult {
-		unhashed::clear_prefix(&Self::final_prefix(), limit, None).into()
+		let counters = unhashed::clear_prefix(&Self::final_prefix(), limit, None);
+		if counters.more {
+			sp_io::KillStorageResult::SomeRemaining(counters.loops)
+		} else {
+			sp_io::KillStorageResult::AllRemoved(counters.loops)
+		}
 	}
 
 	/// Attempt to remove all items from the map.
@@ -1416,7 +1421,10 @@ pub trait StoragePrefixedMap<Value: FullCodec> {
 	/// operating on the same map should always pass `Some`, and this should be equal to the
 	/// previous call result's `maybe_cursor` field.
 	fn clear(limit: u32, maybe_cursor: Option<&[u8]>) -> sp_io::MultiRemovalResults {
-		unhashed::clear_prefix(&Self::final_prefix(), Some(limit), maybe_cursor)
+		let mut cursor = sp_io::MultiRemovalCursor::new();
+		cursor.set_input(maybe_cursor);
+		unhashed::clear_prefix(&Self::final_prefix(), Some(limit), Some(&mut cursor))
+			.into_results(cursor)
 	}
 
 	/// Iter over all value of the storage.

substrate/frame/support/src/storage/unhashed.rs

@@ -20,6 +20,8 @@
 use alloc::vec::Vec;
 use codec::{Decode, Encode};
 
+pub use sp_io::{MultiRemovalCounters, MultiRemovalCursor, MultiRemovalResults};
+
 /// Return the value of the item in storage under `key`, or `None` if there is no explicit entry.
 pub fn get<T: Decode + Sized>(key: &[u8]) -> Option<T> {
 	sp_io::storage::get(key).and_then(|val| {
@@ -110,30 +112,33 @@ pub fn kill(key: &[u8]) {
 ///
 /// # Cursor
 ///
-/// A *cursor* may be passed in to this operation with `maybe_cursor`. `None` should only be
-/// passed once (in the initial call) for any attempt to clear storage. In general, subsequent calls
-/// operating on the same prefix should pass `Some` and this value should be equal to the
-/// previous call result's `maybe_cursor` field. The only exception to this is when you can
-/// guarantee that the subsequent call is in a new block; in this case the previous call's result
-/// cursor need not be passed in and a `None` may be passed instead. This exception may be useful
-/// then making this call solely from a block-hook such as `on_initialize`.
+/// To continue the operation across multiple calls, the caller owns a
+/// [`MultiRemovalCursor`](sp_io::MultiRemovalCursor) and passes `Some(&mut cursor)`: it should be
+/// freshly [created](sp_io::MultiRemovalCursor::new) for the initial call and then passed back in
+/// unchanged on subsequent calls until
+/// [`MultiRemovalCounters::more`](sp_io::MultiRemovalCounters::more) is `false`. Reusing the same
+/// cursor object across iterations avoids re-allocating the cursor buffer on every call.
+///
+/// Pass `None` when you don't need to continue: no cursor is materialized (no allocation), and
+/// [`MultiRemovalCounters::more`](sp_io::MultiRemovalCounters::more) still reports whether keys
+/// remain.
 ///
-/// Returns [`MultiRemovalResults`](sp_io::MultiRemovalResults) to inform about the result. Once the
-/// resultant `maybe_cursor` field is `None`, then no further items remain to be deleted.
+/// Returns [`MultiRemovalCounters`](sp_io::MultiRemovalCounters) to inform about the result. Once
+/// `more` is `false`, no further items remain to be deleted.
 ///
-/// NOTE: After the initial call for any given child storage, it is important that no keys further
-/// keys are inserted. If so, then they may or may not be deleted by subsequent calls.
+/// NOTE: After the initial call for any given prefix, it is important that no further keys are
+/// inserted. If so, then they may or may not be deleted by subsequent calls.
 ///
 /// # Note
 ///
-/// Please note that keys which are residing in the overlay for the child are deleted without
+/// Please note that keys which are residing in the overlay for that prefix are deleted without
 /// counting towards the `limit`.
 pub fn clear_prefix(
 	prefix: &[u8],
 	maybe_limit: Option<u32>,
-	maybe_cursor: Option<&[u8]>,
-) -> sp_io::MultiRemovalResults {
-	sp_io::storage::clear_prefix(prefix, maybe_limit, maybe_cursor)
+	cursor: Option<&mut sp_io::MultiRemovalCursor>,
+) -> sp_io::MultiRemovalCounters {
+	sp_io::storage::clear_prefix(prefix, maybe_limit, cursor)
 }
 
 /// Returns `true` if the storage contains any key, which starts with a certain prefix,