Improve host function usage

cumulus/parachains/common/src/impls.rs

@@ -16,6 +16,7 @@
 //! Auxiliary struct/enums for parachain runtimes.
 //! Taken from polkadot/runtime/common (at a21cd64) and adapted for parachains.
 
+use cumulus_primitives_proof_size_hostfunction::storage_proof_size::storage_proof_size;
 use frame_support::traits::{
 	fungibles::{self, Balanced, Credit},
 	Contains, ContainsPair, Currency, Get, Imbalance, OnUnbalanced,
@@ -25,6 +26,14 @@ use sp_runtime::traits::Zero;
 use sp_std::marker::PhantomData;
 use xcm::latest::{AssetId, Fungibility::Fungible, MultiAsset, MultiLocation};
 
+/// Returns the current storage proof size from the host side.
+///
+/// Returns `None` if proof recording is disabled on the host.
+pub fn get_storage_size() -> Option<u64> {
+	let proof_size = storage_proof_size();
+	(proof_size != u64::MAX).then_some(proof_size)
+}
+
 /// Type alias to conveniently refer to the `Currency::NegativeImbalance` associated type.
 pub type NegativeImbalance<T> = <pallet_balances::Pallet<T> as Currency<
 	<T as frame_system::Config>::AccountId,

cumulus/parachains/common/src/storage_weight_reclaim_extension.rs

@@ -55,7 +55,7 @@ where
 	type AccountId = T::AccountId;
 	type Call = T::RuntimeCall;
 	type AdditionalSigned = ();
-	type Pre = u64;
+	type Pre = Option<u64>;
 
 	fn additional_signed(
 		&self,
@@ -71,7 +71,7 @@ where
 		_info: &sp_runtime::traits::DispatchInfoOf<Self::Call>,
 		_len: usize,
 	) -> Result<Self::Pre, sp_runtime::transaction_validity::TransactionValidityError> {
-		let proof_size = storage_proof_size();
+		let proof_size = crate::impls::get_storage_size();
 		Ok(proof_size)
 	}
 
@@ -82,12 +82,11 @@ where
 		_len: usize,
 		_result: &DispatchResult,
 	) -> Result<(), TransactionValidityError> {
-		if let Some(pre_dispatch_proof_size) = pre {
-			if pre_dispatch_proof_size == 0 {
+		if let Some(Some(pre_dispatch_proof_size)) = pre {
+			let Some(post_dispatch_proof_size) = crate::impls::get_storage_size() else {
+				log::debug!(target: LOG_TARGET, "Proof recording enabled during pre-dispatch, now disabled. This should not happen.");
 				return Ok(())
-			}
-
-			let post_dispatch_proof_size = storage_proof_size();
+			};
 			let benchmarked_weight = info.weight.proof_size();
 			let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);