Contracts: Refactor API to use WeightMeter (paritytech#2943)

Update the Contracts API to use `WeightMeter`, as it simplifies the code
and makes it easier to reason about, rather than taking a mutable weight
or returning a tuple with the weight consumed

---------

Co-authored-by: Alexander Theißen <[email protected]>

substrate/frame/contracts/src/benchmarking/mod.rs

@@ -40,7 +40,7 @@ use frame_support::{
 	self,
 	pallet_prelude::StorageVersion,
 	traits::{fungible::InspectHold, Currency},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 };
 use frame_system::RawOrigin;
 use pallet_balances;
@@ -198,7 +198,7 @@ mod benchmarks {
 	fn on_process_deletion_queue_batch() {
 		#[block]
 		{
-			ContractInfo::<T>::process_deletion_queue_batch(Weight::MAX);
+			ContractInfo::<T>::process_deletion_queue_batch(&mut WeightMeter::new())
 		}
 	}
 
@@ -213,7 +213,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			ContractInfo::<T>::process_deletion_queue_batch(Weight::MAX);
+			ContractInfo::<T>::process_deletion_queue_batch(&mut WeightMeter::new())
 		}
 
 		Ok(())
@@ -226,7 +226,7 @@ mod benchmarks {
 		let mut m = v09::Migration::<T>::default();
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -244,7 +244,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 
 		Ok(())
@@ -259,7 +259,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -276,7 +276,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -291,7 +291,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 		Ok(())
 	}
@@ -307,7 +307,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 	}
 
@@ -322,7 +322,7 @@ mod benchmarks {
 
 		#[block]
 		{
-			m.step();
+			m.step(&mut WeightMeter::new());
 		}
 
 		Ok(())
@@ -335,7 +335,7 @@ mod benchmarks {
 		StorageVersion::new(version).put::<Pallet<T>>();
 		#[block]
 		{
-			Migration::<T>::migrate(Weight::MAX);
+			Migration::<T>::migrate(&mut WeightMeter::new());
 		}
 		assert_eq!(StorageVersion::get::<Pallet<T>>(), version);
 	}
@@ -2776,7 +2776,8 @@ mod benchmarks {
 	#[benchmark(extra, pov_mode = Ignored)]
 	fn print_schedule() -> Result<(), BenchmarkError> {
 		let max_weight = <T as frame_system::Config>::BlockWeights::get().max_block;
-		let (weight_per_key, key_budget) = ContractInfo::<T>::deletion_budget(max_weight);
+		let (weight_per_key, key_budget) =
+			ContractInfo::<T>::deletion_budget(&mut WeightMeter::with_limit(max_weight));
 		let schedule = T::Schedule::get();
 		log::info!(target: LOG_TARGET, "
 		{schedule:#?}

substrate/frame/contracts/src/lib.rs

@@ -123,7 +123,7 @@ use frame_support::{
 		fungible::{Inspect, Mutate, MutateHold},
 		ConstU32, Contains, Get, Randomness, Time,
 	},
-	weights::Weight,
+	weights::{Weight, WeightMeter},
 	BoundedVec, DefaultNoBound, RuntimeDebugNoBound,
 };
 use frame_system::{
@@ -461,17 +461,15 @@ pub mod pallet {
 
 	#[pallet::hooks]
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_idle(_block: BlockNumberFor<T>, mut remaining_weight: Weight) -> Weight {
+		fn on_idle(_block: BlockNumberFor<T>, limit: Weight) -> Weight {
 			use migration::MigrateResult::*;
+			let mut meter = WeightMeter::with_limit(limit);
 
 			loop {
-				let (result, weight) = Migration::<T>::migrate(remaining_weight);
-				remaining_weight.saturating_reduce(weight);
-
-				match result {
-					// There is not enough weight to perform a migration, or make any progress, we
-					// just return the remaining weight.
-					NoMigrationPerformed | InProgress { steps_done: 0 } => return remaining_weight,
+				match Migration::<T>::migrate(&mut meter) {
+					// There is not enough weight to perform a migration.
+					// We can't do anything more, so we return the used weight.
+					NoMigrationPerformed | InProgress { steps_done: 0 } => return meter.consumed(),
 					// Migration is still in progress, we can start the next step.
 					InProgress { .. } => continue,
 					// Either no migration is in progress, or we are done with all migrations, we
@@ -480,8 +478,8 @@ pub mod pallet {
 				}
 			}
 
-			ContractInfo::<T>::process_deletion_queue_batch(remaining_weight)
-				.saturating_add(T::WeightInfo::on_process_deletion_queue_batch())
+			ContractInfo::<T>::process_deletion_queue_batch(&mut meter);
+			meter.consumed()
 		}
 
 		fn integrity_test() {
@@ -924,18 +922,25 @@ pub mod pallet {
 			ensure_signed(origin)?;
 
 			let weight_limit = weight_limit.saturating_add(T::WeightInfo::migrate());
-			let (result, weight) = Migration::<T>::migrate(weight_limit);
+			let mut meter = WeightMeter::with_limit(weight_limit);
+			let result = Migration::<T>::migrate(&mut meter);
 
 			match result {
-				Completed =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::No }),
-				InProgress { steps_done, .. } if steps_done > 0 =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::No }),
-				InProgress { .. } =>
-					Ok(PostDispatchInfo { actual_weight: Some(weight), pays_fee: Pays::Yes }),
+				Completed => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::No,
+				}),
+				InProgress { steps_done, .. } if steps_done > 0 => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::No,
+				}),
+				InProgress { .. } => Ok(PostDispatchInfo {
+					actual_weight: Some(meter.consumed()),
+					pays_fee: Pays::Yes,
+				}),
 				NoMigrationInProgress | NoMigrationPerformed => {
 					let err: DispatchError = <Error<T>>::NoMigrationPerformed.into();
-					Err(err.with_weight(T::WeightInfo::migrate()))
+					Err(err.with_weight(meter.consumed()))
 				},
 			}
 		}