refactor to record cost at the begining

frame/evm/precompile/storage-cleaner/Cargo.toml

@@ -14,6 +14,7 @@ frame-support = { workspace = true }
 frame-system = { workspace = true }
 sp-runtime = { workspace = true }
 sp-std = { workspace = true }
+sp-core = { workspace = true }
 # Frontier
 fp-evm = { workspace = true }
 pallet-evm = { workspace = true }
@@ -41,7 +42,10 @@ std = [
 	"scale-codec/std",
 	# Substrate
 	"frame-support/std",
+	"frame-system/std",
 	"sp-runtime/std",
+	"sp-core/std",
+	"sp-io/std",
 	# Frontier
 	"fp-evm/std",
 	"pallet-evm/std",

frame/evm/precompile/storage-cleaner/src/lib.rs

@@ -18,22 +18,22 @@
 //! Storage cleaner precompile. This precompile is used to clean the storage entries of smart contract that
 //! has been marked as suicided (self-destructed).
 
-extern crate alloc;
-
-pub const ARRAY_LIMIT: u32 = 1_000;
-
 use core::marker::PhantomData;
-use fp_evm::{ACCOUNT_BASIC_PROOF_SIZE, ACCOUNT_STORAGE_PROOF_SIZE};
+use fp_evm::{PrecompileFailure, ACCOUNT_BASIC_PROOF_SIZE, ACCOUNT_STORAGE_PROOF_SIZE};
 use pallet_evm::AddressMapping;
 use precompile_utils::{prelude::*, EvmResult};
+use sp_core::H160;
 use sp_runtime::traits::ConstU32;
 
 #[cfg(test)]
 mod mock;
 #[cfg(test)]
 mod tests;
 
+pub const ARRAY_LIMIT: u32 = 1_000;
 type GetArrayLimit = ConstU32<ARRAY_LIMIT>;
+// Storage key for suicided contracts: Blake2_128(16) + Key (H160(20))
+pub const SUICIDED_STORAGE_KEY: u64 = 36;
 
 #[derive(Debug, Clone)]
 pub struct StorageCleanerPrecompile<Runtime>(PhantomData<Runtime>);
@@ -43,86 +43,164 @@ impl<Runtime> StorageCleanerPrecompile<Runtime>
 where
 	Runtime: pallet_evm::Config,
 {
-	/// Clear Storage entries of smart contracts that has been marked as suicided (self-destructed) up to a certain limit.
-	///
-	/// The function iterates over the addresses, checks if each address is marked as suicided, and then deletes the storage
-	/// entries associated with that address. If there are no remaining entries, the function clears the suicided contract
-	/// by removing the address from the list of suicided contracts and decrementing the sufficients of the associated account.
-	#[precompile::public("clearSuicidedStorage(address[],uint32)")]
+	/// Clear Storage entries of smart contracts that has been marked as suicided (self-destructed). It takes a list of
+	/// addresses and a limit as input. The limit is used to prevent the function from consuming too much gas. The
+	/// maximum number of storage entries that can be removed is limit - 1.
+	#[precompile::public("clearSuicidedStorage(address[],uint64)")]
 	fn clear_suicided_storage(
 		handle: &mut impl PrecompileHandle,
 		addresses: BoundedVec<Address, GetArrayLimit>,
-		limit: u32,
+		limit: u64,
 	) -> EvmResult {
 		let addresses: Vec<_> = addresses.into();
-		let mut deleted_entries = 0;
-
+		let nb_addresses = addresses.len() as u64;
 		if limit == 0 {
 			return Err(revert("Limit should be greater than zero"));
 		}
 
-		for address in &addresses {
-			// Read Suicided storage item
-			// Suicided: Hash (Blake2128(16) + Key (H160(20))
-			handle.record_db_read::<Runtime>(36)?;
-			if !pallet_evm::Pallet::<Runtime>::is_account_suicided(&address.0) {
-				return Err(revert(format!("NotSuicided: {}", address.0)));
+		Self::record_max_cost(handle, nb_addresses, limit)?;
+		let result = Self::clear_suicided_storage_inner(addresses, limit - 1)?;
+		Self::refund_cost(handle, result, nb_addresses, limit);
+
+		Ok(())
+	}
+
+	/// This function iterates over the addresses, checks if each address is marked as suicided, and then deletes the storage
+	/// entries associated with that address. If there are no remaining entries, we clear the suicided contract by calling the
+	/// `clear_suicided_contract` function.
+	fn clear_suicided_storage_inner(
+		addresses: Vec<Address>,
+		limit: u64,
+	) -> Result<RemovalResult, PrecompileFailure> {
+		let mut deleted_entries = 0u64;
+		let mut deleted_contracts = 0u64;
+
+		for Address(address) in addresses {
+			if !pallet_evm::Pallet::<Runtime>::is_account_suicided(&address) {
+				return Err(revert(format!("NotSuicided: {}", address)));
+			}
+
+			let deleted = pallet_evm::AccountStorages::<Runtime>::drain_prefix(address)
+				.take((limit.saturating_sub(deleted_entries)) as usize)
+				.count();
+			deleted_entries = deleted_entries.saturating_add(deleted as u64);
+
+			// Check if the storage of this contract has been completly removed
+			if pallet_evm::AccountStorages::<Runtime>::iter_key_prefix(&address)
+				.next()
+				.is_none()
+			{
+				Self::clear_suicided_contract(address);
+				deleted_contracts = deleted_contracts.saturating_add(1);
 			}
 
-			let mut iter = pallet_evm::AccountStorages::<Runtime>::iter_key_prefix(address.0);
-
-			loop {
-				// Read AccountStorages storage item
-				handle.record_db_read::<Runtime>(ACCOUNT_STORAGE_PROOF_SIZE as usize)?;
-
-				match iter.next() {
-					Some(key) => {
-						// Write AccountStorages storage item
-						handle.record_cost(RuntimeHelper::<Runtime>::db_write_gas_cost())?;
-						pallet_evm::AccountStorages::<Runtime>::remove(address.0, key);
-						deleted_entries += 1;
-						if deleted_entries >= limit {
-							// Check if there are no remaining entries. If there aren't any, clear the contract.
-							// Read AccountStorages storage item
-							handle
-								.record_db_read::<Runtime>(ACCOUNT_STORAGE_PROOF_SIZE as usize)?;
-							// We perform an additional iteration at the end because we cannot predict if there are
-							// remaining entries without checking the next item. If there are no more entries, we clear
-							// the contract and refund the cost of the last empty iteration.
-							if iter.next().is_none() {
-								handle.refund_external_cost(None, Some(ACCOUNT_STORAGE_PROOF_SIZE));
-								Self::clear_suicided_contract(handle, address)?;
-							}
-							return Ok(());
-						}
-					}
-					None => {
-						// No more entries, clear the contract.
-						// Refund the cost of the last iteration.
-						handle.refund_external_cost(None, Some(ACCOUNT_STORAGE_PROOF_SIZE));
-						Self::clear_suicided_contract(handle, address)?;
-						break;
-					}
-				}
+			if deleted_entries >= limit {
+				break;
 			}
 		}
+
+		Ok(RemovalResult {
+			deleted_entries,
+			deleted_contracts,
+		})
+	}
+
+	/// Record the maximum cost (Worst case Scenario) of the clear_suicided_storage function.
+	fn record_max_cost(
+		handle: &mut impl PrecompileHandle,
+		nb_addresses: u64,
+		limit: u64,
+	) -> EvmResult {
+		let read_cost = RuntimeHelper::<Runtime>::db_read_gas_cost();
+		let write_cost = RuntimeHelper::<Runtime>::db_write_gas_cost();
+		let ref_time = 0u64
+			// EVM:: Suicided (reads = nb_addresses)
+			.saturating_add(read_cost.saturating_mul(nb_addresses.into()))
+			// EVM:: Suicided (writes = nb_addresses)
+			.saturating_add(write_cost.saturating_mul(nb_addresses))
+			// System: AccountInfo (reads = nb_addresses) for decrementing sufficients
+			.saturating_add(read_cost.saturating_mul(nb_addresses))
+			// System: AccountInfo (writes = nb_addresses) for decrementing sufficients
+			.saturating_add(write_cost.saturating_mul(nb_addresses))
+			// EVM: AccountStorage (reads = limit)
+			.saturating_add(read_cost.saturating_mul(limit))
+			// EVM: AccountStorage (writes = limit)
+			.saturating_add(write_cost.saturating_mul(limit));
+
+		let proof_size = 0u64
+			// Proof: EVM::Suicided (SUICIDED_STORAGE_KEY) * nb_addresses
+			.saturating_add(SUICIDED_STORAGE_KEY.saturating_mul(nb_addresses))
+			// Proof: EVM::AccountStorage (ACCOUNT_BASIC_PROOF_SIZE) * limit
+			.saturating_add(ACCOUNT_STORAGE_PROOF_SIZE.saturating_mul(limit))
+			// Proof: System::AccountInfo (ACCOUNT_BASIC_PROOF_SIZE) * nb_addresses
+			.saturating_add(ACCOUNT_BASIC_PROOF_SIZE.saturating_mul(nb_addresses));
+
+		handle.record_external_cost(Some(ref_time), Some(proof_size), None)?;
 		Ok(())
 	}
 
+	/// Refund the additional cost recorded for the clear_suicided_storage function.
+	fn refund_cost(
+		handle: &mut impl PrecompileHandle,
+		result: RemovalResult,
+		nb_addresses: u64,
+		limit: u64,
+	) {
+		let read_cost = RuntimeHelper::<Runtime>::db_read_gas_cost();
+		let write_cost = RuntimeHelper::<Runtime>::db_write_gas_cost();
+
+		let extra_entries = limit.saturating_sub(result.deleted_entries);
+		let extra_contracts = nb_addresses.saturating_sub(result.deleted_contracts);
+
+		let mut ref_time = 0u64;
+		let mut proof_size = 0u64;
+
+		// Refund the cost of the remaining entries
+		if extra_entries > 0 {
+			ref_time = ref_time
+				// EVM:: AccountStorage (reads = extra_entries)
+				.saturating_add(read_cost.saturating_mul(extra_entries))
+				// EVM:: AccountStorage (writes = extra_entries)
+				.saturating_add(write_cost.saturating_mul(extra_entries));
+			proof_size = proof_size
+				// Proof: EVM::AccountStorage (ACCOUNT_BASIC_PROOF_SIZE) * extra_entries
+				.saturating_add(ACCOUNT_STORAGE_PROOF_SIZE.saturating_mul(extra_entries));
+		}
+
+		// Refund the cost of the remaining contracts
+		if extra_contracts > 0 {
+			ref_time = ref_time
+				// EVM:: Suicided (reads = extra_contracts)
+				.saturating_add(read_cost.saturating_mul(extra_contracts))
+				// EVM:: Suicided (writes = extra_contracts)
+				.saturating_add(write_cost.saturating_mul(extra_contracts))
+				// System: AccountInfo (reads = extra_contracts) for decrementing sufficients
+				.saturating_add(read_cost.saturating_mul(extra_contracts))
+				// System: AccountInfo (writes = extra_contracts) for decrementing sufficients
+				.saturating_add(write_cost.saturating_mul(extra_contracts));
+			proof_size = proof_size
+				// Proof: EVM::Suicided (SUICIDED_STORAGE_KEY) * extra_contracts
+				.saturating_add(SUICIDED_STORAGE_KEY.saturating_mul(extra_contracts))
+				// Proof: System::AccountInfo (ACCOUNT_BASIC_PROOF_SIZE) * extra_contracts
+				.saturating_add(ACCOUNT_BASIC_PROOF_SIZE.saturating_mul(extra_contracts));
+		}
+
+		handle.refund_external_cost(Some(ref_time), Some(proof_size));
+	}
+
 	/// Clears the storage of a suicided contract.
 	///
 	/// This function will remove the given address from the list of suicided contracts
 	/// and decrement the sufficients of the account associated with the address.
-	fn clear_suicided_contract(handle: &mut impl PrecompileHandle, address: &Address) -> EvmResult {
-		handle.record_cost(RuntimeHelper::<Runtime>::db_write_gas_cost())?;
-		pallet_evm::Suicided::<Runtime>::remove(address.0);
-
-		let account_id = Runtime::AddressMapping::into_account_id(address.0);
-		// dec_sufficients mutates the account, so we need to read it first.
-		// Read Account storage item (AccountBasicProof)
-		handle.record_db_read::<Runtime>(ACCOUNT_BASIC_PROOF_SIZE as usize)?;
-		handle.record_cost(RuntimeHelper::<Runtime>::db_write_gas_cost())?;
+	fn clear_suicided_contract(address: H160) {
+		pallet_evm::Suicided::<Runtime>::remove(address);
+
+		let account_id = Runtime::AddressMapping::into_account_id(address);
 		let _ = frame_system::Pallet::<Runtime>::dec_sufficients(&account_id);
-		Ok(())
 	}
 }
+
+struct RemovalResult {
+	pub deleted_entries: u64,
+	pub deleted_contracts: u64,
+}

frame/evm/precompile/storage-cleaner/src/mock.rs

@@ -88,7 +88,6 @@ impl pallet_balances::Config for Runtime {
 	type FreezeIdentifier = ();
 	type MaxLocks = ();
 	type MaxReserves = ();
-	type MaxHolds = ();
 	type MaxFreezes = ();
 	type RuntimeFreezeReason = ();
 }

frame/evm/precompile/storage-cleaner/src/tests.rs

@@ -81,7 +81,7 @@ fn test_clear_suicided_contract_succesfull() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![suicided_address.into()].into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_returns(());
@@ -117,7 +117,7 @@ fn test_clear_suicided_contract_failed() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![non_suicided_address.into()].into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_reverts(|output| {
@@ -148,7 +148,7 @@ fn test_clear_suicided_empty_input() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![].into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_returns(());
@@ -182,7 +182,7 @@ fn test_clear_suicided_contract_multiple_addresses() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: addresses.clone().into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_returns(());
@@ -217,7 +217,7 @@ fn test_clear_suicided_mixed_suicided_and_non_suicided() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: addresses.clone().into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_reverts(|output| {
@@ -246,7 +246,7 @@ fn test_clear_suicided_no_storage_entries() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: addresses.clone().into(),
-						limit: u32::MAX,
+						limit: u64::MAX,
 					},
 				)
 				.execute_returns(());
@@ -280,7 +280,7 @@ fn test_clear_suicided_contract_limit_works() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: addresses.clone().into(),
-						limit: 4,
+						limit: 5,
 					},
 				)
 				.execute_returns(());
@@ -320,7 +320,7 @@ fn test_clear_suicided_contract_limit_respected() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![suicided_address.into()].into(),
-						limit: 4,
+						limit: 5,
 					},
 				)
 				.execute_returns(());