Revert "FRAME: Create TransactionExtension as a replacement for SignedExtension (#2280)" (#3665)

This PR reverts #2280 which introduced `TransactionExtension` to replace
`SignedExtension`.

As a result of the discussion
[here](https://github.com/paritytech/polkadot-sdk/pull/3623#issuecomment-1986789700),
the changes will be reverted for now with plans to reintroduce the
concept in the future.

---------

Signed-off-by: georgepisaltu <george.pisaltu@parity.io>

cumulus/primitives/storage-weight-reclaim/src/benchmarking.rs

@@ -1,70 +0,0 @@
-// Copyright (C) Parity Technologies (UK) Ltd.
-// SPDX-License-Identifier: Apache-2.0
-
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// 	http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! Benchmarking setup for cumulus-primitives-storage-weight-reclaim
-
-#![cfg(feature = "runtime-benchmarks")]
-
-use super::*;
-
-use frame_benchmarking::{account, v2::*, BenchmarkError};
-use frame_support::pallet_prelude::DispatchClass;
-use frame_system::{BlockWeight, RawOrigin};
-use sp_runtime::traits::{DispatchTransaction, Get};
-use sp_std::{
-	marker::{Send, Sync},
-	prelude::*,
-};
-
-/// Pallet we're benchmarking here.
-pub struct Pallet<T: frame_system::Config>(frame_system::Pallet<T>);
-
-#[benchmarks(where
-    T: Send + Sync,
-    T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>
-)]
-mod benchmarks {
-	use super::*;
-
-	#[benchmark]
-	fn storage_weight_reclaim() -> Result<(), BenchmarkError> {
-		let caller = account("caller", 0, 0);
-		BlockWeight::<T>::mutate(|current_weight| {
-			current_weight.set(Weight::from_parts(0, 1000), DispatchClass::Normal);
-		});
-		let base_extrinsic = <T as frame_system::Config>::BlockWeights::get()
-			.get(DispatchClass::Normal)
-			.base_extrinsic;
-		let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
-		let call: T::RuntimeCall = frame_system::Call::remark { remark: vec![] }.into();
-		let post_info = PostDispatchInfo {
-			actual_weight: Some(Weight::from_parts(0, 200)),
-			pays_fee: Default::default(),
-		};
-		let len = 0_usize;
-		let ext = StorageWeightReclaim::<T>::new();
-
-		#[block]
-		{
-			ext.test_run(RawOrigin::Signed(caller).into(), &call, &info, len, |_| Ok(post_info))
-				.unwrap()
-				.unwrap();
-		}
-
-		assert_eq!(BlockWeight::<T>::get().total().proof_size(), 700 + base_extrinsic.proof_size());
-
-		Ok(())
-	}
-}

cumulus/primitives/storage-weight-reclaim/src/lib.rs

@@ -29,22 +29,12 @@ use frame_support::{
 use frame_system::Config;
 use scale_info::TypeInfo;
 use sp_runtime::{
-	impl_tx_ext_default,
-	traits::{
-		DispatchInfoOf, Dispatchable, PostDispatchInfoOf, TransactionExtension,
-		TransactionExtensionBase,
-	},
+	traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, SignedExtension},
 	transaction_validity::TransactionValidityError,
 	DispatchResult,
 };
 use sp_std::marker::PhantomData;
 
-#[cfg(test)]
-mod tests;
-
-#[cfg(feature = "runtime-benchmarks")]
-mod benchmarking;
-
 const LOG_TARGET: &'static str = "runtime::storage_reclaim";
 
 /// `StorageWeightReclaimer` is a mechanism for manually reclaiming storage weight.
@@ -53,7 +43,7 @@ const LOG_TARGET: &'static str = "runtime::storage_reclaim";
 /// reclaim  it computes the real consumed storage weight and refunds excess weight.
 ///
 /// # Example
-#[doc = docify::embed!("src/tests.rs", simple_reclaimer_example)]
+#[doc = docify::embed!("src/lib.rs", simple_reclaimer_example)]
 pub struct StorageWeightReclaimer {
 	previous_remaining_proof_size: u64,
 	previous_reported_proof_size: Option<u64>,
@@ -129,40 +119,42 @@ impl<T: Config + Send + Sync> core::fmt::Debug for StorageWeightReclaim<T> {
 	}
 }
 
-impl<T: Config + Send + Sync> TransactionExtensionBase for StorageWeightReclaim<T> {
-	const IDENTIFIER: &'static str = "StorageWeightReclaim";
-	type Implicit = ();
-}
-
-impl<T: Config + Send + Sync, Context> TransactionExtension<T::RuntimeCall, Context>
-	for StorageWeightReclaim<T>
+impl<T: Config + Send + Sync> SignedExtension for StorageWeightReclaim<T>
 where
 	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
 {
-	type Val = ();
+	const IDENTIFIER: &'static str = "StorageWeightReclaim";
+
+	type AccountId = T::AccountId;
+	type Call = T::RuntimeCall;
+	type AdditionalSigned = ();
 	type Pre = Option<u64>;
 
-	fn prepare(
+	fn additional_signed(
+		&self,
+	) -> Result<Self::AdditionalSigned, sp_runtime::transaction_validity::TransactionValidityError>
+	{
+		Ok(())
+	}
+
+	fn pre_dispatch(
 		self,
-		_val: Self::Val,
-		_origin: &T::RuntimeOrigin,
-		_call: &T::RuntimeCall,
-		_info: &DispatchInfoOf<T::RuntimeCall>,
+		_who: &Self::AccountId,
+		_call: &Self::Call,
+		_info: &sp_runtime::traits::DispatchInfoOf<Self::Call>,
 		_len: usize,
-		_context: &Context,
-	) -> Result<Self::Pre, TransactionValidityError> {
+	) -> Result<Self::Pre, sp_runtime::transaction_validity::TransactionValidityError> {
 		Ok(get_proof_size())
 	}
 
 	fn post_dispatch(
-		pre: Self::Pre,
-		info: &DispatchInfoOf<T::RuntimeCall>,
-		post_info: &PostDispatchInfoOf<T::RuntimeCall>,
+		pre: Option<Self::Pre>,
+		info: &DispatchInfoOf<Self::Call>,
+		post_info: &PostDispatchInfoOf<Self::Call>,
 		_len: usize,
 		_result: &DispatchResult,
-		_context: &Context,
 	) -> Result<(), TransactionValidityError> {
-		let Some(pre_dispatch_proof_size) = pre else {
+		let Some(Some(pre_dispatch_proof_size)) = pre else {
 			return Ok(());
 		};
 
@@ -202,6 +194,470 @@ where
 		});
 		Ok(())
 	}
-
-	impl_tx_ext_default!(T::RuntimeCall; Context; validate);
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use frame_support::{
+		assert_ok,
+		dispatch::DispatchClass,
+		weights::{Weight, WeightMeter},
+	};
+	use frame_system::{BlockWeight, CheckWeight};
+	use sp_runtime::{AccountId32, BuildStorage};
+	use sp_std::marker::PhantomData;
+	use sp_trie::proof_size_extension::ProofSizeExt;
+
+	type Test = cumulus_test_runtime::Runtime;
+	const CALL: &<Test as Config>::RuntimeCall =
+		&cumulus_test_runtime::RuntimeCall::System(frame_system::Call::set_heap_pages {
+			pages: 0u64,
+		});
+	const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
+	const LEN: usize = 0;
+
+	pub fn new_test_ext() -> sp_io::TestExternalities {
+		let ext: sp_io::TestExternalities = cumulus_test_runtime::RuntimeGenesisConfig::default()
+			.build_storage()
+			.unwrap()
+			.into();
+		ext
+	}
+
+	struct TestRecorder {
+		return_values: Box<[usize]>,
+		counter: std::sync::atomic::AtomicUsize,
+	}
+
+	impl TestRecorder {
+		fn new(values: &[usize]) -> Self {
+			TestRecorder { return_values: values.into(), counter: Default::default() }
+		}
+	}
+
+	impl sp_trie::ProofSizeProvider for TestRecorder {
+		fn estimate_encoded_size(&self) -> usize {
+			let counter = self.counter.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+			self.return_values[counter]
+		}
+	}
+
+	fn setup_test_externalities(proof_values: &[usize]) -> sp_io::TestExternalities {
+		let mut test_ext = new_test_ext();
+		let test_recorder = TestRecorder::new(proof_values);
+		test_ext.register_extension(ProofSizeExt::new(test_recorder));
+		test_ext
+	}
+
+	fn set_current_storage_weight(new_weight: u64) {
+		BlockWeight::<Test>::mutate(|current_weight| {
+			current_weight.set(Weight::from_parts(0, new_weight), DispatchClass::Normal);
+		});
+	}
+
+	#[test]
+	fn basic_refund() {
+		// The real cost will be 100 bytes of storage size
+		let mut test_ext = setup_test_externalities(&[0, 100]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 500
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(0));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			// We expect a refund of 400
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 600);
+		})
+	}
+
+	#[test]
+	fn does_nothing_without_extension() {
+		let mut test_ext = new_test_ext();
+
+		// Proof size extension not registered
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 500
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, None);
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1000);
+		})
+	}
+
+	#[test]
+	fn negative_refund_is_added_to_weight() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 100
+			let info = DispatchInfo { weight: Weight::from_parts(0, 100), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// We expect no refund
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1100);
+		})
+	}
+
+	#[test]
+	fn test_zero_proof_size() {
+		let mut test_ext = setup_test_externalities(&[0, 0]);
+
+		test_ext.execute_with(|| {
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(0));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 0);
+		});
+	}
+
+	#[test]
+	fn test_larger_pre_dispatch_proof_size() {
+		let mut test_ext = setup_test_externalities(&[300, 100]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1300);
+
+			let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(300));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 800);
+		});
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 300
+			let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
+
+			// Actual weight is 50
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 250)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_on_negative() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 50
+			let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
+
+			// Actual weight is 25
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 25)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_reverse_order() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+
+			// Benchmarked storage weight: 300
+			let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
+
+			// Actual weight is 50
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 250)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+			// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
+		})
+	}
+
+	#[test]
+	fn test_incorporates_check_weight_unspent_weight_on_negative_reverse_order() {
+		let mut test_ext = setup_test_externalities(&[100, 300]);
+
+		test_ext.execute_with(|| {
+			set_current_storage_weight(1000);
+			// Benchmarked storage weight: 50
+			let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
+
+			// Actual weight is 25
+			let post_info = PostDispatchInfo {
+				actual_weight: Some(Weight::from_parts(50, 25)),
+				pays_fee: Default::default(),
+			};
+
+			let pre = StorageWeightReclaim::<Test>(PhantomData)
+				.pre_dispatch(&ALICE, CALL, &info, LEN)
+				.unwrap();
+			assert_eq!(pre, Some(100));
+
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
+				Some(pre),
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+			// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+			// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+			// we always need to call `post_dispatch` to verify that they interoperate correctly.
+			assert_ok!(CheckWeight::<Test>::post_dispatch(None, &info, &post_info, 0, &Ok(())));
+
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
+		})
+	}
+
+	#[test]
+	fn storage_size_reported_correctly() {
+		let mut test_ext = setup_test_externalities(&[1000]);
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), Some(1000));
+		});
+
+		let mut test_ext = new_test_ext();
+
+		let test_recorder = TestRecorder::new(&[0]);
+
+		test_ext.register_extension(ProofSizeExt::new(test_recorder));
+
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), Some(0));
+		});
+	}
+
+	#[test]
+	fn storage_size_disabled_reported_correctly() {
+		let mut test_ext = setup_test_externalities(&[PROOF_RECORDING_DISABLED as usize]);
+
+		test_ext.execute_with(|| {
+			assert_eq!(get_proof_size(), None);
+		});
+	}
+
+	#[test]
+	fn test_reclaim_helper() {
+		let mut test_ext = setup_test_externalities(&[1000, 1300, 1800]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 2000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			remaining_weight_meter.consume(Weight::from_parts(0, 500));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 200)));
+
+			remaining_weight_meter.consume(Weight::from_parts(0, 800));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 300)));
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1200));
+		});
+	}
+
+	#[test]
+	fn test_reclaim_helper_does_not_reclaim_negative() {
+		// Benchmarked weight does not change at all
+		let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1000));
+		});
+
+		// Benchmarked weight increases less than storage proof consumes
+		let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+		test_ext.execute_with(|| {
+			let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+			let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+			remaining_weight_meter.consume(Weight::from_parts(0, 0));
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+		});
+	}
+
+	/// Just here for doc purposes
+	fn get_benched_weight() -> Weight {
+		Weight::from_parts(0, 5)
+	}
+
+	/// Just here for doc purposes
+	fn do_work() {}
+
+	#[docify::export_content(simple_reclaimer_example)]
+	fn reclaim_with_weight_meter() {
+		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(10, 10));
+
+		let benched_weight = get_benched_weight();
+
+		// It is important to instantiate the `StorageWeightReclaimer` before we consume the weight
+		// for a piece of work from the weight meter.
+		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+
+		if remaining_weight_meter.try_consume(benched_weight).is_ok() {
+			// Perform some work that takes has `benched_weight` storage weight.
+			do_work();
+
+			// Reclaimer will detect that we only consumed 2 bytes, so 3 bytes are reclaimed.
+			let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+			// We reclaimed 3 bytes of storage size!
+			assert_eq!(reclaimed, Some(Weight::from_parts(0, 3)));
+			assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 10);
+			assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(10, 8));
+		}
+	}
+
+	#[test]
+	fn test_reclaim_helper_works_with_meter() {
+		// The node will report 12 - 10 = 2 consumed storage size between the calls.
+		let mut test_ext = setup_test_externalities(&[10, 12]);
+
+		test_ext.execute_with(|| {
+			// Initial storage size is 10.
+			set_current_storage_weight(10);
+			reclaim_with_weight_meter();
+		});
+	}
 }

cumulus/primitives/storage-weight-reclaim/src/tests.rs

@@ -1,484 +0,0 @@
-// Copyright (C) Parity Technologies (UK) Ltd.
-// SPDX-License-Identifier: Apache-2.0
-
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// 	http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use super::*;
-use frame_support::{
-	assert_ok,
-	dispatch::DispatchClass,
-	weights::{Weight, WeightMeter},
-};
-use frame_system::{BlockWeight, CheckWeight};
-use sp_runtime::{traits::DispatchTransaction, AccountId32, BuildStorage};
-use sp_std::marker::PhantomData;
-use sp_trie::proof_size_extension::ProofSizeExt;
-
-type Test = cumulus_test_runtime::Runtime;
-const CALL: &<Test as Config>::RuntimeCall =
-	&cumulus_test_runtime::RuntimeCall::System(frame_system::Call::set_heap_pages { pages: 0u64 });
-const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
-const LEN: usize = 0;
-
-fn new_test_ext() -> sp_io::TestExternalities {
-	let ext: sp_io::TestExternalities = cumulus_test_runtime::RuntimeGenesisConfig::default()
-		.build_storage()
-		.unwrap()
-		.into();
-	ext
-}
-
-struct TestRecorder {
-	return_values: Box<[usize]>,
-	counter: std::sync::atomic::AtomicUsize,
-}
-
-impl TestRecorder {
-	fn new(values: &[usize]) -> Self {
-		TestRecorder { return_values: values.into(), counter: Default::default() }
-	}
-}
-
-impl sp_trie::ProofSizeProvider for TestRecorder {
-	fn estimate_encoded_size(&self) -> usize {
-		let counter = self.counter.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
-		self.return_values[counter]
-	}
-}
-
-fn setup_test_externalities(proof_values: &[usize]) -> sp_io::TestExternalities {
-	let mut test_ext = new_test_ext();
-	let test_recorder = TestRecorder::new(proof_values);
-	test_ext.register_extension(ProofSizeExt::new(test_recorder));
-	test_ext
-}
-
-fn set_current_storage_weight(new_weight: u64) {
-	BlockWeight::<Test>::mutate(|current_weight| {
-		current_weight.set(Weight::from_parts(0, new_weight), DispatchClass::Normal);
-	});
-}
-
-#[test]
-fn basic_refund() {
-	// The real cost will be 100 bytes of storage size
-	let mut test_ext = setup_test_externalities(&[0, 100]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-
-		// Benchmarked storage weight: 500
-		let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
-		let post_info = PostDispatchInfo::default();
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(0));
-
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		// We expect a refund of 400
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 600);
-	})
-}
-
-#[test]
-fn does_nothing_without_extension() {
-	let mut test_ext = new_test_ext();
-
-	// Proof size extension not registered
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-
-		// Benchmarked storage weight: 500
-		let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
-		let post_info = PostDispatchInfo::default();
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, None);
-
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1000);
-	})
-}
-
-#[test]
-fn negative_refund_is_added_to_weight() {
-	let mut test_ext = setup_test_externalities(&[100, 300]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-		// Benchmarked storage weight: 100
-		let info = DispatchInfo { weight: Weight::from_parts(0, 100), ..Default::default() };
-		let post_info = PostDispatchInfo::default();
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(100));
-
-		// We expect no refund
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1100);
-	})
-}
-
-#[test]
-fn test_zero_proof_size() {
-	let mut test_ext = setup_test_externalities(&[0, 0]);
-
-	test_ext.execute_with(|| {
-		let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
-		let post_info = PostDispatchInfo::default();
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(0));
-
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 0);
-	});
-}
-
-#[test]
-fn test_larger_pre_dispatch_proof_size() {
-	let mut test_ext = setup_test_externalities(&[300, 100]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1300);
-
-		let info = DispatchInfo { weight: Weight::from_parts(0, 500), ..Default::default() };
-		let post_info = PostDispatchInfo::default();
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(300));
-
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 800);
-	});
-}
-
-#[test]
-fn test_incorporates_check_weight_unspent_weight() {
-	let mut test_ext = setup_test_externalities(&[100, 300]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-
-		// Benchmarked storage weight: 300
-		let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
-
-		// Actual weight is 50
-		let post_info = PostDispatchInfo {
-			actual_weight: Some(Weight::from_parts(50, 250)),
-			pays_fee: Default::default(),
-		};
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(100));
-
-		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
-		// we always need to call `post_dispatch` to verify that they interoperate correctly.
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
-	})
-}
-
-#[test]
-fn test_incorporates_check_weight_unspent_weight_on_negative() {
-	let mut test_ext = setup_test_externalities(&[100, 300]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-		// Benchmarked storage weight: 50
-		let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
-
-		// Actual weight is 25
-		let post_info = PostDispatchInfo {
-			actual_weight: Some(Weight::from_parts(50, 25)),
-			pays_fee: Default::default(),
-		};
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(100));
-
-		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
-		// we always need to call `post_dispatch` to verify that they interoperate correctly.
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
-	})
-}
-
-#[test]
-fn test_incorporates_check_weight_unspent_weight_reverse_order() {
-	let mut test_ext = setup_test_externalities(&[100, 300]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-
-		// Benchmarked storage weight: 300
-		let info = DispatchInfo { weight: Weight::from_parts(100, 300), ..Default::default() };
-
-		// Actual weight is 50
-		let post_info = PostDispatchInfo {
-			actual_weight: Some(Weight::from_parts(50, 250)),
-			pays_fee: Default::default(),
-		};
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(100));
-
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-		// `CheckWeight` gets called after `StorageWeightReclaim` this time.
-		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
-		// we always need to call `post_dispatch` to verify that they interoperate correctly.
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 900);
-	})
-}
-
-#[test]
-fn test_incorporates_check_weight_unspent_weight_on_negative_reverse_order() {
-	let mut test_ext = setup_test_externalities(&[100, 300]);
-
-	test_ext.execute_with(|| {
-		set_current_storage_weight(1000);
-		// Benchmarked storage weight: 50
-		let info = DispatchInfo { weight: Weight::from_parts(100, 50), ..Default::default() };
-
-		// Actual weight is 25
-		let post_info = PostDispatchInfo {
-			actual_weight: Some(Weight::from_parts(50, 25)),
-			pays_fee: Default::default(),
-		};
-
-		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
-			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN)
-			.unwrap();
-		assert_eq!(pre, Some(100));
-
-		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch(
-			pre,
-			&info,
-			&post_info,
-			LEN,
-			&Ok(()),
-			&()
-		));
-		// `CheckWeight` gets called after `StorageWeightReclaim` this time.
-		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
-		// we always need to call `post_dispatch` to verify that they interoperate correctly.
-		assert_ok!(CheckWeight::<Test>::post_dispatch((), &info, &post_info, 0, &Ok(()), &()));
-
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 1150);
-	})
-}
-
-#[test]
-fn storage_size_reported_correctly() {
-	let mut test_ext = setup_test_externalities(&[1000]);
-	test_ext.execute_with(|| {
-		assert_eq!(get_proof_size(), Some(1000));
-	});
-
-	let mut test_ext = new_test_ext();
-
-	let test_recorder = TestRecorder::new(&[0]);
-
-	test_ext.register_extension(ProofSizeExt::new(test_recorder));
-
-	test_ext.execute_with(|| {
-		assert_eq!(get_proof_size(), Some(0));
-	});
-}
-
-#[test]
-fn storage_size_disabled_reported_correctly() {
-	let mut test_ext = setup_test_externalities(&[PROOF_RECORDING_DISABLED as usize]);
-
-	test_ext.execute_with(|| {
-		assert_eq!(get_proof_size(), None);
-	});
-}
-
-#[test]
-fn test_reclaim_helper() {
-	let mut test_ext = setup_test_externalities(&[1000, 1300, 1800]);
-
-	test_ext.execute_with(|| {
-		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 2000));
-		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
-		remaining_weight_meter.consume(Weight::from_parts(0, 500));
-		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
-
-		assert_eq!(reclaimed, Some(Weight::from_parts(0, 200)));
-
-		remaining_weight_meter.consume(Weight::from_parts(0, 800));
-		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
-		assert_eq!(reclaimed, Some(Weight::from_parts(0, 300)));
-		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1200));
-	});
-}
-
-#[test]
-fn test_reclaim_helper_does_not_reclaim_negative() {
-	// Benchmarked weight does not change at all
-	let mut test_ext = setup_test_externalities(&[1000, 1300]);
-
-	test_ext.execute_with(|| {
-		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
-		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
-		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
-
-		assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
-		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1000));
-	});
-
-	// Benchmarked weight increases less than storage proof consumes
-	let mut test_ext = setup_test_externalities(&[1000, 1300]);
-
-	test_ext.execute_with(|| {
-		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
-		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
-		remaining_weight_meter.consume(Weight::from_parts(0, 0));
-		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
-
-		assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
-	});
-}
-
-/// Just here for doc purposes
-fn get_benched_weight() -> Weight {
-	Weight::from_parts(0, 5)
-}
-
-/// Just here for doc purposes
-fn do_work() {}
-
-#[docify::export_content(simple_reclaimer_example)]
-fn reclaim_with_weight_meter() {
-	let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(10, 10));
-
-	let benched_weight = get_benched_weight();
-
-	// It is important to instantiate the `StorageWeightReclaimer` before we consume the weight
-	// for a piece of work from the weight meter.
-	let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
-
-	if remaining_weight_meter.try_consume(benched_weight).is_ok() {
-		// Perform some work that takes has `benched_weight` storage weight.
-		do_work();
-
-		// Reclaimer will detect that we only consumed 2 bytes, so 3 bytes are reclaimed.
-		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
-
-		// We reclaimed 3 bytes of storage size!
-		assert_eq!(reclaimed, Some(Weight::from_parts(0, 3)));
-		assert_eq!(BlockWeight::<Test>::get().total().proof_size(), 10);
-		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(10, 8));
-	}
-}
-
-#[test]
-fn test_reclaim_helper_works_with_meter() {
-	// The node will report 12 - 10 = 2 consumed storage size between the calls.
-	let mut test_ext = setup_test_externalities(&[10, 12]);
-
-	test_ext.execute_with(|| {
-		// Initial storage size is 10.
-		set_current_storage_weight(10);
-		reclaim_with_weight_meter();
-	});
-}