diff --git a/cumulus/pallets/dmp-queue/src/lib.rs b/cumulus/pallets/dmp-queue/src/lib.rs index 4b2959a1985a8..c8c7f496f4a35 100644 --- a/cumulus/pallets/dmp-queue/src/lib.rs +++ b/cumulus/pallets/dmp-queue/src/lib.rs @@ -246,16 +246,18 @@ pub mod pallet { let hashed_prefix = twox_128( 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::::put(MigrationState::StartedCleanup { diff --git a/polkadot/runtime/common/src/crowdloan/mod.rs b/polkadot/runtime/common/src/crowdloan/mod.rs index 699ee9a4e9d8b..02307b57a564a 100644 --- a/polkadot/runtime/common/src/crowdloan/mod.rs +++ b/polkadot/runtime/common/src/crowdloan/mod.rs @@ -705,7 +705,13 @@ impl Pallet { } 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( diff --git a/substrate/frame/contracts/src/storage.rs b/substrate/frame/contracts/src/storage.rs index 70a6cc3c9937e..2c14a56634a8e 100644 --- a/substrate/frame/contracts/src/storage.rs +++ b/substrate/frame/contracts/src/storage.rs @@ -327,14 +327,13 @@ impl ContractInfo { 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 { diff --git a/substrate/frame/migrations/src/benchmarking.rs b/substrate/frame/migrations/src/benchmarking.rs index de362145ee86e..55f985f4d9b1b 100644 --- a/substrate/frame/migrations/src/benchmarking.rs +++ b/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"); diff --git a/substrate/frame/migrations/src/migrations.rs b/substrate/frame/migrations/src/migrations.rs index 1a5030bb26c52..c6b5abd27bab2 100644 --- a/substrate/frame/migrations/src/migrations.rs +++ b/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")] diff --git a/substrate/frame/revive/src/storage.rs b/substrate/frame/revive/src/storage.rs index a5af29a48f748..1b62ec76e529d 100644 --- a/substrate/frame/revive/src/storage.rs +++ b/substrate/frame/revive/src/storage.rs @@ -460,7 +460,7 @@ impl ContractInfo { None, ); - if outcome.maybe_cursor.is_some() { + if outcome.more { remaining = remaining .saturating_sub(weight_per_trie_key.saturating_mul(outcome.backend.into())); break; diff --git a/substrate/frame/support/src/migrations.rs b/substrate/frame/support/src/migrations.rs index 7d3297ecd2c8d..e7b43819ace02 100644 --- a/substrate/frame/support/src/migrations.rs +++ b/substrate/frame/support/src/migrations.rs @@ -321,15 +321,14 @@ impl, DbWeight: Get> 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, S: Get<&'static str>, DbWeight: Get> 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()); diff --git a/substrate/frame/support/src/storage/child.rs b/substrate/frame/support/src/storage/child.rs index de7a41b1daf49..a9999261d9581 100644 --- a/substrate/frame/support/src/storage/child.rs +++ b/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(child_info: &ChildInfo, key: &[u8]) -> Option { @@ -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, - 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, ), } } diff --git a/substrate/frame/support/src/storage/generator/double_map.rs b/substrate/frame/support/src/storage/generator/double_map.rs index d76123effb480..d4142714c7e54 100644 --- a/substrate/frame/support/src/storage/generator/double_map.rs +++ b/substrate/frame/support/src/storage/generator/double_map.rs @@ -210,8 +210,16 @@ where where KArg1: EncodeLike, { - 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( @@ -222,12 +230,14 @@ where where KArg1: EncodeLike, { + 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(k1: KArg1) -> bool diff --git a/substrate/frame/support/src/storage/generator/nmap.rs b/substrate/frame/support/src/storage/generator/nmap.rs index 021da750a3764..9e29cbfe505a3 100755 --- a/substrate/frame/support/src/storage/generator/nmap.rs +++ b/substrate/frame/support/src/storage/generator/nmap.rs @@ -186,7 +186,13 @@ where where K: HasKeyPrefix, { - 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( @@ -197,11 +203,14 @@ where where K: HasKeyPrefix, { + 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(partial_key: KP) -> bool diff --git a/substrate/frame/support/src/storage/migration.rs b/substrate/frame/support/src/storage/migration.rs index 63c786d81d598..432fc01255af5 100644 --- a/substrate/frame/support/src/storage/migration.rs +++ b/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`. diff --git a/substrate/frame/support/src/storage/mod.rs b/substrate/frame/support/src/storage/mod.rs index 70e094c9588bb..483b6a03d7bea 100644 --- a/substrate/frame/support/src/storage/mod.rs +++ b/substrate/frame/support/src/storage/mod.rs @@ -1389,7 +1389,12 @@ pub trait StoragePrefixedMap { /// overlay are not taken into account when deleting keys in the backend. #[deprecated = "Use `clear` instead"] fn remove_all(limit: Option) -> 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 { /// 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. diff --git a/substrate/frame/support/src/storage/unhashed.rs b/substrate/frame/support/src/storage/unhashed.rs index e20ca4e235e6f..1cc0d92596945 100644 --- a/substrate/frame/support/src/storage/unhashed.rs +++ b/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(key: &[u8]) -> Option { 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, - 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, diff --git a/substrate/primitives/io/src/lib.rs b/substrate/primitives/io/src/lib.rs index 795b340a2a872..8c8448b767a47 100644 --- a/substrate/primitives/io/src/lib.rs +++ b/substrate/primitives/io/src/lib.rs @@ -139,6 +139,106 @@ use sp_externalities::{Externalities, ExternalitiesExt}; pub use sp_externalities::MultiRemovalResults; +/// A reusable cursor object that holds owned buffers across multiple invocations of a +/// storage-removal operation (e.g. [`storage::clear_prefix`]). +/// +/// The buffers are reused across calls so that subsequent iterations do not allocate +/// (assuming the cursor stays roughly the same size). On the first call no memory is +/// allocated โ€” the buffers grow lazily only when the host actually returns a cursor. +#[derive(Default)] +pub struct MultiRemovalCursor { + /// Input cursor buffer. + in_buf: Vec, + /// Output cursor buffer. + out_buf: Vec, + /// Whether `in_buf` currently contains a valid cursor. + has_cursor: bool, +} + +impl MultiRemovalCursor { + /// Create an empty cursor. No memory is allocated. + pub fn new() -> Self { + Self::default() + } + + /// Returns `true` if the previous call left a cursor, meaning more work is required to + /// complete the operation. + pub fn has_more(&self) -> bool { + self.has_cursor + } + + /// Returns the current cursor as a byte slice, or `None` if there is no cursor. + pub fn as_slice(&self) -> Option<&[u8]> { + self.has_cursor.then(|| self.in_buf.as_slice()) + } + + /// Reset the cursor to the empty state. Allocated buffer capacity is retained for reuse. + pub fn reset(&mut self) { + self.in_buf.clear(); + self.out_buf.clear(); + self.has_cursor = false; + } + + /// Bootstrap the cursor from an externally-provided value. Allocates if the input is `Some`. + /// + /// This is intended for compatibility shims that adapt the legacy `Option<&[u8]>` / + /// `Option>` cursor APIs to the new buffer-reusing wrapper; runtime code that owns its + /// own [`MultiRemovalCursor`] across iterations should not need to call this. + pub fn set_input(&mut self, cursor: Option<&[u8]>) { + self.in_buf.clear(); + if let Some(c) = cursor { + self.in_buf.extend_from_slice(c); + self.has_cursor = true; + } else { + self.has_cursor = false; + } + } + + /// Consume the cursor, returning the current cursor as an owned `Option>`. + /// + /// Compatibility shim for code that still produces the legacy `maybe_cursor: Option>` + /// shape. + pub fn into_inner(self) -> Option> { + if self.has_cursor { + Some(self.in_buf) + } else { + None + } + } +} + +/// Counters returned by storage-removal operations such as [`storage::clear_prefix`]. +/// +/// When the caller passes a [`MultiRemovalCursor`] it also holds the continuation cursor bytes; +/// callers that don't need to continue can pass `None` and rely solely on [`more`](Self::more). +#[derive(Default, Debug, Clone, Copy, Eq, PartialEq)] +pub struct MultiRemovalCounters { + /// The number of items removed from the backend database. + pub backend: u32, + /// The number of unique keys removed, taking into account both the backend and the overlay. + pub unique: u32, + /// The number of iterations (each requiring a storage seek/read) which were done. + pub loops: u32, + /// Whether a continuation cursor remains, i.e. the operation is *not* complete and another + /// call is required to remove the rest. Equivalent to "some keys remain". + pub more: bool, +} + +impl MultiRemovalCounters { + /// Combine these counters with the (consumed) `cursor` into the legacy [`MultiRemovalResults`] + /// shape. + /// + /// Convenience for compatibility shims that still expose the `Option>` cursor. + pub fn into_results(self, cursor: MultiRemovalCursor) -> MultiRemovalResults { + MultiRemovalResults { + maybe_cursor: cursor.into_inner(), + backend: self.backend, + unique: self.unique, + loops: self.loops, + } + } +} + #[cfg(all(not(feature = "disable_allocator"), substrate_runtime, target_family = "wasm"))] mod global_alloc_wasm; @@ -835,36 +935,68 @@ pub trait Storage { /// A convenience wrapper providing a developer-friendly interface for the `clear_prefix` host /// function. + /// + /// Pass `Some` [`MultiRemovalCursor`] only when you need to continue the operation across + /// calls; the cursor holds owned buffers reused across invocations (first call allocates + /// nothing, and the cached cursor is fetched with an exact-size allocation only when it does + /// not fit the reused buffer). Pass `None` when you don't need to continue: no cursor is + /// materialized (no `last_cursor` call, no allocation), and [`MultiRemovalCounters::more`] + /// still reports whether keys remain. #[wrapper] fn clear_prefix( maybe_prefix: impl AsRef<[u8]>, maybe_limit: Option, - maybe_cursor_in: Option<&[u8]>, - ) -> MultiRemovalResults { - let mut result = MultiRemovalResults::default(); - let mut maybe_cursor_out = vec![0u8; 1024]; + cursor: Option<&mut MultiRemovalCursor>, + ) -> MultiRemovalCounters { let mut counters = StorageIterations::default(); + + let Some(cursor) = cursor else { + // No continuation requested: single pass without materializing the output cursor. + let cursor_len = + clear_prefix__raw(maybe_prefix.as_ref(), maybe_limit, None, &mut [], &mut counters) + as usize; + return MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, + }; + }; + + // Make as much as already allocated bytes available for the new cursor + cursor.out_buf.resize(cursor.out_buf.capacity(), 0); + + let cursor_in: Option<&[u8]> = cursor.has_cursor.then(|| cursor.in_buf.as_slice()); + let output_avail = cursor.out_buf.len(); + let cursor_len = clear_prefix__raw( maybe_prefix.as_ref(), maybe_limit, - maybe_cursor_in, - &mut maybe_cursor_out, + cursor_in, + &mut cursor.out_buf, &mut counters, ) as usize; - result.backend = counters.backend; - result.unique = counters.unique; - result.loops = counters.loops; - if cursor_len > 0 { - if maybe_cursor_out.len() < cursor_len { - maybe_cursor_out.resize(cursor_len, 0); - let cached_cursor_len = misc::last_cursor(maybe_cursor_out.as_mut_slice()); - debug_assert!(cached_cursor_len.is_some()); - debug_assert_eq!(cached_cursor_len.unwrap_or(0) as usize, cursor_len); - } - maybe_cursor_out.truncate(cursor_len); - result.maybe_cursor = Some(maybe_cursor_out); + + if cursor_len > output_avail { + // The cursor did not fit in `out_buf`; retrieve the cached cursor from the host. + cursor.out_buf.resize(cursor_len, 0); + let fetched = misc::last_cursor(&mut cursor.out_buf); + debug_assert_eq!(fetched.map(|n| n as usize), Some(cursor_len)); + } else { + cursor.out_buf.truncate(cursor_len); + } + + // Swap: `out_buf` becomes the input for the next call; the previous input buffer + // becomes the next scratch space with its capacity preserved. + core::mem::swap(&mut cursor.in_buf, &mut cursor.out_buf); + cursor.has_cursor = cursor_len > 0; + + MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, } - result } /// Append the encoded `value` to the storage item at `key`. @@ -910,13 +1042,13 @@ pub trait Storage { fn root(&mut self, out: PassFatPointerAndWrite<&mut [u8]>) { let root = self.storage_root(); let encoded = codec::Encode::encode(&root); - out_len = out.len(); - encoded_len = encoded_len(); + let out_len = out.len(); + let encoded_len = encoded.len(); assert!( out_len >= encoded_len, - "Output buffer ({out_len} bytes) provided to store the child storage root hash is not large enough ({encoded_len} bytes needed)" + "Output buffer ({out_len} bytes) provided to store the storage root hash is not large enough ({encoded_len} bytes needed)" ); - out[..encoded.len()].copy_from_slice(&encoded[..]); + out[..encoded_len].copy_from_slice(&encoded[..]); } /// A convenience wrapper providing a developer-friendly interface for the `root` host @@ -1257,37 +1389,61 @@ pub trait DefaultChildStorage { /// A convenience wrapper providing a developer-friendly interface for the `storage_kill` host /// function. + /// + /// See [`MultiRemovalCursor`] for details on how the cursor object is reused across calls. Pass + /// `None` when you don't need to continue the operation: no cursor is materialized and + /// [`MultiRemovalCounters::more`] still reports whether keys remain. #[wrapper] fn storage_kill( storage_key: impl AsRef<[u8]>, maybe_limit: Option, - maybe_cursor: Option<&[u8]>, - ) -> MultiRemovalResults { - let mut result = MultiRemovalResults::default(); - let mut maybe_cursor_out = vec![0u8; 1024]; + cursor: Option<&mut MultiRemovalCursor>, + ) -> MultiRemovalCounters { let mut counters = StorageIterations::default(); + + let Some(cursor) = cursor else { + let cursor_len = + storage_kill__raw(storage_key.as_ref(), maybe_limit, None, &mut [], &mut counters) + as usize; + return MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, + }; + }; + + let out_cap = cursor.out_buf.capacity(); + cursor.out_buf.resize(out_cap, 0); + + let cursor_in: Option<&[u8]> = cursor.has_cursor.then(|| cursor.in_buf.as_slice()); + let output_avail = cursor.out_buf.len(); + let cursor_len = storage_kill__raw( storage_key.as_ref(), maybe_limit, - maybe_cursor, - &mut maybe_cursor_out[..], + cursor_in, + &mut cursor.out_buf, &mut counters, ) as usize; - result.backend = counters.backend; - result.unique = counters.unique; - result.loops = counters.loops; - if cursor_len > 0 { - if maybe_cursor_out.len() < cursor_len { - maybe_cursor_out.resize(cursor_len, 0); - let cached_cursor_len = misc::last_cursor(maybe_cursor_out.as_mut_slice()); - debug_assert!(cached_cursor_len.is_some()); - debug_assert_eq!(cached_cursor_len.unwrap_or(0) as usize, cursor_len); - } - maybe_cursor_out.truncate(cursor_len); - result.maybe_cursor = Some(maybe_cursor_out); + + if cursor_len > output_avail { + cursor.out_buf.resize(cursor_len, 0); + let fetched = misc::last_cursor(&mut cursor.out_buf); + debug_assert_eq!(fetched.map(|n| n as usize), Some(cursor_len)); + } else { + cursor.out_buf.truncate(cursor_len); } - result + core::mem::swap(&mut cursor.in_buf, &mut cursor.out_buf); + cursor.has_cursor = cursor_len > 0; + + MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, + } } /// Check a child storage key. @@ -1386,38 +1542,68 @@ pub trait DefaultChildStorage { /// A convenience wrapper providing a developer-friendly interface for the `clear_prefix` host /// function. + /// + /// See [`MultiRemovalCursor`] for details on how the cursor object is reused across calls. Pass + /// `None` when you don't need to continue the operation: no cursor is materialized and + /// [`MultiRemovalCounters::more`] still reports whether keys remain. #[wrapper] fn clear_prefix( storage_key: impl AsRef<[u8]>, maybe_prefix: impl AsRef<[u8]>, maybe_limit: Option, - maybe_cursor_in: Option<&[u8]>, - ) -> MultiRemovalResults { - let mut result = MultiRemovalResults::default(); - let mut maybe_cursor_out = vec![0u8; 1024]; + cursor: Option<&mut MultiRemovalCursor>, + ) -> MultiRemovalCounters { let mut counters = StorageIterations::default(); + + let Some(cursor) = cursor else { + let cursor_len = clear_prefix__raw( + storage_key.as_ref(), + maybe_prefix.as_ref(), + maybe_limit, + None, + &mut [], + &mut counters, + ) as usize; + return MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, + }; + }; + + let out_cap = cursor.out_buf.capacity(); + cursor.out_buf.resize(out_cap, 0); + + let cursor_in: Option<&[u8]> = cursor.has_cursor.then(|| cursor.in_buf.as_slice()); + let output_avail = cursor.out_buf.len(); + let cursor_len = clear_prefix__raw( storage_key.as_ref(), maybe_prefix.as_ref(), maybe_limit, - maybe_cursor_in, - &mut maybe_cursor_out, + cursor_in, + &mut cursor.out_buf, &mut counters, ) as usize; - result.backend = counters.backend; - result.unique = counters.unique; - result.loops = counters.loops; - if cursor_len > 0 { - if maybe_cursor_out.len() < cursor_len { - maybe_cursor_out.resize(cursor_len, 0); - let cached_cursor_len = misc::last_cursor(maybe_cursor_out.as_mut_slice()); - debug_assert!(cached_cursor_len.is_some()); - debug_assert_eq!(cached_cursor_len.unwrap_or(0) as usize, cursor_len); - } - maybe_cursor_out.truncate(cursor_len); - result.maybe_cursor = Some(maybe_cursor_out); + + if cursor_len > output_avail { + cursor.out_buf.resize(cursor_len, 0); + let fetched = misc::last_cursor(&mut cursor.out_buf); + debug_assert_eq!(fetched.map(|n| n as usize), Some(cursor_len)); + } else { + cursor.out_buf.truncate(cursor_len); + } + + core::mem::swap(&mut cursor.in_buf, &mut cursor.out_buf); + cursor.has_cursor = cursor_len > 0; + + MultiRemovalCounters { + backend: counters.backend, + unique: counters.unique, + loops: counters.loops, + more: cursor_len > 0, } - result } /// Default child root calculation. @@ -2602,7 +2788,8 @@ pub trait Crypto { /// /// Returns the signature. // ERRATA: The RFC gathers all the *_sign_{prehashed} functions under a single definition that - // requires `msg` to be a fat pointer which obviously doesn't make sense for a prehashed message. + // requires `msg` to be a fat pointer which obviously doesn't make sense for a prehashed + // message. fn ecdsa_sign_prehashed( &mut self, id: PassPointerAndReadCopy, @@ -4100,20 +4287,24 @@ mod tests { }); t.execute_with(|| { - let res = storage::clear_prefix(b":abc", None, None); + let mut cursor = MultiRemovalCursor::new(); + let res = storage::clear_prefix(b":abc", None, Some(&mut cursor)); assert_eq!(res.backend, 2); assert_eq!(res.unique, 2); assert_eq!(res.loops, 2); + assert!(!res.more); + assert!(!cursor.has_more()); assert!(storage::get(b":a").is_some()); assert!(storage::get(b":abdd").is_some()); assert!(storage::get(b":abcd").is_none()); assert!(storage::get(b":abc").is_none()); - let res = storage::clear_prefix(b":abc", None, None); + let res = storage::clear_prefix(b":abc", None, Some(&mut cursor)); assert_eq!(res.backend, 0); assert_eq!(res.unique, 0); assert_eq!(res.loops, 0); + assert!(!cursor.has_more()); }); }