RefundRelayerForMessagesFromParachain improvements (#1879)

* RefundRelayerForMessagesFromParachain improvements

* Address code review comments

bridges/bin/runtime-common/src/refund_relayer_extension.rs

@@ -19,37 +19,32 @@
 //! with calls that are: delivering new messsage and all necessary underlying headers
 //! (parachain or relay chain).
 
-use crate::messages::target::FromBridgedChainMessagesProof;
-
-use bp_messages::{target_chain::SourceHeaderChain, LaneId, MessageNonce};
-use bp_polkadot_core::parachains::ParaId;
-use bp_runtime::{Chain, HashOf};
+use crate::messages_call_ext::{
+	MessagesCallSubType, ReceiveMessagesProofHelper, ReceiveMessagesProofInfo,
+};
+use bp_messages::LaneId;
 use codec::{Decode, Encode};
 use frame_support::{
 	dispatch::{CallableCallFor, DispatchInfo, Dispatchable, PostDispatchInfo},
 	traits::IsSubType,
 	CloneNoBound, DefaultNoBound, EqNoBound, PartialEqNoBound, RuntimeDebugNoBound,
 };
-use pallet_bridge_grandpa::{
-	BridgedChain, Call as GrandpaCall, Config as GrandpaConfig, Pallet as GrandpaPallet,
-};
-use pallet_bridge_messages::{
-	Call as MessagesCall, Config as MessagesConfig, Pallet as MessagesPallet,
-};
+use pallet_bridge_grandpa::{CallSubType as GrandpaCallSubType, SubmitFinalityProofHelper};
+use pallet_bridge_messages::Config as MessagesConfig;
 use pallet_bridge_parachains::{
-	Call as ParachainsCall, Config as ParachainsConfig, Pallet as ParachainsPallet, RelayBlockHash,
-	RelayBlockHasher, RelayBlockNumber,
+	BoundedBridgeGrandpaConfig, CallSubType as ParachainsCallSubType, Config as ParachainsConfig,
+	RelayBlockNumber, SubmitParachainHeadsHelper, SubmitParachainHeadsInfo,
 };
 use pallet_bridge_relayers::{Config as RelayersConfig, Pallet as RelayersPallet};
 use pallet_transaction_payment::{Config as TransactionPaymentConfig, OnChargeTransaction};
 use pallet_utility::{Call as UtilityCall, Config as UtilityConfig, Pallet as UtilityPallet};
 use scale_info::TypeInfo;
 use sp_runtime::{
-	traits::{DispatchInfoOf, Get, Header as HeaderT, PostDispatchInfoOf, SignedExtension, Zero},
+	traits::{DispatchInfoOf, Get, PostDispatchInfoOf, SignedExtension, Zero},
 	transaction_validity::{TransactionValidity, TransactionValidityError, ValidTransaction},
 	DispatchResult, FixedPointOperand,
 };
-use sp_std::marker::PhantomData;
+use sp_std::{marker::PhantomData, vec, vec::Vec};
 
 // TODO (https://github.com/paritytech/parity-bridges-common/issues/1667):
 // support multiple bridges in this extension
@@ -81,6 +76,7 @@ where
 		pallet_transaction_payment::Pallet::<R>::compute_actual_fee(len as _, info, post_info, tip)
 	}
 }
+
 /// Signed extension that refunds relayer for new messages coming from the parachain.
 ///
 /// Also refunds relayer for successful finality delivery if it comes in batch (`utility.batchAll`)
@@ -99,18 +95,40 @@ where
 	RuntimeDebugNoBound,
 	TypeInfo,
 )]
-#[scale_info(skip_type_params(RT, GI, PI, MI, BE, PID, LID, FEE))]
-#[allow(clippy::type_complexity)] // TODO: get rid of that in https://github.com/paritytech/parity-bridges-common/issues/1666
-pub struct RefundRelayerForMessagesFromParachain<RT, GI, PI, MI, BE, PID, LID, FEE>(
-	PhantomData<(RT, GI, PI, MI, BE, PID, LID, FEE)>,
+#[scale_info(skip_type_params(RT, GI, PI, MI, PID, LID, FEE))]
+pub struct RefundRelayerForMessagesFromParachain<RT, GI, PI, MI, PID, LID, FEE>(
+	PhantomData<(RT, GI, PI, MI, PID, LID, FEE)>,
 );
 
+impl<R, GI, PI, MI, PID, LID, FEE>
+	RefundRelayerForMessagesFromParachain<R, GI, PI, MI, PID, LID, FEE>
+where
+	R: UtilityConfig<RuntimeCall = CallOf<R>>,
+	CallOf<R>: IsSubType<CallableCallFor<UtilityPallet<R>, R>>,
+{
+	fn expand_call<'a>(&self, call: &'a CallOf<R>) -> Option<Vec<&'a CallOf<R>>> {
+		let calls = match call.is_sub_type() {
+			Some(UtilityCall::<R>::batch_all { ref calls }) => {
+				if calls.len() > 3 {
+					return None
+				}
+
+				calls.iter().collect()
+			},
+			Some(_) => return None,
+			None => vec![call],
+		};
+
+		Some(calls)
+	}
+}
+
 /// Data that is crafted in `pre_dispatch` method and used at `post_dispatch`.
 #[derive(PartialEq)]
 #[cfg_attr(test, derive(Debug))]
 pub struct PreDispatchData<AccountId> {
 	/// Transaction submitter (relayer) account.
-	pub relayer: AccountId,
+	relayer: AccountId,
 	/// Type of the call.
 	pub call_type: CallType,
 }
@@ -119,89 +137,51 @@ pub struct PreDispatchData<AccountId> {
 #[derive(Clone, Copy, PartialEq, RuntimeDebugNoBound)]
 pub enum CallType {
 	/// Relay chain finality + parachain finality + message delivery calls.
-	AllFinalityAndDelivery(ExpectedRelayChainState, ExpectedParachainState, MessagesState),
+	AllFinalityAndDelivery(RelayBlockNumber, SubmitParachainHeadsInfo, ReceiveMessagesProofInfo),
 	/// Parachain finality + message delivery calls.
-	ParachainFinalityAndDelivery(ExpectedParachainState, MessagesState),
+	ParachainFinalityAndDelivery(SubmitParachainHeadsInfo, ReceiveMessagesProofInfo),
 	/// Standalone message delivery call.
-	Delivery(MessagesState),
+	Delivery(ReceiveMessagesProofInfo),
 }
 
 impl CallType {
 	/// Returns the pre-dispatch messages pallet state.
-	fn pre_dispatch_messages_state(&self) -> MessagesState {
+	fn receive_messages_proof_info(&self) -> ReceiveMessagesProofInfo {
 		match *self {
-			Self::AllFinalityAndDelivery(_, _, messages_state) => messages_state,
-			Self::ParachainFinalityAndDelivery(_, messages_state) => messages_state,
-			Self::Delivery(messages_state) => messages_state,
+			Self::AllFinalityAndDelivery(_, _, info) => info,
+			Self::ParachainFinalityAndDelivery(_, info) => info,
+			Self::Delivery(info) => info,
 		}
 	}
 }
 
-/// Expected post-dispatch state of the relay chain pallet.
-#[derive(Clone, Copy, PartialEq, RuntimeDebugNoBound)]
-pub struct ExpectedRelayChainState {
-	/// Best known relay chain block number.
-	pub best_block_number: RelayBlockNumber,
-}
-
-/// Expected post-dispatch state of the parachain pallet.
-#[derive(Clone, Copy, PartialEq, RuntimeDebugNoBound)]
-pub struct ExpectedParachainState {
-	/// At which relay block the parachain head has been updated?
-	pub at_relay_block_number: RelayBlockNumber,
-}
-
-/// Pre-dispatch state of messages pallet.
-///
-/// This struct is for pre-dispatch state of the pallet, not the expected post-dispatch state.
-/// That's because message delivery transaction may deliver some of messages that it brings.
-/// If this happens, we consider it "helpful" and refund its cost. If transaction fails to
-/// deliver at least one message, it is considered wrong and is not refunded.
-#[derive(Clone, Copy, PartialEq, RuntimeDebugNoBound)]
-pub struct MessagesState {
-	/// Best delivered message nonce.
-	pub best_nonce: MessageNonce,
-}
-
 // without this typedef rustfmt fails with internal err
 type BalanceOf<R> =
 	<<R as TransactionPaymentConfig>::OnChargeTransaction as OnChargeTransaction<R>>::Balance;
 type CallOf<R> = <R as frame_system::Config>::RuntimeCall;
 
-impl<R, GI, PI, MI, BE, PID, LID, FEE> SignedExtension
-	for RefundRelayerForMessagesFromParachain<R, GI, PI, MI, BE, PID, LID, FEE>
+impl<R, GI, PI, MI, PID, LID, FEE> SignedExtension
+	for RefundRelayerForMessagesFromParachain<R, GI, PI, MI, PID, LID, FEE>
 where
 	R: 'static
 		+ Send
 		+ Sync
-		+ frame_system::Config
 		+ UtilityConfig<RuntimeCall = CallOf<R>>
-		+ GrandpaConfig<GI>
+		+ BoundedBridgeGrandpaConfig<GI>
 		+ ParachainsConfig<PI, BridgesGrandpaPalletInstance = GI>
 		+ MessagesConfig<MI>
 		+ RelayersConfig,
 	GI: 'static + Send + Sync,
 	PI: 'static + Send + Sync,
 	MI: 'static + Send + Sync,
-	BE: 'static
-		+ Send
-		+ Sync
-		+ Default
-		+ SignedExtension<AccountId = R::AccountId, Call = CallOf<R>>,
 	PID: 'static + Send + Sync + Get<u32>,
 	LID: 'static + Send + Sync + Get<LaneId>,
 	FEE: 'static + Send + Sync + TransactionFeeCalculation<<R as RelayersConfig>::Reward>,
-	<R as frame_system::Config>::RuntimeCall:
-		Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
-	CallOf<R>: IsSubType<CallableCallFor<UtilityPallet<R>, R>>
-		+ IsSubType<CallableCallFor<GrandpaPallet<R, GI>, R>>
-		+ IsSubType<CallableCallFor<ParachainsPallet<R, PI>, R>>
-		+ IsSubType<CallableCallFor<MessagesPallet<R, MI>, R>>,
-	<R as GrandpaConfig<GI>>::BridgedChain:
-		Chain<BlockNumber = RelayBlockNumber, Hash = RelayBlockHash, Hasher = RelayBlockHasher>,
-	<R as MessagesConfig<MI>>::SourceHeaderChain: SourceHeaderChain<
-		MessagesProof = FromBridgedChainMessagesProof<HashOf<BridgedChain<R, GI>>>,
-	>,
+	CallOf<R>: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>
+		+ IsSubType<CallableCallFor<UtilityPallet<R>, R>>
+		+ GrandpaCallSubType<R, GI>
+		+ ParachainsCallSubType<R, PI>
+		+ MessagesCallSubType<R, MI>,
 {
 	const IDENTIFIER: &'static str = "RefundRelayerForMessagesFromParachain";
 	type AccountId = R::AccountId;
@@ -215,17 +195,20 @@ where
 
 	fn validate(
 		&self,
-		who: &Self::AccountId,
+		_who: &Self::AccountId,
 		call: &Self::Call,
-		info: &DispatchInfoOf<Self::Call>,
-		len: usize,
+		_info: &DispatchInfoOf<Self::Call>,
+		_len: usize,
 	) -> TransactionValidity {
-		// reject batch transactions with obsolete headers
-		if let Some(UtilityCall::<R>::batch_all { ref calls }) = call.is_sub_type() {
-			for nested_call in calls {
-				let reject_obsolete_transactions = BE::default();
-				reject_obsolete_transactions.pre_dispatch(who, nested_call, info, len)?;
-			}
+		let calls = match self.expand_call(call) {
+			Some(calls) => calls,
+			None => return Ok(ValidTransaction::default()),
+		};
+
+		for call in calls {
+			call.check_obsolete_submit_finality_proof()?;
+			call.check_obsolete_submit_parachain_heads()?;
+			call.check_obsolete_receive_messages_proof()?;
 		}
 
 		Ok(ValidTransaction::default())
@@ -241,40 +224,37 @@ where
 		// reject batch transactions with obsolete headers
 		self.validate(who, call, info, len).map(drop)?;
 
-		// now try to check if tx matches one of types we support
+		// Try to check if the tx matches one of types we support.
 		let parse_call_type = || {
-			if let Some(UtilityCall::<R>::batch_all { ref calls }) = call.is_sub_type() {
-				if calls.len() == 3 {
-					return Some(CallType::AllFinalityAndDelivery(
-						extract_expected_relay_chain_state::<R, GI>(&calls[0])?,
-						extract_expected_parachain_state::<R, GI, PI, PID>(&calls[1])?,
-						extract_messages_state::<R, GI, MI, LID>(&calls[2])?,
-					))
-				}
-				if calls.len() == 2 {
-					return Some(CallType::ParachainFinalityAndDelivery(
-						extract_expected_parachain_state::<R, GI, PI, PID>(&calls[0])?,
-						extract_messages_state::<R, GI, MI, LID>(&calls[1])?,
-					))
-				}
-				return None
+			let mut calls = self.expand_call(call)?.into_iter();
+			match calls.len() {
+				3 => Some(CallType::AllFinalityAndDelivery(
+					calls.next()?.submit_finality_proof_info()?,
+					calls.next()?.submit_parachain_heads_info_for(PID::get())?,
+					calls.next()?.receive_messages_proof_info_for(LID::get())?,
+				)),
+				2 => Some(CallType::ParachainFinalityAndDelivery(
+					calls.next()?.submit_parachain_heads_info_for(PID::get())?,
+					calls.next()?.receive_messages_proof_info_for(LID::get())?,
+				)),
+				1 => Some(CallType::Delivery(
+					calls.next()?.receive_messages_proof_info_for(LID::get())?,
+				)),
+				_ => None,
 			}
-
-			Some(CallType::Delivery(extract_messages_state::<R, GI, MI, LID>(call)?))
 		};
 
-		Ok(parse_call_type()
-			.map(|call_type| {
-				log::trace!(
-					target: "runtime::bridge",
-					"RefundRelayerForMessagesFromParachain from parachain {} via {:?} parsed bridge transaction in pre-dispatch: {:?}",
-					PID::get(),
-					LID::get(),
-					call_type,
-				);
-				PreDispatchData { relayer: who.clone(), call_type }
-			})
-		)
+		Ok(parse_call_type().map(|call_type| {
+			log::trace!(
+				target: "runtime::bridge",
+				"RefundRelayerForMessagesFromParachain from parachain {} via {:?} \
+				parsed bridge transaction in pre-dispatch: {:?}",
+				PID::get(),
+				LID::get(),
+				call_type,
+			);
+			PreDispatchData { relayer: who.clone(), call_type }
+		}))
 	}
 
 	fn post_dispatch(
@@ -284,40 +264,37 @@ where
 		len: usize,
 		result: &DispatchResult,
 	) -> Result<(), TransactionValidityError> {
-		// we never refund anything if it is not bridge transaction or if it is a bridge
-		// transaction that we do not support here
+		// We don't refund anything if the transaction has failed.
+		if result.is_err() {
+			return Ok(())
+		}
+
+		// We don't refund anything for transactions that we don't support.
 		let (relayer, call_type) = match pre {
 			Some(Some(pre)) => (pre.relayer, pre.call_type),
 			_ => return Ok(()),
 		};
 
-		// we never refund anything if transaction has failed
-		if result.is_err() {
-			return Ok(())
-		}
-
 		// check if relay chain state has been updated
-		if let CallType::AllFinalityAndDelivery(expected_relay_chain_state, _, _) = call_type {
-			let actual_relay_chain_state = relay_chain_state::<R, GI>();
-			if actual_relay_chain_state != Some(expected_relay_chain_state) {
+		if let CallType::AllFinalityAndDelivery(relay_block_number, _, _) = call_type {
+			if !SubmitFinalityProofHelper::<R, GI>::was_successful(relay_block_number) {
 				// we only refund relayer if all calls have updated chain state
 				return Ok(())
 			}
 
-			// there's a conflict between how bridge GRANDPA pallet works and the
-			// `AllFinalityAndDelivery` transaction. If relay chain header is mandatory, the GRANDPA
-			// pallet returns `Pays::No`, because such transaction is mandatory for operating the
-			// bridge. But `utility.batchAll` transaction always requires payment. But in both cases
-			// we'll refund relayer - either explicitly here, or using `Pays::No` if he's choosing
+			// there's a conflict between how bridge GRANDPA pallet works and a `utility.batchAll`
+			// transaction. If relay chain header is mandatory, the GRANDPA pallet returns
+			// `Pays::No`, because such transaction is mandatory for operating the bridge. But
+			// `utility.batchAll` transaction always requires payment. But in both cases we'll
+			// refund relayer - either explicitly here, or using `Pays::No` if he's choosing
 			// to submit dedicated transaction.
 		}
 
 		// check if parachain state has been updated
 		match call_type {
-			CallType::AllFinalityAndDelivery(_, expected_parachain_state, _) |
-			CallType::ParachainFinalityAndDelivery(expected_parachain_state, _) => {
-				let actual_parachain_state = parachain_state::<R, PI, PID>();
-				if actual_parachain_state != Some(expected_parachain_state) {
+			CallType::AllFinalityAndDelivery(_, parachain_heads_info, _) |
+			CallType::ParachainFinalityAndDelivery(parachain_heads_info, _) => {
+				if !SubmitParachainHeadsHelper::<R, PI>::was_successful(&parachain_heads_info) {
 					// we only refund relayer if all calls have updated chain state
 					return Ok(())
 				}
@@ -325,11 +302,10 @@ where
 			_ => (),
 		}
 
-		// check if messages have been delivered
-		let actual_messages_state = messages_state::<R, MI, LID>();
-		let pre_dispatch_messages_state = call_type.pre_dispatch_messages_state();
-		if actual_messages_state == Some(pre_dispatch_messages_state) {
-			// we only refund relayer if all calls have updated chain state
+		// Check if the `ReceiveMessagesProof` call delivered at least some of the messages that
+		// it contained. If this happens, we consider the transaction "helpful" and refund it.
+		let messages_proof_info = call_type.receive_messages_proof_info();
+		if !ReceiveMessagesProofHelper::<R, MI>::was_partially_successful(&messages_proof_info) {
 			return Ok(())
 		}
 
@@ -359,125 +335,22 @@ where
 	}
 }
 
-/// Extracts expected relay chain state from the call.
-fn extract_expected_relay_chain_state<R, GI>(call: &CallOf<R>) -> Option<ExpectedRelayChainState>
-where
-	R: GrandpaConfig<GI>,
-	GI: 'static,
-	<R as GrandpaConfig<GI>>::BridgedChain: Chain<BlockNumber = RelayBlockNumber>,
-	CallOf<R>: IsSubType<CallableCallFor<GrandpaPallet<R, GI>, R>>,
-{
-	if let Some(GrandpaCall::<R, GI>::submit_finality_proof { ref finality_target, .. }) =
-		call.is_sub_type()
-	{
-		return Some(ExpectedRelayChainState { best_block_number: *finality_target.number() })
-	}
-	None
-}
-
-/// Extracts expected parachain state from the call.
-fn extract_expected_parachain_state<R, GI, PI, PID>(
-	call: &CallOf<R>,
-) -> Option<ExpectedParachainState>
-where
-	R: GrandpaConfig<GI> + ParachainsConfig<PI, BridgesGrandpaPalletInstance = GI>,
-	GI: 'static,
-	PI: 'static,
-	PID: Get<u32>,
-	<R as GrandpaConfig<GI>>::BridgedChain:
-		Chain<BlockNumber = RelayBlockNumber, Hash = RelayBlockHash, Hasher = RelayBlockHasher>,
-	CallOf<R>: IsSubType<CallableCallFor<ParachainsPallet<R, PI>, R>>,
-{
-	if let Some(ParachainsCall::<R, PI>::submit_parachain_heads {
-		ref at_relay_block,
-		ref parachains,
-		..
-	}) = call.is_sub_type()
-	{
-		if parachains.len() != 1 || parachains[0].0 != ParaId(PID::get()) {
-			return None
-		}
-
-		return Some(ExpectedParachainState { at_relay_block_number: at_relay_block.0 })
-	}
-	None
-}
-
-/// Extracts messages state from the call.
-fn extract_messages_state<R, GI, MI, LID>(call: &CallOf<R>) -> Option<MessagesState>
-where
-	R: GrandpaConfig<GI> + MessagesConfig<MI>,
-	GI: 'static,
-	MI: 'static,
-	LID: Get<LaneId>,
-	CallOf<R>: IsSubType<CallableCallFor<MessagesPallet<R, MI>, R>>,
-	<R as MessagesConfig<MI>>::SourceHeaderChain: SourceHeaderChain<
-		MessagesProof = FromBridgedChainMessagesProof<HashOf<BridgedChain<R, GI>>>,
-	>,
-{
-	if let Some(MessagesCall::<R, MI>::receive_messages_proof { ref proof, .. }) =
-		call.is_sub_type()
-	{
-		if LID::get() != proof.lane {
-			return None
-		}
-
-		return Some(MessagesState {
-			best_nonce: MessagesPallet::<R, MI>::inbound_lane_data(proof.lane)
-				.last_delivered_nonce(),
-		})
-	}
-	None
-}
-
-/// Returns relay chain state that we are interested in.
-fn relay_chain_state<R, GI>() -> Option<ExpectedRelayChainState>
-where
-	R: GrandpaConfig<GI>,
-	GI: 'static,
-	<R as GrandpaConfig<GI>>::BridgedChain: Chain<BlockNumber = RelayBlockNumber>,
-{
-	GrandpaPallet::<R, GI>::best_finalized_number()
-		.map(|best_block_number| ExpectedRelayChainState { best_block_number })
-}
-
-/// Returns parachain state that we are interested in.
-fn parachain_state<R, PI, PID>() -> Option<ExpectedParachainState>
-where
-	R: ParachainsConfig<PI>,
-	PI: 'static,
-	PID: Get<u32>,
-{
-	ParachainsPallet::<R, PI>::best_parachain_info(ParaId(PID::get())).map(|para_info| {
-		ExpectedParachainState {
-			at_relay_block_number: para_info.best_head_hash.at_relay_block_number,
-		}
-	})
-}
-
-/// Returns messages state that we are interested in.
-fn messages_state<R, MI, LID>() -> Option<MessagesState>
-where
-	R: MessagesConfig<MI>,
-	MI: 'static,
-	LID: Get<LaneId>,
-{
-	Some(MessagesState {
-		best_nonce: MessagesPallet::<R, MI>::inbound_lane_data(LID::get()).last_delivered_nonce(),
-	})
-}
-
 #[cfg(test)]
 mod tests {
 	use super::*;
 	use crate::{messages::target::FromBridgedChainMessagesProof, mock::*};
-	use bp_messages::InboundLaneData;
+	use bp_messages::{InboundLaneData, MessageNonce};
 	use bp_parachains::{BestParaHeadHash, ParaInfo};
-	use bp_polkadot_core::parachains::ParaHeadsProof;
+	use bp_polkadot_core::parachains::{ParaHash, ParaHeadsProof, ParaId};
 	use bp_runtime::HeaderId;
 	use bp_test_utils::make_default_justification;
 	use frame_support::{assert_storage_noop, parameter_types, weights::Weight};
-	use sp_runtime::{transaction_validity::InvalidTransaction, DispatchError};
+	use pallet_bridge_grandpa::Call as GrandpaCall;
+	use pallet_bridge_messages::Call as MessagesCall;
+	use pallet_bridge_parachains::{Call as ParachainsCall, RelayBlockHash};
+	use sp_runtime::{
+		traits::Header as HeaderT, transaction_validity::InvalidTransaction, DispatchError,
+	};
 
 	parameter_types! {
 		pub TestParachain: u32 = 1000;
@@ -489,7 +362,6 @@ mod tests {
 		(),
 		(),
 		(),
-		BridgeRejectObsoleteHeadersAndMessages,
 		TestParachain,
 		TestLaneId,
 		TestRuntime,
@@ -506,6 +378,7 @@ mod tests {
 	fn initialize_environment(
 		best_relay_header_number: RelayBlockNumber,
 		parachain_head_at_relay_header_number: RelayBlockNumber,
+		parachain_head_hash: ParaHash,
 		best_delivered_message: MessageNonce,
 	) {
 		let best_relay_header = HeaderId(best_relay_header_number, RelayBlockHash::default());
@@ -515,7 +388,7 @@ mod tests {
 		let para_info = ParaInfo {
 			best_head_hash: BestParaHeadHash {
 				at_relay_block_number: parachain_head_at_relay_header_number,
-				head_hash: Default::default(),
+				head_hash: parachain_head_hash,
 			},
 			next_imported_hash_position: 0,
 		};
@@ -598,9 +471,17 @@ mod tests {
 		PreDispatchData {
 			relayer: relayer_account_at_this_chain(),
 			call_type: CallType::AllFinalityAndDelivery(
-				ExpectedRelayChainState { best_block_number: 200 },
-				ExpectedParachainState { at_relay_block_number: 200 },
-				MessagesState { best_nonce: 100 },
+				200,
+				SubmitParachainHeadsInfo {
+					at_relay_block_number: 200,
+					para_id: ParaId(TestParachain::get()),
+					para_head_hash: [1u8; 32].into(),
+				},
+				ReceiveMessagesProofInfo {
+					lane_id: TEST_LANE_ID,
+					best_proof_nonce: 200,
+					best_stored_nonce: 100,
+				},
 			),
 		}
 	}
@@ -609,8 +490,16 @@ mod tests {
 		PreDispatchData {
 			relayer: relayer_account_at_this_chain(),
 			call_type: CallType::ParachainFinalityAndDelivery(
-				ExpectedParachainState { at_relay_block_number: 200 },
-				MessagesState { best_nonce: 100 },
+				SubmitParachainHeadsInfo {
+					at_relay_block_number: 200,
+					para_id: ParaId(TestParachain::get()),
+					para_head_hash: [1u8; 32].into(),
+				},
+				ReceiveMessagesProofInfo {
+					lane_id: TEST_LANE_ID,
+					best_proof_nonce: 200,
+					best_stored_nonce: 100,
+				},
 			),
 		}
 	}
@@ -618,7 +507,11 @@ mod tests {
 	fn delivery_pre_dispatch_data() -> PreDispatchData<ThisChainAccountId> {
 		PreDispatchData {
 			relayer: relayer_account_at_this_chain(),
-			call_type: CallType::Delivery(MessagesState { best_nonce: 100 }),
+			call_type: CallType::Delivery(ReceiveMessagesProofInfo {
+				lane_id: TEST_LANE_ID,
+				best_proof_nonce: 200,
+				best_stored_nonce: 100,
+			}),
 		}
 	}
 
@@ -678,7 +571,7 @@ mod tests {
 	#[test]
 	fn validate_allows_non_obsolete_transactions() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(run_validate(message_delivery_call(200)), Ok(ValidTransaction::default()),);
 
@@ -697,7 +590,7 @@ mod tests {
 	#[test]
 	fn ext_rejects_batch_with_obsolete_relay_chain_header() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(all_finality_and_delivery_batch_call(100, 200, 200)),
@@ -714,7 +607,7 @@ mod tests {
 	#[test]
 	fn ext_rejects_batch_with_obsolete_parachain_head() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(all_finality_and_delivery_batch_call(101, 100, 200)),
@@ -741,7 +634,7 @@ mod tests {
 	#[test]
 	fn ext_rejects_batch_with_obsolete_messages() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(all_finality_and_delivery_batch_call(200, 200, 100)),
@@ -768,7 +661,7 @@ mod tests {
 	#[test]
 	fn pre_dispatch_parses_batch_with_relay_chain_and_parachain_headers() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(all_finality_and_delivery_batch_call(200, 200, 200)),
@@ -780,7 +673,7 @@ mod tests {
 	#[test]
 	fn pre_dispatch_parses_batch_with_parachain_header() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(parachain_finality_and_delivery_batch_call(200, 200)),
@@ -792,7 +685,7 @@ mod tests {
 	#[test]
 	fn pre_dispatch_fails_to_parse_batch_with_multiple_parachain_headers() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			let call = RuntimeCall::Utility(UtilityCall::batch_all {
 				calls: vec![
@@ -815,7 +708,7 @@ mod tests {
 	#[test]
 	fn pre_dispatch_parses_message_delivery_transaction() {
 		run_test(|| {
-			initialize_environment(100, 100, 100);
+			initialize_environment(100, 100, Default::default(), 100);
 
 			assert_eq!(
 				run_pre_dispatch(message_delivery_call(200)),
@@ -844,7 +737,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_ignores_transaction_that_has_not_updated_relay_chain_state() {
 		run_test(|| {
-			initialize_environment(100, 200, 200);
+			initialize_environment(100, 200, Default::default(), 200);
 
 			assert_storage_noop!(run_post_dispatch(Some(all_finality_pre_dispatch_data()), Ok(())));
 		});
@@ -853,7 +746,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_ignores_transaction_that_has_not_updated_parachain_state() {
 		run_test(|| {
-			initialize_environment(200, 100, 200);
+			initialize_environment(200, 100, Default::default(), 200);
 
 			assert_storage_noop!(run_post_dispatch(Some(all_finality_pre_dispatch_data()), Ok(())));
 			assert_storage_noop!(run_post_dispatch(
@@ -866,7 +759,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_ignores_transaction_that_has_not_delivered_any_messages() {
 		run_test(|| {
-			initialize_environment(200, 200, 100);
+			initialize_environment(200, 200, Default::default(), 100);
 
 			assert_storage_noop!(run_post_dispatch(Some(all_finality_pre_dispatch_data()), Ok(())));
 			assert_storage_noop!(run_post_dispatch(
@@ -880,7 +773,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_refunds_relayer_in_all_finality_batch() {
 		run_test(|| {
-			initialize_environment(200, 200, 200);
+			initialize_environment(200, 200, [1u8; 32].into(), 200);
 
 			run_post_dispatch(Some(all_finality_pre_dispatch_data()), Ok(()));
 			assert_eq!(
@@ -896,7 +789,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_refunds_relayer_in_parachain_finality_batch() {
 		run_test(|| {
-			initialize_environment(200, 200, 200);
+			initialize_environment(200, 200, [1u8; 32].into(), 200);
 
 			run_post_dispatch(Some(parachain_finality_pre_dispatch_data()), Ok(()));
 			assert_eq!(
@@ -912,7 +805,7 @@ mod tests {
 	#[test]
 	fn post_dispatch_refunds_relayer_in_message_delivery_transaction() {
 		run_test(|| {
-			initialize_environment(200, 200, 200);
+			initialize_environment(200, 200, Default::default(), 200);
 
 			run_post_dispatch(Some(delivery_pre_dispatch_data()), Ok(()));
 			assert_eq!(