Refactoring Checkpoint: (WIP)

pezcumulus/pezpallets/ah-ops/Cargo.toml

@@ -0,0 +1,66 @@
+[package]
+name = "pezpallet-ah-ops"
+description = "Operations cleanup pallet for the post-migration Asset Hub"
+license = "Apache-2.0"
+version = "0.1.0"
+edition.workspace = true
+authors.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezpallet-ah-ops"
+homepage = { workspace = true }
+
+[package.metadata.pezkuwi-sdk]
+exclude-from-umbrella = true # for AHM only
+
+[dependencies]
+codec = { workspace = true, features = ["max-encoded-len"] }
+pezframe-benchmarking = { workspace = true, optional = true }
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+log = { workspace = true }
+pezpallet-balances = { workspace = true }
+pezpallet-timestamp = { workspace = true }
+pezpallet-utility = { workspace = true }
+scale-info = { workspace = true, features = ["derive"] }
+pezsp-application-crypto = { workspace = true }
+pezsp-core = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+pezsp-std = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezframe-benchmarking?/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "log/std",
+  "pezpallet-balances/std",
+  "pezpallet-timestamp/std",
+  "pezpallet-utility/std",
+  "scale-info/std",
+  "pezsp-application-crypto/std",
+  "pezsp-core/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "pezsp-std/std",
+]
+runtime-benchmarks = [
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-balances/runtime-benchmarks",
+  "pezpallet-timestamp/runtime-benchmarks",
+  "pezpallet-utility/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+]
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-balances/try-runtime",
+  "pezpallet-timestamp/try-runtime",
+  "pezpallet-utility/try-runtime",
+  "pezsp-runtime/try-runtime",
+]

pezcumulus/pezpallets/ah-ops/src/benchmarking.rs

@@ -0,0 +1,100 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 crate::*;
+use pezframe_benchmarking::{account, v2::*};
+use pezframe_support::{dispatch::RawOrigin, traits::Currency};
+
+#[benchmarks]
+pub mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn unreserve_lease_deposit() {
+		let sender = account("sender", 0, 0);
+		let ed = <T::Currency as Currency<_>>::minimum_balance();
+		let _ = T::Currency::deposit_creating(&sender, ed + ed);
+		let _ = T::Currency::reserve(&sender, ed);
+		let block = T::RcBlockNumberProvider::current_block_number();
+		let para_id = ParaId::from(1u16);
+		RcLeaseReserve::<T>::insert((block, para_id, &sender), ed);
+
+		assert_eq!(T::Currency::reserved_balance(&sender), ed);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(sender.clone()), block, None, para_id);
+
+		assert_eq!(T::Currency::reserved_balance(&sender), 0);
+		assert_eq!(RcLeaseReserve::<T>::get((block, para_id, &sender)), None);
+	}
+
+	#[benchmark]
+	fn withdraw_crowdloan_contribution() {
+		let pot = account("pot", 0, 0);
+		let ed = <T::Currency as Currency<_>>::minimum_balance();
+		let _ = T::Currency::deposit_creating(&pot, ed + ed);
+		let _ = T::Currency::reserve(&pot, ed);
+		let block = T::RcBlockNumberProvider::current_block_number();
+		let para_id = ParaId::from(1u16);
+		RcLeaseReserve::<T>::insert((block, para_id, &pot), ed);
+
+		let sender = account("sender", 0, 0);
+		RcCrowdloanContribution::<T>::insert((block, para_id, &sender), (pot.clone(), ed));
+
+		assert_eq!(T::Currency::free_balance(&sender), 0);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(sender.clone()), block, None, para_id);
+
+		assert_eq!(RcCrowdloanContribution::<T>::get((block, para_id, &sender)), None);
+		assert_eq!(RcLeaseReserve::<T>::get((block, para_id, &pot)), None);
+		assert_eq!(T::Currency::free_balance(&pot), ed);
+	}
+
+	#[benchmark]
+	fn unreserve_crowdloan_reserve() {
+		let sender = account("sender", 0, 0);
+		let ed = <T::Currency as Currency<_>>::minimum_balance();
+		let _ = T::Currency::deposit_creating(&sender, ed + ed);
+		let _ = T::Currency::reserve(&sender, ed);
+		let block = T::RcBlockNumberProvider::current_block_number();
+		let para_id = ParaId::from(1u16);
+		RcCrowdloanReserve::<T>::insert((block, para_id, &sender), ed);
+
+		assert_eq!(T::Currency::reserved_balance(&sender), ed);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(sender.clone()), block, None, para_id);
+
+		assert_eq!(T::Currency::reserved_balance(&sender), 0);
+		assert_eq!(RcCrowdloanReserve::<T>::get((block, para_id, &sender)), None);
+	}
+
+	#[cfg(feature = "std")]
+	pub fn test_unreserve_lease_deposit<T: Config>() {
+		_unreserve_lease_deposit::<T>(true)
+	}
+
+	#[cfg(feature = "std")]
+	pub fn test_withdraw_crowdloan_contribution<T: Config>() {
+		_withdraw_crowdloan_contribution::<T>(true)
+	}
+
+	#[cfg(feature = "std")]
+	pub fn test_unreserve_crowdloan_reserve<T: Config>() {
+		_unreserve_crowdloan_reserve::<T>(true)
+	}
+}

pezcumulus/pezpallets/ah-ops/src/lib.rs

@@ -0,0 +1,683 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! The operational pallet for the Asset Hub, designed to manage and facilitate the migration of
+//! subsystems such as Governance, Staking, Balances from the Relay Chain to the Asset Hub. This
+//! pallet works alongside its counterpart, `pezpallet_rc_migrator`, which handles migration
+//! processes on the Relay Chain side.
+//!
+//! This pallet is responsible for controlling the initiation, progression, and completion of the
+//! migration process, including managing its various stages and transferring the necessary data.
+//! The pallet directly accesses the storage of other pallets for read/write operations while
+//! maintaining compatibility with their existing APIs.
+//!
+//! To simplify development and avoid the need to edit the original pallets, this pallet may
+//! duplicate private items such as storage entries from the original pallets. This ensures that the
+//! migration logic can be implemented without altering the original implementations.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+#[cfg(feature = "runtime-benchmarks")]
+pub mod benchmarking;
+#[cfg(test)]
+mod mock;
+#[cfg(test)]
+mod tests;
+pub mod weights;
+
+pub use pallet::*;
+pub use weights::WeightInfo;
+
+use codec::DecodeAll;
+use pezframe_support::{
+	pezpallet_prelude::*,
+	traits::{
+		fungible::{Inspect, InspectFreeze, Mutate, MutateFreeze, MutateHold, Unbalanced},
+		tokens::{Fortitude, IdAmount, Precision, Preservation},
+		Defensive, LockableCurrency, ReservableCurrency, WithdrawReasons as LockWithdrawReasons,
+	},
+};
+use pezframe_system::pezpallet_prelude::*;
+use pezpallet_balances::{AccountData, BalanceLock, Reasons as LockReasons};
+use pezsp_application_crypto::ByteArray;
+use pezsp_runtime::{traits::BlockNumberProvider, AccountId32};
+use pezsp_std::prelude::*;
+
+/// The log target of this pallet.
+pub const LOG_TARGET: &str = "runtime::ah-ops";
+
+pub type BalanceOf<T> = <T as pezpallet_balances::Config>::Balance;
+pub type DerivationIndex = u16;
+pub type ParaId = u16;
+
+#[pezframe_support::pallet]
+pub mod pallet {
+	use super::*;
+
+	#[pallet::config]
+	pub trait Config:
+		pezframe_system::Config<AccountData = AccountData<u128>, AccountId = AccountId32>
+		+ pezpallet_balances::Config<Balance = u128>
+		+ pezpallet_timestamp::Config<Moment = u64> // Needed for testing
+	{
+		/// Native asset type.
+		type Currency: Mutate<Self::AccountId, Balance = u128>
+			+ MutateHold<Self::AccountId, Reason = Self::RuntimeHoldReason>
+			+ InspectFreeze<Self::AccountId, Id = Self::FreezeIdentifier>
+			+ MutateFreeze<Self::AccountId>
+			+ Unbalanced<Self::AccountId>
+			+ ReservableCurrency<Self::AccountId, Balance = u128>
+			+ LockableCurrency<Self::AccountId, Balance = u128>;
+
+		/// Access the block number of the Relay Chain.
+		type RcBlockNumberProvider: BlockNumberProvider<BlockNumber = BlockNumberFor<Self>>;
+
+		/// The Weight information for extrinsics in this pallet.
+		type WeightInfo: WeightInfo;
+	}
+
+	/// Amount of balance that was reserved for winning a lease auction.
+	///
+	/// `unreserve_lease_deposit` can be permissionlessly called once the block number passed to
+	/// unreserve the deposit. It is implicitly called by `withdraw_crowdloan_contribution`.
+	///  
+	/// The account here can either be a crowdloan account or a solo bidder. If it is a crowdloan
+	/// account, then the summed up contributions for it in the contributions map will equate the
+	/// reserved balance here.
+	///
+	/// The keys are as follows:
+	/// - Block number after which the deposit can be unreserved.
+	/// - The para_id of the lease slot.
+	/// - The account that will have the balance unreserved.
+	/// - The balance to be unreserved.
+	#[pallet::storage]
+	pub type RcLeaseReserve<T: Config> = StorageNMap<
+		_,
+		(
+			NMapKey<Twox64Concat, BlockNumberFor<T>>,
+			NMapKey<Twox64Concat, ParaId>,
+			NMapKey<Twox64Concat, T::AccountId>,
+		),
+		BalanceOf<T>,
+		OptionQuery,
+	>;
+
+	/// Amount of balance that a contributor made towards a crowdloan.
+	///
+	/// `withdraw_crowdloan_contribution` can be permissionlessly called once the block number
+	/// passed to unlock the balance for a specific account.
+	///
+	/// The keys are as follows:
+	/// - Block number after which the balance can be unlocked.
+	/// - The para_id of the crowdloan.
+	/// - The account that made the contribution.
+	///
+	/// The value is (fund_pot, balance). The contribution pot is the second key in the
+	/// `RcCrowdloanContribution` storage.
+	#[pallet::storage]
+	pub type RcCrowdloanContribution<T: Config> = StorageNMap<
+		_,
+		(
+			NMapKey<Twox64Concat, BlockNumberFor<T>>,
+			NMapKey<Twox64Concat, ParaId>,
+			NMapKey<Twox64Concat, T::AccountId>,
+		),
+		(T::AccountId, BalanceOf<T>),
+		OptionQuery,
+	>;
+
+	/// The reserve that was taken to create a crowdloan.
+	///
+	/// This is normally 500 HEZ and can be refunded as last step after all
+	/// `RcCrowdloanContribution`s of this loan have been withdrawn.
+	///
+	/// Keys:
+	/// - Block number after which this can be unreserved
+	/// - The para_id of the crowdloan
+	/// - The account that will have the balance unreserved
+	#[pallet::storage]
+	pub type RcCrowdloanReserve<T: Config> = StorageNMap<
+		_,
+		(
+			NMapKey<Twox64Concat, BlockNumberFor<T>>,
+			NMapKey<Twox64Concat, ParaId>,
+			NMapKey<Twox64Concat, T::AccountId>,
+		),
+		BalanceOf<T>,
+		OptionQuery,
+	>;
+
+	#[pallet::error]
+	pub enum Error<T> {
+		/// Either no lease deposit or already unreserved.
+		NoLeaseReserve,
+		/// Either no crowdloan contribution or already withdrawn.
+		NoCrowdloanContribution,
+		/// Either no crowdloan reserve or already unreserved.
+		NoCrowdloanReserve,
+		/// Failed to withdraw crowdloan contribution.
+		FailedToWithdrawCrowdloanContribution,
+		/// Block number is not yet reached.
+		NotYet,
+		/// Not all contributions are withdrawn.
+		ContributionsRemaining,
+		/// Passed account IDs are not matching unmigrated child and sibling accounts.
+		WrongSovereignTranslation,
+		/// The account is not a derived account.
+		WrongDerivedTranslation,
+		/// Account cannot be migrated since it is not a sovereign teyrchain account.
+		NotSovereign,
+		/// Internal error, please bug report.
+		InternalError,
+		/// The migrated account would get reaped in the process.
+		WouldReap,
+		/// Failed to put a hold on an account.
+		FailedToPutHold,
+		/// Failed to release a hold from an account.
+		FailedToReleaseHold,
+		/// Failed to thaw a frozen balance.
+		FailedToThaw,
+		/// Failed to set a freeze on an account.
+		FailedToSetFreeze,
+		/// Failed to transfer a balance.
+		FailedToTransfer,
+		/// Failed to reserve a balance.
+		FailedToReserve,
+		/// Failed to unreserve the full balance.
+		CannotUnreserve,
+		/// The from and to accounts are identical.
+		AccountIdentical,
+	}
+
+	#[pallet::event]
+	#[pallet::generate_deposit(pub(crate) fn deposit_event)]
+	pub enum Event<T: Config> {
+		/// Some lease reserve could not be unreserved and needs manual cleanup.
+		LeaseUnreserveRemaining {
+			depositor: T::AccountId,
+			para_id: ParaId,
+			remaining: BalanceOf<T>,
+		},
+
+		/// Some amount for a crowdloan reserve could not be unreserved and needs manual cleanup.
+		CrowdloanUnreserveRemaining {
+			depositor: T::AccountId,
+			para_id: ParaId,
+			remaining: BalanceOf<T>,
+		},
+
+		/// A sovereign teyrchain account has been migrated from its child to sibling
+		/// representation.
+		SovereignMigrated {
+			/// The teyrchain ID that had its account migrated.
+			para_id: ParaId,
+			/// The old account that was migrated out of.
+			from: T::AccountId,
+			/// The new account that was migrated into.
+			to: T::AccountId,
+			/// Set if this account was derived from a para sovereign account.
+			derivation_index: Option<DerivationIndex>,
+		},
+
+		/// An amount of fungible balance was put on hold.
+		HoldPlaced { account: T::AccountId, amount: BalanceOf<T>, reason: T::RuntimeHoldReason },
+
+		/// An amount of fungible balance was released from its hold.
+		HoldReleased { account: T::AccountId, amount: BalanceOf<T>, reason: T::RuntimeHoldReason },
+	}
+
+	#[pallet::pallet]
+	pub struct Pallet<T>(_);
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {
+		/// Unreserve the deposit that was taken for creating a crowdloan.
+		///
+		/// This can be called by any signed origin. It unreserves the lease deposit on the account
+		/// that won the lease auction. It can be unreserved once all leases expired. Note that it
+		/// will be called automatically from `withdraw_crowdloan_contribution` for the matching
+		/// crowdloan account.
+		///
+		/// Solo bidder accounts that won lease auctions can use this to unreserve their amount.
+		#[pallet::call_index(0)]
+		#[pallet::weight(<T as Config>::WeightInfo::unreserve_lease_deposit())]
+		pub fn unreserve_lease_deposit(
+			origin: OriginFor<T>,
+			block: BlockNumberFor<T>,
+			depositor: Option<T::AccountId>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			let depositor = depositor.unwrap_or(sender);
+
+			Self::do_unreserve_lease_deposit(block, depositor, para_id).map_err(Into::into)
+		}
+
+		/// Withdraw the contribution of a finished crowdloan.
+		///
+		/// A crowdloan contribution can be withdrawn if either:
+		/// - The crowdloan failed to in an auction and timed out
+		/// - Won an auction and all leases expired
+		///
+		/// Can be called by any signed origin.
+		#[pallet::call_index(1)]
+		#[pallet::weight(<T as Config>::WeightInfo::withdraw_crowdloan_contribution())]
+		pub fn withdraw_crowdloan_contribution(
+			origin: OriginFor<T>,
+			block: BlockNumberFor<T>,
+			depositor: Option<T::AccountId>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			let depositor = depositor.unwrap_or(sender);
+
+			Self::do_withdraw_crowdloan_contribution(block, depositor, para_id).map_err(Into::into)
+		}
+
+		/// Unreserve the deposit that was taken for creating a crowdloan.
+		///
+		/// This can be called once either:
+		/// - The crowdloan failed to win an auction and timed out
+		/// - Won an auction, all leases expired and all contributions are withdrawn
+		///
+		/// Can be called by any signed origin. The condition that all contributions are withdrawn
+		/// is in place since the reserve acts as a storage deposit.
+		#[pallet::call_index(2)]
+		#[pallet::weight(<T as Config>::WeightInfo::unreserve_crowdloan_reserve())]
+		pub fn unreserve_crowdloan_reserve(
+			origin: OriginFor<T>,
+			block: BlockNumberFor<T>,
+			depositor: Option<T::AccountId>,
+			para_id: ParaId,
+		) -> DispatchResult {
+			let sender = ensure_signed(origin)?;
+			let depositor = depositor.unwrap_or(sender);
+
+			Self::do_unreserve_crowdloan_reserve(block, depositor, para_id).map_err(Into::into)
+		}
+
+		/// Try to migrate a teyrchain sovereign child account to its respective sibling.
+		///
+		/// Takes the old and new account and migrates it only if they are as expected. An event of
+		/// `SovereignMigrated` will be emitted if the account was migrated successfully.
+		///
+		/// Callable by any signed origin.
+		#[pallet::call_index(3)]
+		#[pallet::weight(T::DbWeight::get().reads_writes(15, 15)
+					.saturating_add(Weight::from_parts(0, 50_000)))]
+		pub fn migrate_teyrchain_sovereign_acc(
+			origin: OriginFor<T>,
+			from: T::AccountId,
+			to: T::AccountId,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_migrate_teyrchain_sovereign_derived_acc(&from, &to, None).map_err(Into::into)
+		}
+
+		/// Try to migrate a teyrchain sovereign child account to its respective sibling.
+		///
+		/// Takes the old and new account and migrates it only if they are as expected. An event of
+		/// `SovereignMigrated` will be emitted if the account was migrated successfully.
+		///
+		/// Callable by any signed origin.
+		#[pallet::call_index(5)]
+		#[pallet::weight(T::DbWeight::get().reads_writes(15, 15)
+					.saturating_add(Weight::from_parts(0, 50_000)))]
+		pub fn migrate_teyrchain_sovereign_derived_acc(
+			origin: OriginFor<T>,
+			from: T::AccountId,
+			to: T::AccountId,
+			derivation: (T::AccountId, DerivationIndex),
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_migrate_teyrchain_sovereign_derived_acc(&from, &to, Some(derivation))
+				.map_err(Into::into)
+		}
+
+		/// Force unreserve a named or unnamed reserve.
+		#[pallet::call_index(4)]
+		#[pallet::weight(T::DbWeight::get().reads_writes(10, 10)
+					.saturating_add(Weight::from_parts(0, 50_000)))]
+		pub fn force_unreserve(
+			origin: OriginFor<T>,
+			account: T::AccountId,
+			amount: BalanceOf<T>,
+			reason: Option<T::RuntimeHoldReason>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			Self::do_force_unreserve(account, amount, reason).map_err(Into::into)
+		}
+	}
+
+	impl<T: Config> Pallet<T> {
+		pub fn do_unreserve_lease_deposit(
+			block: BlockNumberFor<T>,
+			depositor: T::AccountId,
+			para_id: ParaId,
+		) -> Result<(), Error<T>> {
+			ensure!(block <= T::RcBlockNumberProvider::current_block_number(), Error::<T>::NotYet);
+			let balance = RcLeaseReserve::<T>::take((block, para_id, &depositor))
+				.ok_or(Error::<T>::NoLeaseReserve)?;
+
+			let remaining = <T as Config>::Currency::unreserve(&depositor, balance);
+			if remaining > 0 {
+				defensive!("Should be able to unreserve all");
+				Self::deposit_event(Event::LeaseUnreserveRemaining {
+					depositor,
+					remaining,
+					para_id,
+				});
+			}
+
+			Ok(())
+		}
+
+		pub fn do_withdraw_crowdloan_contribution(
+			block: BlockNumberFor<T>,
+			depositor: T::AccountId,
+			para_id: ParaId,
+		) -> Result<(), Error<T>> {
+			ensure!(block <= T::RcBlockNumberProvider::current_block_number(), Error::<T>::NotYet);
+			let (pot, contribution) =
+				RcCrowdloanContribution::<T>::take((block, para_id, &depositor))
+					.ok_or(Error::<T>::NoCrowdloanContribution)?;
+
+			// Maybe this is the first one to withdraw and we need to unreserve it from the pot
+			match Self::do_unreserve_lease_deposit(block, pot.clone(), para_id) {
+				Ok(()) => (),
+				Err(Error::<T>::NoLeaseReserve) => (), // fine
+				Err(e) => return Err(e),
+			}
+
+			// Ideally this does not fail. But if it does, then we keep it for manual inspection.
+			let transferred = <T as Config>::Currency::transfer(
+				&pot,
+				&depositor,
+				contribution,
+				Preservation::Preserve,
+			)
+			.defensive()
+			.map_err(|_| Error::<T>::FailedToWithdrawCrowdloanContribution)?;
+			defensive_assert!(transferred == contribution);
+			// Need to reactivate since we deactivated it here https://github.com/pezkuwichain/pezkuwi-sdk/blob/04847d515ef56da4d0801c9b89a4241dfa827b33/polkadot/runtime/common/src/crowdloan/mod.rs#L793
+			<T as Config>::Currency::reactivate(transferred);
+
+			Ok(())
+		}
+
+		pub fn do_unreserve_crowdloan_reserve(
+			block: BlockNumberFor<T>,
+			depositor: T::AccountId,
+			para_id: ParaId,
+		) -> Result<(), Error<T>> {
+			ensure!(block <= T::RcBlockNumberProvider::current_block_number(), Error::<T>::NotYet);
+			ensure!(
+				Self::contributions_withdrawn(block, para_id),
+				Error::<T>::ContributionsRemaining
+			);
+			let amount = RcCrowdloanReserve::<T>::take((block, para_id, &depositor))
+				.ok_or(Error::<T>::NoCrowdloanReserve)?;
+
+			let remaining = <T as Config>::Currency::unreserve(&depositor, amount);
+			if remaining > 0 {
+				defensive!("Should be able to unreserve all");
+				Self::deposit_event(Event::CrowdloanUnreserveRemaining {
+					depositor,
+					remaining,
+					para_id,
+				});
+			}
+
+			Ok(())
+		}
+
+		pub(crate) fn contributions_withdrawn(block: BlockNumberFor<T>, para_id: ParaId) -> bool {
+			let mut contrib_iter = RcCrowdloanContribution::<T>::iter_prefix((block, para_id));
+			contrib_iter.next().is_none()
+		}
+
+		pub fn do_migrate_teyrchain_sovereign_derived_acc(
+			from: &T::AccountId,
+			to: &T::AccountId,
+			derivation: Option<(T::AccountId, DerivationIndex)>,
+		) -> Result<(), Error<T>> {
+			if pezframe_system::Account::<T>::get(from) == Default::default() {
+				// Nothing to do if the account does not exist
+				return Ok(());
+			}
+			if from == to {
+				return Err(Error::<T>::AccountIdentical);
+			}
+			pezpallet_balances::Pallet::<T>::ensure_upgraded(from); // prevent future headache
+
+			let (translated_acc, para_id, index) = if let Some((parent, index)) = derivation {
+				let (parent_translated, para_id) =
+					Self::try_rc_sovereign_derived_to_ah(from, &parent, index)?;
+				(parent_translated, para_id, Some(index))
+			} else {
+				let (translated_acc, para_id) = Self::try_translate_rc_sovereign_to_ah(from)?;
+				(translated_acc, para_id, None)
+			};
+			ensure!(translated_acc == *to, Error::<T>::WrongSovereignTranslation);
+
+			// Release all locks
+			let locks: Vec<BalanceLock<T::Balance>> =
+				pezpallet_balances::Locks::<T>::get(from).into_inner();
+			for lock in &locks {
+				let () = <T as Config>::Currency::remove_lock(lock.id, from);
+			}
+
+			// Thaw all the freezes
+			let freezes: Vec<IdAmount<T::FreezeIdentifier, T::Balance>> =
+				pezpallet_balances::Freezes::<T>::get(from).into();
+
+			for freeze in &freezes {
+				let () = <T as Config>::Currency::thaw(&freeze.id, from)
+					.map_err(|_| Error::<T>::FailedToThaw)?;
+			}
+
+			// Release all holds
+			let holds: Vec<IdAmount<T::RuntimeHoldReason, T::Balance>> =
+				pezpallet_balances::Holds::<T>::get(from).into();
+
+			for IdAmount { id, amount } in &holds {
+				let _ = <T as Config>::Currency::release(id, from, *amount, Precision::Exact)
+					.map_err(|_| Error::<T>::FailedToReleaseHold)?;
+				Self::deposit_event(Event::HoldReleased {
+					account: from.clone(),
+					amount: *amount,
+					reason: *id,
+				});
+			}
+
+			// Unreserve unnamed reserves
+			let unnamed_reserve = <T as Config>::Currency::reserved_balance(from);
+			let missing = <T as Config>::Currency::unreserve(from, unnamed_reserve);
+			defensive_assert!(missing == 0, "Should have unreserved the full amount");
+
+			// Set consumer refs to zero
+			let consumers = pezframe_system::Pallet::<T>::consumers(from);
+			pezframe_system::Account::<T>::mutate(from, |acc| {
+				acc.consumers = 0;
+			});
+			// We dont handle sufficients and there should be none
+			ensure!(pezframe_system::Pallet::<T>::sufficients(from) == 0, Error::<T>::InternalError);
+
+			// Sanity check
+			let total = <T as Config>::Currency::total_balance(from);
+			let reducible = <T as Config>::Currency::reducible_balance(
+				from,
+				Preservation::Expendable,
+				Fortitude::Polite,
+			);
+			defensive_assert!(
+				total >= <T as Config>::Currency::minimum_balance(),
+				"Must have at least ED"
+			);
+			defensive_assert!(total == reducible, "Total balance should be reducible");
+
+			// Now the actual balance transfer to the new account
+			<T as Config>::Currency::transfer(from, to, total, Preservation::Expendable)
+				.defensive()
+				.map_err(|_| Error::<T>::FailedToTransfer)?;
+
+			// Apply consumer refs
+			pezframe_system::Account::<T>::mutate(to, |acc| {
+				acc.consumers += consumers;
+			});
+
+			// Reapply the holds
+			for hold in &holds {
+				<T as Config>::Currency::hold(&hold.id, to, hold.amount)
+					.map_err(|_| Error::<T>::FailedToPutHold)?;
+				// Somehow there are no events for this being emitted... so we emit our own.
+				Self::deposit_event(Event::HoldPlaced {
+					account: to.clone(),
+					amount: hold.amount,
+					reason: hold.id,
+				});
+			}
+
+			// Reapply the reserve
+			<T as Config>::Currency::reserve(to, unnamed_reserve)
+				.defensive()
+				.map_err(|_| Error::<T>::FailedToReserve)?;
+
+			// Reapply the locks
+			for lock in &locks {
+				let reasons = map_lock_reason(lock.reasons);
+				<T as Config>::Currency::set_lock(lock.id, to, lock.amount, reasons);
+			}
+			// Reapply the freezes
+			for freeze in &freezes {
+				<T as Config>::Currency::set_freeze(&freeze.id, to, freeze.amount)
+					.map_err(|_| Error::<T>::FailedToSetFreeze)?;
+			}
+
+			defensive_assert!(
+				pezframe_system::Account::<T>::get(from) == Default::default(),
+				"Must reap old account"
+			);
+			// If new account would die from this, then lets rather not do it and check it manually.
+			ensure!(
+				pezframe_system::Account::<T>::get(to) != Default::default(),
+				Error::<T>::WouldReap
+			);
+
+			Self::deposit_event(Event::SovereignMigrated {
+				para_id,
+				from: from.clone(),
+				to: to.clone(),
+				derivation_index: index,
+			});
+
+			Ok(())
+		}
+
+		pub fn do_force_unreserve(
+			account: T::AccountId,
+			amount: BalanceOf<T>,
+			reason: Option<T::RuntimeHoldReason>,
+		) -> Result<(), Error<T>> {
+			if let Some(reason) = reason {
+				<T as Config>::Currency::release(&reason, &account, amount, Precision::Exact)
+					.map_err(|_| Error::<T>::FailedToReleaseHold)?;
+				Self::deposit_event(Event::HoldReleased {
+					account: account.clone(),
+					amount,
+					reason,
+				});
+			} else {
+				let remaining = <T as Config>::Currency::unreserve(&account, amount);
+				if remaining > 0 {
+					return Err(Error::<T>::CannotUnreserve);
+				}
+			}
+
+			Ok(())
+		}
+
+		/// Try to translate a Teyrchain sovereign account to the Teyrchain AH sovereign account.
+		///
+		/// Returns:
+		/// - `Ok(None)` if the account is not a Teyrchain sovereign account
+		/// - `Ok(Some((ah_account, para_id)))` with the translated account and the para id
+		/// - `Err(())` otherwise
+		///
+		/// The way that this normally works is through the configured
+		/// `SiblingTeyrchainConvertsVia`: <https://github.com/polkadot-fellows/runtimes/blob/7b096c14c2b16cc81ca4e2188eea9103f120b7a4/system-parachains/asset-hubs/asset-hub-polkadot/src/xcm_config.rs#L93-L94>
+		/// it passes the `Sibling` type into it which has type-ID `sibl`:
+		/// <https://github.com/pezkuwichain/pezkuwi-sdk/blob/c10e25aaa8b8afd8665b53f0a0b02e4ea44caa77/polkadot/parachain/src/primitives.rs#L272-L274>
+		/// This type-ID gets used by the converter here:
+		/// <https://github.com/pezkuwichain/pezkuwi-sdk/blob/7ecf3f757a5d6f622309cea7f788e8a547a5dce8/polkadot/xcm/xcm-builder/src/location_conversion.rs#L314>
+		/// and eventually ends up in the encoding here
+		/// <https://github.com/pezkuwichain/pezkuwi-sdk/blob/cdf107de700388a52a17b2fb852c98420c78278e/bizinikiwi/primitives/runtime/src/traits/mod.rs#L1997-L1999>
+		/// The `para` conversion is likewise with `ChildTeyrchainConvertsVia` and the `para`
+		/// type-ID <https://github.com/pezkuwichain/pezkuwi-sdk/blob/c10e25aaa8b8afd8665b53f0a0b02e4ea44caa77/polkadot/parachain/src/primitives.rs#L162-L164>
+		pub fn try_translate_rc_sovereign_to_ah(
+			from: &AccountId32,
+		) -> Result<(AccountId32, ParaId), Error<T>> {
+			let raw = from.to_raw_vec();
+
+			// Must start with "para"
+			let Some(raw) = raw.strip_prefix(b"para") else {
+				return Err(Error::<T>::NotSovereign);
+			};
+			// Must end with 26 zero bytes
+			let Some(raw) = raw.strip_suffix(&[0u8; 26]) else {
+				return Err(Error::<T>::NotSovereign);
+			};
+			let para_id = u16::decode_all(&mut &raw[..]).map_err(|_| Error::<T>::InternalError)?;
+
+			// Translate to AH sibling account
+			let mut ah_raw = [0u8; 32];
+			ah_raw[0..4].copy_from_slice(b"sibl");
+			ah_raw[4..6].copy_from_slice(&para_id.encode());
+
+			Ok((ah_raw.into(), para_id))
+		}
+
+		/// Same as `try_translate_rc_sovereign_to_ah` but for derived accounts.
+		pub fn try_rc_sovereign_derived_to_ah(
+			from: &AccountId32,
+			parent: &AccountId32,
+			index: DerivationIndex,
+		) -> Result<(AccountId32, ParaId), Error<T>> {
+			// check the derivation proof
+			{
+				let derived = pezpallet_utility::derivative_account_id(parent.clone(), index);
+				ensure!(derived == *from, Error::<T>::WrongDerivedTranslation);
+			}
+
+			let (parent_translated, para_id) = Self::try_translate_rc_sovereign_to_ah(parent)?;
+			let parent_translated_derived =
+				pezpallet_utility::derivative_account_id(parent_translated, index);
+			Ok((parent_translated_derived, para_id))
+		}
+	}
+}
+
+/// Backward mapping from <https://github.com/paritytech/polkadot-sdk/blob/74a5e1a242274ddaadac1feb3990fc95c8612079/bizinikiwi/pezframe/balances/src/types.rs#L38>
+pub fn map_lock_reason(reasons: LockReasons) -> LockWithdrawReasons {
+	match reasons {
+		LockReasons::All => LockWithdrawReasons::TRANSACTION_PAYMENT | LockWithdrawReasons::RESERVE,
+		LockReasons::Fee => LockWithdrawReasons::TRANSACTION_PAYMENT,
+		LockReasons::Misc => LockWithdrawReasons::TIP,
+	}
+}

pezcumulus/pezpallets/ah-ops/src/mock.rs

@@ -0,0 +1,78 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 crate as pezpallet_ah_ops;
+use crate::*;
+use pezframe_support::derive_impl;
+use pezsp_core::H256;
+use pezsp_runtime::traits::{BlakeTwo256, IdentityLookup};
+
+type Block = pezframe_system::mocking::MockBlock<Runtime>;
+
+// For testing the pallet, we construct a mock runtime.
+pezframe_support::construct_runtime!(
+	pub enum Runtime {
+		System: pezframe_system,
+		Balances: pezpallet_balances,
+		AhOps: pezpallet_ah_ops,
+		Timestamp: pezpallet_timestamp,
+	}
+);
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl pezframe_system::Config for Runtime {
+	type BaseCallFilter = pezframe_support::traits::Everything;
+	type BlockWeights = ();
+	type BlockLength = ();
+	type DbWeight = ();
+	type RuntimeOrigin = RuntimeOrigin;
+	type Nonce = u64;
+	type Hash = H256;
+	type RuntimeCall = RuntimeCall;
+	type Hashing = BlakeTwo256;
+	type AccountId = AccountId32;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type Block = Block;
+	type RuntimeEvent = RuntimeEvent;
+	type Version = ();
+	type PalletInfo = PalletInfo;
+	type AccountData = pezpallet_balances::AccountData<u128>;
+	type OnNewAccount = ();
+	type OnKilledAccount = ();
+	type SystemWeightInfo = ();
+	type SS58Prefix = ();
+	type OnSetCode = ();
+	type MaxConsumers = pezframe_support::traits::ConstU32<16>;
+}
+
+#[derive_impl(pezpallet_balances::config_preludes::TestDefaultConfig)]
+impl pezpallet_balances::Config for Runtime {
+	type Balance = u128;
+	type AccountStore = System;
+}
+
+impl pezpallet_timestamp::Config for Runtime {
+	type Moment = u64;
+	type OnTimestampSet = ();
+	type MinimumPeriod = ();
+	type WeightInfo = ();
+}
+
+impl Config for Runtime {
+	type Currency = Balances;
+	type RcBlockNumberProvider = System; // Wrong but unused
+	type WeightInfo = ();
+}

pezcumulus/pezpallets/ah-ops/src/tests.rs

@@ -0,0 +1,299 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 pezframe_support::{
+	hypothetically,
+	traits::{fungible::Mutate, InspectLockableCurrency, LockableCurrency, WithdrawReasons},
+};
+use pezsp_runtime::AccountId32;
+use std::str::FromStr;
+
+use crate::mock::Runtime as AssetHub;
+
+#[test]
+fn sovereign_account_translation() {
+	// https://docs.google.com/document/d/1DXYWPXEwi0DkDfG8Fb2ZTI4DQBAz87DBCIW7yQIVrj0
+	let bifrost_cases = [
+		// Bifrost Pezkuwi #1
+		(
+			// para 2030
+			"13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm",
+			// sibl 2030
+			"13cKp89TtYknbyYnqnF6dWN75q5ZosvFSuqzoEVkUAaNR47A",
+			None,
+		),
+		// Bifrost Pezkuwi #2
+		(
+			// para 2030 index 0
+			"14vtfeKAVKh1Jzb3s7e43SqZ3zB5MLsdCxZPoKDxeoCFKLu5",
+			// sibl 2030 index 0
+			"5ETehspFKFNpBbe5DsfuziN6BWq5Qwp1J8qcTQQoAxwa7BsS",
+			// derivation proof (para 2030, index 0)
+			Some(("13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm", 0u16)),
+		),
+		// Bifrost Pezkuwi #3
+		(
+			// para 2030 index 1
+			"14QkQ7wVVDRrhbC1UqHsFwKFUns1SRud94CXMWGHWB8Jhtro",
+			// sibl 2030 index 1
+			"5DNWZkkAxLhqF8tevcbRGyARAVM7abukftmqvoDFUN5dDDDz",
+			// derivation proof (para 2030, index 1)
+			Some(("13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm", 1u16)),
+		),
+		// Bifrost Pezkuwi #4
+		(
+			// para 2030 index 2
+			"13hLwqcVHqjiJMbZhR9LtfdhoxmTdssi7Kp8EJaW2yfk3knK",
+			// sibl 2030 index 2
+			"5EmiwjDYiackJma1GW3aBbQ74rLfWh756UKDb7Cm83XDkUUZ",
+			// derivation proof (para 2030, index 2)
+			Some(("13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm", 2u16)),
+		),
+		// Bifrost Kusama #1
+		(
+			// para 2001
+			"5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E",
+			// sibl 2001
+			"5Eg2fntJDju46yds4uKzu2zuQssqw7JZWohhLMj6mZZjg2pK",
+			None,
+		),
+		// Bifrost Kusama #2
+		(
+			// para 2001 index 0
+			"5E78xTBiaN3nAGYtcNnqTJQJqYAkSDGggKqaDfpNsKyPpbcb",
+			// sibl 2001 index 0
+			"5CzXNqgBZT5yMpMETdfH55saYNKQoJBXsSfnu4d2s1ejYFir",
+			// derivation proof (para 2001, index 0)
+			Some(("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E", 0u16)),
+		),
+		// Bifrost Kusama #3
+		(
+			// para 2001 index 1
+			"5HXi9pzWnTQzk7VKzY6VQn92KfWCcA5NbSm53uKHrYU1VsjP",
+			// sibl 2001 index 1
+			"5GcexD4YNqcKTbW1YWDRczQzpxic61byeNeLaHgqQHk8pxQJ",
+			// derivation proof (para 2001, index 1)
+			Some(("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E", 1u16)),
+		),
+		// Bifrost Kusama #4
+		(
+			// para 2001 index 2
+			"5CkKS3YMx64TguUYrMERc5Bn6Mn2aKMUkcozUFREQDgHS3Tv",
+			// sibl 2001 index 2
+			"5FoYMVucmT552GDMWfYNxcF2XnuuvLbJHt7mU6DfDCpUAS2Y",
+			// derivation proof (para 2001, index 2)
+			Some(("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E", 2u16)),
+		),
+		// Bifrost Kusama #5
+		(
+			// para 2001 index 3
+			"5Crxhmiw5CQq3Mnfcu3dR3yJ3YpjbxjqaeDFtNNtqgmcnN4S",
+			// sibl 2001 index 3
+			"5FP39fgPYhJw3vcLwSMqMnwBuEVGexUMG6JQLPR9yPVhq6Wy",
+			// derivation proof (para 2001, index 3)
+			Some(("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E", 3u16)),
+		),
+		// Bifrost Kusama #5
+		(
+			// para 2001 index 3
+			"5DAZP4gZKZafGv42uoWNTMau4tYuDd2XteJLGL4upermhQpn",
+			// sibl 2001 index 3
+			"5ExtLdYnjHLJbngU1QpumjPieCGaCXwwkH1JrFBQ9GATuNGv",
+			// derivation proof (para 2001, index 4)
+			Some(("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E", 4u16)),
+		),
+	];
+
+	for (from, to, derivation) in bifrost_cases {
+		let from = AccountId32::from_str(from).unwrap();
+		let to = AccountId32::from_str(to).unwrap();
+
+		println!("Translating {from}/{derivation:?} -> {to}");
+		if let Some((parent, index)) = derivation {
+			let parent = AccountId32::from_str(parent).unwrap();
+			let (got_to, _) =
+				crate::Pallet::<AssetHub>::try_rc_sovereign_derived_to_ah(&from, &parent, index)
+					.unwrap();
+			assert_eq!(got_to, to);
+		} else {
+			let (got_to, _) =
+				crate::Pallet::<AssetHub>::try_translate_rc_sovereign_to_ah(&from).unwrap();
+			assert_eq!(got_to, to);
+		}
+	}
+}
+
+#[test]
+fn translate_sovereign_acc_good() {
+	pezsp_io::TestExternalities::new(Default::default()).execute_with(move || {
+		let balance = 1000000000000000000;
+		let lock = balance / 20;
+		const LID: [u8; 8] = *b"lockID00";
+
+		// Test for Para 2030 index 0 (Bifrost Pezkuwi derived 0)
+		let from =
+			AccountId32::from_str("14vtfeKAVKh1Jzb3s7e43SqZ3zB5MLsdCxZPoKDxeoCFKLu5").unwrap();
+		let to = AccountId32::from_str("5ETehspFKFNpBbe5DsfuziN6BWq5Qwp1J8qcTQQoAxwa7BsS").unwrap();
+		let parent =
+			AccountId32::from_str("13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm").unwrap();
+		let derivation_proof = Some((parent, 0u16));
+
+		// Works if the account does not exist
+		hypothetically!({
+			crate::Pallet::<AssetHub>::do_migrate_teyrchain_sovereign_derived_acc(
+				&from,
+				&to,
+				derivation_proof.clone(),
+			)
+			.unwrap();
+			// Also twice
+			crate::Pallet::<AssetHub>::do_migrate_teyrchain_sovereign_derived_acc(
+				&from,
+				&to,
+				derivation_proof.clone(),
+			)
+			.unwrap();
+		});
+
+		// But also if it exists
+		<AssetHub as crate::Config>::Currency::mint_into(&from, balance).unwrap();
+		hypothetically!({
+			crate::Pallet::<AssetHub>::do_migrate_teyrchain_sovereign_derived_acc(
+				&from,
+				&to,
+				derivation_proof.clone(),
+			)
+			.unwrap();
+			// Also twice
+			crate::Pallet::<AssetHub>::do_migrate_teyrchain_sovereign_derived_acc(
+				&from,
+				&to,
+				derivation_proof.clone(),
+			)
+			.unwrap();
+
+			// Balance was moved
+			assert_eq!(<AssetHub as crate::Config>::Currency::free_balance(&from), 0);
+			assert_eq!(<AssetHub as crate::Config>::Currency::free_balance(&to), balance);
+		});
+
+		// Can also have locks
+		<AssetHub as crate::Config>::Currency::set_lock(LID, &from, lock, WithdrawReasons::FEE);
+		hypothetically!({
+			crate::Pallet::<AssetHub>::do_migrate_teyrchain_sovereign_derived_acc(
+				&from,
+				&to,
+				derivation_proof.clone(),
+			)
+			.unwrap();
+
+			// Balance was moved
+			assert_eq!(<AssetHub as crate::Config>::Currency::free_balance(&from), 0);
+			assert_eq!(<AssetHub as crate::Config>::Currency::free_balance(&to), balance);
+
+			// Lock was moved
+			assert_eq!(<AssetHub as crate::Config>::Currency::balance_locked(LID, &from), 0);
+			assert_eq!(<AssetHub as crate::Config>::Currency::balance_locked(LID, &to), lock);
+		});
+	});
+}
+
+#[test]
+fn contributions_withdrawn_works() {
+	pezsp_io::TestExternalities::new(Default::default()).execute_with(|| {
+		let block_number: u64 = 100;
+		let para_id: u16 = 2000;
+		let contributor1 =
+			AccountId32::from_str("13YMK2eeopZtUNpeHnJ1Ws2HqMQG6Ts9PGCZYGyFbSYoZfcm").unwrap();
+		let contributor2 =
+			AccountId32::from_str("14vtfeKAVKh1Jzb3s7e43SqZ3zB5MLsdCxZPoKDxeoCFKLu5").unwrap();
+		let fund_pot =
+			AccountId32::from_str("5Ec4AhPV91i9yNuiWuNunPf6AQCYDhFTTA4G5QCbtqYApH9E").unwrap();
+
+		// Initially no contributions exist, so should return true
+		assert!(
+			crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return true when no contributions exist"
+		);
+
+		// Insert a contribution
+		crate::RcCrowdloanContribution::<AssetHub>::insert(
+			(block_number, para_id, &contributor1),
+			(fund_pot.clone(), 1000u128),
+		);
+
+		// Now should return false since there's a contribution
+		assert!(
+			!crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return false when contributions exist"
+		);
+
+		// Insert another contribution
+		crate::RcCrowdloanContribution::<AssetHub>::insert(
+			(block_number, para_id, &contributor2),
+			(fund_pot.clone(), 2000u128),
+		);
+
+		// Still should return false
+		assert!(
+			!crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return false when multiple contributions exist"
+		);
+
+		// Remove the first contribution
+		crate::RcCrowdloanContribution::<AssetHub>::remove((block_number, para_id, &contributor1));
+
+		// Still should return false (one contribution remains)
+		assert!(
+			!crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return false when one contribution still exists"
+		);
+
+		// Remove the second contribution
+		crate::RcCrowdloanContribution::<AssetHub>::remove((block_number, para_id, &contributor2));
+
+		// Now should return true again
+		assert!(
+			crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return true after all contributions are removed"
+		);
+
+		// Test with different para_id - should still be true (no contributions)
+		let other_para_id: u16 = 2001;
+		assert!(
+			crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, other_para_id),
+			"Should return true for different para_id with no contributions"
+		);
+
+		// Add contribution to original para_id but check different para_id
+		crate::RcCrowdloanContribution::<AssetHub>::insert(
+			(block_number, para_id, &contributor1),
+			(fund_pot, 500u128),
+		);
+
+		// Original para_id should now be false
+		assert!(
+			!crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, para_id),
+			"Should return false for para_id with contribution"
+		);
+
+		// Different para_id should still be true
+		assert!(
+			crate::Pallet::<AssetHub>::contributions_withdrawn(block_number, other_para_id),
+			"Should return true for different para_id"
+		);
+	});
+}

pezcumulus/pezpallets/ah-ops/src/weights.rs

@@ -0,0 +1,175 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Autogenerated weights for `pezpallet_ah_ops`
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE BIZINIKIWI BENCHMARK CLI VERSION 32.0.0
+//! DATE: 2025-04-28, STEPS: `50`, REPEAT: `2`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! HOSTNAME: `cob`, CPU: `<UNKNOWN>`
+//! WASM-EXECUTION: `Compiled`, CHAIN: `None`, DB CACHE: `1024`
+
+// Executed Command:
+// frame-omni-bencher
+// v1
+// benchmark
+// pallet
+// --runtime=target/production/wbuild/asset-hub-pezkuwi-runtime/asset_hub_pezkuwi_runtime.wasm
+// --pallet
+// pezpallet-ah-ops
+// --extrinsic
+// *
+// --output=./pallets/ah-ops/src/weights_1.rs
+// --template=./../../paritytech/pezkuwi-sdk/bizinikiwi/.maintain/frame-weight-template.hbs
+// --no-median-slopes
+// --no-min-squares
+// --repeat=2
+// --steps=50
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+#![allow(missing_docs)]
+#![allow(dead_code)]
+
+use pezframe_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}};
+use core::marker::PhantomData;
+
+/// Weight functions needed for `pezpallet_ah_ops`.
+pub trait WeightInfo {
+	fn unreserve_lease_deposit() -> Weight;
+	fn withdraw_crowdloan_contribution() -> Weight;
+	fn unreserve_crowdloan_reserve() -> Weight;
+}
+
+/// Weights for `pezpallet_ah_ops` using the Bizinikiwi node and recommended hardware.
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcLeaseReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcLeaseReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:1 w:1)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn unreserve_lease_deposit() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `365`
+		//  Estimated: `3593`
+		// Minimum execution time: 30_000_000 picoseconds.
+		Weight::from_parts(33_000_000, 3593)
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcCrowdloanContribution` (r:1 w:1)
+	/// Proof: `AhOps::RcCrowdloanContribution` (`max_values`: None, `max_size`: Some(112), added: 2587, mode: `MaxEncodedLen`)
+	/// Storage: `AhOps::RcLeaseReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcLeaseReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:2 w:2)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn withdraw_crowdloan_contribution() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `538`
+		//  Estimated: `6196`
+		// Minimum execution time: 64_000_000 picoseconds.
+		Weight::from_parts(65_000_000, 6196)
+			.saturating_add(T::DbWeight::get().reads(6_u64))
+			.saturating_add(T::DbWeight::get().writes(4_u64))
+	}
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcCrowdloanContribution` (r:1 w:0)
+	/// Proof: `AhOps::RcCrowdloanContribution` (`max_values`: None, `max_size`: Some(112), added: 2587, mode: `MaxEncodedLen`)
+	/// Storage: `AhOps::RcCrowdloanReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcCrowdloanReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:1 w:1)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn unreserve_crowdloan_reserve() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `365`
+		//  Estimated: `3593`
+		// Minimum execution time: 32_000_000 picoseconds.
+		Weight::from_parts(62_000_000, 3593)
+			.saturating_add(T::DbWeight::get().reads(5_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+}
+
+// For backwards compatibility and tests.
+impl WeightInfo for () {
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcLeaseReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcLeaseReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:1 w:1)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn unreserve_lease_deposit() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `365`
+		//  Estimated: `3593`
+		// Minimum execution time: 30_000_000 picoseconds.
+		Weight::from_parts(33_000_000, 3593)
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+	}
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcCrowdloanContribution` (r:1 w:1)
+	/// Proof: `AhOps::RcCrowdloanContribution` (`max_values`: None, `max_size`: Some(112), added: 2587, mode: `MaxEncodedLen`)
+	/// Storage: `AhOps::RcLeaseReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcLeaseReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:2 w:2)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn withdraw_crowdloan_contribution() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `538`
+		//  Estimated: `6196`
+		// Minimum execution time: 64_000_000 picoseconds.
+		Weight::from_parts(65_000_000, 6196)
+			.saturating_add(RocksDbWeight::get().reads(6_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+	/// Storage: `TeyrchainSystem::ValidationData` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::ValidationData` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `TeyrchainSystem::LastRelayChainBlockNumber` (r:1 w:0)
+	/// Proof: `TeyrchainSystem::LastRelayChainBlockNumber` (`max_values`: Some(1), `max_size`: None, mode: `Measured`)
+	/// Storage: `AhOps::RcCrowdloanContribution` (r:1 w:0)
+	/// Proof: `AhOps::RcCrowdloanContribution` (`max_values`: None, `max_size`: Some(112), added: 2587, mode: `MaxEncodedLen`)
+	/// Storage: `AhOps::RcCrowdloanReserve` (r:1 w:1)
+	/// Proof: `AhOps::RcCrowdloanReserve` (`max_values`: None, `max_size`: Some(80), added: 2555, mode: `MaxEncodedLen`)
+	/// Storage: `System::Account` (r:1 w:1)
+	/// Proof: `System::Account` (`max_values`: None, `max_size`: Some(128), added: 2603, mode: `MaxEncodedLen`)
+	fn unreserve_crowdloan_reserve() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `365`
+		//  Estimated: `3593`
+		// Minimum execution time: 32_000_000 picoseconds.
+		Weight::from_parts(62_000_000, 3593)
+			.saturating_add(RocksDbWeight::get().reads(5_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+	}
+}

pezcumulus/pezpallets/aura-ext/Cargo.toml

@@ -0,0 +1,86 @@
+[package]
+name = "pezcumulus-pezpallet-aura-ext"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+description = "AURA consensus extension pezpallet for teyrchains"
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-aura-ext"
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+# Bizinikiwi
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-aura = { workspace = true }
+pezpallet-timestamp = { workspace = true }
+pezsp-application-crypto = { workspace = true }
+pezsp-consensus-aura = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+# Pezcumulus
+pezcumulus-pezpallet-teyrchain-system = { workspace = true }
+
+[dev-dependencies]
+rstest = { workspace = true }
+
+# Pezcumulus
+pezcumulus-pezpallet-teyrchain-system = { workspace = true, default-features = true }
+pezcumulus-primitives-core = { workspace = true, default-features = true }
+pezcumulus-primitives-proof-size-hostfunction = { workspace = true, default-features = true }
+pezcumulus-test-relay-sproof-builder = { workspace = true, default-features = true }
+
+# Bizinikiwi
+pezsp-core = { workspace = true, default-features = true }
+pezsp-io = { workspace = true, default-features = true }
+pezsp-keyring = { workspace = true, default-features = true }
+pezsp-state-machine = { workspace = true, default-features = true }
+pezsp-trie = { workspace = true, default-features = true }
+pezsp-version = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezcumulus-pezpallet-teyrchain-system/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "pezpallet-aura/std",
+  "pezpallet-timestamp/std",
+  "scale-info/std",
+  "pezsp-application-crypto/std",
+  "pezsp-consensus-aura/std",
+  "pezsp-runtime/std",
+]
+try-runtime = [
+  "pezcumulus-pezpallet-teyrchain-system/try-runtime",
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-aura/try-runtime",
+  "pezpallet-timestamp/try-runtime",
+  "pezsp-runtime/try-runtime",
+]
+runtime-benchmarks = [
+  "pezcumulus-pezpallet-teyrchain-system/runtime-benchmarks",
+  "pezcumulus-primitives-core/runtime-benchmarks",
+  "pezcumulus-primitives-proof-size-hostfunction/runtime-benchmarks",
+  "pezcumulus-test-relay-sproof-builder/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-aura/runtime-benchmarks",
+  "pezpallet-timestamp/runtime-benchmarks",
+  "pezsp-consensus-aura/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-keyring/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "pezsp-state-machine/runtime-benchmarks",
+  "pezsp-trie/runtime-benchmarks",
+  "pezsp-version/runtime-benchmarks",
+]

pezcumulus/pezpallets/aura-ext/src/consensus_hook.rs

@@ -0,0 +1,166 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! The definition of a [`FixedVelocityConsensusHook`] for consensus logic to manage
+//! block velocity.
+use super::{pallet, Aura};
+use core::{marker::PhantomData, num::NonZeroU32};
+use pezcumulus_pezpallet_teyrchain_system::{
+	self as teyrchain_system,
+	consensus_hook::{ConsensusHook, UnincludedSegmentCapacity},
+	relay_state_snapshot::RelayChainStateProof,
+};
+use pezframe_support::pezpallet_prelude::*;
+use pezsp_consensus_aura::{Slot, SlotDuration};
+
+/// A consensus hook that enforces fixed block production velocity and unincluded segment capacity.
+///
+/// It keeps track of relay chain slot information and teyrchain blocks authored per relay chain
+/// slot.
+///
+/// # Type Parameters
+/// - `T` - The runtime configuration trait
+/// - `RELAY_CHAIN_SLOT_DURATION_MILLIS` - Duration of relay chain slots in milliseconds
+/// - `V` - Maximum number of blocks that can be authored per relay chain parent (velocity)
+/// - `C` - Maximum capacity of unincluded segment
+///
+/// # Example Configuration
+/// ```ignore
+/// type ConsensusHook = FixedVelocityConsensusHook<Runtime, 6000, 2, 8>;
+/// ```
+/// This configures:
+/// - 6 second relay chain slots
+/// - Maximum 2 blocks per slot
+/// - Maximum 8 blocks in unincluded segment
+pub struct FixedVelocityConsensusHook<
+	T,
+	const RELAY_CHAIN_SLOT_DURATION_MILLIS: u32,
+	const V: u32,
+	const C: u32,
+>(PhantomData<T>);
+
+impl<
+		T: pallet::Config,
+		const RELAY_CHAIN_SLOT_DURATION_MILLIS: u32,
+		const V: u32,
+		const C: u32,
+	> ConsensusHook for FixedVelocityConsensusHook<T, RELAY_CHAIN_SLOT_DURATION_MILLIS, V, C>
+where
+	<T as pezpallet_timestamp::Config>::Moment: Into<u64>,
+{
+	/// Consensus hook that performs validations on the provided relay chain state
+	/// proof:
+	/// - Ensures blocks are not produced faster than the specified velocity `V`
+	/// - Verifies teyrchain slot alignment with relay chain slot
+	///
+	/// # Panics
+	/// - When the relay chain slot from the state is smaller than the slot from the proof
+	/// - When the number of authored blocks exceeds velocity limit
+	/// - When teyrchain slot is ahead of the calculated slot from relay chain
+	fn on_state_proof(state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		// Ensure velocity is non-zero.
+		let velocity = V.max(1);
+		let relay_chain_slot = state_proof.read_slot().expect("failed to read relay chain slot");
+
+		let (relay_chain_slot, authored_in_relay) = match pallet::RelaySlotInfo::<T>::get() {
+			Some((slot, authored)) if slot == relay_chain_slot => (slot, authored),
+			Some((slot, _)) if slot < relay_chain_slot => (relay_chain_slot, 0),
+			Some((slot, _)) => {
+				panic!("Slot moved backwards: stored_slot={slot:?}, relay_chain_slot={relay_chain_slot:?}")
+			},
+			None => (relay_chain_slot, 0),
+		};
+
+		// We need to allow one additional block to be built to fill the unincluded segment.
+		if authored_in_relay > velocity {
+			panic!("authored blocks limit is reached for the slot: relay_chain_slot={relay_chain_slot:?}, authored={authored_in_relay:?}, velocity={velocity:?}");
+		}
+
+		pallet::RelaySlotInfo::<T>::put((relay_chain_slot, authored_in_relay + 1));
+
+		let para_slot = pezpallet_aura::CurrentSlot::<T>::get();
+
+		// Convert relay chain timestamp.
+		let relay_chain_timestamp =
+			u64::from(RELAY_CHAIN_SLOT_DURATION_MILLIS).saturating_mul(*relay_chain_slot);
+
+		let para_slot_duration = SlotDuration::from_millis(Aura::<T>::slot_duration().into());
+		let para_slot_from_relay =
+			Slot::from_timestamp(relay_chain_timestamp.into(), para_slot_duration);
+
+		if *para_slot != *para_slot_from_relay {
+			panic!(
+				"Teyrchain slot must match relay-derived slot: teyrchain_slot={:?}, derived_from_relay_slot={:?} velocity={:?}, relay_chain_slot={:?}",
+				para_slot,
+				para_slot_from_relay,
+				velocity,
+				relay_chain_slot
+			);
+		}
+
+		let weight = T::DbWeight::get().reads(1);
+
+		(
+			weight,
+			NonZeroU32::new(core::cmp::max(C, 1))
+				.expect("1 is the minimum value and non-zero; qed")
+				.into(),
+		)
+	}
+}
+
+impl<
+		T: pallet::Config + teyrchain_system::Config,
+		const RELAY_CHAIN_SLOT_DURATION_MILLIS: u32,
+		const V: u32,
+		const C: u32,
+	> FixedVelocityConsensusHook<T, RELAY_CHAIN_SLOT_DURATION_MILLIS, V, C>
+{
+	/// Whether it is legal to extend the chain, assuming the given block is the most
+	/// recently included one as-of the relay parent that will be built against, and
+	/// the given slot.
+	///
+	/// This should be consistent with the logic the runtime uses when validating blocks to
+	/// avoid issues.
+	///
+	/// When the unincluded segment is empty, i.e. `included_hash == at`, where at is the block
+	/// whose state we are querying against, this must always return `true` as long as the slot
+	/// is more recent than the included block itself.
+	pub fn can_build_upon(included_hash: T::Hash, new_slot: Slot) -> bool {
+		let velocity = V.max(1);
+		let (last_slot, authored_so_far) = match pallet::RelaySlotInfo::<T>::get() {
+			None => return true,
+			Some(x) => x,
+		};
+
+		let size_after_included =
+			teyrchain_system::Pallet::<T>::unincluded_segment_size_after(included_hash);
+
+		// can never author when the unincluded segment is full.
+		if size_after_included >= C {
+			return false;
+		}
+
+		// Check that we have not authored more than `V + 1` teyrchain blocks in the current relay
+		// chain slot.
+		if last_slot == new_slot {
+			authored_so_far < velocity + 1
+		} else {
+			// disallow slot from moving backwards.
+			last_slot < new_slot
+		}
+	}
+}

pezcumulus/pezpallets/aura-ext/src/lib.rs

@@ -0,0 +1,171 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Pezcumulus extension pallet for AuRa
+//!
+//! This pallet extends the Bizinikiwi AuRa pallet to make it compatible with teyrchains. It
+//! provides the [`Pallet`], the [`Config`] and the [`GenesisConfig`].
+//!
+//! It is also required that the teyrchain runtime uses the provided [`BlockExecutor`] to properly
+//! check the constructed block on the relay chain.
+//!
+//! ```
+//! # struct Runtime;
+//! # struct Executive;
+//! pezcumulus_pezpallet_teyrchain_system::register_validate_block! {
+//!     Runtime = Runtime,
+//!     BlockExecutor = pezcumulus_pezpallet_aura_ext::BlockExecutor::<Runtime, Executive>,
+//! }
+//! ```
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use pezframe_support::traits::{ExecuteBlock, FindAuthor};
+use pezsp_application_crypto::RuntimeAppPublic;
+use pezsp_consensus_aura::{digests::CompatibleDigestItem, Slot};
+use pezsp_runtime::traits::{Block as BlockT, Header as HeaderT, LazyBlock};
+
+pub mod consensus_hook;
+pub mod migration;
+#[cfg(test)]
+mod test;
+
+pub use consensus_hook::FixedVelocityConsensusHook;
+
+type Aura<T> = pezpallet_aura::Pallet<T>;
+
+pub use pallet::*;
+
+#[pezframe_support::pallet]
+pub mod pallet {
+	use super::*;
+	use pezframe_support::pezpallet_prelude::*;
+	use pezframe_system::pezpallet_prelude::*;
+
+	/// The configuration trait.
+	#[pallet::config]
+	pub trait Config: pezpallet_aura::Config + pezframe_system::Config {}
+
+	#[pallet::pallet]
+	#[pallet::storage_version(migration::STORAGE_VERSION)]
+	pub struct Pallet<T>(_);
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn on_finalize(_: BlockNumberFor<T>) {
+			// Update to the latest AuRa authorities.
+			Authorities::<T>::put(pezpallet_aura::Authorities::<T>::get());
+		}
+
+		fn on_initialize(_: BlockNumberFor<T>) -> Weight {
+			// Fetch the authorities once to get them into the storage proof of the PoV.
+			Authorities::<T>::get();
+
+			T::DbWeight::get().reads_writes(1, 0)
+		}
+	}
+
+	/// Serves as cache for the authorities.
+	///
+	/// The authorities in AuRa are overwritten in `on_initialize` when we switch to a new session,
+	/// but we require the old authorities to verify the seal when validating a PoV. This will
+	/// always be updated to the latest AuRa authorities in `on_finalize`.
+	#[pallet::storage]
+	pub(crate) type Authorities<T: Config> = StorageValue<
+		_,
+		BoundedVec<T::AuthorityId, <T as pezpallet_aura::Config>::MaxAuthorities>,
+		ValueQuery,
+	>;
+
+	/// Current relay chain slot paired with a number of authored blocks.
+	///
+	/// This is updated in [`FixedVelocityConsensusHook::on_state_proof`] with the current relay
+	/// chain slot as provided by the relay chain state proof.
+	#[pallet::storage]
+	pub(crate) type RelaySlotInfo<T: Config> = StorageValue<_, (Slot, u32), OptionQuery>;
+
+	#[pallet::genesis_config]
+	#[derive(pezframe_support::DefaultNoBound)]
+	pub struct GenesisConfig<T: Config> {
+		#[serde(skip)]
+		pub _config: core::marker::PhantomData<T>,
+	}
+
+	#[pallet::genesis_build]
+	impl<T: Config> BuildGenesisConfig for GenesisConfig<T> {
+		fn build(&self) {
+			let authorities = pezpallet_aura::Authorities::<T>::get();
+			Authorities::<T>::put(authorities);
+		}
+	}
+}
+
+/// The block executor used when validating a PoV at the relay chain.
+///
+/// When executing the block it will verify the block seal to ensure that the correct author created
+/// the block.
+pub struct BlockExecutor<T, I>(core::marker::PhantomData<(T, I)>);
+
+impl<Block, T, I> ExecuteBlock<Block> for BlockExecutor<T, I>
+where
+	Block: BlockT,
+	T: Config,
+	I: ExecuteBlock<Block>,
+{
+	fn verify_and_remove_seal(block: &mut <Block as BlockT>::LazyBlock) {
+		let header = block.header_mut();
+		// We need to fetch the authorities before we execute the block, to get the authorities
+		// before any potential update.
+		let authorities = Authorities::<T>::get();
+
+		let mut seal = None;
+		header.digest_mut().logs.retain(|s| {
+			let s =
+				CompatibleDigestItem::<<T::AuthorityId as RuntimeAppPublic>::Signature>::as_aura_seal(s);
+			match (s, seal.is_some()) {
+				(Some(_), true) => panic!("Found multiple AuRa seal digests"),
+				(None, _) => true,
+				(Some(s), false) => {
+					seal = Some(s);
+					false
+				},
+			}
+		});
+
+		let seal = seal.expect("Could not find an AuRa seal digest!");
+
+		let author = Aura::<T>::find_author(
+			header.digest().logs().iter().filter_map(|d| d.as_pre_runtime()),
+		)
+		.expect("Could not find AuRa author index!");
+
+		let pre_hash = header.hash();
+
+		if !authorities
+			.get(author as usize)
+			.unwrap_or_else(|| {
+				panic!("Invalid AuRa author index {} for authorities: {:?}", author, authorities)
+			})
+			.verify(&pre_hash, &seal)
+		{
+			panic!("Invalid AuRa seal");
+		}
+	}
+
+	fn execute_verified_block(block: Block::LazyBlock) {
+		I::execute_verified_block(block);
+	}
+}

pezcumulus/pezpallets/aura-ext/src/migration.rs

@@ -0,0 +1,74 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+extern crate alloc;
+
+use crate::{Config, Pallet};
+#[cfg(feature = "try-runtime")]
+use alloc::vec::Vec;
+use pezframe_support::{migrations::VersionedMigration, pezpallet_prelude::StorageVersion};
+
+/// The in-code storage version.
+pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(1);
+
+mod v0 {
+	use super::*;
+	use pezframe_support::{pezpallet_prelude::OptionQuery, storage_alias};
+	use pezsp_consensus_aura::Slot;
+
+	/// Current slot paired with a number of authored blocks.
+	///
+	/// Updated on each block initialization.
+	#[storage_alias]
+	pub(super) type SlotInfo<T: Config> = StorageValue<Pallet<T>, (Slot, u32), OptionQuery>;
+}
+mod v1 {
+	use super::*;
+	use pezframe_support::{pezpallet_prelude::*, traits::UncheckedOnRuntimeUpgrade};
+
+	pub struct UncheckedMigrationToV1<T: Config>(PhantomData<T>);
+
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV1<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let mut weight: Weight = Weight::zero();
+			weight += migrate::<T>();
+			weight
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::TryRuntimeError> {
+			Ok(Vec::new())
+		}
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_state: Vec<u8>) -> Result<(), pezsp_runtime::TryRuntimeError> {
+			ensure!(!v0::SlotInfo::<T>::exists(), "SlotInfo should not exist");
+			Ok(())
+		}
+	}
+
+	pub fn migrate<T: Config>() -> Weight {
+		v0::SlotInfo::<T>::kill();
+		T::DbWeight::get().writes(1)
+	}
+}
+
+/// Migrate `V0` to `V1`.
+pub type MigrateV0ToV1<T> = VersionedMigration<
+	0,
+	1,
+	v1::UncheckedMigrationToV1<T>,
+	Pallet<T>,
+	<T as pezframe_system::Config>::DbWeight,
+>;

pezcumulus/pezpallets/aura-ext/src/test.rs

@@ -0,0 +1,523 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 core::num::NonZeroU32;
+use pezcumulus_pezpallet_teyrchain_system::{
+	consensus_hook::ExpectParentIncluded, Ancestor, AnyRelayNumber, ConsensusHook,
+	RelayChainStateProof, TeyrchainSetCode, UsedBandwidth,
+};
+use pezcumulus_primitives_core::ParaId;
+use pezframe_support::{
+	derive_impl,
+	pezpallet_prelude::ConstU32,
+	parameter_types,
+	traits::{ConstBool, EnqueueWithOrigin, ExecuteBlock, Hooks},
+	BoundedVec,
+};
+use rstest::rstest;
+use pezsp_consensus_aura::{sr25519::AuthorityId, Slot};
+use pezsp_core::{Blake2Hasher, Get, H256};
+use pezsp_io::TestExternalities;
+use pezsp_keyring::Sr25519Keyring::*;
+use pezsp_runtime::{generic::Digest, traits::Block as BlockT};
+use pezsp_trie::{proof_size_extension::ProofSizeExt, recorder::Recorder};
+use pezsp_version::RuntimeVersion;
+use std::cell::RefCell;
+
+// Test pallet that reads storage and calls storage_proof_size
+#[pezframe_support::pallet]
+pub mod test_pallet {
+	use pezframe_support::pezpallet_prelude::*;
+	use pezframe_system::pezpallet_prelude::*;
+
+	#[pallet::config]
+	pub trait Config: pezframe_system::Config {}
+
+	#[pallet::pallet]
+	pub struct Pallet<T>(_);
+
+	#[pallet::storage]
+	pub type TestStorage<T: Config> = StorageValue<_, u64, ValueQuery>;
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn on_initialize(_n: BlockNumberFor<T>) -> Weight {
+			let proof_size =
+				pezcumulus_primitives_proof_size_hostfunction::storage_proof_size::storage_proof_size(
+				);
+			// We need to commit the `proof_size` to ensure that the test is failing if we are
+			// receiving a different proof size later on.
+			TestStorage::<T>::put(proof_size);
+
+			Weight::zero()
+		}
+	}
+}
+
+type Block = pezframe_system::mocking::MockBlock<Test>;
+
+pezframe_support::construct_runtime!(
+	pub enum Test {
+		System: pezframe_system,
+		TeyrchainSystem: pezcumulus_pezpallet_teyrchain_system,
+		Aura: pezpallet_aura,
+		AuraExt: crate,
+		TestPallet: test_pallet,
+	}
+);
+
+parameter_types! {
+	pub Version: RuntimeVersion = RuntimeVersion {
+		spec_name: "test".into(),
+		impl_name: "system-test".into(),
+		authoring_version: 1,
+		spec_version: 1,
+		impl_version: 1,
+		apis: pezsp_version::create_apis_vec!([]),
+		transaction_version: 1,
+		system_version: 1,
+	};
+}
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl pezframe_system::Config for Test {
+	type Block = Block;
+	type Version = Version;
+	type OnSetCode = TeyrchainSetCode<Test>;
+	type RuntimeEvent = ();
+}
+
+impl crate::Config for Test {}
+
+impl test_pallet::Config for Test {}
+
+std::thread_local! {
+	pub static PARA_SLOT_DURATION: RefCell<u64> = RefCell::new(6000);
+}
+
+pub struct TestSlotDuration;
+
+impl TestSlotDuration {
+	pub fn set_slot_duration(slot_duration: u64) {
+		PARA_SLOT_DURATION.with(|v| *v.borrow_mut() = slot_duration);
+	}
+}
+impl Get<u64> for TestSlotDuration {
+	fn get() -> u64 {
+		PARA_SLOT_DURATION.with(|v| v.clone().into_inner())
+	}
+}
+
+impl pezpallet_aura::Config for Test {
+	type AuthorityId = pezsp_consensus_aura::sr25519::AuthorityId;
+	type MaxAuthorities = ConstU32<100_000>;
+	type DisabledValidators = ();
+	type AllowMultipleBlocksPerSlot = ConstBool<true>;
+	type SlotDuration = TestSlotDuration;
+}
+
+impl pezpallet_timestamp::Config for Test {
+	type Moment = u64;
+	type OnTimestampSet = ();
+	type MinimumPeriod = ();
+	type WeightInfo = ();
+}
+
+impl pezcumulus_pezpallet_teyrchain_system::Config for Test {
+	type WeightInfo = ();
+	type RuntimeEvent = ();
+	type OnSystemEvent = ();
+	type SelfParaId = ();
+	type OutboundXcmpMessageSource = ();
+	// Ignore all DMP messages by enqueueing them into `()`:
+	type DmpQueue = EnqueueWithOrigin<(), pezsp_core::ConstU8<0>>;
+	type ReservedDmpWeight = ();
+	type XcmpMessageHandler = ();
+	type ReservedXcmpWeight = ();
+	type CheckAssociatedRelayNumber = AnyRelayNumber;
+	type ConsensusHook = ExpectParentIncluded;
+	type RelayParentOffset = ConstU32<0>;
+}
+
+fn set_ancestors() {
+	let mut ancestors = Vec::new();
+	for i in 0..3 {
+		let mut ancestor = Ancestor::new_unchecked(UsedBandwidth::default(), None);
+		ancestor.replace_para_head_hash(H256::repeat_byte(i + 1));
+		ancestors.push(ancestor);
+	}
+	pezcumulus_pezpallet_teyrchain_system::UnincludedSegment::<Test>::put(ancestors);
+}
+
+fn new_test_ext(para_slot: u64) -> pezsp_io::TestExternalities {
+	let mut ext = TestExternalities::new_empty();
+	ext.execute_with(|| {
+		set_ancestors();
+		// Set initial teyrchain slot
+		pezpallet_aura::CurrentSlot::<Test>::put(Slot::from(para_slot));
+	});
+	ext
+}
+
+fn set_relay_slot(slot: u64, authored: u32) {
+	RelaySlotInfo::<Test>::put((Slot::from(slot), authored))
+}
+
+fn relay_chain_state_proof(relay_slot: u64) -> RelayChainStateProof {
+	let mut builder = pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder::default();
+	builder.current_slot = relay_slot.into();
+
+	let (hash, state_proof) = builder.into_state_root_and_proof();
+
+	RelayChainStateProof::new(ParaId::from(200), hash, state_proof)
+		.expect("Should be able to construct state proof.")
+}
+
+fn assert_slot_info(expected_slot: u64, expected_authored: u32) {
+	let (slot, authored) = pallet::RelaySlotInfo::<Test>::get().unwrap();
+	assert_eq!(slot, Slot::from(expected_slot), "Slot stored in RelaySlotInfo is incorrect.");
+	assert_eq!(
+		authored, expected_authored,
+		"Number of authored blocks stored in RelaySlotInfo is incorrect."
+	);
+}
+
+const DEFAULT_TEST_VELOCITY: u32 = 2;
+
+#[test]
+fn test_velocity() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		let (_, capacity) = Hook::on_state_proof(&state_proof);
+		assert_eq!(capacity, NonZeroU32::new(1).unwrap().into());
+		assert_slot_info(10, 1);
+
+		let (_, capacity) = Hook::on_state_proof(&state_proof);
+		assert_eq!(capacity, NonZeroU32::new(1).unwrap().into());
+		assert_slot_info(10, 2);
+	});
+}
+
+#[test]
+fn test_velocity_2() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 3>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		let (_, capacity) = Hook::on_state_proof(&state_proof);
+		assert_eq!(capacity, NonZeroU32::new(3).unwrap().into());
+		assert_slot_info(10, 1);
+
+		let (_, capacity) = Hook::on_state_proof(&state_proof);
+		assert_eq!(capacity, NonZeroU32::new(3).unwrap().into());
+		assert_slot_info(10, 2);
+	});
+}
+
+#[test]
+#[should_panic(expected = "authored blocks limit is reached for the slot")]
+fn test_exceeding_velocity_limit() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		for authored in 0..=DEFAULT_TEST_VELOCITY + 1 {
+			Hook::on_state_proof(&state_proof);
+			assert_slot_info(10, authored + 1);
+		}
+	});
+}
+
+#[test]
+fn test_para_slot_calculated_from_slot_duration() {
+	type Hook = FixedVelocityConsensusHook<Test, 3000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(5).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		Hook::on_state_proof(&state_proof);
+	});
+}
+
+#[rstest]
+#[case::short_para_slot_okay(2000, 30, 10)]
+#[case::normal_para_slot_okay(6000, 10, 10)]
+// Test boundaries for long teyrchain slots.
+#[case::long_para_slot_okay(24000, 1, 7)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(2), derived_from_relay_slot=Slot(1)"
+)]
+#[case::long_para_slot_mismatch(24000, 2, 7)]
+#[case::long_para_slot_okay(24000, 2, 8)]
+#[case::long_para_slot_okay(24000, 2, 9)]
+#[case::long_para_slot_okay(24000, 2, 10)]
+#[case::long_para_slot_okay(24000, 2, 11)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(2), derived_from_relay_slot=Slot(3)"
+)]
+#[case::long_para_slot_mismatch(24000, 2, 12)]
+#[case::long_para_slot_okay(24000, 3, 12)]
+#[case::short_para_slot(2000, 30, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(31), derived_from_relay_slot=Slot(30)"
+)]
+#[case::short_para_slot_mismatch(2000, 31, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(32), derived_from_relay_slot=Slot(30)"
+)]
+#[case::short_para_slot_mismatch(2000, 32, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(29), derived_from_relay_slot=Slot(30)"
+)]
+#[case::short_para_slot_mismatch(2000, 29, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(1), derived_from_relay_slot=Slot(30)"
+)]
+#[case::short_para_slot_mismatch(2000, 1, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(1), derived_from_relay_slot=Slot(10)"
+)]
+#[case::normal_para_slot_mismatch(6000, 1, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(9), derived_from_relay_slot=Slot(10)"
+)]
+#[case::normal_para_slot_mismatch(6000, 9, 10)]
+#[should_panic(
+	expected = "must match relay-derived slot: teyrchain_slot=Slot(11), derived_from_relay_slot=Slot(10)"
+)]
+#[case::normal_para_slot_mismatch(6000, 11, 10)]
+fn test_para_slot_too_high(
+	#[case] para_slot_duration: u64,
+	#[case] para_slot: u64,
+	#[case] relay_slot: u64,
+) {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 1>;
+
+	TestSlotDuration::set_slot_duration(para_slot_duration);
+	new_test_ext(para_slot).execute_with(|| {
+		let state_proof = relay_chain_state_proof(relay_slot);
+		Hook::on_state_proof(&state_proof);
+	});
+}
+
+#[test]
+fn test_velocity_at_least_one() {
+	// Even though this is 0, one block should always be allowed.
+	const VELOCITY: u32 = 0;
+	type Hook = FixedVelocityConsensusHook<Test, 6000, VELOCITY, 1>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		Hook::on_state_proof(&state_proof);
+	});
+}
+
+#[test]
+#[should_panic(
+	expected = "Teyrchain slot must match relay-derived slot: teyrchain_slot=Slot(8), derived_from_relay_slot=Slot(5) velocity=2"
+)]
+fn test_para_slot_calculated_from_slot_duration_2() {
+	// Note: In contrast to tests below, relay chain slot duration is 3000 here.
+	type Hook = FixedVelocityConsensusHook<Test, 3000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(8).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		let (_, _) = Hook::on_state_proof(&state_proof);
+	});
+}
+
+#[test]
+fn test_velocity_resets_on_new_relay_slot() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		for authored in 0..=DEFAULT_TEST_VELOCITY {
+			Hook::on_state_proof(&state_proof);
+			assert_slot_info(10, authored + 1);
+		}
+
+		// Change teyrchain slot to match the new relay slot
+		pezpallet_aura::CurrentSlot::<Test>::put(Slot::from(11));
+		let state_proof = relay_chain_state_proof(11);
+		for authored in 0..=DEFAULT_TEST_VELOCITY {
+			Hook::on_state_proof(&state_proof);
+			assert_slot_info(11, authored + 1);
+		}
+	});
+}
+
+#[test]
+#[should_panic(expected = "Slot moved backwards: stored_slot=Slot(10), relay_chain_slot=Slot(9)")]
+fn test_backward_relay_slot_not_tolerated() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, 2, 1>;
+
+	new_test_ext(10).execute_with(|| {
+		let state_proof = relay_chain_state_proof(10);
+		Hook::on_state_proof(&state_proof);
+		assert_slot_info(10, 1);
+
+		// Change teyrchain slot to match what would be derived from relay slot 9
+		pezpallet_aura::CurrentSlot::<Test>::put(Slot::from(9));
+		let state_proof = relay_chain_state_proof(9);
+		Hook::on_state_proof(&state_proof);
+	});
+}
+
+#[test]
+fn test_can_build_upon_true_when_empty() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 1>;
+
+	new_test_ext(1).execute_with(|| {
+		let hash = H256::repeat_byte(0x1);
+		assert!(Hook::can_build_upon(hash, Slot::from(1)));
+	});
+}
+
+#[rstest]
+#[case::slot_higher_ok(10, 11, DEFAULT_TEST_VELOCITY, true)]
+#[case::slot_same_ok(10, 10, DEFAULT_TEST_VELOCITY, true)]
+#[case::slot_decrease_illegal(10, 9, DEFAULT_TEST_VELOCITY, false)]
+#[case::velocity_small_ok(10, 10, DEFAULT_TEST_VELOCITY - 1 , true)]
+#[case::velocity_small_ok(10, 10, DEFAULT_TEST_VELOCITY - 2 , true)]
+#[case::velocity_too_high_illegal(10, 10, DEFAULT_TEST_VELOCITY + 1 , false)]
+fn test_can_build_upon_slot_can_not_decrease(
+	#[case] state_relay_slot: u64,
+	#[case] test_relay_slot: u64,
+	#[case] authored_in_slot: u32,
+	#[case] expected_result: bool,
+) {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 10>;
+
+	new_test_ext(1).execute_with(|| {
+		let hash = H256::repeat_byte(0x1);
+
+		set_relay_slot(state_relay_slot, authored_in_slot);
+		// Slot moves backwards
+		assert_eq!(Hook::can_build_upon(hash, Slot::from(test_relay_slot)), expected_result);
+	});
+}
+
+#[test]
+fn test_can_build_upon_unincluded_segment_size() {
+	type Hook = FixedVelocityConsensusHook<Test, 6000, DEFAULT_TEST_VELOCITY, 2>;
+
+	new_test_ext(1).execute_with(|| {
+		let relay_slot = Slot::from(10);
+
+		set_relay_slot(10, DEFAULT_TEST_VELOCITY);
+		// Size after included is two, we can not build
+		assert!(!Hook::can_build_upon(H256::repeat_byte(0x1), relay_slot));
+
+		// Size after included is one, we can build
+		assert!(Hook::can_build_upon(H256::repeat_byte(0x2), relay_slot));
+	});
+}
+
+/// This test ensures that when we call `BlockExecutor::execute_block` in `validate_block`,
+/// it doesn't change the proof size host function return values. Otherwise, it may breaks
+/// logic that is fetching the proof size in `on_initialize`.
+#[test]
+fn block_executor_does_not_influence_proof_size_recordings() {
+	fn build_block(header: <Block as BlockT>::Header) -> <Block as BlockT>::Header {
+		// Initialize the block
+		pezframe_system::Pallet::<Test>::initialize(
+			&header.number,
+			&header.parent_hash,
+			&header.digest(),
+		);
+
+		// We omit `teyrchain-system` as it is not important here.
+		<pezframe_system::Pallet<Test> as Hooks<_>>::on_initialize(header.number);
+		<crate::Pallet<Test> as Hooks<_>>::on_initialize(header.number);
+		<test_pallet::Pallet<Test> as Hooks<_>>::on_initialize(header.number);
+
+		<test_pallet::Pallet<Test> as Hooks<_>>::on_finalize(header.number);
+		<crate::Pallet<Test> as Hooks<_>>::on_finalize(header.number);
+		<pezframe_system::Pallet<Test> as Hooks<_>>::on_finalize(header.number);
+
+		// Finalize the block
+		pezframe_system::Pallet::<Test>::finalize()
+	}
+
+	// Create a simple executive that calls on_initialize and on_finalize
+	struct TestExecutive;
+	impl ExecuteBlock<Block> for TestExecutive {
+		fn verify_and_remove_seal(_: &mut <Block as BlockT>::LazyBlock) {}
+
+		fn execute_verified_block(block: <Block as BlockT>::LazyBlock) {
+			let header = block.header();
+
+			let new_header = build_block(header.clone());
+
+			assert_eq!(*header, new_header);
+		}
+	}
+
+	let mut ext = new_test_ext(10);
+
+	ext.execute_with(|| {
+		// Let's setup some authorities
+		let authority_id = AuthorityId::from(Alice.public());
+		let authorities: BoundedVec<AuthorityId, ConstU32<100_000>> =
+			vec![authority_id.clone()].try_into().unwrap();
+		pezpallet_aura::Authorities::<Test>::put(authorities.clone());
+		Authorities::<Test>::put(authorities.clone());
+	});
+
+	ext.commit_all().unwrap();
+
+	let recorder = Recorder::<Blake2Hasher>::default();
+
+	// Register the ProofSizeExt extension
+	ext.register_extension(ProofSizeExt::new(recorder.clone()));
+
+	let mut header = ext.execute_with_recorder(recorder.clone(), || {
+		let mut digest = Digest::default();
+		digest.push(CompatibleDigestItem::<()>::aura_pre_digest(10u64.into()));
+
+		build_block(HeaderT::new(
+			1,
+			Default::default(),
+			Default::default(),
+			Default::default(),
+			digest,
+		))
+	});
+
+	let sig = Alice.sign(header.hash().as_ref());
+	let seal = CompatibleDigestItem::aura_seal(sig);
+	header.digest_mut().push(seal);
+
+	let mut block = Block::new(header, Default::default()).into();
+
+	ext.reset_overlay();
+	ext.execute_with_recorder(recorder, || {
+		BlockExecutor::<Test, TestExecutive>::verify_and_remove_seal(&mut block);
+	});
+
+	let recorder = Recorder::<Blake2Hasher>::default();
+
+	// Register the ProofSizeExt extension again to overwrite the old one.
+	ext.register_extension(ProofSizeExt::new(recorder.clone()));
+
+	ext.reset_overlay();
+	ext.execute_with_recorder(recorder, || {
+		BlockExecutor::<Test, TestExecutive>::execute_verified_block(block);
+	});
+}

pezcumulus/pezpallets/collator-selection/Cargo.toml

@@ -0,0 +1,83 @@
+[package]
+authors.workspace = true
+description = "Simple pallet to select collators for a teyrchain."
+edition.workspace = true
+homepage.workspace = true
+license = "Apache-2.0"
+name = "pezpallet-collator-selection"
+readme = "README.md"
+repository.workspace = true
+version = "9.0.0"
+documentation = "https://docs.rs/pezpallet-collator-selection"
+
+[lints]
+workspace = true
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+log = { workspace = true }
+rand = { features = ["std_rng"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-authorship = { workspace = true }
+pezpallet-balances = { workspace = true }
+pezpallet-session = { workspace = true }
+pezsp-runtime = { workspace = true }
+pezsp-staking = { workspace = true }
+
+pezframe-benchmarking = { optional = true, workspace = true }
+
+[dev-dependencies]
+pezpallet-aura = { workspace = true, default-features = true }
+pezpallet-timestamp = { workspace = true, default-features = true }
+pezsp-consensus-aura = { workspace = true, default-features = true }
+pezsp-io = { workspace = true, default-features = true }
+pezsp-runtime = { workspace = true, default-features = true }
+pezsp-tracing = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+runtime-benchmarks = [
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-aura/runtime-benchmarks",
+  "pezpallet-authorship/runtime-benchmarks",
+  "pezpallet-balances/runtime-benchmarks",
+  "pezpallet-session/runtime-benchmarks",
+  "pezpallet-timestamp/runtime-benchmarks",
+  "pezsp-consensus-aura/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "pezsp-staking/runtime-benchmarks",
+]
+std = [
+  "codec/std",
+  "pezframe-benchmarking/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "log/std",
+  "pezpallet-authorship/std",
+  "pezpallet-balances/std",
+  "pezpallet-session/std",
+  "rand/std",
+  "scale-info/std",
+  "pezsp-runtime/std",
+  "pezsp-staking/std",
+]
+
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-aura/try-runtime",
+  "pezpallet-authorship/try-runtime",
+  "pezpallet-balances/try-runtime",
+  "pezpallet-session/try-runtime",
+  "pezpallet-timestamp/try-runtime",
+  "pezsp-runtime/try-runtime",
+]

pezcumulus/pezpallets/collator-selection/README.md

@@ -0,0 +1 @@
+License: Apache-2.0

pezcumulus/pezpallets/collator-selection/src/benchmarking.rs

@@ -0,0 +1,464 @@
+// 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 pezpallet-collator-selection
+
+#![cfg(feature = "runtime-benchmarks")]
+
+use super::*;
+
+#[allow(unused)]
+use crate::Pallet as CollatorSelection;
+use alloc::vec::Vec;
+use codec::Decode;
+use core::cmp;
+use pezframe_benchmarking::{account, v2::*, whitelisted_caller, BenchmarkError};
+use pezframe_support::traits::{Currency, EnsureOrigin, Get, ReservableCurrency};
+use pezframe_system::{pezpallet_prelude::BlockNumberFor, EventRecord, RawOrigin};
+use pezpallet_authorship::EventHandler;
+use pezpallet_session::{self as session, SessionManager};
+
+pub type BalanceOf<T> =
+	<<T as Config>::Currency as Currency<<T as pezframe_system::Config>::AccountId>>::Balance;
+
+const SEED: u32 = 0;
+
+fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {
+	let events = pezframe_system::Pallet::<T>::events();
+	let system_event: <T as pezframe_system::Config>::RuntimeEvent = generic_event.into();
+	// compare to the last event record
+	let EventRecord { event, .. } = &events[events.len() - 1];
+	assert_eq!(event, &system_event);
+}
+
+fn create_funded_user<T: Config>(
+	string: &'static str,
+	n: u32,
+	balance_factor: u32,
+) -> T::AccountId {
+	let user = account(string, n, SEED);
+	let balance = <T as pallet::Config>::Currency::minimum_balance() * balance_factor.into();
+	let _ = <T as pallet::Config>::Currency::make_free_balance_be(&user, balance);
+	user
+}
+
+fn keys<T: Config + session::Config>(c: u32) -> <T as session::Config>::Keys {
+	use rand::{RngCore, SeedableRng};
+
+	let keys = {
+		let mut keys = [0u8; 128];
+
+		if c > 0 {
+			let mut rng = rand::rngs::StdRng::seed_from_u64(c as u64);
+			rng.fill_bytes(&mut keys);
+		}
+
+		keys
+	};
+
+	Decode::decode(&mut &keys[..]).unwrap()
+}
+
+fn validator<T: Config + session::Config>(c: u32) -> (T::AccountId, <T as session::Config>::Keys) {
+	(create_funded_user::<T>("candidate", c, 1000), keys::<T>(c))
+}
+
+fn register_validators<T: Config + session::Config>(count: u32) -> Vec<T::AccountId> {
+	let validators = (0..count).map(|c| validator::<T>(c)).collect::<Vec<_>>();
+
+	for (who, keys) in validators.clone() {
+		<session::Pallet<T>>::ensure_can_pay_key_deposit(&who).unwrap();
+		<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new()).unwrap();
+	}
+
+	validators.into_iter().map(|(who, _)| who).collect()
+}
+
+fn register_candidates<T: Config>(count: u32) {
+	let candidates = (0..count).map(|c| account("candidate", c, SEED)).collect::<Vec<_>>();
+	assert!(CandidacyBond::<T>::get() > 0u32.into(), "Bond cannot be zero!");
+
+	for who in candidates {
+		<T as pallet::Config>::Currency::make_free_balance_be(
+			&who,
+			CandidacyBond::<T>::get() * 3u32.into(),
+		);
+		<CollatorSelection<T>>::register_as_candidate(RawOrigin::Signed(who).into()).unwrap();
+	}
+}
+
+fn min_candidates<T: Config>() -> u32 {
+	let min_collators = T::MinEligibleCollators::get();
+	let invulnerable_length = Invulnerables::<T>::get().len();
+	min_collators.saturating_sub(invulnerable_length.try_into().unwrap())
+}
+
+fn min_invulnerables<T: Config>() -> u32 {
+	let min_collators = T::MinEligibleCollators::get();
+	let candidates_length = CandidateList::<T>::decode_len()
+		.unwrap_or_default()
+		.try_into()
+		.unwrap_or_default();
+	min_collators.saturating_sub(candidates_length)
+}
+
+#[benchmarks(where T: pezpallet_authorship::Config + session::Config)]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn set_invulnerables(
+		b: Linear<1, { T::MaxInvulnerables::get() }>,
+	) -> Result<(), BenchmarkError> {
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+
+		let new_invulnerables = register_validators::<T>(b);
+		let mut sorted_new_invulnerables = new_invulnerables.clone();
+		sorted_new_invulnerables.sort();
+
+		#[extrinsic_call]
+		_(origin as T::RuntimeOrigin, new_invulnerables.clone());
+
+		// assert that it comes out sorted
+		assert_last_event::<T>(
+			Event::NewInvulnerables { invulnerables: sorted_new_invulnerables }.into(),
+		);
+		Ok(())
+	}
+
+	#[benchmark]
+	fn add_invulnerable(
+		b: Linear<1, { T::MaxInvulnerables::get() - 1 }>,
+		c: Linear<1, { T::MaxCandidates::get() - 1 }>,
+	) -> Result<(), BenchmarkError> {
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+
+		// need to fill up candidates
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(c);
+		// get accounts and keys for the `c` candidates
+		let mut candidates = (0..c).map(|cc| validator::<T>(cc)).collect::<Vec<_>>();
+		// add one more to the list. should not be in `b` (invulnerables) because it's the account
+		// we will _add_ to invulnerables. we want it to be in `candidates` because we need the
+		// weight associated with removing it.
+		let (new_invulnerable, new_invulnerable_keys) = validator::<T>(b.max(c) + 1);
+		candidates.push((new_invulnerable.clone(), new_invulnerable_keys));
+		// set their keys ...
+		for (who, keys) in candidates.clone() {
+			<session::Pallet<T>>::ensure_can_pay_key_deposit(&who).unwrap();
+			<session::Pallet<T>>::set_keys(RawOrigin::Signed(who).into(), keys, Vec::new())
+				.unwrap();
+		}
+		// ... and register them.
+		for (who, _) in candidates.iter() {
+			let deposit = CandidacyBond::<T>::get();
+			<T as pallet::Config>::Currency::make_free_balance_be(who, deposit * 1000_u32.into());
+			CandidateList::<T>::try_mutate(|list| {
+				list.try_push(CandidateInfo { who: who.clone(), deposit }).unwrap();
+				Ok::<(), BenchmarkError>(())
+			})
+			.unwrap();
+			<T as pallet::Config>::Currency::reserve(who, deposit)?;
+			LastAuthoredBlock::<T>::insert(
+				who.clone(),
+				pezframe_system::Pallet::<T>::block_number() + T::KickThreshold::get(),
+			);
+		}
+
+		// now we need to fill up invulnerables
+		let mut invulnerables = register_validators::<T>(b);
+		invulnerables.sort();
+		let invulnerables: pezframe_support::BoundedVec<_, T::MaxInvulnerables> =
+			pezframe_support::BoundedVec::try_from(invulnerables).unwrap();
+		Invulnerables::<T>::put(invulnerables);
+
+		#[extrinsic_call]
+		_(origin as T::RuntimeOrigin, new_invulnerable.clone());
+
+		assert_last_event::<T>(Event::InvulnerableAdded { account_id: new_invulnerable }.into());
+		Ok(())
+	}
+
+	#[benchmark]
+	fn remove_invulnerable(
+		b: Linear<{ min_invulnerables::<T>() + 1 }, { T::MaxInvulnerables::get() }>,
+	) -> Result<(), BenchmarkError> {
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+		let mut invulnerables = register_validators::<T>(b);
+		invulnerables.sort();
+		let invulnerables: pezframe_support::BoundedVec<_, T::MaxInvulnerables> =
+			pezframe_support::BoundedVec::try_from(invulnerables).unwrap();
+		Invulnerables::<T>::put(invulnerables);
+		let to_remove = Invulnerables::<T>::get().first().unwrap().clone();
+
+		#[extrinsic_call]
+		_(origin as T::RuntimeOrigin, to_remove.clone());
+
+		assert_last_event::<T>(Event::InvulnerableRemoved { account_id: to_remove }.into());
+		Ok(())
+	}
+
+	#[benchmark]
+	fn set_desired_candidates() -> Result<(), BenchmarkError> {
+		let max: u32 = T::MaxCandidates::get();
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+
+		#[extrinsic_call]
+		_(origin as T::RuntimeOrigin, max);
+
+		assert_last_event::<T>(Event::NewDesiredCandidates { desired_candidates: max }.into());
+		Ok(())
+	}
+
+	#[benchmark]
+	fn set_candidacy_bond(
+		c: Linear<0, { T::MaxCandidates::get() }>,
+		k: Linear<0, { T::MaxCandidates::get() }>,
+	) -> Result<(), BenchmarkError> {
+		let initial_bond_amount: BalanceOf<T> =
+			<T as pallet::Config>::Currency::minimum_balance() * 2u32.into();
+		CandidacyBond::<T>::put(initial_bond_amount);
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+		let kicked = cmp::min(k, c);
+		let origin =
+			T::UpdateOrigin::try_successful_origin().map_err(|_| BenchmarkError::Weightless)?;
+		let bond_amount = if k > 0 {
+			CandidateList::<T>::mutate(|candidates| {
+				for info in candidates.iter_mut().skip(kicked as usize) {
+					info.deposit = <T as pallet::Config>::Currency::minimum_balance() * 3u32.into();
+				}
+			});
+			<T as pallet::Config>::Currency::minimum_balance() * 3u32.into()
+		} else {
+			<T as pallet::Config>::Currency::minimum_balance()
+		};
+
+		#[extrinsic_call]
+		_(origin as T::RuntimeOrigin, bond_amount);
+
+		assert_last_event::<T>(Event::NewCandidacyBond { bond_amount }.into());
+		Ok(())
+	}
+
+	#[benchmark]
+	fn update_bond(
+		c: Linear<{ min_candidates::<T>() + 1 }, { T::MaxCandidates::get() }>,
+	) -> Result<(), BenchmarkError> {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(c);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let caller = CandidateList::<T>::get()[0].who.clone();
+		v2::whitelist!(caller);
+
+		let bond_amount: BalanceOf<T> = <T as pallet::Config>::Currency::minimum_balance() +
+			<T as pallet::Config>::Currency::minimum_balance();
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.clone()), bond_amount);
+
+		assert_last_event::<T>(
+			Event::CandidateBondUpdated { account_id: caller, deposit: bond_amount }.into(),
+		);
+		assert!(
+			CandidateList::<T>::get().iter().last().unwrap().deposit ==
+				<T as pallet::Config>::Currency::minimum_balance() * 2u32.into()
+		);
+		Ok(())
+	}
+
+	// worse case is when we have all the max-candidate slots filled except one, and we fill that
+	// one.
+	#[benchmark]
+	fn register_as_candidate(c: Linear<1, { T::MaxCandidates::get() - 1 }>) {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(c + 1);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let caller: T::AccountId = whitelisted_caller();
+		let bond: BalanceOf<T> = <T as pallet::Config>::Currency::minimum_balance() * 2u32.into();
+		<T as pallet::Config>::Currency::make_free_balance_be(&caller, bond);
+
+		<session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
+		<session::Pallet<T>>::set_keys(
+			RawOrigin::Signed(caller.clone()).into(),
+			keys::<T>(c + 1),
+			Vec::new(),
+		)
+		.unwrap();
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.clone()));
+
+		assert_last_event::<T>(
+			Event::CandidateAdded { account_id: caller, deposit: bond / 2u32.into() }.into(),
+		);
+	}
+
+	#[benchmark]
+	fn take_candidate_slot(c: Linear<{ min_candidates::<T>() + 1 }, { T::MaxCandidates::get() }>) {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(1);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let caller: T::AccountId = whitelisted_caller();
+		let bond: BalanceOf<T> = <T as pallet::Config>::Currency::minimum_balance() * 10u32.into();
+		<T as pallet::Config>::Currency::make_free_balance_be(&caller, bond);
+
+		<session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
+		<session::Pallet<T>>::set_keys(
+			RawOrigin::Signed(caller.clone()).into(),
+			keys::<T>(c + 1),
+			Vec::new(),
+		)
+		.unwrap();
+
+		let target = CandidateList::<T>::get().iter().last().unwrap().who.clone();
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller.clone()), bond / 2u32.into(), target.clone());
+
+		assert_last_event::<T>(
+			Event::CandidateReplaced { old: target, new: caller, deposit: bond / 2u32.into() }
+				.into(),
+		);
+	}
+
+	// worse case is the last candidate leaving.
+	#[benchmark]
+	fn leave_intent(c: Linear<{ min_candidates::<T>() + 1 }, { T::MaxCandidates::get() }>) {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(c);
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let leaving = CandidateList::<T>::get().iter().last().unwrap().who.clone();
+		v2::whitelist!(leaving);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(leaving.clone()));
+
+		assert_last_event::<T>(Event::CandidateRemoved { account_id: leaving }.into());
+	}
+
+	// worse case is paying a non-existing candidate account.
+	#[benchmark]
+	fn note_author() {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		<T as pallet::Config>::Currency::make_free_balance_be(
+			&<CollatorSelection<T>>::account_id(),
+			<T as pallet::Config>::Currency::minimum_balance() * 4u32.into(),
+		);
+		let author = account("author", 0, SEED);
+		let new_block: BlockNumberFor<T> = 10u32.into();
+
+		pezframe_system::Pallet::<T>::set_block_number(new_block);
+		assert!(<T as pallet::Config>::Currency::free_balance(&author) == 0u32.into());
+
+		#[block]
+		{
+			<CollatorSelection<T> as EventHandler<_, _>>::note_author(author.clone())
+		}
+
+		assert!(<T as pallet::Config>::Currency::free_balance(&author) > 0u32.into());
+		assert_eq!(pezframe_system::Pallet::<T>::block_number(), new_block);
+	}
+
+	// worst case for new session.
+	#[benchmark]
+	fn new_session(
+		r: Linear<1, { T::MaxCandidates::get() }>,
+		c: Linear<1, { T::MaxCandidates::get() }>,
+	) {
+		CandidacyBond::<T>::put(<T as pallet::Config>::Currency::minimum_balance());
+		DesiredCandidates::<T>::put(c);
+		pezframe_system::Pallet::<T>::set_block_number(0u32.into());
+
+		register_validators::<T>(c);
+		register_candidates::<T>(c);
+
+		let new_block: BlockNumberFor<T> = T::KickThreshold::get();
+		let zero_block: BlockNumberFor<T> = 0u32.into();
+		let candidates: Vec<T::AccountId> = CandidateList::<T>::get()
+			.iter()
+			.map(|candidate_info| candidate_info.who.clone())
+			.collect();
+
+		let non_removals = c.saturating_sub(r);
+
+		for i in 0..c {
+			LastAuthoredBlock::<T>::insert(candidates[i as usize].clone(), zero_block);
+		}
+
+		if non_removals > 0 {
+			for i in 0..non_removals {
+				LastAuthoredBlock::<T>::insert(candidates[i as usize].clone(), new_block);
+			}
+		} else {
+			for i in 0..c {
+				LastAuthoredBlock::<T>::insert(candidates[i as usize].clone(), new_block);
+			}
+		}
+
+		let min_candidates = min_candidates::<T>();
+		let pre_length = CandidateList::<T>::decode_len().unwrap_or_default();
+
+		pezframe_system::Pallet::<T>::set_block_number(new_block);
+
+		let current_length: u32 = CandidateList::<T>::decode_len()
+			.unwrap_or_default()
+			.try_into()
+			.unwrap_or_default();
+		assert!(c == current_length);
+		#[block]
+		{
+			<CollatorSelection<T> as SessionManager<_>>::new_session(0);
+		}
+
+		if c > r && non_removals >= min_candidates {
+			// candidates > removals and remaining candidates > min candidates
+			// => remaining candidates should be shorter than before removal, i.e. some were
+			//    actually removed.
+			assert!(CandidateList::<T>::decode_len().unwrap_or_default() < pre_length);
+		} else if c > r && non_removals < min_candidates {
+			// candidates > removals and remaining candidates would be less than min candidates
+			// => remaining candidates should equal min candidates, i.e. some were removed up to
+			//    the minimum, but then any more were "forced" to stay in candidates.
+			let current_length: u32 = CandidateList::<T>::decode_len()
+				.unwrap_or_default()
+				.try_into()
+				.unwrap_or_default();
+			assert!(min_candidates == current_length);
+		} else {
+			// removals >= candidates, non removals must == 0
+			// can't remove more than exist
+			assert!(CandidateList::<T>::decode_len().unwrap_or_default() == pre_length);
+		}
+	}
+
+	impl_benchmark_test_suite!(CollatorSelection, crate::mock::new_test_ext(), crate::mock::Test);
+}

pezcumulus/pezpallets/collator-selection/src/migration.rs

@@ -0,0 +1,321 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! A module that is responsible for migration of storage for Collator Selection.
+
+use super::*;
+#[cfg(feature = "try-runtime")]
+use alloc::vec::Vec;
+use pezframe_support::traits::{OnRuntimeUpgrade, UncheckedOnRuntimeUpgrade};
+use log;
+
+/// Migrate to v2. Should have been part of <https://github.com/pezkuwichain/kurdistan-sdk/issues/104>.
+pub mod v2 {
+	use super::*;
+	use pezframe_support::{
+		pezpallet_prelude::*,
+		storage_alias,
+		traits::{Currency, ReservableCurrency},
+	};
+	use pezsp_runtime::traits::{Saturating, Zero};
+
+	/// [`UncheckedMigrationToV2`] wrapped in a
+	/// [`VersionedMigration`](pezframe_support::migrations::VersionedMigration), ensuring the
+	/// migration is only performed when on-chain version is 1.
+	pub type MigrationToV2<T> = pezframe_support::migrations::VersionedMigration<
+		1,
+		2,
+		UncheckedMigrationToV2<T>,
+		Pallet<T>,
+		<T as pezframe_system::Config>::DbWeight,
+	>;
+
+	#[storage_alias]
+	pub type Candidates<T: Config> = StorageValue<
+		Pallet<T>,
+		BoundedVec<CandidateInfo<<T as pezframe_system::Config>::AccountId, <<T as Config>::Currency as Currency<<T as pezframe_system::Config>::AccountId>>::Balance>, <T as Config>::MaxCandidates>,
+		ValueQuery,
+	>;
+
+	/// Migrate to V2.
+	pub struct UncheckedMigrationToV2<T>(PhantomData<T>);
+	impl<T: Config + pezpallet_balances::Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV2<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let mut weight = Weight::zero();
+			let mut count: u64 = 0;
+			// candidates who exist under the old `Candidates` key
+			let candidates = Candidates::<T>::take();
+
+			// New candidates who have registered since the upgrade. Under normal circumstances,
+			// this should not exist because the migration should be applied when the upgrade
+			// happens. But in Pezkuwi/Kusama we messed this up, and people registered under
+			// `CandidateList` while their funds were locked in `Candidates`.
+			let new_candidate_list = CandidateList::<T>::get();
+			if new_candidate_list.len().is_zero() {
+				// The new list is empty, so this is essentially being applied correctly. We just
+				// put the candidates into the new storage item.
+				CandidateList::<T>::put(&candidates);
+				// 1 write for the new list
+				weight.saturating_accrue(T::DbWeight::get().reads_writes(0, 1));
+			} else {
+				// Oops, the runtime upgraded without the migration. There are new candidates in
+				// `CandidateList`. So, let's just refund the old ones and assume they have already
+				// started participating in the new system.
+				for candidate in candidates {
+					let err = T::Currency::unreserve(&candidate.who, candidate.deposit);
+					if err > Zero::zero() {
+						log::error!(
+							target: LOG_TARGET,
+							"{:?} balance was unable to be unreserved from {:?}",
+							err, &candidate.who,
+						);
+					}
+					count.saturating_inc();
+				}
+				weight.saturating_accrue(
+					<<T as pezpallet_balances::Config>::WeightInfo as pezpallet_balances::WeightInfo>::force_unreserve().saturating_mul(count.into()),
+				);
+			}
+
+			log::info!(
+				target: LOG_TARGET,
+				"Unreserved locked bond of {} candidates, upgraded storage to version 2",
+				count,
+			);
+
+			weight.saturating_accrue(T::DbWeight::get().reads_writes(3, 2));
+			weight
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::DispatchError> {
+			let number_of_candidates = Candidates::<T>::get().to_vec().len();
+			Ok((number_of_candidates as u32).encode())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(_number_of_candidates: Vec<u8>) -> Result<(), pezsp_runtime::DispatchError> {
+			let new_number_of_candidates = Candidates::<T>::get().to_vec().len();
+			assert_eq!(
+				new_number_of_candidates, 0 as usize,
+				"after migration, the candidates map should be empty"
+			);
+			Ok(())
+		}
+	}
+}
+
+/// Version 1 Migration
+/// This migration ensures that any existing `Invulnerables` storage lists are sorted.
+pub mod v1 {
+	use super::*;
+	use pezframe_support::pezpallet_prelude::*;
+
+	pub struct MigrateToV1<T>(PhantomData<T>);
+	impl<T: Config> OnRuntimeUpgrade for MigrateToV1<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let on_chain_version = Pallet::<T>::on_chain_storage_version();
+			if on_chain_version == 0 {
+				let invulnerables_len = Invulnerables::<T>::get().to_vec().len();
+				Invulnerables::<T>::mutate(|invulnerables| {
+					invulnerables.sort();
+				});
+
+				StorageVersion::new(1).put::<Pallet<T>>();
+				log::info!(
+					target: LOG_TARGET,
+					"Sorted {} Invulnerables, upgraded storage to version 1",
+					invulnerables_len,
+				);
+				// Similar complexity to `set_invulnerables` (put storage value)
+				// Plus 1 read for length, 1 read for `on_chain_version`, 1 write to put version
+				T::WeightInfo::set_invulnerables(invulnerables_len as u32)
+					.saturating_add(T::DbWeight::get().reads_writes(2, 1))
+			} else {
+				log::info!(
+					target: LOG_TARGET,
+					"Migration did not execute. This probably should be removed"
+				);
+				T::DbWeight::get().reads(1)
+			}
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, pezsp_runtime::DispatchError> {
+			let number_of_invulnerables = Invulnerables::<T>::get().to_vec().len();
+			Ok((number_of_invulnerables as u32).encode())
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(number_of_invulnerables: Vec<u8>) -> Result<(), pezsp_runtime::DispatchError> {
+			let stored_invulnerables = Invulnerables::<T>::get().to_vec();
+			let mut sorted_invulnerables = stored_invulnerables.clone();
+			sorted_invulnerables.sort();
+			assert_eq!(
+				stored_invulnerables, sorted_invulnerables,
+				"after migration, the stored invulnerables should be sorted"
+			);
+
+			let number_of_invulnerables: u32 = Decode::decode(
+				&mut number_of_invulnerables.as_slice(),
+			)
+			.expect("the state parameter should be something that was generated by pre_upgrade");
+			let stored_invulnerables_len = stored_invulnerables.len() as u32;
+			assert_eq!(
+				number_of_invulnerables, stored_invulnerables_len,
+				"after migration, there should be the same number of invulnerables"
+			);
+
+			let on_chain_version = Pallet::<T>::on_chain_storage_version();
+			pezframe_support::ensure!(on_chain_version >= 1, "must_upgrade");
+
+			Ok(())
+		}
+	}
+}
+
+#[cfg(all(feature = "try-runtime", test))]
+mod tests {
+	use super::*;
+	use crate::{
+		migration::v2::Candidates,
+		mock::{new_test_ext, Balances, Test},
+	};
+	use pezframe_support::{
+		traits::{Currency, ReservableCurrency, StorageVersion},
+		BoundedVec,
+	};
+	use pezsp_runtime::traits::ConstU32;
+
+	#[test]
+	fn migrate_to_v2_with_new_candidates() {
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(1);
+			storage_version.put::<Pallet<Test>>();
+
+			let one = 1u64;
+			let two = 2u64;
+			let three = 3u64;
+			let deposit = 10u64;
+
+			// Set balance to 100
+			Balances::make_free_balance_be(&one, 100u64);
+			Balances::make_free_balance_be(&two, 100u64);
+			Balances::make_free_balance_be(&three, 100u64);
+
+			// Reservations: 10 for the "old" candidacy and 10 for the "new"
+			Balances::reserve(&one, 10u64).unwrap(); // old
+			Balances::reserve(&two, 20u64).unwrap(); // old + new
+			Balances::reserve(&three, 10u64).unwrap(); // new
+
+			// Candidate info
+			let candidate_one = CandidateInfo { who: one, deposit };
+			let candidate_two = CandidateInfo { who: two, deposit };
+			let candidate_three = CandidateInfo { who: three, deposit };
+
+			// Storage lists
+			let bounded_candidates =
+				BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![
+					candidate_one.clone(),
+					candidate_two.clone(),
+				])
+				.expect("it works");
+			let bounded_candidate_list =
+				BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![
+					candidate_two.clone(),
+					candidate_three.clone(),
+				])
+				.expect("it works");
+
+			// Set storage
+			Candidates::<Test>::put(bounded_candidates);
+			CandidateList::<Test>::put(bounded_candidate_list.clone());
+
+			// Sanity check
+			assert_eq!(Balances::free_balance(one), 90);
+			assert_eq!(Balances::free_balance(two), 80);
+			assert_eq!(Balances::free_balance(three), 90);
+
+			// Run migration
+			v2::MigrationToV2::<Test>::on_runtime_upgrade();
+
+			let new_storage_version = StorageVersion::get::<Pallet<Test>>();
+			assert_eq!(new_storage_version, 2);
+
+			// 10 should have been unreserved from the old candidacy
+			assert_eq!(Balances::free_balance(one), 100);
+			assert_eq!(Balances::free_balance(two), 90);
+			assert_eq!(Balances::free_balance(three), 90);
+			// The storage item should be gone
+			assert!(Candidates::<Test>::get().is_empty());
+			// The new storage item should be preserved
+			assert_eq!(CandidateList::<Test>::get(), bounded_candidate_list);
+		});
+	}
+
+	#[test]
+	fn migrate_to_v2_without_new_candidates() {
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(1);
+			storage_version.put::<Pallet<Test>>();
+
+			let one = 1u64;
+			let two = 2u64;
+			let deposit = 10u64;
+
+			// Set balance to 100
+			Balances::make_free_balance_be(&one, 100u64);
+			Balances::make_free_balance_be(&two, 100u64);
+
+			// Reservations
+			Balances::reserve(&one, 10u64).unwrap(); // old
+			Balances::reserve(&two, 10u64).unwrap(); // old
+
+			// Candidate info
+			let candidate_one = CandidateInfo { who: one, deposit };
+			let candidate_two = CandidateInfo { who: two, deposit };
+
+			// Storage lists
+			let bounded_candidates =
+				BoundedVec::<CandidateInfo<u64, u64>, ConstU32<20>>::try_from(vec![
+					candidate_one.clone(),
+					candidate_two.clone(),
+				])
+				.expect("it works");
+
+			// Set storage
+			Candidates::<Test>::put(bounded_candidates.clone());
+
+			// Sanity check
+			assert_eq!(Balances::free_balance(one), 90);
+			assert_eq!(Balances::free_balance(two), 90);
+
+			// Run migration
+			v2::MigrationToV2::<Test>::on_runtime_upgrade();
+
+			let new_storage_version = StorageVersion::get::<Pallet<Test>>();
+			assert_eq!(new_storage_version, 2);
+
+			// Nothing changes deposit-wise
+			assert_eq!(Balances::free_balance(one), 90);
+			assert_eq!(Balances::free_balance(two), 90);
+			// The storage item should be gone
+			assert!(Candidates::<Test>::get().is_empty());
+			// The new storage item should have the info now
+			assert_eq!(CandidateList::<Test>::get(), bounded_candidates);
+		});
+	}
+}

pezcumulus/pezpallets/collator-selection/src/mock.rs

@@ -0,0 +1,208 @@
+// 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 crate as collator_selection;
+use pezframe_support::{
+	derive_impl, ord_parameter_types, parameter_types,
+	traits::{ConstBool, ConstU32, ConstU64, FindAuthor, ValidatorRegistration},
+	PalletId,
+};
+use pezframe_system as system;
+use pezframe_system::EnsureSignedBy;
+use pezsp_runtime::{testing::UintAuthorityId, traits::OpaqueKeys, BuildStorage, RuntimeAppPublic};
+
+type Block = pezframe_system::mocking::MockBlock<Test>;
+
+// Configure a mock runtime to test the pallet.
+pezframe_support::construct_runtime!(
+	pub enum Test
+	{
+		System: pezframe_system,
+		Timestamp: pezpallet_timestamp,
+		Session: pezpallet_session,
+		Aura: pezpallet_aura,
+		Balances: pezpallet_balances,
+		CollatorSelection: collator_selection,
+		Authorship: pezpallet_authorship,
+	}
+);
+
+parameter_types! {
+	pub const SS58Prefix: u8 = 42;
+}
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl system::Config for Test {
+	type Block = Block;
+	type AccountData = pezpallet_balances::AccountData<u64>;
+	type SS58Prefix = SS58Prefix;
+}
+
+parameter_types! {
+	pub const ExistentialDeposit: u64 = 5;
+}
+
+#[derive_impl(pezpallet_balances::config_preludes::TestDefaultConfig)]
+impl pezpallet_balances::Config for Test {
+	type ExistentialDeposit = ExistentialDeposit;
+	type AccountStore = System;
+}
+
+pub struct Author4;
+impl FindAuthor<u64> for Author4 {
+	fn find_author<'a, I>(_digests: I) -> Option<u64>
+	where
+		I: 'a + IntoIterator<Item = (pezframe_support::ConsensusEngineId, &'a [u8])>,
+	{
+		Some(4)
+	}
+}
+
+impl pezpallet_authorship::Config for Test {
+	type FindAuthor = Author4;
+	type EventHandler = CollatorSelection;
+}
+
+impl pezpallet_timestamp::Config for Test {
+	type Moment = u64;
+	type OnTimestampSet = Aura;
+	type MinimumPeriod = ConstU64<1>;
+	type WeightInfo = ();
+}
+
+impl pezpallet_aura::Config for Test {
+	type AuthorityId = pezsp_consensus_aura::sr25519::AuthorityId;
+	type MaxAuthorities = ConstU32<100_000>;
+	type DisabledValidators = ();
+	type AllowMultipleBlocksPerSlot = ConstBool<false>;
+	type SlotDuration = pezpallet_aura::MinimumPeriodTimesTwo<Self>;
+}
+
+pezsp_runtime::impl_opaque_keys! {
+	pub struct MockSessionKeys {
+		// a key for aura authoring
+		pub aura: UintAuthorityId,
+	}
+}
+
+impl From<UintAuthorityId> for MockSessionKeys {
+	fn from(aura: pezsp_runtime::testing::UintAuthorityId) -> Self {
+		Self { aura }
+	}
+}
+
+parameter_types! {
+	pub static SessionHandlerCollators: Vec<u64> = Vec::new();
+	pub static SessionChangeBlock: u64 = 0;
+}
+
+pub struct TestSessionHandler;
+impl pezpallet_session::SessionHandler<u64> for TestSessionHandler {
+	const KEY_TYPE_IDS: &'static [pezsp_runtime::KeyTypeId] = &[UintAuthorityId::ID];
+	fn on_genesis_session<Ks: OpaqueKeys>(keys: &[(u64, Ks)]) {
+		SessionHandlerCollators::set(keys.iter().map(|(a, _)| *a).collect::<Vec<_>>())
+	}
+	fn on_new_session<Ks: OpaqueKeys>(_: bool, keys: &[(u64, Ks)], _: &[(u64, Ks)]) {
+		SessionChangeBlock::set(System::block_number());
+		SessionHandlerCollators::set(keys.iter().map(|(a, _)| *a).collect::<Vec<_>>())
+	}
+	fn on_before_session_ending() {}
+	fn on_disabled(_: u32) {}
+}
+
+parameter_types! {
+	pub const Offset: u64 = 0;
+	pub const Period: u64 = 10;
+}
+
+impl pezpallet_session::Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type ValidatorId = <Self as pezframe_system::Config>::AccountId;
+	// we don't have stash and controller, thus we don't need the convert as well.
+	type ValidatorIdOf = IdentityCollator;
+	type ShouldEndSession = pezpallet_session::PeriodicSessions<Period, Offset>;
+	type NextSessionRotation = pezpallet_session::PeriodicSessions<Period, Offset>;
+	type SessionManager = CollatorSelection;
+	type SessionHandler = TestSessionHandler;
+	type Keys = MockSessionKeys;
+	type DisablingStrategy = ();
+	type WeightInfo = ();
+	type Currency = Balances;
+	type KeyDeposit = ();
+}
+
+ord_parameter_types! {
+	pub const RootAccount: u64 = 777;
+}
+
+parameter_types! {
+	pub const PotId: PalletId = PalletId(*b"PotStake");
+}
+
+pub struct IsRegistered;
+impl ValidatorRegistration<u64> for IsRegistered {
+	fn is_registered(id: &u64) -> bool {
+		*id != 42u64
+	}
+}
+
+impl Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type Currency = Balances;
+	type UpdateOrigin = EnsureSignedBy<RootAccount, u64>;
+	type PotId = PotId;
+	type MaxCandidates = ConstU32<20>;
+	type MinEligibleCollators = ConstU32<1>;
+	type MaxInvulnerables = ConstU32<20>;
+	type KickThreshold = Period;
+	type ValidatorId = <Self as pezframe_system::Config>::AccountId;
+	type ValidatorIdOf = IdentityCollator;
+	type ValidatorRegistration = IsRegistered;
+	type WeightInfo = ();
+}
+
+pub fn new_test_ext() -> pezsp_io::TestExternalities {
+	pezsp_tracing::try_init_simple();
+	let mut t = pezframe_system::GenesisConfig::<Test>::default().build_storage().unwrap();
+	let invulnerables = vec![2, 1]; // unsorted
+
+	let balances = vec![(1, 100), (2, 100), (3, 100), (4, 100), (5, 100)];
+	let keys = balances
+		.iter()
+		.map(|&(i, _)| (i, i, MockSessionKeys { aura: UintAuthorityId(i) }))
+		.collect::<Vec<_>>();
+	let collator_selection = collator_selection::GenesisConfig::<Test> {
+		desired_candidates: 2,
+		candidacy_bond: 10,
+		invulnerables,
+	};
+	let session = pezpallet_session::GenesisConfig::<Test> { keys, ..Default::default() };
+	pezpallet_balances::GenesisConfig::<Test> { balances, ..Default::default() }
+		.assimilate_storage(&mut t)
+		.unwrap();
+	// collator selection must be initialized before session.
+	collator_selection.assimilate_storage(&mut t).unwrap();
+	session.assimilate_storage(&mut t).unwrap();
+
+	t.into()
+}
+
+pub fn initialize_to_block(n: u64) {
+	for i in System::block_number() + 1..=n {
+		System::set_block_number(i);
+		<AllPalletsWithSystem as pezframe_support::traits::OnInitialize<u64>>::on_initialize(i);
+	}
+}

pezcumulus/pezpallets/collator-selection/src/weights.rs

@@ -0,0 +1,243 @@
+// This file is part of Pezcumulus.
+
+// 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.
+
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+
+use core::marker::PhantomData;
+use pezframe_support::{
+	traits::Get,
+	weights::{constants::RocksDbWeight, Weight},
+};
+
+// The weight info trait for `pezpallet_collator_selection`.
+pub trait WeightInfo {
+	fn set_invulnerables(_b: u32) -> Weight;
+	fn add_invulnerable(_b: u32, _c: u32) -> Weight;
+	fn remove_invulnerable(_b: u32) -> Weight;
+	fn set_desired_candidates() -> Weight;
+	fn set_candidacy_bond(_c: u32, _k: u32) -> Weight;
+	fn register_as_candidate(_c: u32) -> Weight;
+	fn leave_intent(_c: u32) -> Weight;
+	fn update_bond(_c: u32) -> Weight;
+	fn take_candidate_slot(_c: u32) -> Weight;
+	fn note_author() -> Weight;
+	fn new_session(_c: u32, _r: u32) -> Weight;
+}
+
+/// Weights for pezpallet_collator_selection using the Bizinikiwi node and recommended hardware.
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	fn set_invulnerables(b: u32) -> Weight {
+		Weight::from_parts(18_563_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(68_000_u64, 0).saturating_mul(b as u64))
+			.saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	fn set_desired_candidates() -> Weight {
+		Weight::from_parts(16_363_000_u64, 0).saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	fn set_candidacy_bond(_c: u32, _k: u32) -> Weight {
+		Weight::from_parts(16_840_000_u64, 0).saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	fn register_as_candidate(c: u32) -> Weight {
+		Weight::from_parts(71_196_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(198_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	fn leave_intent(c: u32) -> Weight {
+		Weight::from_parts(55_336_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(151_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(1_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	fn update_bond(c: u32) -> Weight {
+		Weight::from_parts(55_336_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(151_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(1_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	fn take_candidate_slot(c: u32) -> Weight {
+		Weight::from_parts(71_196_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(198_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	fn note_author() -> Weight {
+		Weight::from_parts(71_461_000_u64, 0)
+			.saturating_add(T::DbWeight::get().reads(3_u64))
+			.saturating_add(T::DbWeight::get().writes(4_u64))
+	}
+	fn new_session(r: u32, c: u32) -> Weight {
+		Weight::from_parts(0_u64, 0)
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_parts(109_961_000_u64, 0).saturating_mul(r as u64))
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_parts(151_952_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(T::DbWeight::get().reads(1_u64.saturating_mul(r as u64)))
+			.saturating_add(T::DbWeight::get().reads(2_u64.saturating_mul(c as u64)))
+			.saturating_add(T::DbWeight::get().writes(2_u64.saturating_mul(r as u64)))
+			.saturating_add(T::DbWeight::get().writes(2_u64.saturating_mul(c as u64)))
+	}
+	/// Storage: Session NextKeys (r:1 w:0)
+	/// Proof Skipped: Session NextKeys (max_values: None, max_size: None, mode: Measured)
+	/// Storage: CollatorSelection Invulnerables (r:1 w:1)
+	/// Proof: CollatorSelection Invulnerables (max_values: Some(1), max_size: Some(641), added:
+	/// 1136, mode: MaxEncodedLen) Storage: CollatorSelection Candidates (r:1 w:1)
+	/// Proof: CollatorSelection Candidates (max_values: Some(1), max_size: Some(4802), added: 5297,
+	/// mode: MaxEncodedLen) Storage: System Account (r:1 w:1)
+	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode:
+	/// MaxEncodedLen) The range of component `b` is `[1, 19]`.
+	/// The range of component `c` is `[1, 99]`.
+	fn add_invulnerable(b: u32, c: u32) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `757 + b * (32 ±0) + c * (53 ±0)`
+		//  Estimated: `6287 + b * (37 ±0) + c * (53 ±0)`
+		// Minimum execution time: 52_720_000 picoseconds.
+		Weight::from_parts(56_102_459, 0)
+			.saturating_add(Weight::from_parts(0, 6287))
+			// Standard Error: 12_957
+			.saturating_add(Weight::from_parts(26_422, 0).saturating_mul(b.into()))
+			// Standard Error: 2_456
+			.saturating_add(Weight::from_parts(128_528, 0).saturating_mul(c.into()))
+			.saturating_add(T::DbWeight::get().reads(4))
+			.saturating_add(T::DbWeight::get().writes(3))
+			.saturating_add(Weight::from_parts(0, 37).saturating_mul(b.into()))
+			.saturating_add(Weight::from_parts(0, 53).saturating_mul(c.into()))
+	}
+	/// Storage: CollatorSelection Invulnerables (r:1 w:1)
+	/// Proof: CollatorSelection Invulnerables (max_values: Some(1), max_size: Some(3202), added:
+	/// 3697, mode: MaxEncodedLen) The range of component `b` is `[1, 100]`.
+	fn remove_invulnerable(b: u32) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `119 + b * (32 ±0)`
+		//  Estimated: `4687`
+		// Minimum execution time: 183_054_000 picoseconds.
+		Weight::from_parts(197_205_427, 0)
+			.saturating_add(Weight::from_parts(0, 4687))
+			// Standard Error: 13_533
+			.saturating_add(Weight::from_parts(376_231, 0).saturating_mul(b.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+}
+
+// For backwards compatibility and tests
+impl WeightInfo for () {
+	fn set_invulnerables(b: u32) -> Weight {
+		Weight::from_parts(18_563_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(68_000_u64, 0).saturating_mul(b as u64))
+			.saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	fn set_desired_candidates() -> Weight {
+		Weight::from_parts(16_363_000_u64, 0).saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	fn set_candidacy_bond(_c: u32, _k: u32) -> Weight {
+		Weight::from_parts(16_840_000_u64, 0).saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	fn register_as_candidate(c: u32) -> Weight {
+		Weight::from_parts(71_196_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(198_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+	}
+	fn leave_intent(c: u32) -> Weight {
+		Weight::from_parts(55_336_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(151_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(1_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+	}
+	fn note_author() -> Weight {
+		Weight::from_parts(71_461_000_u64, 0)
+			.saturating_add(RocksDbWeight::get().reads(3_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+	fn update_bond(c: u32) -> Weight {
+		Weight::from_parts(55_336_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(151_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(3_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+	fn take_candidate_slot(c: u32) -> Weight {
+		Weight::from_parts(71_196_000_u64, 0)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(198_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(3_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+	fn new_session(r: u32, c: u32) -> Weight {
+		Weight::from_parts(0_u64, 0)
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_parts(109_961_000_u64, 0).saturating_mul(r as u64))
+			// Standard Error: 1_010_000
+			.saturating_add(Weight::from_parts(151_952_000_u64, 0).saturating_mul(c as u64))
+			.saturating_add(RocksDbWeight::get().reads(1_u64.saturating_mul(r as u64)))
+			.saturating_add(RocksDbWeight::get().reads(2_u64.saturating_mul(c as u64)))
+			.saturating_add(RocksDbWeight::get().writes(2_u64.saturating_mul(r as u64)))
+			.saturating_add(RocksDbWeight::get().writes(2_u64.saturating_mul(c as u64)))
+	}
+	/// Storage: Session NextKeys (r:1 w:0)
+	/// Proof Skipped: Session NextKeys (max_values: None, max_size: None, mode: Measured)
+	/// Storage: CollatorSelection Invulnerables (r:1 w:1)
+	/// Proof: CollatorSelection Invulnerables (max_values: Some(1), max_size: Some(641), added:
+	/// 1136, mode: MaxEncodedLen) Storage: CollatorSelection Candidates (r:1 w:1)
+	/// Proof: CollatorSelection Candidates (max_values: Some(1), max_size: Some(4802), added: 5297,
+	/// mode: MaxEncodedLen) Storage: System Account (r:1 w:1)
+	/// Proof: System Account (max_values: None, max_size: Some(128), added: 2603, mode:
+	/// MaxEncodedLen) The range of component `b` is `[1, 19]`.
+	/// The range of component `c` is `[1, 99]`.
+	fn add_invulnerable(b: u32, c: u32) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `757 + b * (32 ±0) + c * (53 ±0)`
+		//  Estimated: `6287 + b * (37 ±0) + c * (53 ±0)`
+		// Minimum execution time: 52_720_000 picoseconds.
+		Weight::from_parts(56_102_459, 0)
+			.saturating_add(Weight::from_parts(0, 6287))
+			// Standard Error: 12_957
+			.saturating_add(Weight::from_parts(26_422, 0).saturating_mul(b.into()))
+			// Standard Error: 2_456
+			.saturating_add(Weight::from_parts(128_528, 0).saturating_mul(c.into()))
+			.saturating_add(RocksDbWeight::get().reads(4))
+			.saturating_add(RocksDbWeight::get().writes(3))
+			.saturating_add(Weight::from_parts(0, 37).saturating_mul(b.into()))
+			.saturating_add(Weight::from_parts(0, 53).saturating_mul(c.into()))
+	}
+	/// Storage: CollatorSelection Invulnerables (r:1 w:1)
+	/// Proof: CollatorSelection Invulnerables (max_values: Some(1), max_size: Some(3202), added:
+	/// 3697, mode: MaxEncodedLen) The range of component `b` is `[1, 100]`.
+	fn remove_invulnerable(b: u32) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `119 + b * (32 ±0)`
+		//  Estimated: `4687`
+		// Minimum execution time: 183_054_000 picoseconds.
+		Weight::from_parts(197_205_427, 0)
+			.saturating_add(Weight::from_parts(0, 4687))
+			// Standard Error: 13_533
+			.saturating_add(Weight::from_parts(376_231, 0).saturating_mul(b.into()))
+			.saturating_add(RocksDbWeight::get().reads(1))
+			.saturating_add(RocksDbWeight::get().writes(1))
+	}
+}

pezcumulus/pezpallets/dmp-queue/Cargo.toml

@@ -0,0 +1,66 @@
+[package]
+name = "pezcumulus-pezpallet-dmp-queue"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+repository.workspace = true
+description = "Migrates messages from the old DMP queue pallet."
+license = "Apache-2.0"
+documentation = "https://docs.rs/pezcumulus-pezpallet-dmp-queue"
+homepage = { workspace = true }
+
+[lints]
+workspace = true
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+log = { workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+pezframe-benchmarking = { optional = true, workspace = true }
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+# Pezkuwi
+xcm = { workspace = true }
+
+# Pezcumulus
+pezcumulus-primitives-core = { workspace = true }
+
+[dev-dependencies]
+pezsp-core = { workspace = true, default-features = true }
+pezsp-tracing = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezcumulus-primitives-core/std",
+  "pezframe-benchmarking?/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "log/std",
+  "scale-info/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "xcm/std",
+]
+runtime-benchmarks = [
+  "pezcumulus-primitives-core/runtime-benchmarks",
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "xcm/runtime-benchmarks",
+]
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezsp-runtime/try-runtime",
+]

pezcumulus/pezpallets/dmp-queue/src/benchmarking.rs

@@ -0,0 +1,108 @@
+// 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 for the `pezcumulus-pezpallet-dmp-queue`.
+
+#![cfg(feature = "runtime-benchmarks")]
+
+use crate::*;
+
+use alloc::vec;
+use pezframe_benchmarking::v2::*;
+use pezframe_support::{pezpallet_prelude::*, traits::Hooks};
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	/// This benchmark uses the proper maximal message length.
+	#[benchmark]
+	fn on_idle_good_msg() {
+		let msg = vec![123; MaxDmpMessageLenOf::<T>::get() as usize];
+
+		Pages::<T>::insert(0, vec![(123, msg.clone())]);
+		PageIndex::<T>::put(PageIndexData { begin_used: 0, end_used: 1, overweight_count: 0 });
+		MigrationStatus::<T>::set(MigrationState::StartedExport { next_begin_used: 0 });
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+
+		assert_last_event::<T>(Event::Exported { page: 0 }.into());
+	}
+
+	/// This benchmark uses 64 KiB messages to emulate a large old message.
+	#[benchmark]
+	fn on_idle_large_msg() {
+		let msg = vec![123; 1 << 16];
+
+		Pages::<T>::insert(0, vec![(123, msg.clone())]);
+		PageIndex::<T>::put(PageIndexData { begin_used: 0, end_used: 1, overweight_count: 0 });
+		MigrationStatus::<T>::set(MigrationState::StartedExport { next_begin_used: 0 });
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+
+		assert_last_event::<T>(Event::Exported { page: 0 }.into());
+	}
+
+	#[benchmark]
+	fn on_idle_overweight_good_msg() {
+		let msg = vec![123; MaxDmpMessageLenOf::<T>::get() as usize];
+
+		Overweight::<T>::insert(0, (123, msg.clone()));
+		PageIndex::<T>::put(PageIndexData { begin_used: 0, end_used: 1, overweight_count: 1 });
+		MigrationStatus::<T>::set(MigrationState::StartedOverweightExport {
+			next_overweight_index: 0,
+		});
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+
+		assert_last_event::<T>(Event::ExportedOverweight { index: 0 }.into());
+	}
+
+	#[benchmark]
+	fn on_idle_overweight_large_msg() {
+		let msg = vec![123; 1 << 16];
+
+		Overweight::<T>::insert(0, (123, msg.clone()));
+		PageIndex::<T>::put(PageIndexData { begin_used: 0, end_used: 1, overweight_count: 1 });
+		MigrationStatus::<T>::set(MigrationState::StartedOverweightExport {
+			next_overweight_index: 0,
+		});
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+
+		assert_last_event::<T>(Event::ExportOverweightFailed { index: 0 }.into());
+	}
+
+	impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Runtime);
+}
+
+fn assert_last_event<T: Config>(generic_event: <T as Config>::RuntimeEvent) {
+	let events = pezframe_system::Pallet::<T>::events();
+	let system_event: <T as pezframe_system::Config>::RuntimeEvent = generic_event.into();
+	let pezframe_system::EventRecord { event, .. } = events.last().expect("Event expected");
+	assert_eq!(event, &system_event.into());
+}

pezcumulus/pezpallets/dmp-queue/src/lib.rs

@@ -0,0 +1,291 @@
+// 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.
+
+//! This pallet used to implement a message queue for downward messages from the relay-chain.
+//!
+//! It is now deprecated and has been refactored to simply drain any remaining messages into
+//! something implementing `HandleMessage`. It proceeds in the state of
+//! [`MigrationState`] one by one by their listing in the source code. The pallet can be removed
+//! from the runtime once `Completed` was emitted.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+#![allow(deprecated)] // The pallet itself is deprecated.
+
+extern crate alloc;
+
+use migration::*;
+pub use pallet::*;
+
+mod benchmarking;
+mod migration;
+mod mock;
+mod tests;
+pub mod weights;
+
+pub use weights::WeightInfo;
+
+/// The maximal length of a DMP message.
+pub type MaxDmpMessageLenOf<T> =
+	<<T as Config>::DmpSink as pezframe_support::traits::HandleMessage>::MaxMessageLen;
+
+#[pezframe_support::pallet]
+#[deprecated(
+	note = "`pezcumulus-pezpallet-dmp-queue` will be removed after November 2024. It can be removed once its lazy migration completed. See <https://github.com/pezkuwichain/kurdistan-sdk/issues/101>."
+)]
+pub mod pallet {
+	use super::*;
+	use pezframe_support::{pezpallet_prelude::*, traits::HandleMessage, weights::WeightMeter};
+	use pezframe_system::pezpallet_prelude::*;
+	use pezsp_io::hashing::twox_128;
+
+	const STORAGE_VERSION: StorageVersion = StorageVersion::new(2);
+
+	#[pallet::pallet]
+	#[pallet::storage_version(STORAGE_VERSION)]
+	pub struct Pallet<T>(_);
+
+	#[pallet::config]
+	pub trait Config: pezframe_system::Config {
+		/// The overarching event type of the runtime.
+		#[allow(deprecated)]
+		type RuntimeEvent: From<Event<Self>> + IsType<<Self as pezframe_system::Config>::RuntimeEvent>;
+
+		/// The sink for all DMP messages that the lazy migration will use.
+		type DmpSink: HandleMessage;
+
+		/// Weight info for this pallet (only needed for the lazy migration).
+		type WeightInfo: WeightInfo;
+	}
+
+	/// The migration state of this pallet.
+	#[pallet::storage]
+	pub type MigrationStatus<T> = StorageValue<_, MigrationState, ValueQuery>;
+
+	/// The lazy-migration state of the pallet.
+	#[derive(
+		codec::Encode, codec::Decode, Debug, PartialEq, Eq, Clone, MaxEncodedLen, TypeInfo,
+	)]
+	pub enum MigrationState {
+		/// Migration has not started yet.
+		NotStarted,
+		/// The export of pages started.
+		StartedExport {
+			/// The next page that should be exported.
+			next_begin_used: PageCounter,
+		},
+		/// The page export completed.
+		CompletedExport,
+		/// The export of overweight messages started.
+		StartedOverweightExport {
+			/// The next overweight index that should be exported.
+			next_overweight_index: u64,
+		},
+		/// The export of overweight messages completed.
+		CompletedOverweightExport,
+		/// The storage cleanup started.
+		StartedCleanup { cursor: Option<BoundedVec<u8, ConstU32<1024>>> },
+		/// The migration finished. The pallet can now be removed from the runtime.
+		Completed,
+	}
+
+	impl Default for MigrationState {
+		fn default() -> Self {
+			Self::NotStarted
+		}
+	}
+
+	#[pallet::event]
+	#[pallet::generate_deposit(pub(super) fn deposit_event)]
+	pub enum Event<T: Config> {
+		/// The export of pages started.
+		StartedExport,
+
+		/// The export of a page completed.
+		Exported { page: PageCounter },
+
+		/// The export of a page failed.
+		///
+		/// This should never be emitted.
+		ExportFailed { page: PageCounter },
+
+		/// The export of pages completed.
+		CompletedExport,
+
+		/// The export of overweight messages started.
+		StartedOverweightExport,
+
+		/// The export of an overweight message completed.
+		ExportedOverweight { index: OverweightIndex },
+
+		/// The export of an overweight message failed.
+		///
+		/// This should never be emitted.
+		ExportOverweightFailed { index: OverweightIndex },
+
+		/// The export of overweight messages completed.
+		CompletedOverweightExport,
+
+		/// The cleanup of remaining pallet storage started.
+		StartedCleanup,
+
+		/// Some debris was cleaned up.
+		CleanedSome { keys_removed: u32 },
+
+		/// The cleanup of remaining pallet storage completed.
+		Completed { error: bool },
+	}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {}
+
+	#[pallet::hooks]
+	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
+		fn integrity_test() {
+			let w = Self::on_idle_weight();
+			assert!(w != Weight::zero());
+			assert!(w.all_lte(T::BlockWeights::get().max_block));
+		}
+
+		fn on_idle(now: BlockNumberFor<T>, limit: Weight) -> Weight {
+			let mut meter = WeightMeter::with_limit(limit);
+
+			if meter.try_consume(Self::on_idle_weight()).is_err() {
+				log::debug!(target: LOG, "Not enough weight for on_idle. {} < {}", Self::on_idle_weight(), limit);
+				return meter.consumed();
+			}
+
+			let state = MigrationStatus::<T>::get();
+			let index = PageIndex::<T>::get();
+			log::debug!(target: LOG, "on_idle: block={:?}, state={:?}, index={:?}", now, state, index);
+
+			match state {
+				MigrationState::NotStarted => {
+					log::debug!(target: LOG, "Init export at page {}", index.begin_used);
+
+					MigrationStatus::<T>::put(MigrationState::StartedExport {
+						next_begin_used: index.begin_used,
+					});
+					Self::deposit_event(Event::StartedExport);
+				},
+				MigrationState::StartedExport { next_begin_used } => {
+					log::debug!(target: LOG, "Exporting page {}", next_begin_used);
+
+					if next_begin_used == index.end_used {
+						MigrationStatus::<T>::put(MigrationState::CompletedExport);
+						log::debug!(target: LOG, "CompletedExport");
+						Self::deposit_event(Event::CompletedExport);
+					} else {
+						let res = migration::migrate_page::<T>(next_begin_used);
+
+						MigrationStatus::<T>::put(MigrationState::StartedExport {
+							next_begin_used: next_begin_used.saturating_add(1),
+						});
+
+						if let Ok(()) = res {
+							log::debug!(target: LOG, "Exported page {}", next_begin_used);
+							Self::deposit_event(Event::Exported { page: next_begin_used });
+						} else {
+							Self::deposit_event(Event::ExportFailed { page: next_begin_used });
+						}
+					}
+				},
+				MigrationState::CompletedExport => {
+					log::debug!(target: LOG, "Init export overweight at index 0");
+
+					MigrationStatus::<T>::put(MigrationState::StartedOverweightExport {
+						next_overweight_index: 0,
+					});
+					Self::deposit_event(Event::StartedOverweightExport);
+				},
+				MigrationState::StartedOverweightExport { next_overweight_index } => {
+					log::debug!(target: LOG, "Exporting overweight index {}", next_overweight_index);
+
+					if next_overweight_index == index.overweight_count {
+						MigrationStatus::<T>::put(MigrationState::CompletedOverweightExport);
+						log::debug!(target: LOG, "CompletedOverweightExport");
+						Self::deposit_event(Event::CompletedOverweightExport);
+					} else {
+						let res = migration::migrate_overweight::<T>(next_overweight_index);
+
+						MigrationStatus::<T>::put(MigrationState::StartedOverweightExport {
+							next_overweight_index: next_overweight_index.saturating_add(1),
+						});
+
+						if let Ok(()) = res {
+							log::debug!(target: LOG, "Exported overweight index {next_overweight_index}");
+							Self::deposit_event(Event::ExportedOverweight {
+								index: next_overweight_index,
+							});
+						} else {
+							Self::deposit_event(Event::ExportOverweightFailed {
+								index: next_overweight_index,
+							});
+						}
+					}
+				},
+				MigrationState::CompletedOverweightExport => {
+					log::debug!(target: LOG, "Init cleanup");
+
+					MigrationStatus::<T>::put(MigrationState::StartedCleanup { cursor: None });
+					Self::deposit_event(Event::StartedCleanup);
+				},
+				MigrationState::StartedCleanup { cursor } => {
+					log::debug!(target: LOG, "Cleaning up");
+					let hashed_prefix =
+						twox_128(<Pallet<T> as PalletInfoAccess>::name().as_bytes());
+
+					let result = pezframe_support::storage::unhashed::clear_prefix(
+						&hashed_prefix,
+						Some(2), // Somehow it does nothing when set to 1, so we set it to 2.
+						cursor.as_ref().map(|c| c.as_ref()),
+					);
+					Self::deposit_event(Event::CleanedSome { keys_removed: result.backend });
+
+					// GOTCHA! We deleted *all* pallet storage; hence we also our own
+					// `MigrationState`. BUT we insert it back:
+					if let Some(unbound_cursor) = result.maybe_cursor {
+						if let Ok(cursor) = unbound_cursor.try_into() {
+							log::debug!(target: LOG, "Next cursor: {:?}", &cursor);
+							MigrationStatus::<T>::put(MigrationState::StartedCleanup {
+								cursor: Some(cursor),
+							});
+						} else {
+							MigrationStatus::<T>::put(MigrationState::Completed);
+							log::error!(target: LOG, "Completed with error: could not bound cursor");
+							Self::deposit_event(Event::Completed { error: true });
+						}
+					} else {
+						MigrationStatus::<T>::put(MigrationState::Completed);
+						log::debug!(target: LOG, "Completed");
+						Self::deposit_event(Event::Completed { error: false });
+					}
+				},
+				MigrationState::Completed => {
+					log::debug!(target: LOG, "Idle; you can remove this pallet");
+				},
+			}
+
+			meter.consumed()
+		}
+	}
+
+	impl<T: Config> Pallet<T> {
+		/// The worst-case weight of [`Self::on_idle`].
+		pub fn on_idle_weight() -> Weight {
+			<T as crate::Config>::WeightInfo::on_idle_good_msg()
+				.max(<T as crate::Config>::WeightInfo::on_idle_large_msg())
+		}
+	}
+}

pezcumulus/pezpallets/dmp-queue/src/migration.rs

@@ -0,0 +1,113 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Migrates the storage from the previously deleted DMP pallet.
+
+use crate::*;
+use alloc::vec::Vec;
+use pezcumulus_primitives_core::relay_chain::BlockNumber as RelayBlockNumber;
+use pezframe_support::{pezpallet_prelude::*, storage_alias, traits::HandleMessage};
+
+pub(crate) const LOG: &str = "runtime::dmp-queue-export-xcms";
+
+/// The old `PageIndexData` struct.
+#[derive(Copy, Clone, Eq, PartialEq, Default, Encode, Decode, RuntimeDebug, TypeInfo)]
+pub struct PageIndexData {
+	/// The lowest used page index.
+	pub begin_used: PageCounter,
+	/// The lowest unused page index.
+	pub end_used: PageCounter,
+	/// The number of overweight messages ever recorded (and thus the lowest free index).
+	pub overweight_count: OverweightIndex,
+}
+
+/// The old `MigrationState` type.
+pub type OverweightIndex = u64;
+/// The old `MigrationState` type.
+pub type PageCounter = u32;
+
+/// The old `PageIndex` storage item.
+#[storage_alias]
+pub type PageIndex<T: Config> = StorageValue<Pallet<T>, PageIndexData, ValueQuery>;
+
+/// The old `Pages` storage item.
+#[storage_alias]
+pub type Pages<T: Config> = StorageMap<
+	Pallet<T>,
+	Blake2_128Concat,
+	PageCounter,
+	Vec<(RelayBlockNumber, Vec<u8>)>,
+	ValueQuery,
+>;
+
+/// The old `Overweight` storage item.
+#[storage_alias]
+pub type Overweight<T: Config> = CountedStorageMap<
+	Pallet<T>,
+	Blake2_128Concat,
+	OverweightIndex,
+	(RelayBlockNumber, Vec<u8>),
+	OptionQuery,
+>;
+
+pub(crate) mod testing_only {
+	use super::*;
+
+	/// This alias is not used by the migration but only for testing.
+	///
+	/// Note that the alias type is wrong on purpose.
+	#[storage_alias]
+	pub type Configuration<T: Config> = StorageValue<Pallet<T>, u32>;
+}
+
+/// Migrates a single page to the `DmpSink`.
+pub(crate) fn migrate_page<T: crate::Config>(p: PageCounter) -> Result<(), ()> {
+	let page = Pages::<T>::take(p);
+	log::debug!(target: LOG, "Migrating page #{p} with {} messages ...", page.len());
+	if page.is_empty() {
+		log::error!(target: LOG, "Page #{p}: EMPTY - storage corrupted?");
+		return Err(());
+	}
+
+	for (m, (block, msg)) in page.iter().enumerate() {
+		let Ok(bound) = BoundedVec::<u8, _>::try_from(msg.clone()) else {
+			log::error!(target: LOG, "[Page {p}] Message #{m}: TOO LONG - dropping");
+			continue;
+		};
+
+		T::DmpSink::handle_message(bound.as_bounded_slice());
+		log::debug!(target: LOG, "[Page {p}] Migrated message #{m} from block {block}");
+	}
+
+	Ok(())
+}
+
+/// Migrates a single overweight message to the `DmpSink`.
+pub(crate) fn migrate_overweight<T: crate::Config>(i: OverweightIndex) -> Result<(), ()> {
+	let Some((block, msg)) = Overweight::<T>::take(i) else {
+		log::error!(target: LOG, "[Overweight {i}] Message: EMPTY - storage corrupted?");
+		return Err(());
+	};
+	let Ok(bound) = BoundedVec::<u8, _>::try_from(msg) else {
+		log::error!(target: LOG, "[Overweight {i}] Message: TOO LONG - dropping");
+		return Err(());
+	};
+
+	T::DmpSink::handle_message(bound.as_bounded_slice());
+	log::debug!(target: LOG, "[Overweight {i}] Migrated message from block {block}");
+
+	Ok(())
+}

pezcumulus/pezpallets/dmp-queue/src/mock.rs

@@ -0,0 +1,76 @@
+// This file is part of Pezcumulus.
+
+// 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.
+
+#![cfg(test)]
+
+use pezframe_support::{derive_impl, parameter_types, traits::HandleMessage};
+use pezsp_core::{bounded_vec::BoundedSlice, ConstU32};
+use pezsp_runtime::traits::IdentityLookup;
+
+type Block = pezframe_system::mocking::MockBlock<Runtime>;
+
+// Configure a mock runtime to test the pallet.
+pezframe_support::construct_runtime!(
+	pub enum Runtime
+	{
+		System: pezframe_system,
+		DmpQueue: crate,
+	}
+);
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl pezframe_system::Config for Runtime {
+	type Block = Block;
+	type RuntimeOrigin = RuntimeOrigin;
+	type RuntimeCall = RuntimeCall;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type RuntimeEvent = RuntimeEvent;
+	type PalletInfo = PalletInfo;
+}
+
+impl crate::Config for Runtime {
+	type RuntimeEvent = RuntimeEvent;
+	type DmpSink = RecordingDmpSink;
+	type WeightInfo = ();
+}
+
+parameter_types! {
+	/// All messages that came into the `DmpSink`.
+	pub static RecordedMessages: Vec<Vec<u8>> = vec![];
+}
+
+/// Can be used as [`Config::DmpSink`] to record all messages that came in.
+pub struct RecordingDmpSink;
+impl HandleMessage for RecordingDmpSink {
+	type MaxMessageLen = ConstU32<16>;
+
+	fn handle_message(msg: BoundedSlice<u8, Self::MaxMessageLen>) {
+		RecordedMessages::mutate(|n| n.push(msg.to_vec()));
+	}
+
+	fn handle_messages<'a>(_: impl Iterator<Item = BoundedSlice<'a, u8, Self::MaxMessageLen>>) {
+		unimplemented!()
+	}
+
+	fn sweep_queue() {
+		unimplemented!()
+	}
+}
+
+pub fn new_test_ext() -> pezsp_io::TestExternalities {
+	pezsp_io::TestExternalities::new(Default::default())
+}

pezcumulus/pezpallets/dmp-queue/src/tests.rs

@@ -0,0 +1,228 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Test the lazy migration.
+
+#![cfg(test)]
+
+use super::{migration::*, mock::*};
+use crate::*;
+
+use pezframe_support::{pezpallet_prelude::*, traits::OnIdle, StorageNoopGuard};
+
+#[test]
+fn migration_works() {
+	let mut ext = new_test_ext();
+	ext.execute_with(|| {
+		pezsp_tracing::try_init_simple();
+		// Insert some storage:
+		PageIndex::<Runtime>::set(PageIndexData {
+			begin_used: 10,
+			end_used: 20,
+			overweight_count: 5,
+		});
+		for p in 10..20 {
+			let msgs = (0..16).map(|i| (p, vec![i as u8; 1])).collect::<Vec<_>>();
+			Pages::<Runtime>::insert(p, msgs);
+		}
+		for i in 0..5 {
+			Overweight::<Runtime>::insert(i, (0, vec![i as u8; 1]));
+		}
+		testing_only::Configuration::<Runtime>::put(123);
+	});
+	// We need to commit, otherwise the keys are removed from the overlay; not the backend.
+	ext.commit_all().unwrap();
+	ext.execute_with(|| {
+		// Run one step of the migration:
+		pre_upgrade_checks::<Runtime>();
+		run_to_block(1);
+		// First we expect a StartedExport event:
+		assert_only_event(Event::StartedExport);
+
+		// Then we expect 10 Exported events:
+		for page in 0..10 {
+			run_to_block(2 + page);
+			assert_only_event(Event::Exported { page: page as u32 + 10 });
+			assert!(!Pages::<Runtime>::contains_key(page as u32), "Page is gone");
+			assert_eq!(
+				MigrationStatus::<Runtime>::get(),
+				MigrationState::StartedExport { next_begin_used: page as u32 + 11 }
+			);
+		}
+
+		// Then we expect a CompletedExport event:
+		run_to_block(12);
+		assert_only_event(Event::CompletedExport);
+		assert_eq!(MigrationStatus::<Runtime>::get(), MigrationState::CompletedExport);
+
+		// Then we expect a StartedOverweightExport event:
+		run_to_block(13);
+		assert_only_event(Event::StartedOverweightExport);
+		assert_eq!(
+			MigrationStatus::<Runtime>::get(),
+			MigrationState::StartedOverweightExport { next_overweight_index: 0 }
+		);
+
+		// Then we expect 5 ExportedOverweight events:
+		for index in 0..5 {
+			run_to_block(14 + index);
+			assert_only_event(Event::ExportedOverweight { index });
+			assert!(!Overweight::<Runtime>::contains_key(index), "Overweight msg is gone");
+			assert_eq!(
+				MigrationStatus::<Runtime>::get(),
+				MigrationState::StartedOverweightExport { next_overweight_index: index + 1 }
+			);
+		}
+
+		// Then we expect a CompletedOverweightExport event:
+		run_to_block(19);
+		assert_only_event(Event::CompletedOverweightExport);
+		assert_eq!(MigrationStatus::<Runtime>::get(), MigrationState::CompletedOverweightExport);
+
+		// Then we expect a StartedCleanup event:
+		run_to_block(20);
+		assert_only_event(Event::StartedCleanup);
+		assert_eq!(
+			MigrationStatus::<Runtime>::get(),
+			MigrationState::StartedCleanup { cursor: None }
+		);
+	});
+	ext.commit_all().unwrap();
+	// Then it cleans up the remaining storage items:
+	ext.execute_with(|| {
+		run_to_block(21);
+		assert_only_event(Event::CleanedSome { keys_removed: 2 });
+	});
+	ext.commit_all().unwrap();
+	ext.execute_with(|| {
+		run_to_block(22);
+		assert_only_event(Event::CleanedSome { keys_removed: 2 });
+	});
+	ext.commit_all().unwrap();
+	ext.execute_with(|| {
+		run_to_block(24);
+		assert_eq!(
+			System::events().into_iter().map(|e| e.event).collect::<Vec<_>>(),
+			vec![
+				Event::CleanedSome { keys_removed: 2 }.into(),
+				Event::Completed { error: false }.into()
+			]
+		);
+		System::reset_events();
+		assert_eq!(MigrationStatus::<Runtime>::get(), MigrationState::Completed);
+
+		post_upgrade_checks::<Runtime>();
+		assert_eq!(RecordedMessages::take().len(), 10 * 16 + 5);
+
+		// Test the storage removal:
+		assert!(!PageIndex::<Runtime>::exists());
+		assert!(!testing_only::Configuration::<Runtime>::exists());
+		assert_eq!(Pages::<Runtime>::iter_keys().count(), 0);
+		assert_eq!(Overweight::<Runtime>::iter_keys().count(), 0);
+
+		// The `MigrationStatus` never disappears and there are no more storage changes:
+		{
+			let _g = StorageNoopGuard::default();
+
+			run_to_block(100);
+			assert_eq!(MigrationStatus::<Runtime>::get(), MigrationState::Completed);
+			assert!(System::events().is_empty());
+			// ... besides the block number
+			System::set_block_number(24);
+		}
+	});
+}
+
+/// Too long messages are dropped by the migration.
+#[test]
+fn migration_too_long_ignored() {
+	new_test_ext().execute_with(|| {
+		// Setup the storage:
+		PageIndex::<Runtime>::set(PageIndexData {
+			begin_used: 10,
+			end_used: 11,
+			overweight_count: 2,
+		});
+
+		let short = vec![1; 16];
+		let long = vec![0; 17];
+		Pages::<Runtime>::insert(10, vec![(10, short.clone()), (10, long.clone())]);
+		// Insert one good and one bad overweight msg:
+		Overweight::<Runtime>::insert(0, (0, short.clone()));
+		Overweight::<Runtime>::insert(1, (0, long.clone()));
+
+		// Run the migration:
+		pre_upgrade_checks::<Runtime>();
+		run_to_block(100);
+		post_upgrade_checks::<Runtime>();
+
+		assert_eq!(RecordedMessages::take(), vec![short.clone(), short]);
+
+		// Test the storage removal:
+		assert!(!PageIndex::<Runtime>::exists());
+		assert_eq!(Pages::<Runtime>::iter_keys().count(), 0);
+		assert_eq!(Overweight::<Runtime>::iter_keys().count(), 0);
+	});
+}
+
+fn run_to_block(n: u64) {
+	System::run_to_block_with::<AllPalletsWithSystem>(
+		n,
+		pezframe_system::RunToBlockHooks::default().after_initialize(|bn| {
+			AllPalletsWithSystem::on_idle(bn, Weight::MAX);
+		}),
+	);
+}
+
+fn assert_only_event(e: Event<Runtime>) {
+	assert_eq!(System::events().pop().expect("Event expected").event, e.clone().into());
+	assert_eq!(System::events().len(), 1, "Got events: {:?} but wanted {:?}", System::events(), e);
+	System::reset_events();
+}
+
+/// TESTING ONLY
+fn pre_upgrade_checks<T: crate::Config>() {
+	let index = PageIndex::<T>::get();
+
+	// Check that all pages are present.
+	assert!(index.begin_used <= index.end_used, "Invalid page index");
+	for p in index.begin_used..index.end_used {
+		assert!(Pages::<T>::contains_key(p), "Missing page");
+		assert!(Pages::<T>::get(p).len() > 0, "Empty page");
+	}
+
+	// Check that all overweight messages are present.
+	for i in 0..index.overweight_count {
+		assert!(Overweight::<T>::contains_key(i), "Missing overweight message");
+	}
+}
+
+/// TESTING ONLY
+fn post_upgrade_checks<T: crate::Config>() {
+	let index = PageIndex::<T>::get();
+
+	// Check that all pages are removed.
+	for p in index.begin_used..index.end_used {
+		assert!(!Pages::<T>::contains_key(p), "Page should be gone");
+	}
+	assert!(Pages::<T>::iter_keys().next().is_none(), "Un-indexed pages");
+
+	// Check that all overweight messages are removed.
+	for i in 0..index.overweight_count {
+		assert!(!Overweight::<T>::contains_key(i), "Overweight message should be gone");
+	}
+	assert!(Overweight::<T>::iter_keys().next().is_none(), "Un-indexed overweight messages");
+}

pezcumulus/pezpallets/dmp-queue/src/weights.rs

@@ -0,0 +1,222 @@
+// 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.
+
+//! Autogenerated weights for `pezcumulus_pezpallet_dmp_queue`
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE BIZINIKIWI BENCHMARK CLI VERSION 4.0.0-dev
+//! DATE: 2023-09-15, STEPS: `50`, REPEAT: `20`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! HOSTNAME: `Olivers-MacBook-Pro.local`, CPU: `<UNKNOWN>`
+//! WASM-EXECUTION: `Compiled`, CHAIN: `Some("asset-hub-kusama-dev")`, DB CACHE: `1024`
+
+// Executed Command:
+// ./target/release/pezkuwi-teyrchain
+// benchmark
+// pallet
+// --pallet
+// pezcumulus-pezpallet-dmp-queue
+// --chain
+// asset-hub-kusama-dev
+// --output
+// pezcumulus/pezpallets/dmp-queue/src/weights.rs
+// --template
+// bizinikiwi/.maintain/frame-weight-template.hbs
+// --extrinsic
+// 
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+#![allow(missing_docs)]
+
+use pezframe_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}};
+use core::marker::PhantomData;
+
+/// Weight functions needed for `pezcumulus_pezpallet_dmp_queue`.
+pub trait WeightInfo {
+	fn on_idle_good_msg() -> Weight;
+	fn on_idle_large_msg() -> Weight;
+	fn on_idle_overweight_good_msg() -> Weight;
+	fn on_idle_overweight_large_msg() -> Weight;
+}
+
+/// Weights for `pezcumulus_pezpallet_dmp_queue` using the Bizinikiwi node and recommended hardware.
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(65585), added: 68060, mode: `MaxEncodedLen`)
+	fn on_idle_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65696`
+		//  Estimated: `69161`
+		// Minimum execution time: 124_651_000 picoseconds.
+		Weight::from_parts(127_857_000, 0)
+			.saturating_add(Weight::from_parts(0, 69161))
+			.saturating_add(T::DbWeight::get().reads(5))
+			.saturating_add(T::DbWeight::get().writes(5))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	fn on_idle_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65659`
+		//  Estimated: `69124`
+		// Minimum execution time: 65_684_000 picoseconds.
+		Weight::from_parts(68_039_000, 0)
+			.saturating_add(Weight::from_parts(0, 69124))
+			.saturating_add(T::DbWeight::get().reads(3))
+			.saturating_add(T::DbWeight::get().writes(2))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(65585), added: 68060, mode: `MaxEncodedLen`)
+	fn on_idle_overweight_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65726`
+		//  Estimated: `69191`
+		// Minimum execution time: 117_657_000 picoseconds.
+		Weight::from_parts(122_035_000, 0)
+			.saturating_add(Weight::from_parts(0, 69191))
+			.saturating_add(T::DbWeight::get().reads(6))
+			.saturating_add(T::DbWeight::get().writes(6))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	fn on_idle_overweight_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65689`
+		//  Estimated: `69154`
+		// Minimum execution time: 59_799_000 picoseconds.
+		Weight::from_parts(61_354_000, 0)
+			.saturating_add(Weight::from_parts(0, 69154))
+			.saturating_add(T::DbWeight::get().reads(4))
+			.saturating_add(T::DbWeight::get().writes(3))
+	}
+}
+
+// For backwards compatibility and tests.
+impl WeightInfo for () {
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(65585), added: 68060, mode: `MaxEncodedLen`)
+	fn on_idle_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65696`
+		//  Estimated: `69161`
+		// Minimum execution time: 124_651_000 picoseconds.
+		Weight::from_parts(127_857_000, 0)
+			.saturating_add(Weight::from_parts(0, 69161))
+			.saturating_add(RocksDbWeight::get().reads(5))
+			.saturating_add(RocksDbWeight::get().writes(5))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca7d95d3e948effbeccff2de2c182672836` (r:1 w:1)
+	fn on_idle_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65659`
+		//  Estimated: `69124`
+		// Minimum execution time: 65_684_000 picoseconds.
+		Weight::from_parts(68_039_000, 0)
+			.saturating_add(Weight::from_parts(0, 69124))
+			.saturating_add(RocksDbWeight::get().reads(3))
+			.saturating_add(RocksDbWeight::get().writes(2))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(65585), added: 68060, mode: `MaxEncodedLen`)
+	fn on_idle_overweight_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65726`
+		//  Estimated: `69191`
+		// Minimum execution time: 117_657_000 picoseconds.
+		Weight::from_parts(122_035_000, 0)
+			.saturating_add(Weight::from_parts(0, 69191))
+			.saturating_add(RocksDbWeight::get().reads(6))
+			.saturating_add(RocksDbWeight::get().writes(6))
+	}
+	/// Storage: `DmpQueue::MigrationStatus` (r:1 w:1)
+	/// Proof: `DmpQueue::MigrationStatus` (`max_values`: Some(1), `max_size`: Some(1028), added: 1523, mode: `MaxEncodedLen`)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca754904d6d8c6fe06c4e5965f9b8397421` (r:1 w:0)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca70f923ef3252d0166429d36d20ed665a8` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0xcd5c1f6df63bc97f4a8ce37f14a50ca772275f64c354954352b71eea39cfaca2` (r:1 w:1)
+	fn on_idle_overweight_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65689`
+		//  Estimated: `69154`
+		// Minimum execution time: 59_799_000 picoseconds.
+		Weight::from_parts(61_354_000, 0)
+			.saturating_add(Weight::from_parts(0, 69154))
+			.saturating_add(RocksDbWeight::get().reads(4))
+			.saturating_add(RocksDbWeight::get().writes(3))
+	}
+}

pezcumulus/pezpallets/session-benchmarking/Cargo.toml

@@ -0,0 +1,43 @@
+[package]
+name = "pezcumulus-pezpallet-session-benchmarking"
+version = "9.0.0"
+authors.workspace = true
+edition.workspace = true
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+description = "FRAME sessions pallet benchmarking"
+readme = "README.md"
+documentation = "https://docs.rs/pezcumulus-pezpallet-session-benchmarking"
+
+[lints]
+workspace = true
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[dependencies]
+codec = { workspace = true }
+pezframe-benchmarking = { optional = true, workspace = true }
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-session = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+[features]
+default = ["std"]
+runtime-benchmarks = [
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-session/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+]
+std = [
+  "codec/std",
+  "pezframe-benchmarking/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "pezpallet-session/std",
+  "pezsp-runtime/std",
+]

pezcumulus/pezpallets/session-benchmarking/README.md

@@ -0,0 +1,3 @@
+Benchmarks for the Session Pallet.
+
+License: Apache-2.0

pezcumulus/pezpallets/session-benchmarking/src/inner.rs

@@ -0,0 +1,64 @@
+// 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 pezpallet-session.
+#![cfg(feature = "runtime-benchmarks")]
+
+use alloc::{vec, vec::Vec};
+
+use codec::Decode;
+use pezframe_benchmarking::v2::*;
+use pezframe_system::RawOrigin;
+use pezpallet_session::*;
+pub struct Pallet<T: Config>(pezpallet_session::Pallet<T>);
+pub trait Config: pezpallet_session::Config {}
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	#[benchmark]
+	fn set_keys() -> Result<(), BenchmarkError> {
+		let caller: T::AccountId = whitelisted_caller();
+		pezframe_system::Pallet::<T>::inc_providers(&caller);
+		let keys = T::Keys::decode(&mut pezsp_runtime::traits::TrailingZeroInput::zeroes()).unwrap();
+		let proof: Vec<u8> = vec![0, 1, 2, 3];
+		<pezpallet_session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller), keys, proof);
+
+		Ok(())
+	}
+
+	#[benchmark]
+	fn purge_keys() -> Result<(), BenchmarkError> {
+		let caller: T::AccountId = whitelisted_caller();
+		pezframe_system::Pallet::<T>::inc_providers(&caller);
+		let keys = T::Keys::decode(&mut pezsp_runtime::traits::TrailingZeroInput::zeroes()).unwrap();
+		let proof: Vec<u8> = vec![0, 1, 2, 3];
+		<pezpallet_session::Pallet<T>>::ensure_can_pay_key_deposit(&caller).unwrap();
+		let _t = pezpallet_session::Pallet::<T>::set_keys(
+			RawOrigin::Signed(caller.clone()).into(),
+			keys,
+			proof,
+		);
+
+		#[extrinsic_call]
+		_(RawOrigin::Signed(caller));
+
+		Ok(())
+	}
+}

pezcumulus/pezpallets/session-benchmarking/src/lib.rs

@@ -0,0 +1,29 @@
+// This file is part of Pezcumulus.
+
+// 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.
+
+//! Benchmarks for the Session Pallet.
+// This is separated into its own crate due to cyclic dependency issues.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+#[cfg(feature = "runtime-benchmarks")]
+pub mod inner;
+
+#[cfg(feature = "runtime-benchmarks")]
+pub use inner::*;

pezcumulus/pezpallets/solo-to-para/Cargo.toml

@@ -0,0 +1,57 @@
+[package]
+name = "pezcumulus-pezpallet-solo-to-para"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+description = "Adds functionality to migrate from a Solo to a Teyrchain"
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-solo-to-para"
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+# Bizinikiwi
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-sudo = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+# Pezkuwi
+pezkuwi-primitives = { workspace = true }
+
+# Pezcumulus
+pezcumulus-pezpallet-teyrchain-system = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezcumulus-pezpallet-teyrchain-system/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "pezpallet-sudo/std",
+  "pezkuwi-primitives/std",
+  "scale-info/std",
+  "pezsp-runtime/std",
+]
+try-runtime = [
+  "pezcumulus-pezpallet-teyrchain-system/try-runtime",
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-sudo/try-runtime",
+  "pezsp-runtime/try-runtime",
+]
+runtime-benchmarks = [
+  "pezcumulus-pezpallet-teyrchain-system/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-sudo/runtime-benchmarks",
+  "pezkuwi-primitives/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+]

pezcumulus/pezpallets/solo-to-para/src/lib.rs

@@ -0,0 +1,107 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+use alloc::vec::Vec;
+use pezcumulus_pezpallet_teyrchain_system as teyrchain_system;
+use pezframe_support::pezpallet_prelude::*;
+use pezframe_system::pezpallet_prelude::*;
+pub use pallet::*;
+use pezkuwi_primitives::PersistedValidationData;
+
+#[pezframe_support::pallet]
+pub mod pallet {
+	use super::*;
+
+	#[pallet::config]
+	pub trait Config:
+		pezframe_system::Config + teyrchain_system::Config + pezpallet_sudo::Config
+	{
+		#[allow(deprecated)]
+		type RuntimeEvent: From<Event> + IsType<<Self as pezframe_system::Config>::RuntimeEvent>;
+	}
+
+	#[pallet::pallet]
+	#[pallet::without_storage_info]
+	pub struct Pallet<T>(_);
+
+	/// In case of a scheduled migration, this storage field contains the custom head data to be
+	/// applied.
+	#[pallet::storage]
+	pub(super) type PendingCustomValidationHeadData<T: Config> =
+		StorageValue<_, Vec<u8>, OptionQuery>;
+
+	#[pallet::event]
+	#[pallet::generate_deposit(pub(super) fn deposit_event)]
+	pub enum Event {
+		/// The custom validation head data has been scheduled to apply.
+		CustomValidationHeadDataStored,
+		/// The custom validation head data was applied as of the contained relay chain block
+		/// number.
+		CustomValidationHeadDataApplied,
+	}
+
+	#[pallet::error]
+	pub enum Error<T> {
+		/// CustomHeadData is not stored in storage.
+		NoCustomHeadData,
+	}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {
+		#[pallet::call_index(0)]
+		#[pallet::weight({0})]
+		pub fn schedule_migration(
+			origin: OriginFor<T>,
+			code: Vec<u8>,
+			head_data: Vec<u8>,
+		) -> DispatchResult {
+			ensure_root(origin)?;
+
+			teyrchain_system::Pallet::<T>::schedule_code_upgrade(code)?;
+			Self::store_pending_custom_validation_head_data(head_data);
+			Ok(())
+		}
+	}
+
+	impl<T: Config> Pallet<T> {
+		/// Set a custom head data that should only be applied when upgradeGoAheadSignal from
+		/// the Relay Chain is GoAhead
+		fn store_pending_custom_validation_head_data(head_data: Vec<u8>) {
+			PendingCustomValidationHeadData::<T>::put(head_data);
+			Self::deposit_event(Event::CustomValidationHeadDataStored);
+		}
+
+		/// Set pending custom head data as head data that will be returned by `validate_block`. on
+		/// the relay chain.
+		fn set_pending_custom_validation_head_data() {
+			if let Some(head_data) = <PendingCustomValidationHeadData<T>>::take() {
+				teyrchain_system::Pallet::<T>::set_custom_validation_head_data(head_data);
+				Self::deposit_event(Event::CustomValidationHeadDataApplied);
+			}
+		}
+	}
+
+	impl<T: Config> teyrchain_system::OnSystemEvent for Pallet<T> {
+		fn on_validation_data(_data: &PersistedValidationData) {}
+		fn on_validation_code_applied() {
+			crate::Pallet::<T>::set_pending_custom_validation_head_data();
+		}
+	}
+}

pezcumulus/pezpallets/teyrchain-system/Cargo.toml

@@ -0,0 +1,141 @@
+[package]
+name = "pezcumulus-pezpallet-teyrchain-system"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+description = "Base pallet for pezcumulus-based teyrchains"
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-teyrchain-system"
+
+[lints]
+workspace = true
+
+[dependencies]
+array-bytes = { workspace = true }
+bytes = { workspace = true }
+codec = { features = ["derive"], workspace = true }
+environmental = { workspace = true }
+hashbrown = { workspace = true }
+impl-trait-for-tuples = { workspace = true }
+log = { workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+trie-db = { workspace = true }
+
+# Bizinikiwi
+pezframe-benchmarking = { optional = true, workspace = true }
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-message-queue = { workspace = true }
+pezsp-consensus-babe = { workspace = true }
+pezsp-core = { workspace = true }
+pezsp-externalities = { workspace = true }
+pezsp-inherents = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+pezsp-state-machine = { workspace = true }
+pezsp-std = { workspace = true }
+pezsp-trie = { workspace = true }
+pezsp-version = { workspace = true }
+
+# Pezkuwi
+pezkuwi-runtime-teyrchains = { workspace = true }
+pezkuwi-teyrchain-primitives = { features = ["wasm-api"], workspace = true }
+xcm = { workspace = true }
+xcm-builder = { workspace = true }
+
+# Pezcumulus
+pezcumulus-pezpallet-teyrchain-system-proc-macro = { workspace = true }
+pezcumulus-primitives-core = { workspace = true }
+pezcumulus-primitives-proof-size-hostfunction = { workspace = true }
+pezcumulus-primitives-teyrchain-inherent = { workspace = true }
+
+[dev-dependencies]
+assert_matches = { workspace = true }
+futures = { workspace = true }
+hex-literal = { workspace = true, default-features = true }
+rand = { workspace = true, default-features = true }
+rstest = { workspace = true }
+trie-standardmap = { workspace = true }
+
+
+# Bizinikiwi
+pezsc-consensus = { workspace = true }
+pezsp-api = { workspace = true, default-features = true }
+pezsp-consensus-slots = { workspace = true, default-features = true }
+pezsp-crypto-hashing = { workspace = true, default-features = true }
+pezsp-keyring = { workspace = true, default-features = true }
+pezsp-tracing = { workspace = true, default-features = true }
+pezsp-version = { workspace = true, default-features = true }
+# Pezcumulus
+pezcumulus-test-client = { workspace = true }
+pezcumulus-test-relay-sproof-builder = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+std = [
+  "bytes/std",
+  "codec/std",
+  "pezcumulus-pezpallet-teyrchain-system-proc-macro/std",
+  "pezcumulus-primitives-core/std",
+  "pezcumulus-primitives-proof-size-hostfunction/std",
+  "pezcumulus-primitives-teyrchain-inherent/std",
+  "environmental/std",
+  "pezframe-benchmarking/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "log/std",
+  "pezpallet-message-queue/std",
+  "pezkuwi-runtime-teyrchains/std",
+  "pezkuwi-teyrchain-primitives/std",
+  "scale-info/std",
+  "pezsp-consensus-babe/std",
+  "pezsp-core/std",
+  "pezsp-externalities/std",
+  "pezsp-inherents/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "pezsp-state-machine/std",
+  "pezsp-std/std",
+  "pezsp-tracing/std",
+  "pezsp-trie/std",
+  "trie-db/std",
+  "xcm-builder/std",
+  "xcm/std",
+]
+
+runtime-benchmarks = [
+  "pezcumulus-primitives-core/runtime-benchmarks",
+  "pezcumulus-primitives-proof-size-hostfunction/runtime-benchmarks",
+  "pezcumulus-primitives-teyrchain-inherent/runtime-benchmarks",
+  "pezcumulus-test-client/runtime-benchmarks",
+  "pezcumulus-test-relay-sproof-builder/runtime-benchmarks",
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-message-queue/runtime-benchmarks",
+  "pezkuwi-runtime-teyrchains/runtime-benchmarks",
+  "pezkuwi-teyrchain-primitives/runtime-benchmarks",
+  "pezsc-consensus/runtime-benchmarks",
+  "pezsp-api/runtime-benchmarks",
+  "pezsp-consensus-babe/runtime-benchmarks",
+  "pezsp-consensus-slots/runtime-benchmarks",
+  "pezsp-inherents/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-keyring/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "pezsp-state-machine/runtime-benchmarks",
+  "pezsp-trie/runtime-benchmarks",
+  "pezsp-version/runtime-benchmarks",
+  "xcm-builder/runtime-benchmarks",
+  "xcm/runtime-benchmarks",
+]
+
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-message-queue/try-runtime",
+  "pezkuwi-runtime-teyrchains/try-runtime",
+  "pezsp-runtime/try-runtime",
+]

pezcumulus/pezpallets/teyrchain-system/proc-macro/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "pezcumulus-pezpallet-teyrchain-system-proc-macro"
+version = "0.6.0"
+authors.workspace = true
+edition.workspace = true
+description = "Proc macros provided by the teyrchain-system pallet"
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-teyrchain-system-proc-macro"
+
+[lints]
+workspace = true
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro-crate = { workspace = true }
+proc-macro2 = { workspace = true }
+quote = { workspace = true }
+syn = { workspace = true }
+
+[features]
+default = ["std"]
+std = []

pezcumulus/pezpallets/teyrchain-system/proc-macro/src/lib.rs

@@ -0,0 +1,138 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 proc_macro2::Span;
+use proc_macro_crate::{crate_name, FoundCrate};
+use syn::{
+	parse::{Parse, ParseStream},
+	spanned::Spanned,
+	token, Error, Ident, Path,
+};
+
+mod keywords {
+	syn::custom_keyword!(Runtime);
+	syn::custom_keyword!(BlockExecutor);
+	syn::custom_keyword!(CheckInherents);
+}
+
+struct Input {
+	runtime: Path,
+	block_executor: Path,
+}
+
+impl Parse for Input {
+	fn parse(input: ParseStream) -> Result<Self, Error> {
+		let mut runtime = None;
+		let mut block_executor = None;
+
+		fn parse_inner<KW: Parse + Spanned>(
+			input: ParseStream,
+			result: &mut Option<Path>,
+		) -> Result<(), Error> {
+			let kw = input.parse::<KW>()?;
+
+			if result.is_none() {
+				input.parse::<token::Eq>()?;
+				*result = Some(input.parse::<Path>()?);
+				if input.peek(token::Comma) {
+					input.parse::<token::Comma>()?;
+				}
+
+				Ok(())
+			} else {
+				Err(Error::new(kw.span(), "Is only allowed to be passed once"))
+			}
+		}
+
+		while !input.is_empty() || runtime.is_none() || block_executor.is_none() {
+			let lookahead = input.lookahead1();
+
+			if lookahead.peek(keywords::Runtime) {
+				parse_inner::<keywords::Runtime>(input, &mut runtime)?;
+			} else if lookahead.peek(keywords::BlockExecutor) {
+				parse_inner::<keywords::BlockExecutor>(input, &mut block_executor)?;
+			} else if lookahead.peek(keywords::CheckInherents) {
+				return Err(Error::new(input.span(), "`CheckInherents` is not supported anymore!"));
+			} else {
+				return Err(lookahead.error());
+			}
+		}
+
+		Ok(Self {
+			runtime: runtime.expect("Everything is parsed before; qed"),
+			block_executor: block_executor.expect("Everything is parsed before; qed"),
+		})
+	}
+}
+
+fn crate_() -> Result<Ident, Error> {
+	match crate_name("pezcumulus-pezpallet-teyrchain-system") {
+		Ok(FoundCrate::Itself) =>
+			Ok(syn::Ident::new("pezcumulus_pezpallet_teyrchain_system", Span::call_site())),
+		Ok(FoundCrate::Name(name)) => Ok(Ident::new(&name, Span::call_site())),
+		Err(e) => Err(Error::new(Span::call_site(), e)),
+	}
+}
+
+#[proc_macro]
+pub fn register_validate_block(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+	let Input { runtime, block_executor } = match syn::parse(input) {
+		Ok(t) => t,
+		Err(e) => return e.into_compile_error().into(),
+	};
+
+	let crate_ = match crate_() {
+		Ok(c) => c,
+		Err(e) => return e.into_compile_error().into(),
+	};
+
+	if cfg!(not(feature = "std")) {
+		quote::quote! {
+			#[doc(hidden)]
+			mod teyrchain_validate_block {
+				use super::*;
+
+				#[no_mangle]
+				unsafe fn validate_block(arguments: *mut u8, arguments_len: usize) -> u64 {
+					// We convert the `arguments` into a boxed slice and then into `Bytes`.
+					let args = #crate_::validate_block::Box::from_raw(
+						#crate_::validate_block::slice::from_raw_parts_mut(
+							arguments,
+							arguments_len,
+						)
+					);
+					let args = #crate_::validate_block::bytes::Bytes::from(args);
+
+					// Then we decode from these bytes the `MemoryOptimizedValidationParams`.
+					let params = #crate_::validate_block::decode_from_bytes::<
+						#crate_::validate_block::MemoryOptimizedValidationParams
+					>(args).expect("Invalid arguments to `validate_block`.");
+
+					let res = #crate_::validate_block::implementation::validate_block::<
+						<#runtime as #crate_::validate_block::GetRuntimeBlockType>::RuntimeBlock,
+						#block_executor,
+						#runtime,
+					>(params);
+
+					#crate_::validate_block::pezkuwi_teyrchain_primitives::write_result(&res)
+				}
+			}
+		}
+	} else {
+		quote::quote!()
+	}
+	.into()
+}

pezcumulus/pezpallets/teyrchain-system/src/benchmarking.rs

@@ -0,0 +1,72 @@
+// This file is part of Pezcumulus.
+
+// 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 for the teyrchain-system pallet.
+
+#![cfg(feature = "runtime-benchmarks")]
+
+use super::*;
+use crate::teyrchain_inherent::InboundDownwardMessages;
+use pezcumulus_primitives_core::{relay_chain::Hash as RelayHash, InboundDownwardMessage};
+use pezframe_benchmarking::v2::*;
+use pezsp_runtime::traits::BlakeTwo256;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	/// Enqueue `n` messages via `enqueue_inbound_downward_messages`.
+	///
+	/// The limit is set to `1000` for benchmarking purposes as the actual limit is only known at
+	/// runtime. However, the limit (and default) for Dotsama are magnitudes smaller.
+	#[benchmark]
+	fn enqueue_inbound_downward_messages(n: Linear<0, 1000>) {
+		let msg = InboundDownwardMessage {
+			sent_at: n, // The block number does not matter.
+			msg: vec![0u8; MaxDmpMessageLenOf::<T>::get() as usize],
+		};
+		let msgs = vec![msg; n as usize];
+		let head = mqp_head(&msgs);
+
+		#[block]
+		{
+			Pallet::<T>::enqueue_inbound_downward_messages(
+				head,
+				InboundDownwardMessages::new(msgs).into_abridged(&mut usize::MAX.clone()),
+			);
+		}
+
+		assert_eq!(ProcessedDownwardMessages::<T>::get(), n);
+		assert_eq!(LastDmqMqcHead::<T>::get().head(), head);
+	}
+
+	/// Re-implements an easy version of the `MessageQueueChain` for testing purposes.
+	fn mqp_head(msgs: &Vec<InboundDownwardMessage>) -> RelayHash {
+		let mut head = Default::default();
+		for msg in msgs.iter() {
+			let msg_hash = BlakeTwo256::hash_of(&msg.msg);
+			head = BlakeTwo256::hash_of(&(head, msg.sent_at, msg_hash));
+		}
+		head
+	}
+
+	impl_benchmark_test_suite! {
+		Pallet,
+		crate::mock::new_test_ext(),
+		crate::mock::Test
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/consensus_hook.rs

@@ -0,0 +1,109 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! The definition of a [`ConsensusHook`] trait for consensus logic to manage the backlog
+//! of teyrchain blocks ready to submit to the relay chain, as well as some basic implementations.
+
+use super::relay_state_snapshot::RelayChainStateProof;
+use core::num::NonZeroU32;
+use pezframe_support::weights::Weight;
+
+/// The possible capacity of the unincluded segment.
+#[derive(Clone, Debug, PartialEq)]
+pub struct UnincludedSegmentCapacity(UnincludedSegmentCapacityInner);
+
+impl UnincludedSegmentCapacity {
+	pub(crate) fn get(&self) -> u32 {
+		match self.0 {
+			UnincludedSegmentCapacityInner::ExpectParentIncluded => 1,
+			UnincludedSegmentCapacityInner::Value(v) => v.get(),
+		}
+	}
+
+	pub(crate) fn is_expecting_included_parent(&self) -> bool {
+		match self.0 {
+			UnincludedSegmentCapacityInner::ExpectParentIncluded => true,
+			UnincludedSegmentCapacityInner::Value(_) => false,
+		}
+	}
+}
+
+#[derive(Clone, Debug, PartialEq)]
+pub(crate) enum UnincludedSegmentCapacityInner {
+	ExpectParentIncluded,
+	Value(NonZeroU32),
+}
+
+impl From<NonZeroU32> for UnincludedSegmentCapacity {
+	fn from(value: NonZeroU32) -> Self {
+		UnincludedSegmentCapacity(UnincludedSegmentCapacityInner::Value(value))
+	}
+}
+
+/// The consensus hook for dealing with the unincluded segment.
+///
+/// Higher-level and user-configurable consensus logic is more informed about the
+/// desired unincluded segment length, as well as any rules for adapting it dynamically
+/// according to the relay-chain state.
+pub trait ConsensusHook {
+	/// This hook is called partway through the `set_validation_data` inherent in teyrchain-system.
+	///
+	/// The hook is allowed to panic if customized consensus rules aren't met and is required
+	/// to return a maximum capacity for the unincluded segment with weight consumed.
+	fn on_state_proof(state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity);
+}
+
+/// A special consensus hook for handling the migration to asynchronous backing gracefully,
+/// even if collators haven't been updated to provide the last included parent in the state
+/// proof yet.
+///
+/// This behaves as though the parent is included, even if the relay chain state proof doesn't
+/// contain the included para head. If the para head is present in the state proof, this does ensure
+/// the parent is included.
+pub struct ExpectParentIncluded;
+
+impl ConsensusHook for ExpectParentIncluded {
+	fn on_state_proof(_state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		(
+			Weight::zero(),
+			UnincludedSegmentCapacity(UnincludedSegmentCapacityInner::ExpectParentIncluded),
+		)
+	}
+}
+
+/// A consensus hook for a fixed unincluded segment length. This hook does nothing but
+/// set the capacity of the unincluded segment to the constant N.
+///
+/// Since it is illegal to provide an unincluded segment length of 0, this sets a minimum of
+/// 1.
+pub struct FixedCapacityUnincludedSegment<const N: u32>;
+
+impl<const N: u32> ConsensusHook for FixedCapacityUnincludedSegment<N> {
+	fn on_state_proof(_state_proof: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		(
+			Weight::zero(),
+			NonZeroU32::new(core::cmp::max(N, 1))
+				.expect("1 is the minimum value and non-zero; qed")
+				.into(),
+		)
+	}
+}
+
+/// A fixed-capacity unincluded segment hook, which requires that the parent block is
+/// included prior to the current block being authored.
+///
+/// This is a simple type alias around a fixed-capacity unincluded segment with a size of 1.
+pub type RequireParentIncluded = FixedCapacityUnincludedSegment<1>;

pezcumulus/pezpallets/teyrchain-system/src/descendant_validation.rs

@@ -0,0 +1,617 @@
+// 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 crate::{
+	descendant_validation::RelayParentVerificationError::InvalidNumberOfDescendants,
+	RelayChainStateProof,
+};
+use alloc::vec::Vec;
+use pezsp_consensus_babe::{
+	digests::{CompatibleDigestItem, NextEpochDescriptor},
+	AuthorityIndex,
+};
+use pezsp_runtime::{traits::Header, RuntimeAppPublic};
+
+/// Verifies that the provided relay parent descendants form a valid chain
+/// and are signed by relay chain authorities. If relay chain descendants shall be checked,
+/// a set of authorities for the epoch of the relay parent must be provided in the
+/// relay chain state proof. If any of the descendants indicate the beginning of a new epoch,
+/// the authority set for the next relay chain epoch must be included in the state proof too.
+///
+/// # Parameters
+///
+/// - `relay_state_proof`: The proof of the relay chain state, which contains details about the
+///   authority sets and other chain data.
+/// - `relay_parent_descendants`: A vector of relay chain headers representing the descendants of
+///   the relay parent that need to be validated. The first item in this vector must be the relay
+///   parent itself.
+/// - `relay_parent_state_root`: The state root hash of the relay parent. This
+///   will be matched with the first relay parent header from the descendants.
+///   **Note:** This parameter can be removed once the hash of the relay parent is available
+///   to the runtime. https://github.com/pezkuwichain/pezkuwi-sdk/issues/99
+/// - `expected_rp_descendants_num`: The expected number of headers in the
+///   `relay_parent_descendants`. A mismatch will cause the function to return an error.
+///
+/// # Errors
+///
+/// This function will error under the following scenarios:
+///
+/// - The number of headers in `relay_parent_descendants` does not match
+///   `expected_rp_descendants_num`.
+/// - No authorities are provided in the state proof.
+/// - The state root of the provided relay parent does not match the expected value.
+/// - A relay header does not contain a BABE pre-digest.
+/// - A header with an invalid seal signature is found, or the authorities required to verify the
+///   signature are missing (current or next epoch).
+pub(crate) fn verify_relay_parent_descendants<H: Header>(
+	relay_state_proof: &RelayChainStateProof,
+	relay_parent_descendants: Vec<H>,
+	relay_parent_state_root: H::Hash,
+	expected_rp_descendants_num: u32,
+) -> Result<(), RelayParentVerificationError<H>> {
+	if relay_parent_descendants.len() != (expected_rp_descendants_num + 1) as usize {
+		return Err(InvalidNumberOfDescendants {
+			expected: expected_rp_descendants_num + 1,
+			received: relay_parent_descendants.len(),
+		});
+	}
+
+	let Ok(mut current_authorities) = relay_state_proof.read_authorities() else {
+		return Err(RelayParentVerificationError::MissingAuthorities);
+	};
+	let mut maybe_next_authorities = relay_state_proof.read_next_authorities().ok().flatten();
+
+	let mut next_expected_parent_hash = None;
+
+	// Verify that the state root of the first block is the same as the one
+	// from the relay parent. In the PVF, we don't have the relay parent header hash
+	// available, so we need to use the storage root here to establish a chain.
+	if let Some(relay_parent) = relay_parent_descendants.get(0) {
+		if *relay_parent.state_root() != relay_parent_state_root {
+			return Err(RelayParentVerificationError::InvalidStateRoot {
+				expected: relay_parent_state_root,
+				found: *relay_parent.state_root(),
+			});
+		}
+	};
+
+	for (index, mut current_header) in relay_parent_descendants.into_iter().enumerate() {
+		// Hash calculated while seal is intact
+		let sealed_header_hash = current_header.hash();
+		let relay_number = *current_header.number();
+
+		// Verify that the blocks actually form a chain
+		if let Some(expected_hash) = next_expected_parent_hash {
+			if *current_header.parent_hash() != expected_hash {
+				return Err(RelayParentVerificationError::InvalidChainSequence {
+					expected: expected_hash,
+					found: *current_header.parent_hash(),
+					number: relay_number,
+				});
+			}
+		}
+		next_expected_parent_hash = Some(sealed_header_hash);
+
+		log::debug!(target: crate::LOG_TARGET, "Validating header #{relay_number:?} ({sealed_header_hash:?})");
+		let (authority_index, next_epoch_descriptor) =
+			find_authority_idx_epoch_digest(&current_header).ok_or_else(|| {
+				RelayParentVerificationError::MissingPredigest { hash: sealed_header_hash }
+			})?;
+
+		// Once we have seen a next epoch descriptor, we must always use the authorities of the
+		// next epoch. If the relay parent contains epoch descriptor, we shall not rotate
+		// authorities. As in that case the authorities in the state proof reflect the
+		// new authorities already.
+		if let Some(descriptor) = next_epoch_descriptor {
+			// If the relay parent itself contains the epoch change, we must _not_ use the next
+			// authorities, as they have already been rotated in storage.
+			if index != 0 {
+				let Some(next_authorities) = maybe_next_authorities else {
+					return Err(RelayParentVerificationError::MissingNextEpochAuthorities {
+						number: relay_number,
+						hash: sealed_header_hash,
+					});
+				};
+				log::debug!(
+					target: crate::LOG_TARGET,
+					"Header {sealed_header_hash:?} contains epoch change! \
+					Using next authority set to verify signatures going forward."
+				);
+				// Rotate authorities, all headers following are to
+				// be verified against the new authorities. The authorities for the next epoch
+				// have been signed by a current authority, we can use it for further epochs.
+				current_authorities = next_authorities;
+				maybe_next_authorities = Some(descriptor.authorities);
+			}
+		}
+
+		let Some(authority_id) = current_authorities.get(authority_index as usize) else {
+			return Err(RelayParentVerificationError::MissingAuthorityId);
+		};
+
+		let Some(seal) = current_header.digest_mut().pop() else {
+			return Err(RelayParentVerificationError::MissingSeal { hash: sealed_header_hash });
+		};
+		let Some(signature) = seal.as_babe_seal() else {
+			return Err(RelayParentVerificationError::InvalidSeal { hash: sealed_header_hash });
+		};
+
+		if !authority_id.0.verify(&current_header.hash(), &signature) {
+			return Err(RelayParentVerificationError::InvalidSignature {
+				number: relay_number,
+				hash: sealed_header_hash,
+			});
+		}
+		log::debug!(target: crate::LOG_TARGET, "Validated header #{relay_number:?} ({sealed_header_hash:?})");
+	}
+
+	Ok(())
+}
+
+/// Extract babe digest items from the header.
+/// - [AuthorityIndex]: We extract the authority index from the babe predigest. We use it to verify
+///   that a correct authority from the authority set signed the header.
+/// - [NextEpochDescriptor]: We extract it because we need to know which block starts a new epoch on
+///   the relay chain. Epoch change indicates a switch in authority set, so we need to verify
+///   following signatures against the new authorities.
+pub fn find_authority_idx_epoch_digest<H: Header>(
+	header: &H,
+) -> Option<(AuthorityIndex, Option<NextEpochDescriptor>)> {
+	let mut babe_pre_digest = None;
+	let mut next_epoch_digest = None;
+	for log in header.digest().logs() {
+		if let Some(digest) = log.as_babe_pre_digest() {
+			babe_pre_digest = Some(digest);
+		}
+
+		if let Some(digest) = log.as_next_epoch_descriptor() {
+			next_epoch_digest = Some(digest);
+		}
+	}
+
+	babe_pre_digest.map(|pd| (pd.authority_index(), next_epoch_digest))
+}
+
+/// Errors that can occur during descendant validation
+#[derive(Debug, PartialEq)]
+pub(crate) enum RelayParentVerificationError<H: Header> {
+	/// The number of descendants provided doesn't match the expected count
+	InvalidNumberOfDescendants { expected: u32, received: usize },
+	/// No authorities were provided in the state proof
+	MissingAuthorities,
+	/// The state root of the relay parent doesn't match
+	InvalidStateRoot { expected: H::Hash, found: H::Hash },
+	/// The chain sequence is invalid (parent hash doesn't match expected hash)
+	InvalidChainSequence { expected: H::Hash, found: H::Hash, number: H::Number },
+	/// The header is missing the required pre-digest
+	MissingPredigest { hash: H::Hash },
+	/// The header is missing the required seal
+	MissingSeal { hash: H::Hash },
+	/// The header contains an invalid seal
+	InvalidSeal { hash: H::Hash },
+	/// Next epoch authorities are missing when they are required
+	MissingNextEpochAuthorities { number: H::Number, hash: H::Hash },
+	/// Unable to find the authority ID at the expected index
+	MissingAuthorityId,
+	/// The signature verification failed
+	InvalidSignature { number: H::Number, hash: H::Hash },
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use codec::{Decode, Encode};
+	use pezcumulus_primitives_core::relay_chain;
+	use pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+	use rstest::rstest;
+	use pezsp_consensus_babe::{
+		digests::{CompatibleDigestItem, NextEpochDescriptor, PreDigest, PrimaryPreDigest},
+		AuthorityId, AuthorityPair, BabeAuthorityWeight, ConsensusLog, BABE_ENGINE_ID,
+	};
+	use pezsp_core::{
+		sr25519::vrf::{VrfPreOutput, VrfProof, VrfSignature},
+		Pair, H256,
+	};
+	use pezsp_keyring::Sr25519Keyring;
+	use pezsp_runtime::{testing::Header as TestHeader, DigestItem};
+	const PARA_ID: u32 = 2000;
+
+	/// Verify a header chain with different lengths and different number of authorities included in
+	/// the storage proof.
+	#[rstest]
+	fn test_verify_relay_parent_descendants_happy_case(
+		#[values(1, 2, 3, 4, 100)] num_headers: u64,
+		#[values(1, 3, 100, 1000)] num_authorities: u64,
+	) {
+		let (relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(num_headers, num_authorities, None);
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		assert!(verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		)
+		.is_ok());
+	}
+
+	#[rstest]
+	fn test_verify_relay_parent_broken_state_root() {
+		let (relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(10, 10, None);
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		// Set a erroneous state root
+		let relay_parent_state_root = H256::repeat_byte(0x9);
+
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+
+		assert_eq!(
+			result,
+			Err(RelayParentVerificationError::<TestHeader>::InvalidStateRoot {
+				expected: H256::repeat_byte(0x9),
+				found: H256::repeat_byte(0x0),
+			})
+		);
+	}
+
+	#[rstest]
+	#[case::too_few_1(1)]
+	#[case::too_few_2(8)]
+	// 9 would be just right, but we want to panic
+	#[case::too_many_1(10)]
+	#[case::too_many_2(100)]
+	fn test_incorrect_number_of_headers(#[case] expected_number_of_descendants: u32) {
+		let (relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(10, 10, None);
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+
+		assert_eq!(
+			result,
+			Err(InvalidNumberOfDescendants::<TestHeader> {
+				expected: expected_number_of_descendants + 1,
+				received: 10,
+			})
+		);
+	}
+
+	#[rstest]
+	fn test_authorities_missing() {
+		let (relay_parent_descendants, _, _) = build_relay_parent_descendants(10, 10, None);
+		// No authorities, this is bad!
+		let relay_state_proof = build_relay_chain_storage_proof(None, None);
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+
+		assert_eq!(result, Err(RelayParentVerificationError::<TestHeader>::MissingAuthorities));
+	}
+
+	#[rstest]
+	fn test_relay_parents_do_not_form_chain() {
+		let (mut relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(10, 10, None);
+		let header_to_modify = relay_parent_descendants.get_mut(2).expect("Parent is available");
+		let expected_hash = header_to_modify.parent_hash;
+		// Parent hash does not point to the proper parent, incomplete chain
+		header_to_modify.parent_hash = H256::repeat_byte(0x9);
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+
+		assert_eq!(
+			result,
+			Err(RelayParentVerificationError::<TestHeader>::InvalidChainSequence {
+				number: 2,
+				expected: expected_hash,
+				found: H256::repeat_byte(0x9),
+			})
+		);
+	}
+
+	#[rstest]
+	fn test_relay_parent_with_wrong_signature() {
+		let (mut relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(10, 10, None);
+
+		// Pop the seal of the last descendant and put some invalid signature into the digests
+		let rp_to_modify = relay_parent_descendants.last_mut().expect("Parent is available");
+		rp_to_modify.digest_mut().logs.pop();
+		let invalid_signature =
+			Sr25519Keyring::Alice.sign(b"Not the signature you are looking for.");
+		rp_to_modify.digest_mut().push(DigestItem::babe_seal(invalid_signature.into()));
+		let expected_hash = rp_to_modify.hash();
+
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+
+		assert_eq!(
+			result,
+			Err(RelayParentVerificationError::<TestHeader>::InvalidSignature {
+				number: 9,
+				hash: expected_hash,
+			})
+		);
+	}
+
+	#[rstest]
+	fn test_verify_relay_parent_descendants_missing_next_authorities_with_epoch_change() {
+		pezsp_tracing::try_init_simple();
+		let (relay_parent_descendants, authorities, _) =
+			build_relay_parent_descendants(10, 10, Some(5));
+		let relay_state_proof = build_relay_chain_storage_proof(Some(authorities), None);
+
+		let expected_hash = relay_parent_descendants[5].hash();
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		let result = verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		);
+		assert_eq!(
+			result,
+			Err(RelayParentVerificationError::<TestHeader>::MissingNextEpochAuthorities {
+				number: 5,
+				hash: expected_hash,
+			})
+		);
+	}
+
+	#[rstest]
+	fn test_verify_relay_parent_descendants_happy_case_with_epoch_change(
+		#[values(1, 2, 3, 4, 100)] num_headers: u64,
+		#[values(1, 3, 100, 1000)] num_authorities: u64,
+	) {
+		pezsp_tracing::try_init_simple();
+		let (relay_parent_descendants, authorities, next_authorities) =
+			build_relay_parent_descendants(num_headers, num_authorities, Some(5));
+		let relay_state_proof =
+			build_relay_chain_storage_proof(Some(authorities), Some(next_authorities));
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		assert!(verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		)
+		.is_ok());
+	}
+
+	/// Test some interesting epoch change positions, like epoch change on RP directly, and last
+	/// block.
+	#[rstest]
+	fn test_verify_relay_parent_with_epoch_change_at_positions(
+		#[values(0, 5, 10)] epoch_change_position: u64,
+	) {
+		pezsp_tracing::try_init_simple();
+		let (relay_parent_descendants, authorities, next_authorities) =
+			build_relay_parent_descendants(10, 10, Some(epoch_change_position));
+		let relay_state_proof =
+			build_relay_chain_storage_proof(Some(authorities), Some(next_authorities));
+
+		// Make sure that the first relay parent has the correct state root set
+		let relay_parent_state_root = relay_parent_descendants.get(0).unwrap().state_root;
+		// Expected number of parents passed to the function does not include actual relay parent
+		let expected_number_of_descendants = (relay_parent_descendants.len() - 1) as u32;
+
+		assert!(verify_relay_parent_descendants(
+			&relay_state_proof,
+			relay_parent_descendants,
+			relay_parent_state_root,
+			expected_number_of_descendants,
+		)
+		.is_ok());
+	}
+
+	/// Helper function to create a mock `RelayChainStateProof`.
+	fn build_relay_chain_storage_proof(
+		authorities: Option<Vec<(AuthorityId, BabeAuthorityWeight)>>,
+		next_authorities: Option<Vec<(AuthorityId, BabeAuthorityWeight)>>,
+	) -> RelayChainStateProof {
+		// Create a mock implementation or structure, adjust this to match the proof's definition
+		let mut proof_builder = RelayStateSproofBuilder::default();
+		if let Some(authorities) = authorities {
+			proof_builder
+				.additional_key_values
+				.push((relay_chain::well_known_keys::AUTHORITIES.to_vec(), authorities.encode()));
+		}
+
+		if let Some(next_authorities) = next_authorities {
+			proof_builder.additional_key_values.push((
+				relay_chain::well_known_keys::NEXT_AUTHORITIES.to_vec(),
+				next_authorities.encode(),
+			));
+		}
+		let (hash, relay_storage_proof) = proof_builder.into_state_root_and_proof();
+		RelayChainStateProof::new(PARA_ID.into(), hash, relay_storage_proof).unwrap()
+	}
+
+	/// This method generates some vrf data, but only to make the compiler happy.
+	/// This data is not verified and we don't care :).
+	fn generate_testing_vrf() -> VrfSignature {
+		let vrf_proof_bytes = [0u8; 64];
+		let proof: VrfProof = VrfProof::decode(&mut vrf_proof_bytes.as_slice()).unwrap();
+		let vrf_pre_out_bytes = [0u8; 32];
+		let pre_output: VrfPreOutput =
+			VrfPreOutput::decode(&mut vrf_pre_out_bytes.as_slice()).unwrap();
+		VrfSignature { pre_output, proof }
+	}
+
+	/// Build a chain of relay parent descendants.
+	///
+	/// Returns the relay parent header as well as the current and next epoch authorities.
+	fn build_relay_parent_descendants(
+		num_headers: u64,
+		num_authorities: u64,
+		epoch_change_at: Option<u64>,
+	) -> (
+		Vec<TestHeader>,
+		Vec<(AuthorityId, BabeAuthorityWeight)>,
+		Vec<(AuthorityId, BabeAuthorityWeight)>,
+	) {
+		// Generate initial authorities
+		let (authorities, next_authorities) = generate_authority_pairs(num_authorities);
+		let authorities_for_storage = convert_to_authority_weight_pair(&authorities);
+		let next_authorities_for_storage = convert_to_authority_weight_pair(&next_authorities);
+
+		// Generate headers chain
+		let mut headers = Vec::with_capacity(num_headers as usize);
+		let mut current_authorities = authorities.clone();
+		let mut previous_hash = None;
+
+		for block_number in 0..=num_headers - 1 {
+			let mut header = create_header(block_number, previous_hash);
+			let authority_index = (block_number as u32) % (num_authorities as u32);
+
+			// Add pre-digest
+			add_pre_digest(&mut header, authority_index, block_number);
+
+			// Handle epoch change if needed
+			if epoch_change_at.map_or(false, |change_at| block_number == change_at) {
+				add_epoch_change_digest(&mut header, num_authorities);
+				if block_number > 0 {
+					current_authorities = next_authorities.clone();
+				}
+			}
+
+			// Sign and seal header
+			let signature =
+				current_authorities[authority_index as usize].sign(header.hash().as_bytes());
+			header.digest_mut().push(DigestItem::babe_seal(signature.into()));
+
+			previous_hash = Some(header.hash());
+			headers.push(header);
+		}
+
+		(headers, authorities_for_storage, next_authorities_for_storage)
+	}
+
+	// Helper functions
+	fn generate_authority_pairs(num_authorities: u64) -> (Vec<AuthorityPair>, Vec<AuthorityPair>) {
+		let authorities: Vec<_> = (0..num_authorities).map(|_| Pair::generate().0).collect();
+		let next_authorities: Vec<_> = (0..num_authorities).map(|_| Pair::generate().0).collect();
+		(authorities, next_authorities)
+	}
+
+	fn convert_to_authority_weight_pair(
+		authorities: &[AuthorityPair],
+	) -> Vec<(AuthorityId, BabeAuthorityWeight)> {
+		authorities
+			.iter()
+			.map(|auth| (auth.public().into(), Default::default()))
+			.collect()
+	}
+
+	fn create_header(block_number: u64, previous_hash: Option<H256>) -> TestHeader {
+		let mut header = TestHeader::new_from_number(block_number);
+		if let Some(parent_hash) = previous_hash {
+			header.parent_hash = parent_hash;
+		}
+		header
+	}
+
+	fn add_pre_digest(header: &mut TestHeader, authority_index: u32, block_number: u64) {
+		let pre_digest = PrimaryPreDigest {
+			authority_index,
+			slot: block_number.into(),
+			vrf_signature: generate_testing_vrf(),
+		};
+		header
+			.digest_mut()
+			.push(DigestItem::babe_pre_digest(PreDigest::Primary(pre_digest)));
+	}
+
+	fn add_epoch_change_digest(header: &mut TestHeader, num_authorities: u64) {
+		let digest_authorities: Vec<(AuthorityId, BabeAuthorityWeight)> = (0..num_authorities)
+			.map(|_| {
+				let authority_pair: AuthorityPair = Pair::generate().0;
+				(authority_pair.public().into(), Default::default())
+			})
+			.collect();
+
+		header.digest_mut().push(DigestItem::Consensus(
+			BABE_ENGINE_ID,
+			ConsensusLog::NextEpochData(NextEpochDescriptor {
+				authorities: digest_authorities,
+				randomness: [0; 32],
+			})
+			.encode(),
+		));
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/migration.rs

@@ -0,0 +1,90 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 crate::{Config, Pallet, ReservedDmpWeightOverride, ReservedXcmpWeightOverride};
+use pezframe_support::{
+	pezpallet_prelude::*,
+	traits::{Get, OnRuntimeUpgrade, StorageVersion},
+	weights::Weight,
+};
+
+/// The in-code storage version.
+pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(2);
+
+/// Migrates the pallet storage to the most recent version.
+pub struct Migration<T: Config>(PhantomData<T>);
+
+impl<T: Config> OnRuntimeUpgrade for Migration<T> {
+	fn on_runtime_upgrade() -> Weight {
+		let mut weight: Weight = T::DbWeight::get().reads(2);
+
+		if StorageVersion::get::<Pallet<T>>() == 0 {
+			weight = weight
+				.saturating_add(v1::migrate::<T>())
+				.saturating_add(T::DbWeight::get().writes(1));
+			StorageVersion::new(1).put::<Pallet<T>>();
+		}
+
+		if StorageVersion::get::<Pallet<T>>() == 1 {
+			weight = weight
+				.saturating_add(v2::migrate::<T>())
+				.saturating_add(T::DbWeight::get().writes(1));
+			StorageVersion::new(2).put::<Pallet<T>>();
+		}
+
+		weight
+	}
+}
+
+/// V2: Migrate to 2D weights for ReservedXcmpWeightOverride and ReservedDmpWeightOverride.
+mod v2 {
+	use super::*;
+	const DEFAULT_POV_SIZE: u64 = 64 * 1024; // 64 KB
+
+	pub fn migrate<T: Config>() -> Weight {
+		let translate = |pre: u64| -> Weight { Weight::from_parts(pre, DEFAULT_POV_SIZE) };
+
+		if ReservedXcmpWeightOverride::<T>::translate(|pre| pre.map(translate)).is_err() {
+			log::error!(
+				target: "teyrchain_system",
+				"unexpected error when performing translation of the ReservedXcmpWeightOverride type during storage upgrade to v2"
+			);
+		}
+
+		if ReservedDmpWeightOverride::<T>::translate(|pre| pre.map(translate)).is_err() {
+			log::error!(
+				target: "teyrchain_system",
+				"unexpected error when performing translation of the ReservedDmpWeightOverride type during storage upgrade to v2"
+			);
+		}
+
+		T::DbWeight::get().reads_writes(2, 2)
+	}
+}
+
+/// V1: `LastUpgrade` block number is removed from the storage since the upgrade
+/// mechanism now uses signals instead of block offsets.
+mod v1 {
+	use crate::{Config, Pallet};
+	#[allow(deprecated)]
+	use pezframe_support::{migration::remove_storage_prefix, pezpallet_prelude::*};
+
+	pub fn migrate<T: Config>() -> Weight {
+		#[allow(deprecated)]
+		remove_storage_prefix(<Pallet<T>>::name().as_bytes(), b"LastUpgrade", b"");
+		T::DbWeight::get().writes(1)
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/mock.rs

@@ -0,0 +1,492 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Test setup and helpers.
+
+#![cfg(test)]
+
+use super::*;
+
+use alloc::collections::vec_deque::VecDeque;
+use codec::Encode;
+use core::num::NonZeroU32;
+use pezcumulus_primitives_core::{
+	relay_chain::BlockNumber as RelayBlockNumber, AggregateMessageOrigin, InboundDownwardMessage,
+	InboundHrmpMessage, PersistedValidationData,
+};
+use pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+use pezframe_support::{
+	derive_impl,
+	inherent::{InherentData, ProvideInherent},
+	parameter_types,
+	traits::{
+		OnFinalize, OnInitialize, ProcessMessage, ProcessMessageError, UnfilteredDispatchable,
+	},
+	weights::{Weight, WeightMeter},
+};
+use pezframe_system::{limits::BlockWeights, pezpallet_prelude::BlockNumberFor, RawOrigin};
+use pezsp_core::ConstU32;
+use pezsp_runtime::{traits::BlakeTwo256, BuildStorage};
+use pezsp_version::RuntimeVersion;
+use std::cell::RefCell;
+
+use crate as teyrchain_system;
+use crate::consensus_hook::UnincludedSegmentCapacity;
+
+type Block = pezframe_system::mocking::MockBlock<Test>;
+
+pezframe_support::construct_runtime!(
+	pub enum Test {
+		System: pezframe_system,
+		TeyrchainSystem: teyrchain_system,
+		MessageQueue: pezpallet_message_queue,
+	}
+);
+
+parameter_types! {
+	pub Version: RuntimeVersion = RuntimeVersion {
+		spec_name: alloc::borrow::Cow::Borrowed("test"),
+		impl_name: alloc::borrow::Cow::Borrowed("system-test"),
+		authoring_version: 1,
+		spec_version: 1,
+		impl_version: 1,
+		apis: pezsp_version::create_apis_vec!([]),
+		transaction_version: 1,
+		system_version: 1,
+	};
+	pub RuntimeBlockWeights: BlockWeights =
+		BlockWeights::simple_max(Weight::from_parts(1024, 6 * 1024 * 1024));
+	pub const TeyrchainId: ParaId = ParaId::new(200);
+	pub const ReservedXcmpWeight: Weight = Weight::zero();
+	pub const ReservedDmpWeight: Weight = Weight::zero();
+}
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl pezframe_system::Config for Test {
+	type BlockWeights = RuntimeBlockWeights;
+	type Block = Block;
+	type Version = Version;
+	type OnSetCode = TeyrchainSetCode<Self>;
+}
+
+parameter_types! {
+	pub const RelayOrigin: AggregateMessageOrigin = AggregateMessageOrigin::Parent;
+}
+
+impl Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type OnSystemEvent = ();
+	type SelfParaId = TeyrchainId;
+	type OutboundXcmpMessageSource = FromThreadLocal;
+	type DmpQueue = pezframe_support::traits::EnqueueWithOrigin<MessageQueue, RelayOrigin>;
+	type ReservedDmpWeight = ReservedDmpWeight;
+	type XcmpMessageHandler = SaveIntoThreadLocal;
+	type ReservedXcmpWeight = ReservedXcmpWeight;
+	type CheckAssociatedRelayNumber = AnyRelayNumber;
+	type ConsensusHook = TestConsensusHook;
+	type WeightInfo = ();
+	type RelayParentOffset = ConstU32<0>;
+}
+
+std::thread_local! {
+	pub static CONSENSUS_HOOK: RefCell<Box<dyn Fn(&RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity)>>
+		= RefCell::new(Box::new(|_| (Weight::zero(), NonZeroU32::new(1).unwrap().into())));
+}
+
+pub struct TestConsensusHook;
+
+impl ConsensusHook for TestConsensusHook {
+	fn on_state_proof(s: &RelayChainStateProof) -> (Weight, UnincludedSegmentCapacity) {
+		CONSENSUS_HOOK.with(|f| f.borrow_mut()(s))
+	}
+}
+
+parameter_types! {
+	pub const MaxWeight: Weight = Weight::MAX;
+}
+
+impl pezpallet_message_queue::Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	// NOTE that normally for benchmarking we should use the No-OP message processor, but in this
+	// case its a mocked runtime and will only be used to generate insecure default weights.
+	type MessageProcessor = SaveIntoThreadLocal;
+	type Size = u32;
+	type QueueChangeHandler = ();
+	type QueuePausedQuery = ();
+	type HeapSize = pezsp_core::ConstU32<{ 103 * 1024 }>;
+	type MaxStale = pezsp_core::ConstU32<8>;
+	type ServiceWeight = MaxWeight;
+	type IdleMaxServiceWeight = ();
+	type WeightInfo = ();
+}
+
+/// A `XcmpMessageSource` that takes messages from thread-local.
+pub struct FromThreadLocal;
+
+/// A `MessageProcessor` that stores all messages in thread-local.
+pub struct SaveIntoThreadLocal;
+
+std::thread_local! {
+	pub static HANDLED_DMP_MESSAGES: RefCell<Vec<Vec<u8>>> = RefCell::new(Vec::new());
+	pub static HANDLED_XCMP_MESSAGES: RefCell<Vec<(ParaId, relay_chain::BlockNumber, Vec<u8>)>> = RefCell::new(Vec::new());
+	pub static SENT_MESSAGES: RefCell<Vec<(ParaId, Vec<u8>)>> = RefCell::new(Vec::new());
+}
+
+pub fn send_message(dest: ParaId, message: Vec<u8>) {
+	SENT_MESSAGES.with(|m| m.borrow_mut().push((dest, message)));
+}
+
+impl XcmpMessageSource for FromThreadLocal {
+	fn take_outbound_messages(maximum_channels: usize) -> Vec<(ParaId, Vec<u8>)> {
+		let mut ids = std::collections::BTreeSet::<ParaId>::new();
+		let mut taken_messages = 0;
+		let mut taken_bytes = 0;
+		let mut result = Vec::new();
+		SENT_MESSAGES.with(|ms| {
+			ms.borrow_mut().retain(|m| {
+				let status = <Pallet<Test> as GetChannelInfo>::get_channel_status(m.0);
+				let (max_size_now, max_size_ever) = match status {
+					ChannelStatus::Ready(now, ever) => (now, ever),
+					ChannelStatus::Closed => return false, // drop message
+					ChannelStatus::Full => return true,    // keep message queued.
+				};
+
+				let msg_len = m.1.len();
+
+				if !ids.contains(&m.0) &&
+					taken_messages < maximum_channels &&
+					msg_len <= max_size_ever &&
+					taken_bytes + msg_len <= max_size_now
+				{
+					ids.insert(m.0);
+					taken_messages += 1;
+					taken_bytes += msg_len;
+					result.push(m.clone());
+					false
+				} else {
+					true
+				}
+			})
+		});
+		result
+	}
+}
+
+impl ProcessMessage for SaveIntoThreadLocal {
+	type Origin = AggregateMessageOrigin;
+
+	fn process_message(
+		message: &[u8],
+		origin: Self::Origin,
+		_meter: &mut WeightMeter,
+		_id: &mut [u8; 32],
+	) -> Result<bool, ProcessMessageError> {
+		assert_eq!(origin, Self::Origin::Parent);
+
+		HANDLED_DMP_MESSAGES.with(|m| {
+			m.borrow_mut().push(message.to_vec());
+			Weight::zero()
+		});
+		Ok(true)
+	}
+}
+
+impl XcmpMessageHandler for SaveIntoThreadLocal {
+	fn handle_xcmp_messages<'a, I: Iterator<Item = (ParaId, RelayBlockNumber, &'a [u8])>>(
+		iter: I,
+		_max_weight: Weight,
+	) -> Weight {
+		HANDLED_XCMP_MESSAGES.with(|m| {
+			for (sender, sent_at, message) in iter {
+				m.borrow_mut().push((sender, sent_at, message.to_vec()));
+			}
+			Weight::zero()
+		})
+	}
+}
+
+// This function basically just builds a genesis storage key/value store according to
+// our desired mockup.
+pub fn new_test_ext() -> pezsp_io::TestExternalities {
+	HANDLED_DMP_MESSAGES.with(|m| m.borrow_mut().clear());
+	HANDLED_XCMP_MESSAGES.with(|m| m.borrow_mut().clear());
+
+	pezframe_system::GenesisConfig::<Test>::default().build_storage().unwrap().into()
+}
+
+#[allow(dead_code)]
+pub fn mk_dmp(sent_at: u8, size: usize) -> InboundDownwardMessage {
+	InboundDownwardMessage { sent_at: sent_at as u32, msg: vec![sent_at; size] }
+}
+
+pub fn mk_hrmp(sent_at: u8, size: usize) -> InboundHrmpMessage {
+	InboundHrmpMessage { sent_at: sent_at as u32, data: vec![sent_at; size] }
+}
+
+pub struct ReadRuntimeVersion(pub Vec<u8>);
+
+impl pezsp_core::traits::ReadRuntimeVersion for ReadRuntimeVersion {
+	fn read_runtime_version(
+		&self,
+		_wasm_code: &[u8],
+		_ext: &mut dyn pezsp_externalities::Externalities,
+	) -> Result<Vec<u8>, String> {
+		Ok(self.0.clone())
+	}
+}
+
+pub fn wasm_ext() -> pezsp_io::TestExternalities {
+	let version = RuntimeVersion {
+		spec_name: "test".into(),
+		spec_version: 2,
+		impl_version: 1,
+		..Default::default()
+	};
+
+	let mut ext = new_test_ext();
+	ext.register_extension(pezsp_core::traits::ReadRuntimeVersionExt::new(ReadRuntimeVersion(
+		version.encode(),
+	)));
+	ext
+}
+
+pub struct BlockTest {
+	n: BlockNumberFor<Test>,
+	within_block: Box<dyn Fn()>,
+	after_block: Option<Box<dyn Fn()>>,
+}
+
+/// BlockTests exist to test blocks with some setup: we have to assume that
+/// `validate_block` will mutate and check storage in certain predictable
+/// ways, for example, and we want to always ensure that tests are executed
+/// in the context of some particular block number.
+#[derive(Default)]
+pub struct BlockTests {
+	tests: Vec<BlockTest>,
+	without_externalities: bool,
+	pending_upgrade: Option<RelayChainBlockNumber>,
+	ran: bool,
+	relay_sproof_builder_hook:
+		Option<Box<dyn Fn(&BlockTests, RelayChainBlockNumber, &mut RelayStateSproofBuilder)>>,
+	inherent_data_hook:
+		Option<Box<dyn Fn(&BlockTests, RelayChainBlockNumber, &mut TeyrchainInherentData)>>,
+	inclusion_delay: Option<usize>,
+	relay_block_number: Option<Box<dyn Fn(&BlockNumberFor<Test>) -> RelayChainBlockNumber>>,
+
+	included_para_head: Option<relay_chain::HeadData>,
+	pending_blocks: VecDeque<relay_chain::HeadData>,
+}
+
+impl BlockTests {
+	pub fn new() -> BlockTests {
+		Default::default()
+	}
+
+	pub fn new_without_externalities() -> BlockTests {
+		let mut tests = BlockTests::new();
+		tests.without_externalities = true;
+		tests
+	}
+
+	pub fn add_raw(mut self, test: BlockTest) -> Self {
+		self.tests.push(test);
+		self
+	}
+
+	pub fn add<F>(self, n: BlockNumberFor<Test>, within_block: F) -> Self
+	where
+		F: 'static + Fn(),
+	{
+		self.add_raw(BlockTest { n, within_block: Box::new(within_block), after_block: None })
+	}
+
+	pub fn add_with_post_test<F1, F2>(
+		self,
+		n: BlockNumberFor<Test>,
+		within_block: F1,
+		after_block: F2,
+	) -> Self
+	where
+		F1: 'static + Fn(),
+		F2: 'static + Fn(),
+	{
+		self.add_raw(BlockTest {
+			n,
+			within_block: Box::new(within_block),
+			after_block: Some(Box::new(after_block)),
+		})
+	}
+
+	pub fn with_relay_sproof_builder<F>(mut self, f: F) -> Self
+	where
+		F: 'static + Fn(&BlockTests, RelayChainBlockNumber, &mut RelayStateSproofBuilder),
+	{
+		self.relay_sproof_builder_hook = Some(Box::new(f));
+		self
+	}
+
+	pub fn with_relay_block_number<F>(mut self, f: F) -> Self
+	where
+		F: 'static + Fn(&BlockNumberFor<Test>) -> RelayChainBlockNumber,
+	{
+		self.relay_block_number = Some(Box::new(f));
+		self
+	}
+
+	pub fn with_inherent_data<F>(mut self, f: F) -> Self
+	where
+		F: 'static + Fn(&BlockTests, RelayChainBlockNumber, &mut TeyrchainInherentData),
+	{
+		self.inherent_data_hook = Some(Box::new(f));
+		self
+	}
+
+	pub fn with_inclusion_delay(mut self, inclusion_delay: usize) -> Self {
+		self.inclusion_delay.replace(inclusion_delay);
+		self
+	}
+
+	pub fn run(&mut self) {
+		wasm_ext().execute_with(|| {
+			self.run_without_ext();
+		});
+	}
+
+	pub fn run_without_ext(&mut self) {
+		self.ran = true;
+
+		let mut parent_head_data = {
+			let header = HeaderFor::<Test>::new_from_number(0);
+			relay_chain::HeadData(header.encode())
+		};
+
+		self.included_para_head = Some(parent_head_data.clone());
+
+		for BlockTest { n, within_block, after_block } in self.tests.iter() {
+			let relay_parent_number = self
+				.relay_block_number
+				.as_ref()
+				.map(|f| f(n))
+				.unwrap_or(*n as RelayChainBlockNumber);
+			// clear pending updates, as applicable
+			if let Some(upgrade_block) = self.pending_upgrade {
+				if n >= &upgrade_block.into() {
+					self.pending_upgrade = None;
+				}
+			}
+
+			// begin initialization
+			let parent_hash = BlakeTwo256::hash(&parent_head_data.0);
+			System::reset_events();
+			System::initialize(n, &parent_hash, &Default::default());
+
+			// now mess with the storage the way validate_block does
+			let mut sproof_builder = RelayStateSproofBuilder::default();
+			sproof_builder.included_para_head = self
+				.included_para_head
+				.clone()
+				.unwrap_or_else(|| parent_head_data.clone())
+				.into();
+			if let Some(ref hook) = self.relay_sproof_builder_hook {
+				hook(self, relay_parent_number, &mut sproof_builder);
+			}
+			let (relay_parent_storage_root, relay_chain_state) =
+				sproof_builder.into_state_root_and_proof();
+			let vfp = PersistedValidationData {
+				relay_parent_number,
+				relay_parent_storage_root,
+				..Default::default()
+			};
+
+			<ValidationData<Test>>::put(&vfp);
+			NewValidationCode::<Test>::kill();
+
+			// It is insufficient to push the validation function params
+			// to storage; they must also be included in the inherent data.
+			let inherent_data = {
+				let mut inherent_data = InherentData::default();
+				let mut system_inherent_data = TeyrchainInherentData {
+					validation_data: vfp.clone(),
+					relay_chain_state,
+					downward_messages: Default::default(),
+					horizontal_messages: Default::default(),
+					relay_parent_descendants: Default::default(),
+					collator_peer_id: None,
+				};
+				if let Some(ref hook) = self.inherent_data_hook {
+					hook(self, relay_parent_number, &mut system_inherent_data);
+				}
+				inherent_data
+					.put_data(
+						pezcumulus_primitives_teyrchain_inherent::INHERENT_IDENTIFIER,
+						&system_inherent_data,
+					)
+					.expect("failed to put VFP inherent");
+				inherent_data
+			};
+
+			// execute the block
+			TeyrchainSystem::on_initialize(*n);
+			TeyrchainSystem::create_inherent(&inherent_data)
+				.expect("got an inherent")
+				.dispatch_bypass_filter(RawOrigin::None.into())
+				.expect("dispatch succeeded");
+			MessageQueue::on_initialize(*n);
+			within_block();
+			MessageQueue::on_finalize(*n);
+			TeyrchainSystem::on_finalize(*n);
+
+			// did block execution set new validation code?
+			if NewValidationCode::<Test>::exists() && self.pending_upgrade.is_some() {
+				panic!("attempted to set validation code while upgrade was pending");
+			}
+
+			// clean up
+			let header = System::finalize();
+			let head_data = relay_chain::HeadData(header.encode());
+			parent_head_data = head_data.clone();
+			match self.inclusion_delay {
+				Some(delay) if delay > 0 => {
+					self.pending_blocks.push_back(head_data);
+					if self.pending_blocks.len() > delay {
+						let included = self.pending_blocks.pop_front().unwrap();
+
+						self.included_para_head.replace(included);
+					}
+				},
+				_ => {
+					self.included_para_head.replace(head_data);
+				},
+			}
+
+			if let Some(after_block) = after_block {
+				after_block();
+			}
+		}
+	}
+}
+
+impl Drop for BlockTests {
+	fn drop(&mut self) {
+		if !self.ran {
+			if self.without_externalities {
+				self.run_without_ext();
+			} else {
+				self.run();
+			}
+		}
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/relay_state_snapshot.rs

@@ -0,0 +1,386 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Relay chain state proof provides means for accessing part of relay chain storage for reads.
+
+use alloc::vec::Vec;
+use codec::{Decode, Encode};
+use pezcumulus_primitives_core::{
+	relay_chain, AbridgedHostConfiguration, AbridgedHrmpChannel, ParaId,
+};
+use scale_info::TypeInfo;
+use pezsp_runtime::traits::HashingFor;
+use pezsp_state_machine::{Backend, TrieBackend, TrieBackendBuilder};
+use pezsp_trie::{HashDBT, MemoryDB, StorageProof, EMPTY_PREFIX};
+
+/// The capacity of the upward message queue of a teyrchain on the relay chain.
+// The field order should stay the same as the data can be found in the proof to ensure both are
+// have the same encoded representation.
+#[derive(Clone, Encode, Decode, TypeInfo, Default)]
+pub struct RelayDispatchQueueRemainingCapacity {
+	/// The number of additional messages that can be enqueued.
+	pub remaining_count: u32,
+	/// The total size of additional messages that can be enqueued.
+	pub remaining_size: u32,
+}
+
+/// A snapshot of some messaging related state of relay chain pertaining to the current teyrchain.
+///
+/// This data is essential for making sure that the teyrchain is aware of current resource use on
+/// the relay chain and that the candidates produced for this teyrchain do not exceed any of these
+/// limits.
+#[derive(Clone, Encode, Decode, TypeInfo)]
+pub struct MessagingStateSnapshot {
+	/// The current message queue chain head for downward message queue.
+	///
+	/// If the value is absent on the relay chain this will be set to all zeros.
+	pub dmq_mqc_head: relay_chain::Hash,
+
+	/// The current capacity of the upward message queue of the current teyrchain on the relay
+	/// chain.
+	pub relay_dispatch_queue_remaining_capacity: RelayDispatchQueueRemainingCapacity,
+
+	/// Information about all the inbound HRMP channels.
+	///
+	/// These are structured as a list of tuples. The para id in the tuple specifies the sender
+	/// of the channel. Obviously, the recipient is the current teyrchain.
+	///
+	/// The channels are sorted by the sender para id ascension.
+	pub ingress_channels: Vec<(ParaId, AbridgedHrmpChannel)>,
+
+	/// Information about all the outbound HRMP channels.
+	///
+	/// These are structured as a list of tuples. The para id in the tuple specifies the recipient
+	/// of the channel. Obviously, the sender is the current teyrchain.
+	///
+	/// The channels are sorted by the recipient para id ascension.
+	pub egress_channels: Vec<(ParaId, AbridgedHrmpChannel)>,
+}
+
+#[derive(Debug)]
+pub enum Error {
+	/// The provided proof was created against unexpected storage root.
+	RootMismatch,
+	/// The entry cannot be read.
+	ReadEntry(ReadEntryErr),
+	/// The optional entry cannot be read.
+	ReadOptionalEntry(ReadEntryErr),
+	/// The slot cannot be extracted.
+	Slot(ReadEntryErr),
+	/// The upgrade go-ahead signal cannot be read.
+	UpgradeGoAhead(ReadEntryErr),
+	/// The upgrade restriction signal cannot be read.
+	UpgradeRestriction(ReadEntryErr),
+	/// The host configuration cannot be extracted.
+	Config(ReadEntryErr),
+	/// The DMQ MQC head cannot be extracted.
+	DmqMqcHead(ReadEntryErr),
+	/// Relay dispatch queue cannot be extracted.
+	RelayDispatchQueueRemainingCapacity(ReadEntryErr),
+	/// The hrmp ingress channel index cannot be extracted.
+	HrmpIngressChannelIndex(ReadEntryErr),
+	/// The hrmp egress channel index cannot be extracted.
+	HrmpEgressChannelIndex(ReadEntryErr),
+	/// The channel identified by the sender and receiver cannot be extracted.
+	HrmpChannel(ParaId, ParaId, ReadEntryErr),
+	/// The latest included teyrchain head cannot be extracted.
+	ParaHead(ReadEntryErr),
+	/// The relay chain authorities cannot be extracted
+	Authorities(ReadEntryErr),
+	/// The relay chain authorities for the next epoch cannot be extracted
+	NextAuthorities(ReadEntryErr),
+}
+
+#[derive(Debug)]
+pub enum ReadEntryErr {
+	/// The value cannot be extracted from the proof.
+	Proof,
+	/// The value cannot be decoded.
+	Decode,
+	/// The value is expected to be present on the relay chain, but it doesn't exist.
+	Absent,
+}
+
+/// Read an entry given by the key and try to decode it. If the value specified by the key according
+/// to the proof is empty, the `fallback` value will be returned.
+///
+/// Returns `Err` in case the backend can't return the value under the specific key (likely due to
+/// a malformed proof), in case the decoding fails, or in case where the value is empty in the relay
+/// chain state and no fallback was provided.
+fn read_entry<T, B>(backend: &B, key: &[u8], fallback: Option<T>) -> Result<T, ReadEntryErr>
+where
+	T: Decode,
+	B: Backend<HashingFor<relay_chain::Block>>,
+{
+	backend
+		.storage(key)
+		.map_err(|_| ReadEntryErr::Proof)?
+		.map(|raw_entry| T::decode(&mut &raw_entry[..]).map_err(|_| ReadEntryErr::Decode))
+		.transpose()?
+		.or(fallback)
+		.ok_or(ReadEntryErr::Absent)
+}
+
+/// Read an optional entry given by the key and try to decode it.
+/// Returns `None` if the value specified by the key according to the proof is empty.
+///
+/// Returns `Err` in case the backend can't return the value under the specific key (likely due to
+/// a malformed proof) or if the value couldn't be decoded.
+fn read_optional_entry<T, B>(backend: &B, key: &[u8]) -> Result<Option<T>, ReadEntryErr>
+where
+	T: Decode,
+	B: Backend<HashingFor<relay_chain::Block>>,
+{
+	match read_entry(backend, key, None) {
+		Ok(v) => Ok(Some(v)),
+		Err(ReadEntryErr::Absent) => Ok(None),
+		Err(err) => Err(err),
+	}
+}
+
+/// A state proof extracted from the relay chain.
+///
+/// This state proof is extracted from the relay chain block we are building on top of.
+pub struct RelayChainStateProof {
+	para_id: ParaId,
+	trie_backend:
+		TrieBackend<MemoryDB<HashingFor<relay_chain::Block>>, HashingFor<relay_chain::Block>>,
+}
+
+impl RelayChainStateProof {
+	/// Create a new instance of `Self`.
+	///
+	/// Returns an error if the given `relay_parent_storage_root` is not the root of the given
+	/// `proof`.
+	pub fn new(
+		para_id: ParaId,
+		relay_parent_storage_root: relay_chain::Hash,
+		proof: StorageProof,
+	) -> Result<Self, Error> {
+		let db = proof.into_memory_db::<HashingFor<relay_chain::Block>>();
+		if !db.contains(&relay_parent_storage_root, EMPTY_PREFIX) {
+			return Err(Error::RootMismatch);
+		}
+		let trie_backend = TrieBackendBuilder::new(db, relay_parent_storage_root).build();
+
+		Ok(Self { para_id, trie_backend })
+	}
+
+	/// Read the [`MessagingStateSnapshot`] from the relay chain state proof.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_messaging_state_snapshot(
+		&self,
+		host_config: &AbridgedHostConfiguration,
+	) -> Result<MessagingStateSnapshot, Error> {
+		let dmq_mqc_head: relay_chain::Hash = read_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::dmq_mqc_head(self.para_id),
+			Some(Default::default()),
+		)
+		.map_err(Error::DmqMqcHead)?;
+
+		let relay_dispatch_queue_remaining_capacity = read_optional_entry::<
+			RelayDispatchQueueRemainingCapacity,
+			_,
+		>(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::relay_dispatch_queue_remaining_capacity(self.para_id)
+				.key,
+		);
+
+		// TODO paritytech/pezkuwi#6283: Remove all usages of `relay_dispatch_queue_size`
+		//
+		// When the relay chain and all teyrchains support
+		// `relay_dispatch_queue_remaining_capacity`, this code here needs to be removed and above
+		// needs to be changed to `read_entry` that returns an error if
+		// `relay_dispatch_queue_remaining_capacity` can not be found/decoded.
+		//
+		// For now we just fallback to the old dispatch queue size on `ReadEntryErr::Absent`.
+		// `ReadEntryErr::Decode` and `ReadEntryErr::Proof` are potentially subject to meddling
+		// by malicious collators, so we reject the block in those cases.
+		let relay_dispatch_queue_remaining_capacity = match relay_dispatch_queue_remaining_capacity
+		{
+			Ok(Some(r)) => r,
+			Ok(None) => {
+				let res = read_entry::<(u32, u32), _>(
+					&self.trie_backend,
+					#[allow(deprecated)]
+					&relay_chain::well_known_keys::relay_dispatch_queue_size(self.para_id),
+					Some((0, 0)),
+				)
+				.map_err(Error::RelayDispatchQueueRemainingCapacity)?;
+
+				let remaining_count = host_config.max_upward_queue_count.saturating_sub(res.0);
+				let remaining_size = host_config.max_upward_queue_size.saturating_sub(res.1);
+				RelayDispatchQueueRemainingCapacity { remaining_count, remaining_size }
+			},
+			Err(e) => return Err(Error::RelayDispatchQueueRemainingCapacity(e)),
+		};
+
+		let ingress_channel_index: Vec<ParaId> = read_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::hrmp_ingress_channel_index(self.para_id),
+			Some(Vec::new()),
+		)
+		.map_err(Error::HrmpIngressChannelIndex)?;
+
+		let egress_channel_index: Vec<ParaId> = read_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::hrmp_egress_channel_index(self.para_id),
+			Some(Vec::new()),
+		)
+		.map_err(Error::HrmpEgressChannelIndex)?;
+
+		let mut ingress_channels = Vec::with_capacity(ingress_channel_index.len());
+		for sender in ingress_channel_index {
+			let channel_id = relay_chain::HrmpChannelId { sender, recipient: self.para_id };
+			let hrmp_channel: AbridgedHrmpChannel = read_entry(
+				&self.trie_backend,
+				&relay_chain::well_known_keys::hrmp_channels(channel_id),
+				None,
+			)
+			.map_err(|read_err| Error::HrmpChannel(sender, self.para_id, read_err))?;
+			ingress_channels.push((sender, hrmp_channel));
+		}
+
+		let mut egress_channels = Vec::with_capacity(egress_channel_index.len());
+		for recipient in egress_channel_index {
+			let channel_id = relay_chain::HrmpChannelId { sender: self.para_id, recipient };
+			let hrmp_channel: AbridgedHrmpChannel = read_entry(
+				&self.trie_backend,
+				&relay_chain::well_known_keys::hrmp_channels(channel_id),
+				None,
+			)
+			.map_err(|read_err| Error::HrmpChannel(self.para_id, recipient, read_err))?;
+			egress_channels.push((recipient, hrmp_channel));
+		}
+
+		// NOTE that ingress_channels and egress_channels promise to be sorted. We satisfy this
+		// property by relying on the fact that `ingress_channel_index` and `egress_channel_index`
+		// are themselves sorted.
+		Ok(MessagingStateSnapshot {
+			dmq_mqc_head,
+			relay_dispatch_queue_remaining_capacity,
+			ingress_channels,
+			egress_channels,
+		})
+	}
+
+	/// Read the [`AbridgedHostConfiguration`] from the relay chain state proof.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_abridged_host_configuration(&self) -> Result<AbridgedHostConfiguration, Error> {
+		read_entry(&self.trie_backend, relay_chain::well_known_keys::ACTIVE_CONFIG, None)
+			.map_err(Error::Config)
+	}
+
+	/// Read latest included teyrchain [head data](`relay_chain::HeadData`) from the relay chain
+	/// state proof.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_included_para_head(&self) -> Result<relay_chain::HeadData, Error> {
+		read_entry(&self.trie_backend, &relay_chain::well_known_keys::para_head(self.para_id), None)
+			.map_err(Error::ParaHead)
+	}
+
+	/// Read relay chain authorities.
+	pub fn read_authorities(
+		&self,
+	) -> Result<Vec<(pezsp_consensus_babe::AuthorityId, pezsp_consensus_babe::BabeAuthorityWeight)>, Error>
+	{
+		read_entry(&self.trie_backend, &relay_chain::well_known_keys::AUTHORITIES, None)
+			.map_err(Error::Authorities)
+	}
+
+	/// Read relay chain authorities for the next epoch.
+	pub fn read_next_authorities(
+		&self,
+	) -> Result<
+		Option<Vec<(pezsp_consensus_babe::AuthorityId, pezsp_consensus_babe::BabeAuthorityWeight)>>,
+		Error,
+	> {
+		read_optional_entry(&self.trie_backend, &relay_chain::well_known_keys::NEXT_AUTHORITIES)
+			.map_err(Error::NextAuthorities)
+	}
+
+	/// Read the [`Slot`](relay_chain::Slot) from the relay chain state proof.
+	///
+	/// The slot is slot of the relay chain block this state proof was extracted from.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_slot(&self) -> Result<relay_chain::Slot, Error> {
+		read_entry(&self.trie_backend, relay_chain::well_known_keys::CURRENT_SLOT, None)
+			.map_err(Error::Slot)
+	}
+
+	/// Read the go-ahead signal for the upgrade from the relay chain state proof.
+	///
+	/// The go-ahead specifies whether the teyrchain can apply the upgrade or should abort it. If
+	/// the value is absent then there is either no judgment by the relay chain yet or no upgrade
+	/// is pending.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_upgrade_go_ahead_signal(
+		&self,
+	) -> Result<Option<relay_chain::UpgradeGoAhead>, Error> {
+		read_optional_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::upgrade_go_ahead_signal(self.para_id),
+		)
+		.map_err(Error::UpgradeGoAhead)
+	}
+
+	/// Read the upgrade restriction signal for the upgrade from the relay chain state proof.
+	///
+	/// If the upgrade restriction is not `None`, then the teyrchain cannot signal an upgrade at
+	/// this block.
+	///
+	/// Returns an error if anything failed at reading or decoding.
+	pub fn read_upgrade_restriction_signal(
+		&self,
+	) -> Result<Option<relay_chain::UpgradeRestriction>, Error> {
+		read_optional_entry(
+			&self.trie_backend,
+			&relay_chain::well_known_keys::upgrade_restriction_signal(self.para_id),
+		)
+		.map_err(Error::UpgradeRestriction)
+	}
+
+	/// Read an entry given by the key and try to decode it. If the value specified by the key
+	/// according to the proof is empty, the `fallback` value will be returned.
+	///
+	/// Returns `Err` in case the backend can't return the value under the specific key (likely due
+	/// to a malformed proof), in case the decoding fails, or in case where the value is empty in
+	/// the relay chain state and no fallback was provided.
+	pub fn read_entry<T>(&self, key: &[u8], fallback: Option<T>) -> Result<T, Error>
+	where
+		T: Decode,
+	{
+		read_entry(&self.trie_backend, key, fallback).map_err(Error::ReadEntry)
+	}
+
+	/// Read an optional entry given by the key and try to decode it.
+	///
+	/// Returns `Err` in case the backend can't return the value under the specific key (likely due
+	/// to a malformed proof) or if the value couldn't be decoded.
+	pub fn read_optional_entry<T>(&self, key: &[u8]) -> Result<Option<T>, Error>
+	where
+		T: Decode,
+	{
+		read_optional_entry(&self.trie_backend, key).map_err(Error::ReadOptionalEntry)
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/teyrchain_inherent.rs

@@ -0,0 +1,538 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Pezcumulus teyrchain inherent related structures.
+
+use alloc::{collections::btree_map::BTreeMap, vec, vec::Vec};
+use core::fmt::Debug;
+use pezcumulus_primitives_core::{
+	relay_chain::{
+		ApprovedPeerId, BlockNumber as RelayChainBlockNumber, BlockNumber, Header as RelayHeader,
+	},
+	InboundDownwardMessage, InboundHrmpMessage, ParaId, PersistedValidationData,
+};
+use pezcumulus_primitives_teyrchain_inherent::{HashedMessage, TeyrchainInherentData};
+use pezframe_support::{
+	defensive,
+	pezpallet_prelude::{Decode, DecodeWithMemTracking, Encode},
+};
+use scale_info::TypeInfo;
+use pezsp_core::{bounded::BoundedSlice, Get};
+
+/// A structure that helps identify a message inside a collection of messages sorted by `sent_at`.
+///
+/// This structure contains a `sent_at` field and a reverse index. Using this information, we can
+/// identify a message inside a sorted collection by walking back `reverse_idx` positions starting
+/// from the last message that has the provided `sent_at`.
+///
+/// We use a reverse index instead of a normal index because sometimes the messages at the
+/// beginning of the collection are being pruned.
+///
+/// # Example
+///
+///
+/// For the collection
+/// `msgs = [{sent_at: 0}, {sent_at: 1}, {sent_at: 1}, {sent_at: 1}, {sent_at: 1}, {sent_at: 3}]`
+///
+/// `InboundMessageId {sent_at: 1, reverse_idx: 0}` points to `msgs[4]`
+/// `InboundMessageId {sent_at: 1, reverse_idx: 3}` points to `msgs[1]`
+/// `InboundMessageId {sent_at: 1, reverse_idx: 4}` points to `msgs[0]`
+#[derive(
+	Encode,
+	Decode,
+	DecodeWithMemTracking,
+	Clone,
+	Default,
+	pezsp_runtime::RuntimeDebug,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct InboundMessageId {
+	/// The block number at which this message was added to the message passing queue
+	/// on the relay chain.
+	pub sent_at: BlockNumber,
+	/// The reverse index of the message in the collection of messages sent at `sent_at`.
+	pub reverse_idx: u32,
+}
+
+/// A message that was received by the teyrchain.
+pub trait InboundMessage {
+	/// The corresponding compressed message.
+	/// This should be an equivalent message that stores the same metadata as the current message,
+	/// but stores only a hash of the message data.
+	type CompressedMessage: Debug;
+
+	/// Gets the message data.
+	fn data(&self) -> &[u8];
+
+	/// Gets the relay chain number where the current message was pushed to the corresponding
+	/// relay chain queue.
+	fn sent_at(&self) -> RelayChainBlockNumber;
+
+	/// Converts the current message into a `CompressedMessage`
+	fn to_compressed(&self) -> Self::CompressedMessage;
+}
+
+/// A collection of inbound messages.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct InboundMessagesCollection<Message: InboundMessage> {
+	messages: Vec<Message>,
+}
+
+impl<Message: InboundMessage> InboundMessagesCollection<Message> {
+	/// Creates a new instance of `InboundMessagesCollection` that contains the provided `messages`.
+	pub fn new(messages: Vec<Message>) -> Self {
+		Self { messages }
+	}
+
+	/// Drop all the messages up to `last_processed_msg`.
+	pub fn drop_processed_messages(&mut self, last_processed_msg: &InboundMessageId) {
+		let mut last_processed_msg_idx = None;
+		let messages = &mut self.messages;
+		for (idx, message) in messages.iter().enumerate().rev() {
+			let sent_at = message.sent_at();
+			if sent_at == last_processed_msg.sent_at {
+				last_processed_msg_idx = idx.checked_sub(last_processed_msg.reverse_idx as usize);
+				break;
+			}
+			// If we build on the same relay parent twice, we will receive the same messages again
+			// while `last_processed_msg` may have been increased. We need this check to make sure
+			// that the old messages are dropped.
+			if sent_at < last_processed_msg.sent_at {
+				last_processed_msg_idx = Some(idx);
+				break;
+			}
+		}
+		if let Some(last_processed_msg_idx) = last_processed_msg_idx {
+			messages.drain(..=last_processed_msg_idx);
+		}
+	}
+
+	/// Converts `self` into an [`AbridgedInboundMessagesCollection`].
+	///
+	/// The first messages in `self` (up to the provided `size_limit`) are kept in their current
+	/// form (they will contain the full message data).
+	/// The messages that exceed that limit are hashed.
+	pub fn into_abridged(
+		self,
+		size_limit: &mut usize,
+	) -> AbridgedInboundMessagesCollection<Message> {
+		let mut messages = self.messages;
+
+		let mut split_off_pos = messages.len();
+		for (idx, message) in messages.iter().enumerate() {
+			if *size_limit < message.data().len() {
+				break;
+			}
+			*size_limit -= message.data().len();
+
+			split_off_pos = idx + 1;
+		}
+
+		let extra_messages = messages.split_off(split_off_pos);
+		let hashed_messages = extra_messages.iter().map(|msg| msg.to_compressed()).collect();
+
+		AbridgedInboundMessagesCollection { full_messages: messages, hashed_messages }
+	}
+}
+
+/// A compressed collection of inbound messages.
+///
+/// The first messages in the collection (up to a limit) contain the full message data.
+/// The messages that exceed that limit are hashed.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct AbridgedInboundMessagesCollection<Message: InboundMessage> {
+	full_messages: Vec<Message>,
+	hashed_messages: Vec<Message::CompressedMessage>,
+}
+
+impl<Message: InboundMessage> AbridgedInboundMessagesCollection<Message> {
+	/// Gets a tuple containing both the full messages and the hashed messages
+	/// stored by the current collection.
+	pub fn messages(&self) -> (&[Message], &[Message::CompressedMessage]) {
+		(&self.full_messages, &self.hashed_messages)
+	}
+
+	/// Check that the current collection contains as many full messages as possible.
+	///
+	/// The `AbridgedInboundMessagesCollection` is provided to the runtime by a collator.
+	/// A malicious collator can provide a collection that contains no full messages or fewer
+	/// full messages than possible, leading to censorship.
+	pub fn check_enough_messages_included(&self, collection_name: &str) {
+		if self.hashed_messages.is_empty() {
+			return;
+		}
+
+		// Ideally, we should check that the collection contains as many full messages as possible
+		// without exceeding the max expected size. The worst case scenario is that were the first
+		// message that had to be hashed is a max size message. So in this case, the min expected
+		// size would be `max_expected_size - max_msg_size`. However, there are multiple issues:
+		// 1. The max message size config can change while we still have to process messages with
+		//    the old max message size.
+		// 2. We can't access the max downward message size from the teyrchain runtime.
+		//
+		// So the safest approach is to check that there is at least 1 full message.
+		assert!(
+			self.full_messages.len() >= 1,
+			"[{}] Advancement rule violation: mandatory messages missing",
+			collection_name,
+		);
+	}
+}
+
+impl<Message: InboundMessage> Default for AbridgedInboundMessagesCollection<Message> {
+	fn default() -> Self {
+		Self { full_messages: vec![], hashed_messages: vec![] }
+	}
+}
+
+impl InboundMessage for InboundDownwardMessage<RelayChainBlockNumber> {
+	type CompressedMessage = HashedMessage;
+
+	fn data(&self) -> &[u8] {
+		&self.msg
+	}
+
+	fn sent_at(&self) -> RelayChainBlockNumber {
+		self.sent_at
+	}
+
+	fn to_compressed(&self) -> Self::CompressedMessage {
+		self.into()
+	}
+}
+
+pub type InboundDownwardMessages =
+	InboundMessagesCollection<InboundDownwardMessage<RelayChainBlockNumber>>;
+
+pub type AbridgedInboundDownwardMessages =
+	AbridgedInboundMessagesCollection<InboundDownwardMessage<RelayChainBlockNumber>>;
+
+impl AbridgedInboundDownwardMessages {
+	/// Returns an iterator over the messages that maps them to `BoundedSlices`.
+	pub fn bounded_msgs_iter<MaxMessageLen: Get<u32>>(
+		&self,
+	) -> impl Iterator<Item = BoundedSlice<'_, u8, MaxMessageLen>> {
+		self.full_messages
+			.iter()
+			// Note: we are not using `.defensive()` here since that prints the whole value to
+			// console. In case that the message is too long, this clogs up the log quite badly.
+			.filter_map(|m| match BoundedSlice::try_from(&m.msg[..]) {
+				Ok(bounded) => Some(bounded),
+				Err(_) => {
+					defensive!("Inbound Downward message was too long; dropping");
+					None
+				},
+			})
+	}
+}
+
+impl InboundMessage for (ParaId, InboundHrmpMessage) {
+	type CompressedMessage = (ParaId, HashedMessage);
+
+	fn data(&self) -> &[u8] {
+		&self.1.data
+	}
+
+	fn sent_at(&self) -> RelayChainBlockNumber {
+		self.1.sent_at
+	}
+
+	fn to_compressed(&self) -> Self::CompressedMessage {
+		let (sender, message) = self;
+		(*sender, message.into())
+	}
+}
+
+pub type InboundHrmpMessages = InboundMessagesCollection<(ParaId, InboundHrmpMessage)>;
+
+impl InboundHrmpMessages {
+	// Prepare horizontal messages for a more convenient processing:
+	//
+	// Instead of a mapping from a para to a list of inbound HRMP messages, we will have a
+	// list of tuples `(sender, message)` first ordered by `sent_at` (the relay chain block
+	// number in which the message hit the relay-chain) and second ordered by para id
+	// ascending.
+	pub fn from_map(messages_map: BTreeMap<ParaId, Vec<InboundHrmpMessage>>) -> Self {
+		let mut messages = messages_map
+			.into_iter()
+			.flat_map(|(sender, channel_contents)| {
+				channel_contents.into_iter().map(move |message| (sender, message))
+			})
+			.collect::<Vec<_>>();
+		messages.sort_by(|(sender_a, msg_a), (sender_b, msg_b)| {
+			// first sort by sent-at and then by the para id
+			(msg_a.sent_at, sender_a).cmp(&(msg_b.sent_at, sender_b))
+		});
+
+		Self { messages }
+	}
+}
+
+pub type AbridgedInboundHrmpMessages =
+	AbridgedInboundMessagesCollection<(ParaId, InboundHrmpMessage)>;
+
+impl AbridgedInboundHrmpMessages {
+	/// Returns an iterator over the deconstructed messages.
+	pub fn flat_msgs_iter(&self) -> impl Iterator<Item = (ParaId, RelayChainBlockNumber, &[u8])> {
+		self.full_messages
+			.iter()
+			.map(|&(sender, ref message)| (sender, message.sent_at, &message.data[..]))
+	}
+}
+
+/// The basic inherent data that is passed by the collator to the teyrchain runtime.
+/// This data doesn't contain any messages.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct BasicTeyrchainInherentData {
+	pub validation_data: PersistedValidationData,
+	pub relay_chain_state: pezsp_trie::StorageProof,
+	pub relay_parent_descendants: Vec<RelayHeader>,
+	pub collator_peer_id: Option<ApprovedPeerId>,
+}
+
+/// The messages that are passed by the collator to the teyrchain runtime as part of the
+/// inherent data.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct InboundMessagesData {
+	pub downward_messages: AbridgedInboundDownwardMessages,
+	pub horizontal_messages: AbridgedInboundHrmpMessages,
+}
+
+impl InboundMessagesData {
+	/// Creates a new instance of `InboundMessagesData` with the provided messages.
+	pub fn new(
+		dmq_msgs: AbridgedInboundDownwardMessages,
+		hrmp_msgs: AbridgedInboundHrmpMessages,
+	) -> Self {
+		Self { downward_messages: dmq_msgs, horizontal_messages: hrmp_msgs }
+	}
+}
+
+/// Deconstructs a `TeyrchainInherentData` instance.
+pub fn deconstruct_teyrchain_inherent_data(
+	data: TeyrchainInherentData,
+) -> (BasicTeyrchainInherentData, InboundDownwardMessages, InboundHrmpMessages) {
+	(
+		BasicTeyrchainInherentData {
+			validation_data: data.validation_data,
+			relay_chain_state: data.relay_chain_state,
+			relay_parent_descendants: data.relay_parent_descendants,
+			collator_peer_id: data.collator_peer_id,
+		},
+		InboundDownwardMessages::new(data.downward_messages),
+		InboundHrmpMessages::from_map(data.horizontal_messages),
+	)
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	fn build_inbound_dm_vec(
+		info: &[(RelayChainBlockNumber, usize)],
+	) -> Vec<InboundDownwardMessage<RelayChainBlockNumber>> {
+		let mut messages = vec![];
+		for (sent_at, size) in info.iter() {
+			let data = vec![1; *size];
+			messages.push(InboundDownwardMessage { sent_at: *sent_at, msg: data })
+		}
+		messages
+	}
+
+	#[test]
+	fn drop_processed_messages_works() {
+		let msgs_vec =
+			build_inbound_dm_vec(&[(0, 0), (0, 0), (2, 0), (2, 0), (2, 0), (2, 0), (3, 0)]);
+
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 3, reverse_idx: 0 });
+		assert_eq!(msgs.messages, []);
+
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 2, reverse_idx: 0 });
+		assert_eq!(msgs.messages, msgs_vec[6..]);
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 2, reverse_idx: 1 });
+		assert_eq!(msgs.messages, msgs_vec[5..]);
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 2, reverse_idx: 4 });
+		assert_eq!(msgs.messages, msgs_vec[2..]);
+
+		// Go back starting from the last message sent at block 2, with 1 more message than the
+		// total number of messages sent at 2.
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 2, reverse_idx: 5 });
+		assert_eq!(msgs.messages, msgs_vec[1..]);
+
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 0, reverse_idx: 1 });
+		assert_eq!(msgs.messages, msgs_vec[1..]);
+		// Go back starting from the last message sent at block 0, with 1 more message than the
+		// total number of messages sent at 0.
+		let mut msgs = InboundDownwardMessages::new(msgs_vec.clone());
+		msgs.drop_processed_messages(&InboundMessageId { sent_at: 0, reverse_idx: 3 });
+		assert_eq!(msgs.messages, msgs_vec);
+	}
+
+	#[test]
+	fn into_abridged_works() {
+		let msgs = InboundDownwardMessages::new(vec![]);
+		let mut size_limit = 0;
+		let abridged_msgs = msgs.into_abridged(&mut size_limit);
+		assert_eq!(size_limit, 0);
+		assert_eq!(&abridged_msgs.full_messages, &vec![]);
+		assert_eq!(abridged_msgs.hashed_messages, vec![]);
+
+		let msgs_vec = build_inbound_dm_vec(&[(0, 100), (0, 100), (0, 150), (0, 50)]);
+		let msgs = InboundDownwardMessages::new(msgs_vec.clone());
+
+		let mut size_limit = 150;
+		let abridged_msgs = msgs.clone().into_abridged(&mut size_limit);
+		assert_eq!(size_limit, 50);
+		assert_eq!(&abridged_msgs.full_messages, &msgs_vec[..1]);
+		assert_eq!(
+			abridged_msgs.hashed_messages,
+			vec![(&msgs_vec[1]).into(), (&msgs_vec[2]).into(), (&msgs_vec[3]).into()]
+		);
+
+		let mut size_limit = 200;
+		let abridged_msgs = msgs.clone().into_abridged(&mut size_limit);
+		assert_eq!(size_limit, 0);
+		assert_eq!(&abridged_msgs.full_messages, &msgs_vec[..2]);
+		assert_eq!(
+			abridged_msgs.hashed_messages,
+			vec![(&msgs_vec[2]).into(), (&msgs_vec[3]).into()]
+		);
+
+		let mut size_limit = 399;
+		let abridged_msgs = msgs.clone().into_abridged(&mut size_limit);
+		assert_eq!(size_limit, 49);
+		assert_eq!(&abridged_msgs.full_messages, &msgs_vec[..3]);
+		assert_eq!(abridged_msgs.hashed_messages, vec![(&msgs_vec[3]).into()]);
+
+		let mut size_limit = 400;
+		let abridged_msgs = msgs.clone().into_abridged(&mut size_limit);
+		assert_eq!(size_limit, 0);
+		assert_eq!(&abridged_msgs.full_messages, &msgs_vec);
+		assert_eq!(abridged_msgs.hashed_messages, vec![]);
+	}
+
+	#[test]
+	fn from_map_works() {
+		let mut messages_map: BTreeMap<ParaId, Vec<InboundHrmpMessage>> = BTreeMap::new();
+		messages_map.insert(
+			1000.into(),
+			vec![
+				InboundHrmpMessage { sent_at: 0, data: vec![0] },
+				InboundHrmpMessage { sent_at: 0, data: vec![1] },
+				InboundHrmpMessage { sent_at: 1, data: vec![2] },
+			],
+		);
+		messages_map.insert(
+			2000.into(),
+			vec![
+				InboundHrmpMessage { sent_at: 0, data: vec![3] },
+				InboundHrmpMessage { sent_at: 0, data: vec![4] },
+				InboundHrmpMessage { sent_at: 1, data: vec![5] },
+			],
+		);
+		messages_map.insert(
+			3000.into(),
+			vec![
+				InboundHrmpMessage { sent_at: 0, data: vec![6] },
+				InboundHrmpMessage { sent_at: 1, data: vec![7] },
+				InboundHrmpMessage { sent_at: 2, data: vec![8] },
+				InboundHrmpMessage { sent_at: 3, data: vec![9] },
+				InboundHrmpMessage { sent_at: 4, data: vec![10] },
+			],
+		);
+
+		let msgs = InboundHrmpMessages::from_map(messages_map);
+		assert_eq!(
+			msgs.messages,
+			[
+				(1000.into(), InboundHrmpMessage { sent_at: 0, data: vec![0] }),
+				(1000.into(), InboundHrmpMessage { sent_at: 0, data: vec![1] }),
+				(2000.into(), InboundHrmpMessage { sent_at: 0, data: vec![3] }),
+				(2000.into(), InboundHrmpMessage { sent_at: 0, data: vec![4] }),
+				(3000.into(), InboundHrmpMessage { sent_at: 0, data: vec![6] }),
+				(1000.into(), InboundHrmpMessage { sent_at: 1, data: vec![2] }),
+				(2000.into(), InboundHrmpMessage { sent_at: 1, data: vec![5] }),
+				(3000.into(), InboundHrmpMessage { sent_at: 1, data: vec![7] }),
+				(3000.into(), InboundHrmpMessage { sent_at: 2, data: vec![8] }),
+				(3000.into(), InboundHrmpMessage { sent_at: 3, data: vec![9] }),
+				(3000.into(), InboundHrmpMessage { sent_at: 4, data: vec![10] })
+			]
+		)
+	}
+
+	#[test]
+	fn check_enough_messages_included_works() {
+		let mut messages = AbridgedInboundHrmpMessages {
+			full_messages: vec![(
+				1000.into(),
+				InboundHrmpMessage { sent_at: 0, data: vec![1; 100] },
+			)],
+			hashed_messages: vec![(
+				2000.into(),
+				HashedMessage { sent_at: 1, msg_hash: Default::default() },
+			)],
+		};
+
+		messages.check_enough_messages_included("Test");
+
+		messages.full_messages = vec![];
+		let result = std::panic::catch_unwind(|| messages.check_enough_messages_included("Test"));
+		assert!(result.is_err());
+
+		messages.hashed_messages = vec![];
+		messages.check_enough_messages_included("Test");
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/unincluded_segment.rs

@@ -0,0 +1,857 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Primitives used for tracking message queues constraints in an unincluded block segment
+//! of the teyrchain.
+//!
+//! Unincluded segment describes a chain of latest included block descendants, which are not yet
+//! sent to relay chain.
+
+use super::relay_state_snapshot::{MessagingStateSnapshot, RelayDispatchQueueRemainingCapacity};
+use alloc::collections::btree_map::BTreeMap;
+use codec::{Decode, Encode};
+use core::marker::PhantomData;
+use pezcumulus_primitives_core::{relay_chain, ParaId};
+use scale_info::TypeInfo;
+use pezsp_runtime::RuntimeDebug;
+
+/// Constraints on outbound HRMP channel.
+#[derive(Clone, RuntimeDebug)]
+pub struct HrmpOutboundLimits {
+	/// The maximum bytes that can be written to the channel.
+	pub bytes_remaining: u32,
+	/// The maximum messages that can be written to the channel.
+	pub messages_remaining: u32,
+}
+
+/// Limits on outbound message bandwidth.
+#[derive(Clone, RuntimeDebug)]
+pub struct OutboundBandwidthLimits {
+	/// The amount of UMP messages remaining.
+	pub ump_messages_remaining: u32,
+	/// The amount of UMP bytes remaining.
+	pub ump_bytes_remaining: u32,
+	/// The limitations of all registered outbound HRMP channels.
+	pub hrmp_outgoing: BTreeMap<ParaId, HrmpOutboundLimits>,
+}
+
+impl OutboundBandwidthLimits {
+	/// Creates new limits from the messaging state and upward message queue maximums fetched
+	/// from the host configuration.
+	///
+	/// These will be the total bandwidth limits across the entire unincluded segment.
+	pub fn from_relay_chain_state(messaging_state: &MessagingStateSnapshot) -> Self {
+		let RelayDispatchQueueRemainingCapacity { remaining_count, remaining_size } =
+			messaging_state.relay_dispatch_queue_remaining_capacity;
+
+		let hrmp_outgoing = messaging_state
+			.egress_channels
+			.iter()
+			.map(|(id, channel)| {
+				(
+					*id,
+					HrmpOutboundLimits {
+						bytes_remaining: channel.max_total_size.saturating_sub(channel.total_size),
+						messages_remaining: channel.max_capacity.saturating_sub(channel.msg_count),
+					},
+				)
+			})
+			.collect();
+
+		Self {
+			ump_messages_remaining: remaining_count,
+			ump_bytes_remaining: remaining_size,
+			hrmp_outgoing,
+		}
+	}
+}
+
+/// The error type for updating bandwidth used by a segment.
+#[derive(RuntimeDebug)]
+#[cfg_attr(test, derive(PartialEq))]
+pub enum BandwidthUpdateError {
+	/// Too many messages submitted to HRMP channel.
+	HrmpMessagesOverflow {
+		/// Teyrchain id of the recipient.
+		recipient: ParaId,
+		/// The amount of remaining messages in the capacity of the channel.
+		messages_remaining: u32,
+		/// The amount of messages submitted to the channel.
+		messages_submitted: u32,
+	},
+	/// Too many bytes submitted to HRMP channel.
+	HrmpBytesOverflow {
+		/// Teyrchain id of the recipient.
+		recipient: ParaId,
+		/// The amount of remaining bytes in the capacity of the channel.
+		bytes_remaining: u32,
+		/// The amount of bytes submitted to the channel.
+		bytes_submitted: u32,
+	},
+	/// Too many messages submitted to UMP queue.
+	UmpMessagesOverflow {
+		/// The amount of remaining messages in the capacity of UMP.
+		messages_remaining: u32,
+		/// The amount of messages submitted to UMP.
+		messages_submitted: u32,
+	},
+	/// Too many bytes submitted to UMP.
+	UmpBytesOverflow {
+		/// The amount of remaining bytes in the capacity of UMP.
+		bytes_remaining: u32,
+		/// The amount of bytes submitted to UMP.
+		bytes_submitted: u32,
+	},
+	/// Invalid HRMP watermark.
+	InvalidHrmpWatermark {
+		/// HRMP watermark submitted by the candidate.
+		submitted: relay_chain::BlockNumber,
+		/// Latest tracked HRMP watermark.
+		latest: relay_chain::BlockNumber,
+	},
+	/// Upgrade signal sent by relay chain was already processed by
+	/// some ancestor from the segment.
+	UpgradeGoAheadAlreadyProcessed,
+}
+
+/// The number of messages and size in bytes submitted to HRMP channel.
+#[derive(RuntimeDebug, Default, Copy, Clone, Encode, Decode, TypeInfo)]
+pub struct HrmpChannelUpdate {
+	/// The amount of messages submitted to the channel.
+	pub msg_count: u32,
+	/// The amount of bytes submitted to the channel.
+	pub total_bytes: u32,
+}
+
+impl HrmpChannelUpdate {
+	/// Returns `true` if the update is empty, `false` otherwise.
+	fn is_empty(&self) -> bool {
+		self.msg_count == 0 && self.total_bytes == 0
+	}
+
+	/// Tries to append another update, respecting given bandwidth limits.
+	fn append(
+		&self,
+		other: &Self,
+		recipient: ParaId,
+		limits: &OutboundBandwidthLimits,
+	) -> Result<Self, BandwidthUpdateError> {
+		let limits = limits
+			.hrmp_outgoing
+			.get(&recipient)
+			.expect("limit for declared hrmp channel must be present; qed");
+
+		let mut new = *self;
+
+		new.msg_count = new.msg_count.saturating_add(other.msg_count);
+		if new.msg_count > limits.messages_remaining {
+			return Err(BandwidthUpdateError::HrmpMessagesOverflow {
+				recipient,
+				messages_remaining: limits.messages_remaining,
+				messages_submitted: new.msg_count,
+			});
+		}
+		new.total_bytes = new.total_bytes.saturating_add(other.total_bytes);
+		if new.total_bytes > limits.bytes_remaining {
+			return Err(BandwidthUpdateError::HrmpBytesOverflow {
+				recipient,
+				bytes_remaining: limits.bytes_remaining,
+				bytes_submitted: new.total_bytes,
+			});
+		}
+
+		Ok(new)
+	}
+
+	/// Subtracts previously added channel update.
+	fn subtract(&mut self, other: &Self) {
+		self.msg_count -= other.msg_count;
+		self.total_bytes -= other.total_bytes;
+	}
+}
+
+/// Bandwidth used by a teyrchain block(s).
+///
+/// This struct can be created with pub items, however, it should
+/// never hit the storage directly to avoid bypassing limitations checks.
+#[derive(Default, Clone, Encode, Decode, TypeInfo, RuntimeDebug)]
+pub struct UsedBandwidth {
+	/// The amount of UMP messages sent.
+	pub ump_msg_count: u32,
+	/// The amount of UMP bytes sent.
+	pub ump_total_bytes: u32,
+	/// Outbound HRMP channels updates.
+	pub hrmp_outgoing: BTreeMap<ParaId, HrmpChannelUpdate>,
+}
+
+impl UsedBandwidth {
+	/// Tries to append another update, respecting given bandwidth limits.
+	fn append(
+		&self,
+		other: &Self,
+		limits: &OutboundBandwidthLimits,
+	) -> Result<Self, BandwidthUpdateError> {
+		let mut new = self.clone();
+
+		new.ump_msg_count = new.ump_msg_count.saturating_add(other.ump_msg_count);
+		if new.ump_msg_count > limits.ump_messages_remaining {
+			return Err(BandwidthUpdateError::UmpMessagesOverflow {
+				messages_remaining: limits.ump_messages_remaining,
+				messages_submitted: new.ump_msg_count,
+			});
+		}
+		new.ump_total_bytes = new.ump_total_bytes.saturating_add(other.ump_total_bytes);
+		if new.ump_total_bytes > limits.ump_bytes_remaining {
+			return Err(BandwidthUpdateError::UmpBytesOverflow {
+				bytes_remaining: limits.ump_bytes_remaining,
+				bytes_submitted: new.ump_total_bytes,
+			});
+		}
+
+		for (id, channel) in other.hrmp_outgoing.iter() {
+			let current = new.hrmp_outgoing.entry(*id).or_default();
+			*current = current.append(channel, *id, limits)?;
+		}
+
+		Ok(new)
+	}
+
+	/// Subtracts previously added bandwidth update.
+	fn subtract(&mut self, other: &Self) {
+		self.ump_msg_count -= other.ump_msg_count;
+		self.ump_total_bytes -= other.ump_total_bytes;
+
+		for (id, channel) in other.hrmp_outgoing.iter() {
+			let entry = self
+				.hrmp_outgoing
+				.get_mut(id)
+				.expect("entry's been inserted earlier with `append`; qed");
+			entry.subtract(channel);
+		}
+
+		self.hrmp_outgoing.retain(|_, channel| !channel.is_empty());
+	}
+}
+
+/// Ancestor of the block being currently executed, not yet included
+/// into the relay chain.
+#[derive(Encode, Decode, TypeInfo, RuntimeDebug)]
+pub struct Ancestor<H> {
+	/// Bandwidth used by this block.
+	used_bandwidth: UsedBandwidth,
+	/// Output head data hash of this block. This may be optional in case the head data has not
+	/// yet been posted on chain, but should be updated during initialization of the next block.
+	para_head_hash: Option<H>,
+	/// Optional go-ahead signal sent by the relay-chain this ancestor has processed.
+	consumed_go_ahead_signal: Option<relay_chain::UpgradeGoAhead>,
+}
+
+impl<H> Ancestor<H> {
+	/// Creates new ancestor without validating the bandwidth used.
+	pub fn new_unchecked(
+		used_bandwidth: UsedBandwidth,
+		consumed_go_ahead_signal: Option<relay_chain::UpgradeGoAhead>,
+	) -> Self {
+		Self { used_bandwidth, para_head_hash: None, consumed_go_ahead_signal }
+	}
+
+	/// Returns [`UsedBandwidth`] of this block.
+	pub fn used_bandwidth(&self) -> &UsedBandwidth {
+		&self.used_bandwidth
+	}
+
+	/// Returns hashed [output head data](`relay_chain::HeadData`) of this block.
+	pub fn para_head_hash(&self) -> Option<&H> {
+		self.para_head_hash.as_ref()
+	}
+
+	/// Set para head hash of this block.
+	pub fn replace_para_head_hash(&mut self, para_head_hash: H) {
+		self.para_head_hash.replace(para_head_hash);
+	}
+}
+
+/// An update to the HRMP watermark. This is always a relay-chain block number,
+/// but the two variants have different semantic meanings.
+pub enum HrmpWatermarkUpdate {
+	/// An update to the HRMP watermark where the new value is set to be equal to the
+	/// relay-parent's block number, i.e. the "head" of the relay chain.
+	/// This is always legal.
+	Head(relay_chain::BlockNumber),
+	/// An update to the HRMP watermark where the new value falls into the "trunk" of the
+	/// relay-chain. In this case, the watermark must be greater than the previous value.
+	Trunk(relay_chain::BlockNumber),
+}
+
+impl HrmpWatermarkUpdate {
+	/// Create a new update based on the desired watermark value and the current
+	/// relay-parent number.
+	pub fn new(
+		watermark: relay_chain::BlockNumber,
+		relay_parent_number: relay_chain::BlockNumber,
+	) -> Self {
+		// Hard constrain the watermark to the relay-parent number.
+		if watermark >= relay_parent_number {
+			HrmpWatermarkUpdate::Head(relay_parent_number)
+		} else {
+			HrmpWatermarkUpdate::Trunk(watermark)
+		}
+	}
+}
+
+/// Struct that keeps track of bandwidth used by the unincluded part of the chain
+/// along with the latest HRMP watermark.
+#[derive(Default, Encode, Decode, TypeInfo, RuntimeDebug)]
+pub struct SegmentTracker<H> {
+	/// Bandwidth used by the segment.
+	used_bandwidth: UsedBandwidth,
+	/// The mark which specifies the block number up to which all inbound HRMP messages are
+	/// processed.
+	hrmp_watermark: Option<relay_chain::BlockNumber>,
+	/// Optional go-ahead signal sent by the relay-chain some ancestor from the segment has
+	/// processed. Only single block is allowed to have this set within the whole segment.
+	consumed_go_ahead_signal: Option<relay_chain::UpgradeGoAhead>,
+	/// `H` is the type of para head hash.
+	phantom_data: PhantomData<H>,
+}
+
+impl<H> SegmentTracker<H> {
+	/// Tries to append another block to the tracker, respecting given bandwidth limits.
+	/// In practice, the bandwidth limits supplied should be the total allowed within the
+	/// block.
+	pub fn append(
+		&mut self,
+		block: &Ancestor<H>,
+		new_watermark: HrmpWatermarkUpdate,
+		limits: &OutboundBandwidthLimits,
+	) -> Result<(), BandwidthUpdateError> {
+		if self.consumed_go_ahead_signal.is_some() && block.consumed_go_ahead_signal.is_some() {
+			return Err(BandwidthUpdateError::UpgradeGoAheadAlreadyProcessed);
+		}
+		if let Some(watermark) = self.hrmp_watermark.as_ref() {
+			if let HrmpWatermarkUpdate::Trunk(new) = new_watermark {
+				if &new < watermark {
+					return Err(BandwidthUpdateError::InvalidHrmpWatermark {
+						submitted: new,
+						latest: *watermark,
+					});
+				}
+			}
+		}
+
+		self.used_bandwidth = self.used_bandwidth.append(block.used_bandwidth(), limits)?;
+
+		if let Some(consumed) = block.consumed_go_ahead_signal.as_ref() {
+			self.consumed_go_ahead_signal.replace(*consumed);
+		}
+		self.hrmp_watermark.replace(match new_watermark {
+			HrmpWatermarkUpdate::Trunk(w) | HrmpWatermarkUpdate::Head(w) => w,
+		});
+
+		Ok(())
+	}
+
+	/// Removes previously added block from the tracker.
+	pub fn subtract(&mut self, block: &Ancestor<H>) {
+		self.used_bandwidth.subtract(block.used_bandwidth());
+		if let Some(consumed) = block.consumed_go_ahead_signal.as_ref() {
+			// This is the same signal stored in the tracker.
+			let signal_in_segment = self.consumed_go_ahead_signal.take();
+			assert_eq!(signal_in_segment, Some(*consumed));
+		}
+		// Watermark doesn't need to be updated since the is always dropped
+		// from the tail of the segment.
+	}
+
+	/// Return a reference to the used bandwidth across the entire segment.
+	pub fn used_bandwidth(&self) -> &UsedBandwidth {
+		&self.used_bandwidth
+	}
+
+	/// Return go ahead signal consumed by some ancestor in a segment, if any.
+	pub fn consumed_go_ahead_signal(&self) -> Option<relay_chain::UpgradeGoAhead> {
+		self.consumed_go_ahead_signal
+	}
+}
+
+pub(crate) fn size_after_included<H: PartialEq>(included_hash: H, segment: &[Ancestor<H>]) -> u32 {
+	let pivot = segment
+		.iter()
+		.position(|ancestor| ancestor.para_head_hash() == Some(&included_hash))
+		.map(|p| p + 1)
+		.unwrap_or(0);
+
+	(segment.len() - pivot) as u32
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use alloc::{vec, vec::Vec};
+	use assert_matches::assert_matches;
+
+	#[test]
+	fn outbound_limits_constructed_correctly() {
+		let para_a = ParaId::from(0);
+		let para_a_channel = relay_chain::AbridgedHrmpChannel {
+			max_message_size: 15,
+
+			// Msg count capacity left is 2.
+			msg_count: 5,
+			max_capacity: 7,
+
+			// Bytes capacity left is 10.
+			total_size: 50,
+			max_total_size: 60,
+			mqc_head: None,
+		};
+
+		let para_b = ParaId::from(1);
+		let para_b_channel = relay_chain::AbridgedHrmpChannel {
+			max_message_size: 15,
+
+			// Msg count capacity left is 10.
+			msg_count: 40,
+			max_capacity: 50,
+
+			// Bytes capacity left is 0.
+			total_size: 500,
+			max_total_size: 500,
+			mqc_head: None,
+		};
+		let relay_dispatch_queue_remaining_capacity =
+			RelayDispatchQueueRemainingCapacity { remaining_count: 1, remaining_size: 50 };
+		let messaging_state = MessagingStateSnapshot {
+			dmq_mqc_head: relay_chain::Hash::zero(),
+			relay_dispatch_queue_remaining_capacity,
+			ingress_channels: Vec::new(),
+			egress_channels: vec![(para_a, para_a_channel), (para_b, para_b_channel)],
+		};
+
+		let limits = OutboundBandwidthLimits::from_relay_chain_state(&messaging_state);
+
+		// UMP.
+		assert_eq!(limits.ump_messages_remaining, 1);
+		assert_eq!(limits.ump_bytes_remaining, 50);
+
+		// HRMP.
+		let para_a_limits = limits.hrmp_outgoing.get(&para_a).expect("channel must be present");
+		let para_b_limits = limits.hrmp_outgoing.get(&para_b).expect("channel must be present");
+		assert_eq!(para_a_limits.bytes_remaining, 10);
+		assert_eq!(para_a_limits.messages_remaining, 2);
+		assert_eq!(para_b_limits.bytes_remaining, 0);
+		assert_eq!(para_b_limits.messages_remaining, 10);
+	}
+
+	#[test]
+	fn hrmp_msg_count_limits() {
+		let para_0 = ParaId::from(0);
+		let para_0_limits = HrmpOutboundLimits { bytes_remaining: u32::MAX, messages_remaining: 5 };
+
+		let para_1 = ParaId::from(1);
+		let para_1_limits = HrmpOutboundLimits { bytes_remaining: u32::MAX, messages_remaining: 3 };
+		let hrmp_outgoing = [(para_0, para_0_limits), (para_1, para_1_limits)].into();
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing,
+		};
+
+		let mut hrmp_update = HrmpChannelUpdate::default();
+		assert!(hrmp_update.is_empty());
+
+		for _ in 0..5 {
+			hrmp_update = hrmp_update
+				.append(&HrmpChannelUpdate { msg_count: 1, total_bytes: 10 }, para_0, &limits)
+				.expect("update is within the limits");
+		}
+		assert_matches!(
+			hrmp_update.append(
+				&HrmpChannelUpdate { msg_count: 1, total_bytes: 10 },
+				para_0,
+				&limits,
+			),
+			Err(BandwidthUpdateError::HrmpMessagesOverflow {
+				recipient,
+				messages_remaining,
+				messages_submitted,
+			}) if recipient == para_0 && messages_remaining == 5 && messages_submitted == 6
+		);
+
+		let mut hrmp_update = HrmpChannelUpdate::default();
+		hrmp_update = hrmp_update
+			.append(&HrmpChannelUpdate { msg_count: 2, total_bytes: 10 }, para_1, &limits)
+			.expect("update is within the limits");
+		assert_matches!(
+			hrmp_update.append(
+				&HrmpChannelUpdate { msg_count: 3, total_bytes: 10 },
+				para_1,
+				&limits,
+			),
+			Err(BandwidthUpdateError::HrmpMessagesOverflow {
+				recipient,
+				messages_remaining,
+				messages_submitted,
+			}) if recipient == para_1 && messages_remaining == 3 && messages_submitted == 5
+		);
+	}
+
+	#[test]
+	fn hrmp_bytes_limits() {
+		let para_0 = ParaId::from(0);
+		let para_0_limits =
+			HrmpOutboundLimits { bytes_remaining: 25, messages_remaining: u32::MAX };
+
+		let hrmp_outgoing = [(para_0, para_0_limits)].into();
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing,
+		};
+
+		let mut hrmp_update = HrmpChannelUpdate::default();
+		assert!(hrmp_update.is_empty());
+
+		for _ in 0..5 {
+			hrmp_update = hrmp_update
+				.append(&HrmpChannelUpdate { msg_count: 1, total_bytes: 4 }, para_0, &limits)
+				.expect("update is within the limits");
+		}
+		assert_matches!(
+			hrmp_update.append(
+				&HrmpChannelUpdate { msg_count: 1, total_bytes: 6 },
+				para_0,
+				&limits,
+			),
+			Err(BandwidthUpdateError::HrmpBytesOverflow {
+				recipient,
+				bytes_remaining,
+				bytes_submitted,
+			}) if recipient == para_0 && bytes_remaining == 25 && bytes_submitted == 26
+		);
+	}
+
+	#[test]
+	fn hrmp_limits_with_segment() {
+		let create_used_hrmp =
+			|hrmp_outgoing| UsedBandwidth { ump_msg_count: 0, ump_total_bytes: 0, hrmp_outgoing };
+
+		let para_0 = ParaId::from(0);
+		let para_0_limits = HrmpOutboundLimits { bytes_remaining: 30, messages_remaining: 10 };
+
+		let para_1 = ParaId::from(1);
+		let para_1_limits = HrmpOutboundLimits { bytes_remaining: 20, messages_remaining: 3 };
+		let hrmp_outgoing = [(para_0, para_0_limits), (para_1, para_1_limits)].into();
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing,
+		};
+
+		let mut segment = SegmentTracker::default();
+
+		let para_0_update = HrmpChannelUpdate { msg_count: 1, total_bytes: 6 };
+		let ancestor_0 = Ancestor {
+			used_bandwidth: create_used_hrmp([(para_0, para_0_update)].into()),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor_0, HrmpWatermarkUpdate::Trunk(0), &limits)
+			.expect("update is within the limits");
+
+		for watermark in 1..5 {
+			let ancestor = Ancestor {
+				used_bandwidth: create_used_hrmp([(para_0, para_0_update)].into()),
+				para_head_hash: None::<relay_chain::Hash>,
+				consumed_go_ahead_signal: None,
+			};
+			segment
+				.append(&ancestor, HrmpWatermarkUpdate::Trunk(watermark), &limits)
+				.expect("update is within the limits");
+		}
+
+		let para_0_update = HrmpChannelUpdate { msg_count: 1, total_bytes: 1 };
+		let ancestor_5 = Ancestor {
+			used_bandwidth: create_used_hrmp([(para_0, para_0_update)].into()),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		assert_matches!(
+			segment.append(&ancestor_5, HrmpWatermarkUpdate::Trunk(5), &limits),
+			Err(BandwidthUpdateError::HrmpBytesOverflow {
+				recipient,
+				bytes_remaining,
+				bytes_submitted,
+			}) if recipient == para_0 && bytes_remaining == 30 && bytes_submitted == 31
+		);
+		// Remove the first ancestor from the segment to make space.
+		segment.subtract(&ancestor_0);
+		segment
+			.append(&ancestor_5, HrmpWatermarkUpdate::Trunk(5), &limits)
+			.expect("update is within the limits");
+
+		let para_1_update = HrmpChannelUpdate { msg_count: 3, total_bytes: 10 };
+		let ancestor = Ancestor {
+			used_bandwidth: create_used_hrmp([(para_1, para_1_update)].into()),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor, HrmpWatermarkUpdate::Trunk(6), &limits)
+			.expect("update is within the limits");
+
+		assert_matches!(
+			segment.append(&ancestor, HrmpWatermarkUpdate::Trunk(7), &limits),
+			Err(BandwidthUpdateError::HrmpMessagesOverflow {
+				recipient,
+				messages_remaining,
+				messages_submitted,
+			}) if recipient == para_1 && messages_remaining == 3 && messages_submitted == 6
+		);
+	}
+
+	#[test]
+	fn ump_limits_with_segment() {
+		let create_used_ump = |(ump_msg_count, ump_total_bytes)| UsedBandwidth {
+			ump_msg_count,
+			ump_total_bytes,
+			hrmp_outgoing: BTreeMap::default(),
+		};
+
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 5,
+			ump_bytes_remaining: 50,
+			hrmp_outgoing: BTreeMap::default(),
+		};
+
+		let mut segment = SegmentTracker::default();
+
+		let ancestor_0 = Ancestor {
+			used_bandwidth: create_used_ump((1, 10)),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor_0, HrmpWatermarkUpdate::Trunk(0), &limits)
+			.expect("update is within the limits");
+
+		for watermark in 1..4 {
+			let ancestor = Ancestor {
+				used_bandwidth: create_used_ump((1, 10)),
+				para_head_hash: None::<relay_chain::Hash>,
+				consumed_go_ahead_signal: None,
+			};
+			segment
+				.append(&ancestor, HrmpWatermarkUpdate::Trunk(watermark), &limits)
+				.expect("update is within the limits");
+		}
+
+		let ancestor_4 = Ancestor {
+			used_bandwidth: create_used_ump((1, 30)),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		assert_matches!(
+			segment.append(&ancestor_4, HrmpWatermarkUpdate::Trunk(4), &limits),
+			Err(BandwidthUpdateError::UmpBytesOverflow {
+				bytes_remaining,
+				bytes_submitted,
+			}) if bytes_remaining == 50 && bytes_submitted == 70
+		);
+
+		let ancestor = Ancestor {
+			used_bandwidth: create_used_ump((1, 5)),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor, HrmpWatermarkUpdate::Trunk(4), &limits)
+			.expect("update is within the limits");
+		assert_matches!(
+			segment.append(&ancestor, HrmpWatermarkUpdate::Trunk(5), &limits),
+			Err(BandwidthUpdateError::UmpMessagesOverflow {
+				messages_remaining,
+				messages_submitted,
+			}) if messages_remaining == 5 && messages_submitted == 6
+		);
+	}
+
+	#[test]
+	fn segment_hrmp_watermark() {
+		let mut segment = SegmentTracker::default();
+
+		let ancestor = Ancestor {
+			used_bandwidth: UsedBandwidth::default(),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing: BTreeMap::default(),
+		};
+
+		segment
+			.append(&ancestor, HrmpWatermarkUpdate::Head(0), &limits)
+			.expect("nothing to compare the watermark with in default segment");
+		assert_matches!(segment.append(&ancestor, HrmpWatermarkUpdate::Trunk(0), &limits), Ok(()));
+
+		// Trunk updates are allowed when the new watermark is greater than the current one
+		for watermark in 1..5 {
+			segment
+				.append(&ancestor, HrmpWatermarkUpdate::Trunk(watermark), &limits)
+				.expect("hrmp watermark is valid");
+		}
+		// Trunk updates are not allowed when the new watermark is smaller than the current one
+		for watermark in 0..4 {
+			assert_eq!(
+				segment.append(&ancestor, HrmpWatermarkUpdate::Trunk(watermark), &limits),
+				Err(BandwidthUpdateError::InvalidHrmpWatermark { submitted: watermark, latest: 4 }),
+			);
+		}
+		// The current watermark is still valid.
+		segment
+			.append(&ancestor, HrmpWatermarkUpdate::Trunk(5), &limits)
+			.expect("hrmp watermark is valid");
+
+		// Head updates are allowed even if the new watermark is smaller than the current one
+		segment
+			.append(&ancestor, HrmpWatermarkUpdate::Head(4), &limits)
+			.expect("head updates always valid");
+	}
+
+	#[test]
+	fn segment_drops_empty_hrmp_channels() {
+		let create_used_hrmp =
+			|hrmp_outgoing| UsedBandwidth { ump_msg_count: 0, ump_total_bytes: 0, hrmp_outgoing };
+
+		let para_0 = ParaId::from(0);
+		let para_0_limits =
+			HrmpOutboundLimits { bytes_remaining: u32::MAX, messages_remaining: u32::MAX };
+
+		let para_1 = ParaId::from(1);
+		let para_1_limits =
+			HrmpOutboundLimits { bytes_remaining: u32::MAX, messages_remaining: u32::MAX };
+		let hrmp_outgoing = [(para_0, para_0_limits), (para_1, para_1_limits)].into();
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing,
+		};
+
+		let mut segment = SegmentTracker::default();
+
+		let para_0_update = HrmpChannelUpdate { msg_count: 1, total_bytes: 1 };
+		let ancestor_0 = Ancestor {
+			used_bandwidth: create_used_hrmp([(para_0, para_0_update)].into()),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor_0, HrmpWatermarkUpdate::Head(0), &limits)
+			.expect("update is within the limits");
+		let para_1_update = HrmpChannelUpdate { msg_count: 3, total_bytes: 10 };
+		let ancestor_1 = Ancestor {
+			used_bandwidth: create_used_hrmp([(para_1, para_1_update)].into()),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor_1, HrmpWatermarkUpdate::Head(1), &limits)
+			.expect("update is within the limits");
+
+		assert_eq!(segment.used_bandwidth.hrmp_outgoing.len(), 2);
+
+		segment.subtract(&ancestor_0);
+		assert_eq!(segment.used_bandwidth.hrmp_outgoing.len(), 1);
+
+		segment.subtract(&ancestor_1);
+		assert_eq!(segment.used_bandwidth.hrmp_outgoing.len(), 0);
+	}
+
+	#[test]
+	fn segment_go_ahead_signal_is_unique() {
+		let limits = OutboundBandwidthLimits {
+			ump_messages_remaining: 0,
+			ump_bytes_remaining: 0,
+			hrmp_outgoing: BTreeMap::default(),
+		};
+
+		let mut segment = SegmentTracker::default();
+
+		let ancestor_0 = Ancestor {
+			used_bandwidth: UsedBandwidth::default(),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: Some(relay_chain::UpgradeGoAhead::GoAhead),
+		};
+		segment
+			.append(&ancestor_0, HrmpWatermarkUpdate::Head(0), &limits)
+			.expect("update is within the limits");
+
+		let ancestor_1 = Ancestor {
+			used_bandwidth: UsedBandwidth::default(),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: None,
+		};
+		segment
+			.append(&ancestor_1, HrmpWatermarkUpdate::Head(1), &limits)
+			.expect("update is within the limits");
+
+		let ancestor_2 = Ancestor {
+			used_bandwidth: UsedBandwidth::default(),
+			para_head_hash: None::<relay_chain::Hash>,
+			consumed_go_ahead_signal: Some(relay_chain::UpgradeGoAhead::Abort),
+		};
+		assert_matches!(
+			segment.append(&ancestor_2, HrmpWatermarkUpdate::Head(2), &limits),
+			Err(BandwidthUpdateError::UpgradeGoAheadAlreadyProcessed)
+		);
+
+		segment.subtract(&ancestor_0);
+		segment
+			.append(&ancestor_2, HrmpWatermarkUpdate::Head(1), &limits)
+			.expect("update is within the limits");
+	}
+
+	#[test]
+	fn size_after_included_works() {
+		let segment = vec![
+			Ancestor {
+				used_bandwidth: Default::default(),
+				para_head_hash: Some("a"),
+				consumed_go_ahead_signal: None,
+			},
+			Ancestor {
+				used_bandwidth: Default::default(),
+				para_head_hash: Some("b"),
+				consumed_go_ahead_signal: None,
+			},
+			Ancestor {
+				used_bandwidth: Default::default(),
+				para_head_hash: Some("c"),
+				consumed_go_ahead_signal: None,
+			},
+		];
+
+		assert_eq!(size_after_included("a", &segment), 2,);
+		assert_eq!(size_after_included("b", &segment), 1,);
+		assert_eq!(size_after_included("c", &segment), 0,);
+		assert_eq!(size_after_included("d", &segment), 3,);
+
+		assert_eq!(size_after_included("x", &[]), 0,);
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/implementation.rs

@@ -0,0 +1,550 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! The actual implementation of the validate block functionality.
+
+use super::{trie_cache, trie_recorder, MemoryOptimizedValidationParams};
+use alloc::vec::Vec;
+use codec::{Decode, Encode};
+use pezcumulus_primitives_core::{
+	relay_chain::{BlockNumber as RNumber, Hash as RHash, UMPSignal, UMP_SEPARATOR},
+	ClaimQueueOffset, CoreSelector, PersistedValidationData, TeyrchainBlockData,
+};
+use pezframe_support::{
+	traits::{ExecuteBlock, Get, IsSubType},
+	BoundedVec,
+};
+use pezkuwi_teyrchain_primitives::primitives::{HeadData, ValidationResult};
+use pezsp_core::storage::{well_known_keys, ChildInfo, StateVersion};
+use pezsp_externalities::{set_and_run_with_externalities, Externalities};
+use pezsp_io::{hashing::blake2_128, KillStorageResult};
+use pezsp_runtime::traits::{
+	Block as BlockT, ExtrinsicCall, Hash as HashT, HashingFor, Header as HeaderT, LazyBlock,
+};
+use pezsp_state_machine::OverlayedChanges;
+use pezsp_trie::{HashDBT, ProofSizeProvider, EMPTY_PREFIX};
+use trie_recorder::{SeenNodes, SizeOnlyRecorderProvider};
+
+type Ext<'a, Block, Backend> = pezsp_state_machine::Ext<'a, HashingFor<Block>, Backend>;
+
+fn with_externalities<F: FnOnce(&mut dyn Externalities) -> R, R>(f: F) -> R {
+	pezsp_externalities::with_externalities(f).expect("Environmental externalities not set.")
+}
+
+// Recorder instance to be used during this validate_block call.
+environmental::environmental!(recorder: trait ProofSizeProvider);
+
+/// Validate the given teyrchain block.
+///
+/// This function is doing roughly the following:
+///
+/// 1. We decode the [`TeyrchainBlockData`] from the `block_data` in `params`.
+///
+/// 2. We are doing some security checks like checking that the `parent_head` in `params`
+/// is the parent of the block we are going to check. We also ensure that the `set_validation_data`
+/// inherent is present in the block and that the validation data matches the values in `params`.
+///
+/// 3. We construct the sparse in-memory database from the storage proof inside the block data and
+/// then ensure that the storage root matches the storage root in the `parent_head`.
+///
+/// 4. We replace all the storage related host functions with functions inside the wasm blob.
+/// This means instead of calling into the host, we will stay inside the wasm execution. This is
+/// very important as the relay chain validator hasn't the state required to verify the block. But
+/// we have the in-memory database that contains all the values from the state of the teyrchain
+/// that we require to verify the block.
+///
+/// 5. The last step is to execute the entire block in the machinery we just have setup. Executing
+/// the blocks include running all transactions in the block against our in-memory database and
+/// ensuring that the final storage root matches the storage root in the header of the block. In the
+/// end we return back the [`ValidationResult`] with all the required information for the validator.
+#[doc(hidden)]
+pub fn validate_block<B: BlockT, E: ExecuteBlock<B>, PSC: crate::Config>(
+	MemoryOptimizedValidationParams {
+		block_data,
+		parent_head: teyrchain_head,
+		relay_parent_number,
+		relay_parent_storage_root,
+	}: MemoryOptimizedValidationParams,
+) -> ValidationResult
+where
+	B::Extrinsic: ExtrinsicCall,
+	<B::Extrinsic as ExtrinsicCall>::Call: IsSubType<crate::Call<PSC>>,
+{
+	let _guard = (
+		// Replace storage calls with our own implementations
+		pezsp_io::storage::host_read.replace_implementation(host_storage_read),
+		pezsp_io::storage::host_set.replace_implementation(host_storage_set),
+		pezsp_io::storage::host_get.replace_implementation(host_storage_get),
+		pezsp_io::storage::host_exists.replace_implementation(host_storage_exists),
+		pezsp_io::storage::host_clear.replace_implementation(host_storage_clear),
+		pezsp_io::storage::host_root.replace_implementation(host_storage_root),
+		pezsp_io::storage::host_clear_prefix.replace_implementation(host_storage_clear_prefix),
+		pezsp_io::storage::host_append.replace_implementation(host_storage_append),
+		pezsp_io::storage::host_next_key.replace_implementation(host_storage_next_key),
+		pezsp_io::storage::host_start_transaction
+			.replace_implementation(host_storage_start_transaction),
+		pezsp_io::storage::host_rollback_transaction
+			.replace_implementation(host_storage_rollback_transaction),
+		pezsp_io::storage::host_commit_transaction
+			.replace_implementation(host_storage_commit_transaction),
+		pezsp_io::default_child_storage::host_get
+			.replace_implementation(host_default_child_storage_get),
+		pezsp_io::default_child_storage::host_read
+			.replace_implementation(host_default_child_storage_read),
+		pezsp_io::default_child_storage::host_set
+			.replace_implementation(host_default_child_storage_set),
+		pezsp_io::default_child_storage::host_clear
+			.replace_implementation(host_default_child_storage_clear),
+		pezsp_io::default_child_storage::host_storage_kill
+			.replace_implementation(host_default_child_storage_storage_kill),
+		pezsp_io::default_child_storage::host_exists
+			.replace_implementation(host_default_child_storage_exists),
+		pezsp_io::default_child_storage::host_clear_prefix
+			.replace_implementation(host_default_child_storage_clear_prefix),
+		pezsp_io::default_child_storage::host_root
+			.replace_implementation(host_default_child_storage_root),
+		pezsp_io::default_child_storage::host_next_key
+			.replace_implementation(host_default_child_storage_next_key),
+		pezsp_io::offchain_index::host_set.replace_implementation(host_offchain_index_set),
+		pezsp_io::offchain_index::host_clear.replace_implementation(host_offchain_index_clear),
+		pezcumulus_primitives_proof_size_hostfunction::storage_proof_size::host_storage_proof_size
+			.replace_implementation(host_storage_proof_size),
+	);
+
+	let block_data = codec::decode_from_bytes::<TeyrchainBlockData<B::LazyBlock>>(block_data)
+		.expect("Invalid teyrchain block data");
+
+	// Initialize hashmaps randomness.
+	pezsp_trie::add_extra_randomness(build_seed_from_head_data::<B>(
+		&block_data,
+		relay_parent_storage_root,
+	));
+
+	let mut parent_header =
+		codec::decode_from_bytes::<B::Header>(teyrchain_head.clone()).expect("Invalid parent head");
+
+	let (blocks, proof) = block_data.into_inner();
+
+	assert_eq!(
+		*blocks
+			.first()
+			.expect("BlockData should have at least one block")
+			.header()
+			.parent_hash(),
+		parent_header.hash(),
+		"Teyrchain head needs to be the parent of the first block"
+	);
+
+	blocks.iter().fold(parent_header.hash(), |p, b| {
+		assert_eq!(
+			p,
+			*b.header().parent_hash(),
+			"Not a valid chain of blocks :(; {:?} not a parent of {:?}?",
+			array_bytes::bytes2hex("0x", p.as_ref()),
+			array_bytes::bytes2hex("0x", b.header().parent_hash().as_ref()),
+		);
+		b.header().hash()
+	});
+
+	let mut processed_downward_messages = 0;
+	let mut upward_messages = BoundedVec::default();
+	let mut upward_message_signals = Vec::<Vec<_>>::new();
+	let mut horizontal_messages = BoundedVec::default();
+	let mut hrmp_watermark = Default::default();
+	let mut head_data = None;
+	let mut new_validation_code = None;
+	let num_blocks = blocks.len();
+
+	// Create the db
+	let mut db = match proof.to_memory_db(Some(parent_header.state_root())) {
+		Ok((db, _)) => db,
+		Err(_) => panic!("Compact proof decoding failure."),
+	};
+
+	core::mem::drop(proof);
+
+	let cache_provider = trie_cache::CacheProvider::new();
+	let seen_nodes = SeenNodes::<HashingFor<B>>::default();
+
+	for (block_index, mut block) in blocks.into_iter().enumerate() {
+		// We use the storage root of the `parent_head` to ensure that it is the correct root.
+		// This is already being done above while creating the in-memory db, but let's be paranoid!!
+		let backend = pezsp_state_machine::TrieBackendBuilder::new_with_cache(
+			&db,
+			*parent_header.state_root(),
+			&cache_provider,
+		)
+		.build();
+
+		// We use the same recorder when executing all blocks. So, each node only contributes once
+		// to the total size of the storage proof. This recorder should only be used for
+		// `execute_block`.
+		let mut execute_recorder = SizeOnlyRecorderProvider::with_seen_nodes(seen_nodes.clone());
+		// `backend` with the `execute_recorder`. As the `execute_recorder`, this should only be
+		// used for `execute_block`.
+		let execute_backend = pezsp_state_machine::TrieBackendBuilder::wrap(&backend)
+			.with_recorder(execute_recorder.clone())
+			.build();
+
+		// We let all blocks contribute to the same overlay. Data written by a previous block will
+		// be directly accessible without going to the db.
+		let mut overlay = OverlayedChanges::default();
+
+		parent_header = block.header().clone();
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&backend,
+			&mut Default::default(),
+			&mut Default::default(),
+			|| {
+				E::verify_and_remove_seal(&mut block);
+			},
+		);
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&execute_backend,
+			// Here is the only place where we want to use the recorder.
+			// We want to ensure that we not accidentally read something from the proof, that
+			// was not yet read and thus, alter the proof size. Otherwise, we end up with
+			// mismatches in later blocks.
+			&mut execute_recorder,
+			&mut overlay,
+			|| {
+				E::execute_verified_block(block);
+			},
+		);
+
+		if overlay.storage(well_known_keys::CODE).is_some() && num_blocks > 1 {
+			panic!("When applying a runtime upgrade, only one block per PoV is allowed. Received {num_blocks}.")
+		}
+
+		run_with_externalities_and_recorder::<B, _, _>(
+			&backend,
+			&mut Default::default(),
+			// We are only reading here, but need to know what the old block has written. Thus, we
+			// are passing here the overlay.
+			&mut overlay,
+			|| {
+				// Ensure the validation data is correct.
+				validate_validation_data(
+					crate::ValidationData::<PSC>::get()
+						.expect("`ValidationData` must be set after executing a block; qed"),
+					&teyrchain_head,
+					relay_parent_number,
+					relay_parent_storage_root,
+				);
+
+				new_validation_code =
+					new_validation_code.take().or(crate::NewValidationCode::<PSC>::get());
+
+				let mut found_separator = false;
+				crate::UpwardMessages::<PSC>::get()
+					.into_iter()
+					.filter_map(|m| {
+						// Filter out the `UMP_SEPARATOR` and the `UMPSignals`.
+						if m == UMP_SEPARATOR {
+							found_separator = true;
+							None
+						} else if found_separator {
+							if upward_message_signals.iter().all(|s| *s != m) {
+								upward_message_signals.push(m);
+							}
+							None
+						} else {
+							// No signal or separator
+							Some(m)
+						}
+					})
+					.for_each(|m| {
+						upward_messages.try_push(m)
+							.expect(
+								"Number of upward messages should not be greater than `MAX_UPWARD_MESSAGE_NUM`",
+							)
+					});
+
+				processed_downward_messages += crate::ProcessedDownwardMessages::<PSC>::get();
+				horizontal_messages.try_extend(crate::HrmpOutboundMessages::<PSC>::get().into_iter()).expect(
+					"Number of horizontal messages should not be greater than `MAX_HORIZONTAL_MESSAGE_NUM`",
+				);
+				hrmp_watermark = crate::HrmpWatermark::<PSC>::get();
+
+				if block_index + 1 == num_blocks {
+					head_data = Some(
+						crate::CustomValidationHeadData::<PSC>::get()
+							.map_or_else(|| HeadData(parent_header.encode()), HeadData),
+					);
+				}
+			},
+		);
+
+		if block_index + 1 != num_blocks {
+			let mut changes = overlay
+				.drain_storage_changes(
+					&backend,
+					<PSC as pezframe_system::Config>::Version::get().state_version(),
+				)
+				.expect("Failed to get drain storage changes from the overlay.");
+
+			drop(backend);
+
+			// We just forward the changes directly to our db.
+			changes.transaction.drain().into_iter().for_each(|(_, (value, count))| {
+				// We only care about inserts and not deletes.
+				if count > 0 {
+					db.insert(EMPTY_PREFIX, &value);
+
+					let hash = HashingFor::<B>::hash(&value);
+					seen_nodes.borrow_mut().insert(hash);
+				}
+			});
+		}
+	}
+
+	if !upward_message_signals.is_empty() {
+		let mut selected_core: Option<(CoreSelector, ClaimQueueOffset)> = None;
+		let mut approved_peer = None;
+
+		upward_message_signals.iter().for_each(|s| {
+			match UMPSignal::decode(&mut &s[..]).expect("Failed to decode `UMPSignal`") {
+				UMPSignal::SelectCore(selector, offset) => match &selected_core {
+					Some(selected_core) if *selected_core != (selector, offset) => {
+						panic!(
+							"All `SelectCore` signals need to select the same core: {selected_core:?} vs {:?}",
+							(selector, offset),
+						)
+					},
+					Some(_) => {},
+					None => {
+						selected_core = Some((selector, offset));
+					},
+				},
+				UMPSignal::ApprovedPeer(new_approved_peer) => match &approved_peer {
+					Some(approved_peer) if *approved_peer != new_approved_peer => {
+						panic!(
+							"All `ApprovedPeer` signals need to select the same peer_id: {new_approved_peer:?} vs {approved_peer:?}",
+						)
+					},
+					Some(_) => {},
+					None => {
+						approved_peer = Some(new_approved_peer);
+					},
+				},
+			}
+		});
+
+		upward_messages
+			.try_push(UMP_SEPARATOR)
+			.expect("UMPSignals does not fit in UMPMessages");
+		upward_messages
+			.try_extend(upward_message_signals.into_iter())
+			.expect("UMPSignals does not fit in UMPMessages");
+	}
+
+	ValidationResult {
+		head_data: head_data.expect("HeadData not set"),
+		new_validation_code: new_validation_code.map(Into::into),
+		upward_messages,
+		processed_downward_messages,
+		horizontal_messages,
+		hrmp_watermark,
+	}
+}
+
+/// Validates the given [`PersistedValidationData`] against the data from the relay chain.
+fn validate_validation_data(
+	validation_data: PersistedValidationData,
+	parent_header: &[u8],
+	relay_parent_number: RNumber,
+	relay_parent_storage_root: RHash,
+) {
+	assert_eq!(parent_header, &validation_data.parent_head.0, "Parent head doesn't match");
+	assert_eq!(
+		relay_parent_number, validation_data.relay_parent_number,
+		"Relay parent number doesn't match",
+	);
+	assert_eq!(
+		relay_parent_storage_root, validation_data.relay_parent_storage_root,
+		"Relay parent storage root doesn't match",
+	);
+}
+
+/// Build a seed from the head data of the teyrchain block.
+///
+/// Uses both the relay parent storage root and the hash of the blocks
+/// in the block data, to make sure the seed changes every block and that
+/// the user cannot find about it ahead of time.
+fn build_seed_from_head_data<B: BlockT>(
+	block_data: &TeyrchainBlockData<B::LazyBlock>,
+	relay_parent_storage_root: crate::relay_chain::Hash,
+) -> [u8; 16] {
+	let mut bytes_to_hash = Vec::with_capacity(
+		block_data.blocks().len() * size_of::<B::Hash>() + size_of::<crate::relay_chain::Hash>(),
+	);
+
+	bytes_to_hash.extend_from_slice(relay_parent_storage_root.as_ref());
+	block_data.blocks().iter().for_each(|block| {
+		bytes_to_hash.extend_from_slice(block.header().hash().as_ref());
+	});
+
+	blake2_128(&bytes_to_hash)
+}
+
+/// Run the given closure with the externalities and recorder set.
+fn run_with_externalities_and_recorder<Block: BlockT, R, F: FnOnce() -> R>(
+	backend: &impl pezsp_state_machine::Backend<HashingFor<Block>>,
+	recorder: &mut SizeOnlyRecorderProvider<HashingFor<Block>>,
+	overlay: &mut OverlayedChanges<HashingFor<Block>>,
+	execute: F,
+) -> R {
+	let mut ext = Ext::<Block, _>::new(overlay, backend);
+
+	recorder::using(recorder, || set_and_run_with_externalities(&mut ext, || execute()))
+}
+
+fn host_storage_read(key: &[u8], value_out: &mut [u8], value_offset: u32) -> Option<u32> {
+	match with_externalities(|ext| ext.storage(key)) {
+		Some(value) => {
+			let value_offset = value_offset as usize;
+			let data = &value[value_offset.min(value.len())..];
+			let written = core::cmp::min(data.len(), value_out.len());
+			value_out[..written].copy_from_slice(&data[..written]);
+			Some(value.len() as u32)
+		},
+		None => None,
+	}
+}
+
+fn host_storage_set(key: &[u8], value: &[u8]) {
+	with_externalities(|ext| ext.place_storage(key.to_vec(), Some(value.to_vec())))
+}
+
+fn host_storage_get(key: &[u8]) -> Option<bytes::Bytes> {
+	with_externalities(|ext| ext.storage(key).map(|value| value.into()))
+}
+
+fn host_storage_exists(key: &[u8]) -> bool {
+	with_externalities(|ext| ext.exists_storage(key))
+}
+
+fn host_storage_clear(key: &[u8]) {
+	with_externalities(|ext| ext.place_storage(key.to_vec(), None))
+}
+
+fn host_storage_proof_size() -> u64 {
+	recorder::with(|rec| rec.estimate_encoded_size()).expect("Recorder is always set; qed") as _
+}
+
+fn host_storage_root(version: StateVersion) -> Vec<u8> {
+	with_externalities(|ext| ext.storage_root(version))
+}
+
+fn host_storage_clear_prefix(prefix: &[u8], limit: Option<u32>) -> KillStorageResult {
+	with_externalities(|ext| ext.clear_prefix(prefix, limit, None).into())
+}
+
+fn host_storage_append(key: &[u8], value: Vec<u8>) {
+	with_externalities(|ext| ext.storage_append(key.to_vec(), value))
+}
+
+fn host_storage_next_key(key: &[u8]) -> Option<Vec<u8>> {
+	with_externalities(|ext| ext.next_storage_key(key))
+}
+
+fn host_storage_start_transaction() {
+	with_externalities(|ext| ext.storage_start_transaction())
+}
+
+fn host_storage_rollback_transaction() {
+	with_externalities(|ext| ext.storage_rollback_transaction().ok())
+		.expect("No open transaction that can be rolled back.");
+}
+
+fn host_storage_commit_transaction() {
+	with_externalities(|ext| ext.storage_commit_transaction().ok())
+		.expect("No open transaction that can be committed.");
+}
+
+fn host_default_child_storage_get(storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.child_storage(&child_info, key))
+}
+
+fn host_default_child_storage_read(
+	storage_key: &[u8],
+	key: &[u8],
+	value_out: &mut [u8],
+	value_offset: u32,
+) -> Option<u32> {
+	let child_info = ChildInfo::new_default(storage_key);
+	match with_externalities(|ext| ext.child_storage(&child_info, key)) {
+		Some(value) => {
+			let value_offset = value_offset as usize;
+			let data = &value[value_offset.min(value.len())..];
+			let written = core::cmp::min(data.len(), value_out.len());
+			value_out[..written].copy_from_slice(&data[..written]);
+			Some(value.len() as u32)
+		},
+		None => None,
+	}
+}
+
+fn host_default_child_storage_set(storage_key: &[u8], key: &[u8], value: &[u8]) {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| {
+		ext.place_child_storage(&child_info, key.to_vec(), Some(value.to_vec()))
+	})
+}
+
+fn host_default_child_storage_clear(storage_key: &[u8], key: &[u8]) {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.place_child_storage(&child_info, key.to_vec(), None))
+}
+
+fn host_default_child_storage_storage_kill(
+	storage_key: &[u8],
+	limit: Option<u32>,
+) -> KillStorageResult {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.kill_child_storage(&child_info, limit, None).into())
+}
+
+fn host_default_child_storage_exists(storage_key: &[u8], key: &[u8]) -> bool {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.exists_child_storage(&child_info, key))
+}
+
+fn host_default_child_storage_clear_prefix(
+	storage_key: &[u8],
+	prefix: &[u8],
+	limit: Option<u32>,
+) -> KillStorageResult {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.clear_child_prefix(&child_info, prefix, limit, None).into())
+}
+
+fn host_default_child_storage_root(storage_key: &[u8], version: StateVersion) -> Vec<u8> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.child_storage_root(&child_info, version))
+}
+
+fn host_default_child_storage_next_key(storage_key: &[u8], key: &[u8]) -> Option<Vec<u8>> {
+	let child_info = ChildInfo::new_default(storage_key);
+	with_externalities(|ext| ext.next_child_storage_key(&child_info, key))
+}
+
+fn host_offchain_index_set(_key: &[u8], _value: &[u8]) {}
+
+fn host_offchain_index_clear(_key: &[u8]) {}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/mod.rs

@@ -0,0 +1,70 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! A module that enables a runtime to work as teyrchain.
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub mod implementation;
+#[cfg(test)]
+mod tests;
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub mod trie_cache;
+
+#[cfg(any(test, not(feature = "std")))]
+#[doc(hidden)]
+pub mod trie_recorder;
+
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use alloc::{boxed::Box, slice};
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use bytes;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use codec::decode_from_bytes;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezkuwi_teyrchain_primitives;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezsp_runtime::traits::GetRuntimeBlockType;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use pezsp_std;
+
+/// Basically the same as
+/// [`ValidationParams`](pezkuwi_teyrchain_primitives::primitives::ValidationParams), but a little
+/// bit optimized for our use case here.
+///
+/// `block_data` and `head_data` are represented as [`bytes::Bytes`] to make them reuse
+/// the memory of the input parameter of the exported `validate_blocks` function.
+///
+/// The layout of this type must match exactly the layout of
+/// [`ValidationParams`](pezkuwi_teyrchain_primitives::primitives::ValidationParams) to have the
+/// same SCALE encoding.
+#[derive(codec::Decode)]
+#[cfg_attr(feature = "std", derive(codec::Encode))]
+#[doc(hidden)]
+pub struct MemoryOptimizedValidationParams {
+	pub parent_head: bytes::Bytes,
+	pub block_data: bytes::Bytes,
+	pub relay_parent_number: pezcumulus_primitives_core::relay_chain::BlockNumber,
+	pub relay_parent_storage_root: pezcumulus_primitives_core::relay_chain::Hash,
+}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/tests.rs

@@ -0,0 +1,777 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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 crate::{validate_block::MemoryOptimizedValidationParams, *};
+use codec::{Decode, DecodeAll, Encode};
+use pezcumulus_primitives_core::{relay_chain, PersistedValidationData, TeyrchainBlockData};
+use pezcumulus_test_client::{
+	generate_extrinsic, generate_extrinsic_with_pair,
+	runtime::{
+		self as test_runtime, Block, Hash, Header, SudoCall, SystemCall, TestPalletCall,
+		UncheckedExtrinsic, WASM_BINARY,
+	},
+	seal_block, transfer, BlockData, BlockOrigin, BuildTeyrchainBlockData, Client,
+	DefaultTestClientBuilderExt, HeadData, InitBlockBuilder,
+	Sr25519Keyring::{Alice, Bob, Charlie},
+	TestClientBuilder, TestClientBuilderExt, ValidationParams,
+};
+use pezcumulus_test_relay_sproof_builder::RelayStateSproofBuilder;
+use pezkuwi_teyrchain_primitives::primitives::ValidationResult;
+use pezsc_consensus::{BlockImport, BlockImportParams, ForkChoiceStrategy};
+use pezsp_api::{ApiExt, Core, ProofRecorder, ProvideRuntimeApi};
+use pezsp_consensus_slots::SlotDuration;
+use pezsp_core::{Hasher, H256};
+use pezsp_runtime::{
+	traits::{BlakeTwo256, Block as BlockT, Header as HeaderT},
+	DigestItem,
+};
+use pezsp_trie::{proof_size_extension::ProofSizeExt, recorder::IgnoredNodes, StorageProof};
+use std::{env, process::Command};
+
+fn call_validate_block_validation_result(
+	validation_code: &[u8],
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<ValidationResult> {
+	pezcumulus_test_client::validate_block(
+		ValidationParams {
+			block_data: BlockData(block_data.encode()),
+			parent_head: HeadData(parent_head.encode()),
+			relay_parent_number: 1,
+			relay_parent_storage_root,
+		},
+		validation_code,
+	)
+}
+
+fn call_validate_block(
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<Header> {
+	call_validate_block_validation_result(
+		WASM_BINARY.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block_data,
+		relay_parent_storage_root,
+	)
+	.map(|v| Header::decode(&mut &v.head_data.0[..]).expect("Decodes `Header`."))
+}
+
+/// Call `validate_block` in the runtime with `elastic-scaling` activated.
+fn call_validate_block_elastic_scaling(
+	parent_head: Header,
+	block_data: TeyrchainBlockData<Block>,
+	relay_parent_storage_root: Hash,
+) -> pezcumulus_test_client::ExecutorResult<Header> {
+	call_validate_block_validation_result(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block_data,
+		relay_parent_storage_root,
+	)
+	.map(|v| Header::decode(&mut &v.head_data.0[..]).expect("Decodes `Header`."))
+}
+
+fn create_test_client() -> (Client, Header) {
+	let client = TestClientBuilder::new().enable_import_proof_recording().build();
+
+	let genesis_header = client
+		.header(client.chain_info().genesis_hash)
+		.ok()
+		.flatten()
+		.expect("Genesis header exists; qed");
+
+	(client, genesis_header)
+}
+
+/// Create test client using the runtime with `elastic-scaling` feature enabled.
+fn create_elastic_scaling_test_client() -> (Client, Header) {
+	let mut builder = TestClientBuilder::new();
+	builder.genesis_init_mut().wasm = Some(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!")
+			.to_vec(),
+	);
+	let client = builder.enable_import_proof_recording().build();
+
+	let genesis_header = client
+		.header(client.chain_info().genesis_hash)
+		.ok()
+		.flatten()
+		.expect("Genesis header exists; qed");
+
+	(client, genesis_header)
+}
+
+fn pop_seal(mut block: Block) -> Block {
+	assert!(block.header.digest.pop().unwrap().as_seal().is_some());
+	block
+}
+
+struct TestBlockData {
+	block: TeyrchainBlockData<Block>,
+	validation_data: PersistedValidationData,
+}
+
+fn build_block_with_witness(
+	client: &Client,
+	extra_extrinsics: Vec<UncheckedExtrinsic>,
+	parent_head: Header,
+	mut sproof_builder: RelayStateSproofBuilder,
+	pre_digests: Vec<DigestItem>,
+) -> TestBlockData {
+	sproof_builder.para_id = test_runtime::TEYRCHAIN_ID.into();
+	sproof_builder.included_para_head = Some(HeadData(parent_head.encode()));
+
+	let validation_data = PersistedValidationData {
+		relay_parent_number: 1,
+		parent_head: parent_head.encode().into(),
+		..Default::default()
+	};
+
+	let pezcumulus_test_client::BlockBuilderAndSupportData {
+		mut block_builder,
+		persisted_validation_data,
+	} = client.init_block_builder_with_pre_digests(Some(validation_data), sproof_builder, pre_digests);
+
+	extra_extrinsics.into_iter().for_each(|e| block_builder.push(e).unwrap());
+
+	let mut block = block_builder.build_teyrchain_block(*parent_head.state_root());
+
+	block.blocks_mut()[0] = seal_block(block.blocks()[0].clone(), client);
+
+	TestBlockData { block, validation_data: persisted_validation_data }
+}
+
+fn build_multiple_blocks_with_witness(
+	client: &Client,
+	mut parent_head: Header,
+	mut sproof_builder: RelayStateSproofBuilder,
+	num_blocks: u32,
+	extra_extrinsics: impl Fn(u32) -> Vec<UncheckedExtrinsic>,
+) -> TestBlockData {
+	let parent_head_root = *parent_head.state_root();
+	sproof_builder.para_id = test_runtime::TEYRCHAIN_ID.into();
+	sproof_builder.included_para_head = Some(HeadData(parent_head.encode()));
+
+	let timestamp = if sproof_builder.current_slot == 0u64 {
+		let timestamp = std::time::SystemTime::now()
+			.duration_since(std::time::SystemTime::UNIX_EPOCH)
+			.expect("Time is always after UNIX_EPOCH; qed")
+			.as_millis() as u64;
+		sproof_builder.current_slot = (timestamp / 6000).into();
+
+		timestamp
+	} else {
+		sproof_builder
+			.current_slot
+			.timestamp(SlotDuration::from_millis(6000))
+			.unwrap()
+			.as_millis()
+	};
+
+	let validation_data = PersistedValidationData {
+		relay_parent_number: 1,
+		parent_head: parent_head.encode().into(),
+		..Default::default()
+	};
+
+	let mut persisted_validation_data = None;
+	let mut blocks = Vec::new();
+	let mut proof = StorageProof::empty();
+	let mut ignored_nodes = IgnoredNodes::<H256>::default();
+
+	for i in 0..num_blocks {
+		let pezcumulus_test_client::BlockBuilderAndSupportData {
+			mut block_builder,
+			persisted_validation_data: p_v_data,
+		} = client.init_block_builder_with_ignored_nodes(
+			parent_head.hash(),
+			Some(validation_data.clone()),
+			sproof_builder.clone(),
+			timestamp,
+			ignored_nodes.clone(),
+		);
+
+		persisted_validation_data = Some(p_v_data);
+
+		for ext in (extra_extrinsics)(i) {
+			block_builder.push(ext).unwrap();
+		}
+
+		let mut built_block = block_builder.build().unwrap();
+		built_block.block = seal_block(built_block.block, &client);
+
+		futures::executor::block_on({
+			let parent_hash = *built_block.block.header.parent_hash();
+			let state = client.state_at(parent_hash).unwrap();
+
+			let mut api = client.runtime_api();
+			let proof_recorder = ProofRecorder::<Block>::with_ignored_nodes(ignored_nodes.clone());
+			api.record_proof_with_recorder(proof_recorder.clone());
+			api.register_extension(ProofSizeExt::new(proof_recorder));
+			api.execute_block(parent_hash, pop_seal(built_block.block.clone()).into())
+				.unwrap();
+
+			let (mut header, extrinsics) = built_block.block.clone().deconstruct();
+			let seal = header.digest.pop().unwrap();
+
+			let mut import = BlockImportParams::new(BlockOrigin::Own, header);
+			import.body = Some(extrinsics);
+			import.post_digests.push(seal);
+			import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+			import.state_action = api.into_storage_changes(&state, parent_hash).unwrap().into();
+
+			BlockImport::import_block(&client, import)
+		})
+		.unwrap();
+
+		ignored_nodes.extend(IgnoredNodes::from_storage_proof::<BlakeTwo256>(
+			&built_block.proof.clone().unwrap(),
+		));
+		ignored_nodes.extend(IgnoredNodes::from_memory_db(built_block.storage_changes.transaction));
+		proof = StorageProof::merge([proof, built_block.proof.unwrap()]);
+
+		parent_head = built_block.block.header.clone();
+
+		blocks.push(built_block.block);
+	}
+
+	let proof = proof.into_compact_proof::<BlakeTwo256>(parent_head_root).unwrap();
+
+	TestBlockData {
+		block: TeyrchainBlockData::new(blocks, proof),
+		validation_data: persisted_validation_data.unwrap(),
+	}
+}
+
+#[test]
+fn validate_block_works() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		Vec::new(),
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let header = block.blocks()[0].header().clone();
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_multiple_blocks_work() {
+	pezsp_tracing::try_init_simple();
+
+	let blocks_per_pov = 4;
+	let (client, parent_head) = create_elastic_scaling_test_client();
+	let TestBlockData { block, validation_data } = build_multiple_blocks_with_witness(
+		&client,
+		parent_head.clone(),
+		Default::default(),
+		blocks_per_pov,
+		|i| {
+			vec![generate_extrinsic_with_pair(
+				&client,
+				Charlie.into(),
+				TestPalletCall::read_and_write_big_value {},
+				Some(i),
+			)]
+		},
+	);
+
+	assert!(block.proof().encoded_size() < 3 * 1024 * 1024);
+
+	let header = block.blocks().last().unwrap().header().clone();
+	let res_header = call_validate_block_elastic_scaling(
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_block_with_extra_extrinsics() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let extra_extrinsics = vec![
+		transfer(&client, Alice, Bob, 69),
+		transfer(&client, Bob, Charlie, 100),
+		transfer(&client, Charlie, Alice, 500),
+	];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+	let header = block.blocks()[0].header().clone();
+
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn validate_block_returns_custom_head_data() {
+	pezsp_tracing::try_init_simple();
+
+	let expected_header = vec![1, 3, 3, 7, 4, 5, 6];
+
+	let (client, parent_head) = create_test_client();
+	let extra_extrinsics = vec![
+		transfer(&client, Alice, Bob, 69),
+		generate_extrinsic(
+			&client,
+			Charlie,
+			TestPalletCall::set_custom_validation_head_data {
+				custom_header: expected_header.clone(),
+			},
+		),
+		transfer(&client, Bob, Charlie, 100),
+	];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+	let header = block.blocks()[0].header().clone();
+	assert_ne!(expected_header, header.encode());
+
+	let res_header = call_validate_block_validation_result(
+		WASM_BINARY.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`")
+	.head_data
+	.0;
+	assert_eq!(expected_header, res_header);
+}
+
+#[test]
+fn validate_block_rejects_invalid_seal() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { mut block, validation_data, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+		let (id, data) =
+			block.blocks_mut()[0].header.digest.logs.last().unwrap().as_seal().unwrap();
+		let mut data = data.to_vec();
+		let random = BlakeTwo256::hash(&data);
+		data[..random.as_ref().len()].copy_from_slice(random.as_ref());
+
+		*block.blocks_mut()[0].header.digest.logs.last_mut().unwrap() = DigestItem::Seal(id, data);
+
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_rejects_invalid_seal", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap()).contains("Invalid AuRa seal"));
+	}
+}
+#[test]
+fn validate_block_invalid_parent_hash() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { mut block, validation_data, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+		block.blocks_mut()[0].header.set_parent_hash(Hash::from_low_u64_be(1));
+
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_invalid_parent_hash", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Teyrchain head needs to be the parent of the first block"));
+	}
+}
+
+#[test]
+fn validate_block_fails_on_invalid_validation_data() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_test_client();
+		let TestBlockData { block, .. } = build_block_with_witness(
+			&client,
+			Vec::new(),
+			parent_head.clone(),
+			Default::default(),
+			Default::default(),
+		);
+
+		call_validate_block(parent_head, block, Hash::random()).unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["validate_block_fails_on_invalid_validation_data", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Relay parent storage root doesn't match"));
+	}
+}
+
+/// Test that ensures that `ValidationParams` and `MemoryOptimizedValidationParams`
+/// are encoding/decoding.
+#[test]
+fn validation_params_and_memory_optimized_validation_params_encode_and_decode() {
+	const BLOCK_DATA: &[u8] = &[1, 2, 3, 4, 5];
+	const PARENT_HEAD: &[u8] = &[1, 3, 4, 5, 6, 7, 9];
+
+	let validation_params = ValidationParams {
+		block_data: BlockData(BLOCK_DATA.encode()),
+		parent_head: HeadData(PARENT_HEAD.encode()),
+		relay_parent_number: 1,
+		relay_parent_storage_root: Hash::random(),
+	};
+
+	let encoded = validation_params.encode();
+
+	let decoded = MemoryOptimizedValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded.relay_parent_number, validation_params.relay_parent_number);
+	assert_eq!(decoded.relay_parent_storage_root, validation_params.relay_parent_storage_root);
+	assert_eq!(decoded.block_data, validation_params.block_data.0);
+	assert_eq!(decoded.parent_head, validation_params.parent_head.0);
+
+	let encoded = decoded.encode();
+
+	let decoded = ValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded, validation_params);
+}
+
+/// Test for ensuring that we are differentiating in the `validation::trie_cache` between different
+/// child tries.
+///
+/// This is achieved by first building a block using `read_and_write_child_tries` that should set
+/// the values in the child tries. In the second step we are building a second block with the same
+/// extrinsic that reads the values from the child tries and it asserts that we read the correct
+/// data from the state.
+#[test]
+fn validate_block_works_with_child_tries() {
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_test_client();
+	let TestBlockData { block, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Charlie, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let (mut header, extrinsics) = block.blocks()[0].clone().deconstruct();
+	let seal = header.digest.pop().unwrap();
+
+	let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+	import.body = Some(extrinsics);
+	import.post_digests.push(seal);
+	import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+	futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+
+	let parent_head = block.blocks()[0].header.clone();
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic(&client, Alice, TestPalletCall::read_and_write_child_tries {})],
+		parent_head.clone(),
+		Default::default(),
+		Default::default(),
+	);
+
+	let header = block.blocks()[0].header().clone();
+	let res_header =
+		call_validate_block(parent_head, block, validation_data.relay_parent_storage_root)
+			.expect("Calls `validate_block`");
+	assert_eq!(header, res_header);
+}
+
+#[test]
+fn state_changes_in_multiple_blocks_are_applied_in_exact_order() {
+	pezsp_tracing::try_init_simple();
+
+	let blocks_per_pov = 12;
+	// disable the core selection logic
+	let (client, genesis_head) = create_elastic_scaling_test_client();
+
+	// 1. Build the initial block that stores values in the map.
+	let TestBlockData { block: initial_block_data, .. } = build_block_with_witness(
+		&client,
+		vec![generate_extrinsic_with_pair(
+			&client,
+			Alice.into(),
+			TestPalletCall::store_values_in_map { max_key: 4095 },
+			Some(0),
+		)],
+		genesis_head.clone(),
+		RelayStateSproofBuilder { current_slot: 1.into(), ..Default::default() },
+		Vec::new(),
+	);
+
+	let initial_block = initial_block_data.blocks()[0].clone();
+	let (mut header, extrinsics) = initial_block.clone().deconstruct();
+	let seal = header.digest.pop().unwrap();
+
+	let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+	import.body = Some(extrinsics);
+	import.post_digests.push(seal);
+	import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+	futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+	let initial_block_header = initial_block.header().clone();
+
+	// 2. Build the PoV block that removes values from the map.
+	let TestBlockData { block: pov_block_data, validation_data: pov_validation_data } =
+		build_multiple_blocks_with_witness(
+			&client,
+			initial_block_header.clone(), // Start building PoV from the initial block's header
+			RelayStateSproofBuilder { current_slot: 2.into(), ..Default::default() },
+			blocks_per_pov,
+			|i| {
+				// Each block `i` (0-11) removes key `116 + i`.
+				let key_to_remove = 116 + i;
+				vec![generate_extrinsic_with_pair(
+					&client,
+					Bob.into(), // Use Bob to avoid nonce conflicts with Alice
+					TestPalletCall::remove_value_from_map { key: key_to_remove },
+					Some(i),
+				)]
+			},
+		);
+
+	// 3. Validate the PoV.
+	let final_pov_header = pov_block_data.blocks().last().unwrap().header().clone();
+	let res_header = call_validate_block_elastic_scaling(
+		initial_block_header, // The parent is the head of the initial block before the PoV
+		pov_block_data,
+		pov_validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block` after building the PoV");
+	assert_eq!(final_pov_header, res_header);
+}
+
+#[test]
+fn validate_block_handles_ump_signal() {
+	use pezcumulus_primitives_core::{
+		relay_chain::{UMPSignal, UMP_SEPARATOR},
+		ClaimQueueOffset, CoreInfo, CoreSelector,
+	};
+
+	pezsp_tracing::try_init_simple();
+
+	let (client, parent_head) = create_elastic_scaling_test_client();
+	let extra_extrinsics =
+		vec![transfer(&client, Alice, Bob, 69), transfer(&client, Bob, Charlie, 100)];
+
+	let TestBlockData { block, validation_data } = build_block_with_witness(
+		&client,
+		extra_extrinsics,
+		parent_head.clone(),
+		Default::default(),
+		vec![CumulusDigestItem::CoreInfo(CoreInfo {
+			selector: CoreSelector(0),
+			claim_queue_offset: ClaimQueueOffset(0),
+			number_of_cores: 1.into(),
+		})
+		.to_digest_item()],
+	);
+
+	let upward_messages = call_validate_block_validation_result(
+		test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!"),
+		parent_head,
+		block,
+		validation_data.relay_parent_storage_root,
+	)
+	.expect("Calls `validate_block`")
+	.upward_messages;
+
+	assert_eq!(
+		upward_messages,
+		vec![UMP_SEPARATOR, UMPSignal::SelectCore(CoreSelector(0), ClaimQueueOffset(0)).encode()]
+	);
+}
+
+#[test]
+fn ensure_we_only_like_blockchains() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, parent_head) = create_elastic_scaling_test_client();
+		let TestBlockData { mut block, validation_data } = build_multiple_blocks_with_witness(
+			&client,
+			parent_head.clone(),
+			Default::default(),
+			4,
+			|_| Default::default(),
+		);
+
+		// Reference some non existing parent.
+		block.blocks_mut()[2].header.parent_hash = Hash::default();
+
+		call_validate_block_elastic_scaling(
+			parent_head,
+			block,
+			validation_data.relay_parent_storage_root,
+		)
+		.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args(["ensure_we_only_like_blockchains", "--", "--nocapture"])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("Not a valid chain of blocks :("));
+	}
+}
+
+#[test]
+fn rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade() {
+	pezsp_tracing::try_init_simple();
+
+	if env::var("RUN_TEST").is_ok() {
+		let (client, genesis_head) = create_elastic_scaling_test_client();
+
+		let code = test_runtime::elastic_scaling_500ms::WASM_BINARY
+			.expect("You need to build the WASM binaries to run the tests!")
+			.to_vec();
+		let code_len = code.len() as u32;
+
+		let mut proof_builder =
+			RelayStateSproofBuilder { current_slot: 1.into(), ..Default::default() };
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// Build the block that send the runtime upgrade.
+		let TestBlockData { block: initial_block_data, .. } = build_block_with_witness(
+			&client,
+			vec![generate_extrinsic_with_pair(
+				&client,
+				Alice.into(),
+				SudoCall::sudo {
+					call: Box::new(SystemCall::set_code_without_checks { code }.into()),
+				},
+				Some(0),
+			)],
+			genesis_head.clone(),
+			proof_builder,
+			Vec::new(),
+		);
+
+		let initial_block = initial_block_data.blocks()[0].clone();
+		let (mut header, extrinsics) = initial_block.clone().deconstruct();
+		let seal = header.digest.pop().unwrap();
+
+		let mut import = BlockImportParams::new(BlockOrigin::Own, header.clone());
+		import.body = Some(extrinsics);
+		import.post_digests.push(seal);
+		import.fork_choice = Some(ForkChoiceStrategy::Custom(true));
+
+		futures::executor::block_on(BlockImport::import_block(&client, import)).unwrap();
+		let initial_block_header = initial_block.header().clone();
+
+		let mut proof_builder = RelayStateSproofBuilder {
+			current_slot: 2.into(),
+			upgrade_go_ahead: Some(relay_chain::UpgradeGoAhead::GoAhead),
+			..Default::default()
+		};
+		proof_builder.host_config.max_code_size = code_len * 2;
+
+		// 2. Build a PoV that consists of multiple blocks.
+		let TestBlockData { block: pov_block_data, validation_data: pov_validation_data } =
+			build_multiple_blocks_with_witness(
+				&client,
+				initial_block_header.clone(), // Start building PoV from the initial block's header
+				proof_builder,
+				4,
+				|_| Vec::new(),
+			);
+
+		// 3. Validate the PoV.
+		call_validate_block_elastic_scaling(
+			initial_block_header, // The parent is the head of the initial block before the PoV
+			pov_block_data,
+			pov_validation_data.relay_parent_storage_root,
+		)
+		.unwrap_err();
+	} else {
+		let output = Command::new(env::current_exe().unwrap())
+			.args([
+				"rejects_multiple_blocks_per_pov_when_applying_runtime_upgrade",
+				"--",
+				"--nocapture",
+			])
+			.env("RUN_TEST", "1")
+			.output()
+			.expect("Runs the test");
+
+		assert!(output.status.success());
+
+		assert!(dbg!(String::from_utf8(output.stderr).unwrap())
+			.contains("only one block per PoV is allowed"));
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/trie_cache.rs

@@ -0,0 +1,158 @@
+// This file is part of Pezcumulus.
+
+// 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 alloc::boxed::Box;
+
+use core::cell::{RefCell, RefMut};
+use hashbrown::{hash_map::Entry, HashMap};
+use pezsp_state_machine::TrieCacheProvider;
+use pezsp_trie::{NodeCodec, RandomState};
+use trie_db::{node::NodeOwned, Hasher};
+
+/// Special purpose trie cache implementation that is able to cache an unlimited number
+/// of values. To be used in `validate_block` to serve values and nodes that
+/// have already been loaded and decoded from the storage proof.
+pub struct TrieCache<'a, H: Hasher> {
+	node_cache: RefMut<'a, HashMap<H::Out, NodeOwned<H::Out>, RandomState>>,
+	value_cache: Option<RefMut<'a, HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>, RandomState>>>,
+}
+
+impl<'a, H: Hasher> trie_db::TrieCache<NodeCodec<H>> for TrieCache<'a, H> {
+	fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&trie_db::CachedValue<H::Out>> {
+		self.value_cache.as_ref().and_then(|cache| cache.get(key))
+	}
+
+	fn cache_value_for_key(&mut self, key: &[u8], value: trie_db::CachedValue<H::Out>) {
+		self.value_cache.as_mut().and_then(|cache| cache.insert(key.into(), value));
+	}
+
+	fn get_or_insert_node(
+		&mut self,
+		hash: <NodeCodec<H> as trie_db::NodeCodec>::HashOut,
+		fetch_node: &mut dyn FnMut() -> trie_db::Result<
+			NodeOwned<H::Out>,
+			H::Out,
+			<NodeCodec<H> as trie_db::NodeCodec>::Error,
+		>,
+	) -> trie_db::Result<&NodeOwned<H::Out>, H::Out, <NodeCodec<H> as trie_db::NodeCodec>::Error> {
+		match self.node_cache.entry(hash) {
+			Entry::Occupied(entry) => Ok(entry.into_mut()),
+			Entry::Vacant(entry) => Ok(entry.insert(fetch_node()?)),
+		}
+	}
+
+	fn get_node(
+		&mut self,
+		hash: &H::Out,
+	) -> Option<&NodeOwned<<NodeCodec<H> as trie_db::NodeCodec>::HashOut>> {
+		self.node_cache.get(hash)
+	}
+}
+
+/// Provider of [`TrieCache`] instances.
+pub struct CacheProvider<H: Hasher> {
+	node_cache: RefCell<HashMap<H::Out, NodeOwned<H::Out>, RandomState>>,
+	/// Cache: `storage_root` => `storage_key` => `value`.
+	///
+	/// One `block` can for example use multiple tries (child tries) and we need to distinguish the
+	/// cached (`storage_key`, `value`) between them. For this we are using the `storage_root` to
+	/// distinguish them (even if the storage root is the same for two child tries, it just means
+	/// that both are exactly the same trie and there would happen no collision).
+	value_cache: RefCell<
+		HashMap<H::Out, HashMap<Box<[u8]>, trie_db::CachedValue<H::Out>, RandomState>, RandomState>,
+	>,
+}
+
+impl<H: Hasher> CacheProvider<H> {
+	/// Constructs a new instance of [`CacheProvider`] with an uninitialized state
+	/// and empty node and value caches.
+	pub fn new() -> Self {
+		CacheProvider { node_cache: Default::default(), value_cache: Default::default() }
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for &&CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		Self: 'a,
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_cache(**self, storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_mut_cache(**self)
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: <H as Hasher>::Out) {
+		TrieCacheProvider::merge(**self, other, new_root)
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for &CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		Self: 'a,
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_cache(*self, storage_root)
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		TrieCacheProvider::<H>::as_trie_db_mut_cache(*self)
+	}
+
+	fn merge<'a>(&'a self, other: Self::Cache<'a>, new_root: <H as Hasher>::Out) {
+		TrieCacheProvider::merge(*self, other, new_root)
+	}
+}
+
+impl<H: Hasher> TrieCacheProvider<H> for CacheProvider<H> {
+	type Cache<'a>
+		= TrieCache<'a, H>
+	where
+		H: 'a;
+
+	fn as_trie_db_cache(&self, storage_root: <H as Hasher>::Out) -> Self::Cache<'_> {
+		TrieCache {
+			value_cache: Some(RefMut::map(self.value_cache.borrow_mut(), |c| {
+				c.entry(storage_root).or_default()
+			})),
+			node_cache: self.node_cache.borrow_mut(),
+		}
+	}
+
+	fn as_trie_db_mut_cache(&self) -> Self::Cache<'_> {
+		// This method is called when we calculate the storage root.
+		// We are not interested in caching new values (as we would throw them away directly after a
+		// block is validated) and thus, we don't pass any `value_cache`.
+		TrieCache { value_cache: None, node_cache: self.node_cache.borrow_mut() }
+	}
+
+	fn merge<'a>(&'a self, _other: Self::Cache<'a>, _new_root: <H as Hasher>::Out) {
+		// This is called to merge the `value_cache` from `as_trie_db_mut_cache`, which is not
+		// activated, so we don't need to do anything here.
+	}
+}
+
+// This is safe here since we are single-threaded in WASM
+unsafe impl<H: Hasher> Send for CacheProvider<H> {}
+unsafe impl<H: Hasher> Sync for CacheProvider<H> {}

pezcumulus/pezpallets/teyrchain-system/src/validate_block/trie_recorder.rs

@@ -0,0 +1,296 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Provide a specialized trie-recorder and provider for use in validate-block.
+//!
+//! This file defines two main structs, [`SizeOnlyRecorder`] and
+//! [`SizeOnlyRecorderProvider`]. They are used to track the current
+//! proof-size without actually recording the accessed nodes themselves.
+
+use alloc::rc::Rc;
+use codec::Encode;
+use core::cell::{RefCell, RefMut};
+use hashbrown::{HashMap, HashSet};
+use pezsp_trie::{NodeCodec, ProofSizeProvider, RandomState, StorageProof};
+use trie_db::{Hasher, RecordedForKey, TrieAccess};
+
+pub(crate) type SeenNodes<H> = Rc<RefCell<HashSet<<H as Hasher>::Out, RandomState>>>;
+
+/// A trie recorder that only keeps track of the proof size.
+///
+/// The internal size counting logic should align
+/// with ['pezsp_trie::recorder::Recorder'].
+pub struct SizeOnlyRecorder<'a, H: Hasher> {
+	seen_nodes: RefMut<'a, HashSet<H::Out, RandomState>>,
+	encoded_size: RefMut<'a, usize>,
+	recorded_keys: RefMut<'a, HashMap<Rc<[u8]>, RecordedForKey, RandomState>>,
+}
+
+impl<'a, H: trie_db::Hasher> trie_db::TrieRecorder<H::Out> for SizeOnlyRecorder<'a, H> {
+	fn record(&mut self, access: TrieAccess<'_, H::Out>) {
+		let mut encoded_size_update = 0;
+		match access {
+			TrieAccess::NodeOwned { hash, node_owned } =>
+				if self.seen_nodes.insert(hash) {
+					let node = node_owned.to_encoded::<NodeCodec<H>>();
+					encoded_size_update += node.encoded_size();
+				},
+			TrieAccess::EncodedNode { hash, encoded_node } =>
+				if self.seen_nodes.insert(hash) {
+					encoded_size_update += encoded_node.encoded_size();
+				},
+			TrieAccess::Value { hash, value, full_key } => {
+				if self.seen_nodes.insert(hash) {
+					encoded_size_update += value.encoded_size();
+				}
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+			TrieAccess::Hash { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.or_insert_with(|| RecordedForKey::Hash);
+			},
+			TrieAccess::NonExisting { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+			TrieAccess::InlineValue { full_key } => {
+				self.recorded_keys
+					.entry(full_key.into())
+					.and_modify(|e| *e = RecordedForKey::Value)
+					.or_insert_with(|| RecordedForKey::Value);
+			},
+		};
+
+		*self.encoded_size += encoded_size_update;
+	}
+
+	fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey {
+		self.recorded_keys.get(key).copied().unwrap_or(RecordedForKey::None)
+	}
+}
+
+#[derive(Clone)]
+pub struct SizeOnlyRecorderProvider<H: Hasher> {
+	seen_nodes: SeenNodes<H>,
+	encoded_size: Rc<RefCell<usize>>,
+	recorded_keys: Rc<RefCell<HashMap<Rc<[u8]>, RecordedForKey, RandomState>>>,
+}
+
+impl<H: Hasher> Default for SizeOnlyRecorderProvider<H> {
+	fn default() -> Self {
+		Self {
+			seen_nodes: Default::default(),
+			encoded_size: Default::default(),
+			recorded_keys: Default::default(),
+		}
+	}
+}
+
+impl<H: Hasher> SizeOnlyRecorderProvider<H> {
+	/// Use the given `seen_nodes` to populate the internal state.
+	#[cfg(not(feature = "std"))]
+	pub(crate) fn with_seen_nodes(seen_nodes: SeenNodes<H>) -> Self {
+		Self { seen_nodes, ..Default::default() }
+	}
+}
+
+impl<H: trie_db::Hasher> pezsp_trie::TrieRecorderProvider<H> for SizeOnlyRecorderProvider<H> {
+	type Recorder<'a>
+		= SizeOnlyRecorder<'a, H>
+	where
+		H: 'a;
+
+	fn drain_storage_proof(self) -> Option<StorageProof> {
+		None
+	}
+
+	fn as_trie_recorder(&self, _storage_root: H::Out) -> Self::Recorder<'_> {
+		SizeOnlyRecorder {
+			encoded_size: self.encoded_size.borrow_mut(),
+			seen_nodes: self.seen_nodes.borrow_mut(),
+			recorded_keys: self.recorded_keys.borrow_mut(),
+		}
+	}
+}
+
+impl<H: trie_db::Hasher> ProofSizeProvider for SizeOnlyRecorderProvider<H> {
+	fn estimate_encoded_size(&self) -> usize {
+		*self.encoded_size.borrow()
+	}
+}
+
+// This is safe here since we are single-threaded in WASM
+unsafe impl<H: Hasher> Send for SizeOnlyRecorderProvider<H> {}
+unsafe impl<H: Hasher> Sync for SizeOnlyRecorderProvider<H> {}
+
+#[cfg(test)]
+mod tests {
+	use rand::Rng;
+	use pezsp_trie::{
+		cache::{CacheSize, SharedTrieCache},
+		MemoryDB, ProofSizeProvider, TrieRecorderProvider,
+	};
+	use trie_db::{Trie, TrieDBBuilder, TrieDBMutBuilder, TrieHash, TrieMut, TrieRecorder};
+	use trie_standardmap::{Alphabet, StandardMap, ValueMode};
+
+	use super::*;
+
+	type Recorder = pezsp_trie::recorder::Recorder<pezsp_core::Blake2Hasher>;
+
+	fn create_trie() -> (
+		pezsp_trie::MemoryDB<pezsp_core::Blake2Hasher>,
+		TrieHash<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>,
+		Vec<(Vec<u8>, Vec<u8>)>,
+	) {
+		let mut db = MemoryDB::default();
+		let mut root = Default::default();
+
+		let mut seed = Default::default();
+		let test_data: Vec<(Vec<u8>, Vec<u8>)> = StandardMap {
+			alphabet: Alphabet::Low,
+			min_key: 16,
+			journal_key: 0,
+			value_mode: ValueMode::Random,
+			count: 1000,
+		}
+		.make_with(&mut seed)
+		.into_iter()
+		.map(|(k, v)| {
+			// Double the length so we end up with some values of 2 bytes and some of 64
+			let v = [v.clone(), v].concat();
+			(k, v)
+		})
+		.collect();
+
+		// Fill database with values
+		{
+			let mut trie = TrieDBMutBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(
+				&mut db, &mut root,
+			)
+			.build();
+			for (k, v) in &test_data {
+				trie.insert(k, v).expect("Inserts data");
+			}
+		}
+
+		(db, root, test_data)
+	}
+
+	#[test]
+	fn recorder_equivalence_cache() {
+		let (db, root, test_data) = create_trie();
+
+		let mut rng = rand::thread_rng();
+		for _ in 1..10 {
+			let reference_recorder = Recorder::default();
+			let recorder_for_test: SizeOnlyRecorderProvider<pezsp_core::Blake2Hasher> =
+				SizeOnlyRecorderProvider::default();
+			let reference_cache: SharedTrieCache<pezsp_core::Blake2Hasher> =
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
+			let cache_for_test: SharedTrieCache<pezsp_core::Blake2Hasher> =
+				SharedTrieCache::new(CacheSize::new(1024 * 5), None);
+			{
+				let local_cache = cache_for_test.local_cache_untrusted();
+				let mut trie_cache_for_reference = local_cache.as_trie_db_cache(root);
+				let mut reference_trie_recorder = reference_recorder.as_trie_recorder(root);
+				let reference_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut reference_trie_recorder)
+						.with_cache(&mut trie_cache_for_reference)
+						.build();
+
+				let local_cache_for_test = reference_cache.local_cache_untrusted();
+				let mut trie_cache_for_test = local_cache_for_test.as_trie_db_cache(root);
+				let mut trie_recorder_under_test = recorder_for_test.as_trie_recorder(root);
+				let test_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut trie_recorder_under_test)
+						.with_cache(&mut trie_cache_for_test)
+						.build();
+
+				// Access random values from the test data
+				for _ in 0..100 {
+					let index: usize = rng.gen_range(0..test_data.len());
+					test_trie.get(&test_data[index].0).unwrap().unwrap();
+					reference_trie.get(&test_data[index].0).unwrap().unwrap();
+				}
+
+				// Check that we have the same nodes recorded for both recorders
+				for (key, _) in test_data.iter() {
+					let reference = reference_trie_recorder.trie_nodes_recorded_for_key(key);
+					let test_value = trie_recorder_under_test.trie_nodes_recorded_for_key(key);
+					assert_eq!(format!("{:?}", reference), format!("{:?}", test_value));
+				}
+			}
+
+			// Check that we have the same size recorded for both recorders
+			assert_eq!(
+				reference_recorder.estimate_encoded_size(),
+				recorder_for_test.estimate_encoded_size()
+			);
+		}
+	}
+
+	#[test]
+	fn recorder_equivalence_no_cache() {
+		let (db, root, test_data) = create_trie();
+
+		let mut rng = rand::thread_rng();
+		for _ in 1..10 {
+			let reference_recorder = Recorder::default();
+			let recorder_for_test: SizeOnlyRecorderProvider<pezsp_core::Blake2Hasher> =
+				SizeOnlyRecorderProvider::default();
+			{
+				let mut reference_trie_recorder = reference_recorder.as_trie_recorder(root);
+				let reference_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut reference_trie_recorder)
+						.build();
+
+				let mut trie_recorder_under_test = recorder_for_test.as_trie_recorder(root);
+				let test_trie =
+					TrieDBBuilder::<pezsp_trie::LayoutV1<pezsp_core::Blake2Hasher>>::new(&db, &root)
+						.with_recorder(&mut trie_recorder_under_test)
+						.build();
+
+				for _ in 0..200 {
+					let index: usize = rng.gen_range(0..test_data.len());
+					test_trie.get(&test_data[index].0).unwrap().unwrap();
+					reference_trie.get(&test_data[index].0).unwrap().unwrap();
+				}
+
+				// Check that we have the same nodes recorded for both recorders
+				for (key, _) in test_data.iter() {
+					let reference = reference_trie_recorder.trie_nodes_recorded_for_key(key);
+					let test_value = trie_recorder_under_test.trie_nodes_recorded_for_key(key);
+					assert_eq!(format!("{:?}", reference), format!("{:?}", test_value));
+				}
+			}
+
+			// Check that we have the same size recorded for both recorders
+			assert_eq!(
+				reference_recorder.estimate_encoded_size(),
+				recorder_for_test.estimate_encoded_size()
+			);
+		}
+	}
+}

pezcumulus/pezpallets/teyrchain-system/src/weights.rs

@@ -0,0 +1,115 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Autogenerated weights for pezcumulus_pezpallet_teyrchain_system
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE BIZINIKIWI BENCHMARK CLI VERSION 4.0.0-dev
+//! DATE: 2023-03-28, STEPS: `50`, REPEAT: `20`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! HOSTNAME: `i9`, CPU: `13th Gen Intel(R) Core(TM) i9-13900K`
+//! EXECUTION: Some(Wasm), WASM-EXECUTION: Compiled, CHAIN: Some("westmint-dev"), DB CACHE: 1024
+
+// Executed Command:
+// ./target/release/pezkuwi-teyrchain
+// benchmark
+// pallet
+// --chain
+// westmint-dev
+// --pallet
+// pezcumulus_pezpallet_teyrchain_system
+// --extrinsic
+// *
+// --execution
+// wasm
+// --wasm-execution
+// compiled
+// --output
+// pezpallets/teyrchain-system/src/weights.rs
+// --steps
+// 50
+// --repeat
+// 20
+// --template
+// ../bizinikiwi/.maintain/frame-weight-template.hbs
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+
+use pezframe_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}};
+use core::marker::PhantomData;
+
+/// Weight functions needed for pezcumulus_pezpallet_teyrchain_system.
+pub trait WeightInfo {
+	fn enqueue_inbound_downward_messages(n: u32, ) -> Weight;
+}
+
+/// Weights for pezcumulus_pezpallet_teyrchain_system using the Bizinikiwi node and recommended hardware.
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	/// Storage: TeyrchainSystem LastDmqMqcHead (r:1 w:1)
+	/// Proof Skipped: TeyrchainSystem LastDmqMqcHead (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: TeyrchainSystem ReservedDmpWeightOverride (r:1 w:0)
+	/// Proof Skipped: TeyrchainSystem ReservedDmpWeightOverride (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: MessageQueue BookStateFor (r:1 w:1)
+	/// Proof: MessageQueue BookStateFor (max_values: None, max_size: Some(52), added: 2527, mode: MaxEncodedLen)
+	/// Storage: MessageQueue ServiceHead (r:1 w:1)
+	/// Proof: MessageQueue ServiceHead (max_values: Some(1), max_size: Some(5), added: 500, mode: MaxEncodedLen)
+	/// Storage: TeyrchainSystem ProcessedDownwardMessages (r:0 w:1)
+	/// Proof Skipped: TeyrchainSystem ProcessedDownwardMessages (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: MessageQueue Pages (r:0 w:16)
+	/// Proof: MessageQueue Pages (max_values: None, max_size: Some(65585), added: 68060, mode: MaxEncodedLen)
+	/// The range of component `n` is `[0, 1000]`.
+	fn enqueue_inbound_downward_messages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `12`
+		//  Estimated: `8013`
+		// Minimum execution time: 1_625_000 picoseconds.
+		Weight::from_parts(1_735_000, 8013)
+			// Standard Error: 14_563
+			.saturating_add(Weight::from_parts(25_300_108, 0).saturating_mul(n.into()))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(4_u64))
+	}
+}
+
+// For backwards compatibility and tests
+impl WeightInfo for () {
+	/// Storage: TeyrchainSystem LastDmqMqcHead (r:1 w:1)
+	/// Proof Skipped: TeyrchainSystem LastDmqMqcHead (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: TeyrchainSystem ReservedDmpWeightOverride (r:1 w:0)
+	/// Proof Skipped: TeyrchainSystem ReservedDmpWeightOverride (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: MessageQueue BookStateFor (r:1 w:1)
+	/// Proof: MessageQueue BookStateFor (max_values: None, max_size: Some(52), added: 2527, mode: MaxEncodedLen)
+	/// Storage: MessageQueue ServiceHead (r:1 w:1)
+	/// Proof: MessageQueue ServiceHead (max_values: Some(1), max_size: Some(5), added: 500, mode: MaxEncodedLen)
+	/// Storage: TeyrchainSystem ProcessedDownwardMessages (r:0 w:1)
+	/// Proof Skipped: TeyrchainSystem ProcessedDownwardMessages (max_values: Some(1), max_size: None, mode: Measured)
+	/// Storage: MessageQueue Pages (r:0 w:16)
+	/// Proof: MessageQueue Pages (max_values: None, max_size: Some(65585), added: 68060, mode: MaxEncodedLen)
+	/// The range of component `n` is `[0, 1000]`.
+	fn enqueue_inbound_downward_messages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `12`
+		//  Estimated: `8013`
+		// Minimum execution time: 1_625_000 picoseconds.
+		Weight::from_parts(1_735_000, 8013)
+			// Standard Error: 14_563
+			.saturating_add(Weight::from_parts(25_300_108, 0).saturating_mul(n.into()))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(4_u64))
+	}
+}

pezcumulus/pezpallets/weight-reclaim/Cargo.toml

@@ -0,0 +1,68 @@
+[package]
+name = "pezcumulus-pezpallet-weight-reclaim"
+version = "1.0.0"
+authors.workspace = true
+edition.workspace = true
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+description = "pallet and transaction extensions for accurate proof size reclaim"
+documentation = "https://docs.rs/pezcumulus-pezpallet-weight-reclaim"
+
+[lints]
+workspace = true
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]
+
+[dependencies]
+# Bizinikiwi dependencies
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+pezsp-trie = { workspace = true }
+
+pezcumulus-primitives-storage-weight-reclaim = { workspace = true }
+pezframe-benchmarking = { optional = true, workspace = true }
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+
+# Other dependencies
+codec = { features = ["derive"], workspace = true }
+derive-where = { workspace = true }
+docify = { workspace = true }
+log = { workspace = true, default-features = true }
+scale-info = { features = ["derive"], workspace = true }
+
+[dev-dependencies]
+pezcumulus-primitives-proof-size-hostfunction = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezcumulus-primitives-proof-size-hostfunction/std",
+  "pezcumulus-primitives-storage-weight-reclaim/std",
+  "pezframe-benchmarking?/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "log/std",
+  "scale-info/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "pezsp-trie/std",
+]
+runtime-benchmarks = [
+  "pezcumulus-primitives-proof-size-hostfunction/runtime-benchmarks",
+  "pezcumulus-primitives-storage-weight-reclaim/runtime-benchmarks",
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "pezsp-trie/runtime-benchmarks",
+]
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezsp-runtime/try-runtime",
+]

pezcumulus/pezpallets/weight-reclaim/src/benchmarks.rs

@@ -0,0 +1,71 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+#![cfg(feature = "runtime-benchmarks")]
+
+use super::*;
+use pezframe_support::pezpallet_prelude::{DispatchClass, Pays};
+use pezframe_system::RawOrigin;
+use pezsp_runtime::traits::{AsTransactionAuthorizedOrigin, DispatchTransaction};
+
+#[pezframe_benchmarking::v2::benchmarks(
+	where T: Send + Sync,
+		<T as pezframe_system::Config>::RuntimeCall:
+			Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
+		<T as pezframe_system::Config>::RuntimeOrigin: AsTransactionAuthorizedOrigin,
+)]
+mod bench {
+	use super::*;
+	use pezframe_benchmarking::impl_test_function;
+
+	#[benchmark]
+	fn storage_weight_reclaim() {
+		let ext = StorageWeightReclaim::<T, ()>::new(());
+
+		let origin = RawOrigin::Root.into();
+		let call = T::RuntimeCall::from(pezframe_system::Call::remark { remark: alloc::vec![] });
+
+		let overestimate = 10_000;
+		let info = DispatchInfo {
+			call_weight: Weight::zero().add_proof_size(overestimate),
+			extension_weight: Weight::zero(),
+			class: DispatchClass::Normal,
+			pays_fee: Pays::No,
+		};
+
+		let post_info = PostDispatchInfo { actual_weight: None, pays_fee: Pays::No };
+
+		let mut block_weight = pezframe_system::ConsumedWeight::default();
+		block_weight.accrue(Weight::from_parts(0, overestimate), info.class);
+
+		pezframe_system::BlockWeight::<T>::put(block_weight);
+
+		#[block]
+		{
+			assert!(ext.test_run(origin, &call, &info, 0, 0, |_| Ok(post_info)).unwrap().is_ok());
+		}
+
+		let final_block_proof_size =
+			pezframe_system::BlockWeight::<T>::get().get(info.class).proof_size();
+
+		assert!(
+			final_block_proof_size < overestimate,
+			"The proof size measured should be less than {overestimate}"
+		);
+	}
+
+	impl_benchmark_test_suite!(Pallet, crate::tests::setup_test_ext_default(), crate::tests::Test);
+}

pezcumulus/pezpallets/weight-reclaim/src/lib.rs

@@ -0,0 +1,306 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Pallet and transaction extensions to reclaim PoV proof size weight after an extrinsic has been
+//! applied.
+//!
+//! This crate provides:
+//! * [`StorageWeightReclaim`] transaction extension: it must wrap the whole transaction extension
+//!   pipeline.
+//! * The pallet required for the transaction extensions weight information and benchmarks.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+#[cfg(not(feature = "std"))]
+use alloc::vec::Vec;
+use codec::{Decode, DecodeWithMemTracking, Encode};
+use pezcumulus_primitives_storage_weight_reclaim::get_proof_size;
+use derive_where::derive_where;
+use pezframe_support::{
+	dispatch::{DispatchInfo, PostDispatchInfo},
+	pezpallet_prelude::Weight,
+	traits::Defensive,
+};
+use scale_info::TypeInfo;
+use pezsp_runtime::{
+	traits::{DispatchInfoOf, Dispatchable, Implication, PostDispatchInfoOf, TransactionExtension},
+	transaction_validity::{TransactionSource, TransactionValidityError, ValidTransaction},
+	DispatchResult,
+};
+
+#[cfg(feature = "runtime-benchmarks")]
+pub mod benchmarks;
+#[cfg(test)]
+mod tests;
+mod weights;
+
+pub use pallet::*;
+pub use weights::WeightInfo;
+
+const LOG_TARGET: &'static str = "runtime::storage_reclaim_pallet";
+
+/// Pallet to use alongside the transaction extension [`StorageWeightReclaim`], the pallet provides
+/// weight information and benchmarks.
+#[pezframe_support::pallet]
+pub mod pallet {
+	use super::*;
+
+	#[pallet::pallet]
+	pub struct Pallet<T>(_);
+
+	#[pallet::config]
+	pub trait Config: pezframe_system::Config {
+		type WeightInfo: WeightInfo;
+	}
+}
+
+/// Storage weight reclaim mechanism.
+///
+/// This extension must wrap all the transaction extensions:
+#[doc = docify::embed!("./src/tests.rs", Tx)]
+///
+/// This extension checks the size of the node-side storage proof before and after executing a given
+/// extrinsic using the proof size host function. The difference between benchmarked and used weight
+/// is reclaimed.
+///
+/// If the benchmark was underestimating the proof size, then it is added to the block weight.
+///
+/// For the time part of the weight, it does same as system `WeightReclaim` extension, it
+/// calculates the unused weight using the post information and reclaim the unused weight.
+/// So this extension can be used as a drop-in replacement for `WeightReclaim` extension for
+/// teyrchains.
+#[derive(Encode, Decode, DecodeWithMemTracking, TypeInfo)]
+#[derive_where(Clone, Eq, PartialEq, Default; S)]
+#[scale_info(skip_type_params(T))]
+pub struct StorageWeightReclaim<T, S>(pub S, core::marker::PhantomData<T>);
+
+impl<T, S> StorageWeightReclaim<T, S> {
+	/// Create a new `StorageWeightReclaim` instance.
+	pub fn new(s: S) -> Self {
+		Self(s, Default::default())
+	}
+}
+
+impl<T, S> From<S> for StorageWeightReclaim<T, S> {
+	fn from(s: S) -> Self {
+		Self::new(s)
+	}
+}
+
+impl<T, S: core::fmt::Debug> core::fmt::Debug for StorageWeightReclaim<T, S> {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
+		#[cfg(feature = "std")]
+		let _ = write!(f, "StorageWeightReclaim<{:?}>", self.0);
+
+		#[cfg(not(feature = "std"))]
+		let _ = write!(f, "StorageWeightReclaim<wasm-stripped>");
+
+		Ok(())
+	}
+}
+
+impl<T: Config + Send + Sync, S: TransactionExtension<T::RuntimeCall>>
+	TransactionExtension<T::RuntimeCall> for StorageWeightReclaim<T, S>
+where
+	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
+{
+	const IDENTIFIER: &'static str = "StorageWeightReclaim<Use `metadata()`!>";
+
+	type Implicit = S::Implicit;
+
+	// Initial proof size and inner extension value.
+	type Val = (Option<u64>, S::Val);
+
+	// Initial proof size and inner extension pre.
+	type Pre = (Option<u64>, S::Pre);
+
+	fn implicit(&self) -> Result<Self::Implicit, TransactionValidityError> {
+		self.0.implicit()
+	}
+
+	fn metadata() -> Vec<pezsp_runtime::traits::TransactionExtensionMetadata> {
+		let mut inner = S::metadata();
+		inner.push(pezsp_runtime::traits::TransactionExtensionMetadata {
+			identifier: "StorageWeightReclaim",
+			ty: scale_info::meta_type::<()>(),
+			implicit: scale_info::meta_type::<()>(),
+		});
+		inner
+	}
+
+	fn weight(&self, call: &T::RuntimeCall) -> Weight {
+		T::WeightInfo::storage_weight_reclaim().saturating_add(self.0.weight(call))
+	}
+
+	fn validate(
+		&self,
+		origin: T::RuntimeOrigin,
+		call: &T::RuntimeCall,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+		self_implicit: Self::Implicit,
+		inherited_implication: &impl Implication,
+		source: TransactionSource,
+	) -> Result<(ValidTransaction, Self::Val, T::RuntimeOrigin), TransactionValidityError> {
+		let proof_size = get_proof_size();
+
+		self.0
+			.validate(origin, call, info, len, self_implicit, inherited_implication, source)
+			.map(|(validity, val, origin)| (validity, (proof_size, val), origin))
+	}
+
+	fn prepare(
+		self,
+		val: Self::Val,
+		origin: &T::RuntimeOrigin,
+		call: &T::RuntimeCall,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+	) -> Result<Self::Pre, TransactionValidityError> {
+		let (proof_size, inner_val) = val;
+		self.0.prepare(inner_val, origin, call, info, len).map(|pre| (proof_size, pre))
+	}
+
+	fn post_dispatch_details(
+		pre: Self::Pre,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		post_info: &PostDispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+		result: &DispatchResult,
+	) -> Result<Weight, TransactionValidityError> {
+		let (proof_size_before_dispatch, inner_pre) = pre;
+
+		let mut post_info_with_inner = *post_info;
+		S::post_dispatch(inner_pre, info, &mut post_info_with_inner, len, result)?;
+
+		let inner_refund = if let (Some(before_weight), Some(after_weight)) =
+			(post_info.actual_weight, post_info_with_inner.actual_weight)
+		{
+			before_weight.saturating_sub(after_weight)
+		} else {
+			Weight::zero()
+		};
+
+		let Some(proof_size_before_dispatch) = proof_size_before_dispatch else {
+			// We have no proof size information, there is nothing we can do.
+			return Ok(inner_refund);
+		};
+
+		let Some(proof_size_after_dispatch) = get_proof_size().defensive_proof(
+			"Proof recording enabled during prepare, now disabled. This should not happen.",
+		) else {
+			return Ok(inner_refund);
+		};
+
+		// The consumed proof size as measured by the host.
+		let measured_proof_size =
+			proof_size_after_dispatch.saturating_sub(proof_size_before_dispatch);
+
+		// The consumed weight as benchmarked. Calculated from post info and info.
+		// NOTE: `calc_actual_weight` will take the minimum of `post_info` and `info` weights.
+		// This means any underestimation of compute time in the pre dispatch info will not be
+		// taken into account.
+		let benchmarked_actual_weight = post_info_with_inner.calc_actual_weight(info);
+
+		let benchmarked_actual_proof_size = benchmarked_actual_weight.proof_size();
+		if benchmarked_actual_proof_size < measured_proof_size {
+			log::error!(
+				target: LOG_TARGET,
+				"Benchmarked storage weight smaller than consumed storage weight. \
+				benchmarked: {benchmarked_actual_proof_size} consumed: {measured_proof_size}"
+			);
+		} else {
+			log::trace!(
+				target: LOG_TARGET,
+				"Reclaiming storage weight. benchmarked: {benchmarked_actual_proof_size},
+				consumed: {measured_proof_size}"
+			);
+		}
+
+		let accurate_weight = benchmarked_actual_weight.set_proof_size(measured_proof_size);
+
+		let pov_size_missing_from_node = pezframe_system::BlockWeight::<T>::mutate(|current_weight| {
+			let already_reclaimed = pezframe_system::ExtrinsicWeightReclaimed::<T>::get();
+			current_weight.accrue(already_reclaimed, info.class);
+			current_weight.reduce(info.total_weight(), info.class);
+			current_weight.accrue(accurate_weight, info.class);
+
+			// If we encounter a situation where the node-side proof size is already higher than
+			// what we have in the runtime bookkeeping, we add the difference to the `BlockWeight`.
+			// This prevents that the proof size grows faster than the runtime proof size.
+			let extrinsic_len = pezframe_system::AllExtrinsicsLen::<T>::get().unwrap_or(0);
+			let node_side_pov_size = proof_size_after_dispatch.saturating_add(extrinsic_len.into());
+			let block_weight_proof_size = current_weight.total().proof_size();
+			let pov_size_missing_from_node =
+				node_side_pov_size.saturating_sub(block_weight_proof_size);
+			if pov_size_missing_from_node > 0 {
+				log::warn!(
+					target: LOG_TARGET,
+					"Node-side PoV size higher than runtime proof size weight. node-side: \
+					{node_side_pov_size} extrinsic_len: {extrinsic_len} runtime: \
+					{block_weight_proof_size}, missing: {pov_size_missing_from_node}. Setting to \
+					node-side proof size."
+				);
+				current_weight
+					.accrue(Weight::from_parts(0, pov_size_missing_from_node), info.class);
+			}
+
+			pov_size_missing_from_node
+		});
+
+		// The saturation will happen if the pre-dispatch weight is underestimating the proof
+		// size or if the node-side proof size is higher than expected.
+		// In this case the extrinsic proof size weight reclaimed is 0 and not a negative reclaim.
+		let accurate_unspent = info
+			.total_weight()
+			.saturating_sub(accurate_weight)
+			.saturating_sub(Weight::from_parts(0, pov_size_missing_from_node));
+		pezframe_system::ExtrinsicWeightReclaimed::<T>::put(accurate_unspent);
+
+		// Call have already returned their unspent amount.
+		// (also transaction extension prior in the pipeline, but there shouldn't be any.)
+		let already_unspent_in_tx_ext_pipeline = post_info.calc_unspent(info);
+		Ok(accurate_unspent.saturating_sub(already_unspent_in_tx_ext_pipeline))
+	}
+
+	fn bare_validate(
+		call: &T::RuntimeCall,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+	) -> pezframe_support::pezpallet_prelude::TransactionValidity {
+		S::bare_validate(call, info, len)
+	}
+
+	fn bare_validate_and_prepare(
+		call: &T::RuntimeCall,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+	) -> Result<(), TransactionValidityError> {
+		S::bare_validate_and_prepare(call, info, len)
+	}
+
+	fn bare_post_dispatch(
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		post_info: &mut PostDispatchInfoOf<T::RuntimeCall>,
+		len: usize,
+		result: &DispatchResult,
+	) -> Result<(), TransactionValidityError> {
+		S::bare_post_dispatch(info, post_info, len, result)?;
+
+		pezframe_system::Pallet::<T>::reclaim_weight(info, post_info)
+	}
+}

pezcumulus/pezpallets/weight-reclaim/src/weights.rs

@@ -0,0 +1,75 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Autogenerated weights for `pezcumulus_pezpallet_weight_reclaim`
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE BIZINIKIWI BENCHMARK CLI VERSION 32.0.0
+//! DATE: 2024-08-10, STEPS: `50`, REPEAT: `20`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! HOSTNAME: `fedora`, CPU: `13th Gen Intel(R) Core(TM) i7-1360P`
+//! WASM-EXECUTION: `Compiled`, CHAIN: `Some("dev")`, DB CACHE: `1024`
+
+// Executed Command:
+// ./target/release/teyrchain-template-node
+// benchmark
+// pallet
+// --pallet
+// pezcumulus-pezpallet-weight-reclaim
+// --chain
+// dev
+// --output
+// pezcumulus/pezpallets/weight-reclaim/src/weights.rs
+// --template
+// bizinikiwi/.maintain/frame-weight-template.hbs
+// --extrinsic
+// *
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+#![allow(missing_docs)]
+
+use pezframe_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}};
+use core::marker::PhantomData;
+
+/// Weight functions needed for `pezcumulus_pezpallet_weight_reclaim`.
+pub trait WeightInfo {
+	fn storage_weight_reclaim() -> Weight;
+}
+
+/// Weights for `pezcumulus_pezpallet_weight_reclaim` using the Bizinikiwi node and recommended hardware.
+#[allow(dead_code)]
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	fn storage_weight_reclaim() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `0`
+		//  Estimated: `0`
+		// Minimum execution time: 2_247_000 picoseconds.
+		Weight::from_parts(2_466_000, 0)
+	}
+}
+
+// For backwards compatibility and tests.
+impl WeightInfo for () {
+	fn storage_weight_reclaim() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `0`
+		//  Estimated: `0`
+		// Minimum execution time: 2_247_000 picoseconds.
+		Weight::from_parts(2_466_000, 0)
+	}
+}

pezcumulus/pezpallets/xcm/Cargo.toml

@@ -0,0 +1,52 @@
+[package]
+authors.workspace = true
+edition.workspace = true
+name = "pezcumulus-pezpallet-xcm"
+version = "0.7.0"
+license = "Apache-2.0"
+description = "Pallet for stuff specific to teyrchains' usage of XCM"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-xcm"
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+xcm = { workspace = true }
+
+pezcumulus-primitives-core = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+  "codec/std",
+  "pezcumulus-primitives-core/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "scale-info/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "xcm/std",
+]
+try-runtime = [
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezsp-runtime/try-runtime",
+]
+runtime-benchmarks = [
+  "pezcumulus-primitives-core/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "xcm/runtime-benchmarks",
+]

pezcumulus/pezpallets/xcm/src/lib.rs

@@ -0,0 +1,118 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Pallet for stuff specific to teyrchains' usage of XCM. Right now that's just the origin
+//! used by teyrchains when receiving `Transact` messages from other teyrchains or the Relay chain
+//! which must be natively represented.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use codec::{Decode, Encode};
+use pezcumulus_primitives_core::ParaId;
+pub use pallet::*;
+use scale_info::TypeInfo;
+use pezsp_runtime::{traits::BadOrigin, RuntimeDebug};
+use xcm::latest::{ExecuteXcm, Outcome};
+
+#[pezframe_support::pallet]
+pub mod pallet {
+	use super::*;
+	use pezframe_support::pezpallet_prelude::*;
+
+	#[pallet::pallet]
+	pub struct Pallet<T>(_);
+
+	/// The module configuration trait.
+	#[pallet::config]
+	pub trait Config: pezframe_system::Config {
+		/// The overarching event type.
+		#[allow(deprecated)]
+		type RuntimeEvent: From<Event<Self>> + IsType<<Self as pezframe_system::Config>::RuntimeEvent>;
+
+		type XcmExecutor: ExecuteXcm<Self::RuntimeCall>;
+	}
+
+	#[pallet::event]
+	pub enum Event<T: Config> {
+		/// Downward message is invalid XCM.
+		/// \[ id \]
+		InvalidFormat([u8; 32]),
+		/// Downward message is unsupported version of XCM.
+		/// \[ id \]
+		UnsupportedVersion([u8; 32]),
+		/// Downward message executed with the given outcome.
+		/// \[ id, outcome \]
+		ExecutedDownward([u8; 32], Outcome),
+	}
+
+	/// Origin for the teyrchains module.
+	#[derive(
+		PartialEq,
+		Eq,
+		Clone,
+		Encode,
+		Decode,
+		DecodeWithMemTracking,
+		TypeInfo,
+		RuntimeDebug,
+		MaxEncodedLen,
+	)]
+	#[pallet::origin]
+	pub enum Origin {
+		/// It comes from the (parent) relay chain.
+		Relay,
+		/// It comes from a (sibling) teyrchain.
+		SiblingTeyrchain(ParaId),
+	}
+
+	#[pallet::call]
+	impl<T: Config> Pallet<T> {}
+
+	impl From<ParaId> for Origin {
+		fn from(id: ParaId) -> Origin {
+			Origin::SiblingTeyrchain(id)
+		}
+	}
+	impl From<u32> for Origin {
+		fn from(id: u32) -> Origin {
+			Origin::SiblingTeyrchain(id.into())
+		}
+	}
+}
+
+/// Ensure that the origin `o` represents a sibling teyrchain.
+/// Returns `Ok` with the teyrchain ID of the sibling or an `Err` otherwise.
+pub fn ensure_sibling_para<OuterOrigin>(o: OuterOrigin) -> Result<ParaId, BadOrigin>
+where
+	OuterOrigin: Into<Result<Origin, OuterOrigin>>,
+{
+	match o.into() {
+		Ok(Origin::SiblingTeyrchain(id)) => Ok(id),
+		_ => Err(BadOrigin),
+	}
+}
+
+/// Ensure that the origin `o` represents is the relay chain.
+/// Returns `Ok` if it does or an `Err` otherwise.
+pub fn ensure_relay<OuterOrigin>(o: OuterOrigin) -> Result<(), BadOrigin>
+where
+	OuterOrigin: Into<Result<Origin, OuterOrigin>>,
+{
+	match o.into() {
+		Ok(Origin::Relay) => Ok(()),
+		_ => Err(BadOrigin),
+	}
+}

pezcumulus/pezpallets/xcmp-queue/Cargo.toml

@@ -0,0 +1,108 @@
+[package]
+name = "pezcumulus-pezpallet-xcmp-queue"
+version = "0.7.1"
+authors.workspace = true
+edition.workspace = true
+description = "Pallet to queue outbound and inbound XCMP messages."
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+documentation = "https://docs.rs/pezcumulus-pezpallet-xcmp-queue"
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+tracing = { workspace = true }
+
+# Bizinikiwi
+pezframe-support = { workspace = true }
+pezframe-system = { workspace = true }
+pezpallet-message-queue = { workspace = true }
+pezsp-core = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+# Pezkuwi
+pezkuwi-runtime-common = { workspace = true }
+pezkuwi-runtime-teyrchains = { workspace = true }
+xcm = { workspace = true }
+xcm-builder = { workspace = true }
+xcm-executor = { workspace = true }
+
+# Pezcumulus
+pezcumulus-primitives-core = { workspace = true }
+
+# Optional import for weight accuracy testing
+approx = { workspace = true }
+# Optional import for benchmarking
+bounded-collections = { workspace = true }
+pezframe-benchmarking = { optional = true, workspace = true }
+
+# Bridges
+bp-xcm-bridge-hub-router = { optional = true, workspace = true }
+
+[dev-dependencies]
+# Bizinikiwi
+pezframe-support = { features = [
+  "experimental",
+], workspace = true, default-features = true }
+pezpallet-balances = { workspace = true, default-features = true }
+pezsp-core = { workspace = true, default-features = true }
+
+# Pezcumulus
+pezcumulus-pezpallet-teyrchain-system = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+std = [
+  "approx/std",
+  "bounded-collections/std",
+  "bp-xcm-bridge-hub-router?/std",
+  "codec/std",
+  "pezcumulus-primitives-core/std",
+  "pezframe-benchmarking?/std",
+  "pezframe-support/std",
+  "pezframe-system/std",
+  "pezpallet-message-queue/std",
+  "pezkuwi-runtime-common/std",
+  "pezkuwi-runtime-teyrchains/std",
+  "scale-info/std",
+  "pezsp-core/std",
+  "pezsp-io/std",
+  "pezsp-runtime/std",
+  "tracing/std",
+  "xcm-builder/std",
+  "xcm-executor/std",
+  "xcm/std",
+]
+runtime-benchmarks = [
+  "bp-xcm-bridge-hub-router?/runtime-benchmarks",
+  "pezcumulus-pezpallet-teyrchain-system/runtime-benchmarks",
+  "pezcumulus-primitives-core/runtime-benchmarks",
+  "pezframe-benchmarking/runtime-benchmarks",
+  "pezframe-support/runtime-benchmarks",
+  "pezframe-system/runtime-benchmarks",
+  "pezpallet-balances/runtime-benchmarks",
+  "pezpallet-message-queue/runtime-benchmarks",
+  "pezkuwi-runtime-common/runtime-benchmarks",
+  "pezkuwi-runtime-teyrchains/runtime-benchmarks",
+  "pezsp-io/runtime-benchmarks",
+  "pezsp-runtime/runtime-benchmarks",
+  "xcm-builder/runtime-benchmarks",
+  "xcm-executor/runtime-benchmarks",
+  "xcm/runtime-benchmarks",
+]
+try-runtime = [
+  "pezcumulus-pezpallet-teyrchain-system/try-runtime",
+  "pezframe-support/try-runtime",
+  "pezframe-system/try-runtime",
+  "pezpallet-balances/try-runtime",
+  "pezpallet-message-queue/try-runtime",
+  "pezkuwi-runtime-common/try-runtime",
+  "pezkuwi-runtime-teyrchains/try-runtime",
+  "pezsp-runtime/try-runtime",
+]
+bridging = ["bp-xcm-bridge-hub-router"]

pezcumulus/pezpallets/xcmp-queue/src/benchmarking.rs

@@ -0,0 +1,341 @@
+// 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 pezcumulus-pezpallet-xcmp-queue
+
+use crate::{weights_ext::get_average_page_pos, *};
+
+use alloc::vec;
+use codec::DecodeAll;
+use pezframe_benchmarking::v2::*;
+use pezframe_support::{assert_ok, traits::Hooks};
+use pezframe_system::RawOrigin;
+use xcm::MAX_INSTRUCTIONS_TO_DECODE;
+
+#[benchmarks]
+mod benchmarks {
+	use super::*;
+
+	/// Modify any of the `QueueConfig` fields with a new `u32` value.
+	///
+	/// Used as weight for:
+	/// - update_suspend_threshold
+	/// - update_drop_threshold
+	/// - update_resume_threshold
+	#[benchmark]
+	fn set_config_with_u32() {
+		#[extrinsic_call]
+		Pallet::<T>::update_resume_threshold(RawOrigin::Root, 1);
+	}
+
+	/// Add a XCMP message of `n` bytes to the message queue.
+	///
+	/// The message will be added on a new page and also, the `BookState` will be added
+	/// to the ready ring.
+	#[benchmark]
+	fn enqueue_n_bytes_xcmp_message(n: Linear<0, { MaxXcmpMessageLenOf::<T>::get() }>) {
+		#[cfg(test)]
+		{
+			mock::EnqueuedMessages::set(vec![]);
+		}
+
+		let msg = BoundedVec::try_from(vec![0; n as usize]).unwrap();
+
+		#[cfg(not(test))]
+		let fp_before = T::XcmpQueue::footprint(0.into());
+		#[block]
+		{
+			assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+				0.into(),
+				&[msg.as_bounded_slice()],
+				true,
+				&mut WeightMeter::new()
+			));
+		}
+		#[cfg(not(test))]
+		{
+			let fp_after = T::XcmpQueue::footprint(0.into());
+			assert_eq!(fp_after.ready_pages, fp_before.ready_pages + 1);
+		}
+	}
+
+	/// Add `n` XCMP message of 0 bytes to the message queue.
+	///
+	/// Only for the first message a new page will be created and the `BookState` will be added
+	/// to the ready ring.
+	#[benchmark]
+	fn enqueue_n_empty_xcmp_messages(n: Linear<0, 1000>) {
+		#[cfg(test)]
+		{
+			mock::EnqueuedMessages::set(vec![]);
+			<QueueConfig<T>>::set(QueueConfigData {
+				suspend_threshold: 1100,
+				drop_threshold: 1100,
+				resume_threshold: 1100,
+			});
+		}
+
+		let msg = BoundedVec::new();
+		let msgs = vec![msg.as_bounded_slice(); n as usize];
+
+		#[cfg(not(test))]
+		let fp_before = T::XcmpQueue::footprint(0.into());
+		#[block]
+		{
+			assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+				0.into(),
+				&msgs,
+				true,
+				&mut WeightMeter::new()
+			));
+		}
+		#[cfg(not(test))]
+		{
+			let fp_after = T::XcmpQueue::footprint(0.into());
+			if !msgs.is_empty() {
+				assert_eq!(fp_after.ready_pages, fp_before.ready_pages + 1);
+			}
+		}
+	}
+
+	/// Add an XCMP message of 0 bytes to the message queue at the provided position
+	/// on an existing page.
+	#[benchmark(pov_mode = Measured)]
+	fn enqueue_empty_xcmp_message_at(
+		n: Linear<0, { crate::MaxXcmpMessageLenOf::<T>::get() - 10 }>,
+	) {
+		#[cfg(test)]
+		{
+			mock::EnqueuedMessages::set(vec![]);
+		}
+
+		assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+			0.into(),
+			&[BoundedVec::try_from(vec![0; n as usize]).unwrap().as_bounded_slice()],
+			true,
+			&mut WeightMeter::new()
+		));
+
+		#[cfg(not(test))]
+		let fp_before = T::XcmpQueue::footprint(0.into());
+		#[block]
+		{
+			assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+				0.into(),
+				&[BoundedVec::new().as_bounded_slice()],
+				true,
+				&mut WeightMeter::new()
+			));
+		}
+		#[cfg(not(test))]
+		{
+			let fp_after = T::XcmpQueue::footprint(0.into());
+			assert_eq!(fp_after.ready_pages, fp_before.ready_pages);
+		}
+	}
+
+	/// Add `n` pages to the message queue.
+	///
+	/// We add one page by enqueueing a maximal size message which fills it.
+	#[benchmark]
+	fn enqueue_n_full_pages(n: Linear<0, 100>) {
+		#[cfg(test)]
+		{
+			mock::EnqueuedMessages::set(vec![]);
+		}
+		<QueueConfig<T>>::set(QueueConfigData {
+			suspend_threshold: 200,
+			drop_threshold: 200,
+			resume_threshold: 200,
+		});
+
+		let max_msg_len = MaxXcmpMessageLenOf::<T>::get() as usize;
+		let mut msgs = vec![];
+		for _i in 0..n {
+			let msg = BoundedVec::try_from(vec![0; max_msg_len]).unwrap();
+			msgs.push(msg);
+		}
+
+		#[cfg(not(test))]
+		let fp_before = T::XcmpQueue::footprint(0.into());
+		#[block]
+		{
+			assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+				0.into(),
+				&msgs.iter().map(|msg| msg.as_bounded_slice()).collect::<Vec<_>>(),
+				true,
+				&mut WeightMeter::new()
+			));
+		}
+		#[cfg(not(test))]
+		{
+			let fp_after = T::XcmpQueue::footprint(0.into());
+			assert_eq!(fp_after.ready_pages, fp_before.ready_pages + n);
+		}
+	}
+
+	#[benchmark(pov_mode = Measured)]
+	fn enqueue_1000_small_xcmp_messages() {
+		#[cfg(test)]
+		{
+			<QueueConfig<T>>::set(QueueConfigData {
+				suspend_threshold: 1100,
+				drop_threshold: 1100,
+				resume_threshold: 1100,
+			});
+		}
+
+		assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+			0.into(),
+			&[BoundedVec::try_from(vec![
+				0;
+				get_average_page_pos(MaxXcmpMessageLenOf::<T>::get())
+					as usize
+			])
+			.unwrap()
+			.as_bounded_slice()],
+			true,
+			&mut WeightMeter::new()
+		));
+
+		let mut msgs = vec![];
+		for _i in 0..1000 {
+			msgs.push(BoundedVec::try_from(vec![0; 3]).unwrap());
+		}
+
+		#[cfg(not(test))]
+		let fp_before = T::XcmpQueue::footprint(0.into());
+		#[block]
+		{
+			assert_ok!(Pallet::<T>::enqueue_xcmp_messages(
+				0.into(),
+				&msgs.iter().map(|msg| msg.as_bounded_slice()).collect::<Vec<_>>(),
+				true,
+				&mut WeightMeter::new()
+			));
+		}
+		#[cfg(not(test))]
+		{
+			let fp_after = T::XcmpQueue::footprint(0.into());
+			assert_eq!(fp_after.ready_pages, fp_before.ready_pages);
+		}
+	}
+
+	#[benchmark]
+	fn suspend_channel() {
+		let para = 123.into();
+		let data = ChannelSignal::Suspend.encode();
+
+		#[block]
+		{
+			ChannelSignal::decode_all(&mut &data[..]).unwrap();
+			Pallet::<T>::suspend_channel(para);
+		}
+
+		assert_eq!(
+			OutboundXcmpStatus::<T>::get()
+				.iter()
+				.find(|p| p.recipient == para)
+				.unwrap()
+				.state,
+			OutboundState::Suspended
+		);
+	}
+
+	#[benchmark]
+	fn resume_channel() {
+		let para = 123.into();
+		let data = ChannelSignal::Resume.encode();
+
+		Pallet::<T>::suspend_channel(para);
+
+		#[block]
+		{
+			ChannelSignal::decode_all(&mut &data[..]).unwrap();
+			Pallet::<T>::resume_channel(para);
+		}
+
+		assert!(
+			OutboundXcmpStatus::<T>::get().iter().all(|p| p.recipient != para),
+			"No messages in the channel; therefore removed."
+		);
+	}
+
+	/// Split a singular XCM.
+	#[benchmark]
+	fn take_first_concatenated_xcm(
+		n: Linear<0, { MAX_INSTRUCTIONS_TO_DECODE as u32 - MAX_XCM_DECODE_DEPTH }>,
+	) {
+		let mut xcm = Xcm::<T>(vec![ClearOrigin; n as usize]);
+		for _ in 0..MAX_XCM_DECODE_DEPTH - 1 {
+			xcm = Xcm::<T>(vec![Instruction::SetAppendix(xcm)]);
+		}
+		let data = VersionedXcm::<T>::from(xcm).encode();
+
+		#[block]
+		{
+			Pallet::<T>::take_first_concatenated_xcm(&mut &data[..], &mut WeightMeter::new())
+				.unwrap();
+		}
+	}
+
+	/// Benchmark the migration for a maximal sized message.
+	#[benchmark]
+	fn on_idle_good_msg() {
+		use migration::v3;
+
+		let block = 5;
+		let para = ParaId::from(4);
+		let message = vec![123u8; MaxXcmpMessageLenOf::<T>::get() as usize];
+		let message_metadata = vec![(block, XcmpMessageFormat::ConcatenatedVersionedXcm)];
+
+		v3::InboundXcmpMessages::<T>::insert(para, block, message);
+		v3::InboundXcmpStatus::<T>::set(Some(vec![v3::InboundChannelDetails {
+			sender: para,
+			state: v3::InboundState::Ok,
+			message_metadata,
+		}]));
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+	}
+
+	/// Benchmark the migration with a 64 KiB message that will not be possible to enqueue.
+	#[benchmark]
+	fn on_idle_large_msg() {
+		use migration::v3;
+
+		let block = 5;
+		let para = ParaId::from(4);
+		let message = vec![123u8; 1 << 16]; // 64 KiB message
+		let message_metadata = vec![(block, XcmpMessageFormat::ConcatenatedVersionedXcm)];
+
+		v3::InboundXcmpMessages::<T>::insert(para, block, message);
+		v3::InboundXcmpStatus::<T>::set(Some(vec![v3::InboundChannelDetails {
+			sender: para,
+			state: v3::InboundState::Ok,
+			message_metadata,
+		}]));
+
+		#[block]
+		{
+			Pallet::<T>::on_idle(0u32.into(), Weight::MAX);
+		}
+	}
+
+	impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test);
+}

pezcumulus/pezpallets/xcmp-queue/src/bridging.rs

@@ -0,0 +1,101 @@
+// 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 crate::{pallet, OutboundState};
+use pezcumulus_primitives_core::ParaId;
+use xcm::latest::prelude::*;
+
+/// Adapter implementation for `bp_xcm_bridge_hub_router::XcmChannelStatusProvider` which checks
+/// both `OutboundXcmpStatus` and `InboundXcmpStatus` for defined `Location` if any of those is
+/// suspended.
+pub struct InAndOutXcmpChannelStatusProvider<Runtime>(core::marker::PhantomData<Runtime>);
+impl<Runtime: crate::Config> bp_xcm_bridge_hub_router::XcmChannelStatusProvider
+	for InAndOutXcmpChannelStatusProvider<Runtime>
+{
+	fn is_congested(with: &Location) -> bool {
+		// handle congestion only for a sibling teyrchain locations.
+		let sibling_para_id: ParaId = match with.unpack() {
+			(_, [Teyrchain(para_id)]) => (*para_id).into(),
+			_ => return false,
+		};
+
+		// if the inbound channel with recipient is suspended, it means that we are unable to
+		// receive congestion reports from the `with` location. So we assume the pipeline is
+		// congested too.
+		if pallet::Pallet::<Runtime>::is_inbound_channel_suspended(sibling_para_id) {
+			return true;
+		}
+
+		// if the outbound channel with recipient is suspended, it means that one of further
+		// queues (e.g. bridge queue between two bridge hubs) is overloaded, so we shall
+		// take larger fee for our outbound messages
+		OutXcmpChannelStatusProvider::<Runtime>::is_congested(with)
+	}
+}
+
+/// Adapter implementation for `bp_xcm_bridge::ChannelStatusProvider` and/or
+/// `bp_xcm_bridge_hub_router::XcmChannelStatusProvider` which checks only `OutboundXcmpStatus`
+/// for defined `Location` if is suspended.
+pub struct OutXcmpChannelStatusProvider<Runtime>(core::marker::PhantomData<Runtime>);
+impl<Runtime: crate::Config> OutXcmpChannelStatusProvider<Runtime> {
+	fn is_congested(with: &Location) -> bool {
+		// handle congestion only for a sibling teyrchain locations.
+		let sibling_para_id: ParaId = match with.unpack() {
+			(_, [Teyrchain(para_id)]) => (*para_id).into(),
+			_ => return false,
+		};
+
+		// let's find the channel's state with the sibling teyrchain,
+		let Some((outbound_state, queued_pages)) =
+			pallet::Pallet::<Runtime>::outbound_channel_state(sibling_para_id)
+		else {
+			return false;
+		};
+		// suspended channel => it is congested
+		if outbound_state == OutboundState::Suspended {
+			return true;
+		}
+
+		// It takes some time for target teyrchain to suspend inbound channel with the target BH and
+		// during that we will keep accepting new message delivery transactions. Let's also reject
+		// new deliveries if there are too many "pages" (concatenated XCM messages) in the target BH
+		// -> target teyrchain queue.
+
+		// If the outbound channel has at least `N` pages enqueued, let's assume it is congested.
+		// Normally, the chain with a few opened HRMP channels, will "send" pages at every block.
+		// Having `N` pages means that for last `N` blocks we either have not sent any messages,
+		// or have sent signals.
+
+		const MAX_QUEUED_PAGES_BEFORE_DEACTIVATION: u16 = 4;
+		if queued_pages > MAX_QUEUED_PAGES_BEFORE_DEACTIVATION {
+			return true;
+		}
+
+		false
+	}
+}
+
+impl<Runtime: crate::Config> bp_xcm_bridge_hub_router::XcmChannelStatusProvider
+	for OutXcmpChannelStatusProvider<Runtime>
+{
+	fn is_congested(with: &Location) -> bool {
+		Self::is_congested(with)
+	}
+}
+
+#[cfg(feature = "runtime-benchmarks")]
+pub fn suspend_channel_for_benchmarks<T: crate::Config>(target: ParaId) {
+	pallet::Pallet::<T>::suspend_channel(target)
+}

pezcumulus/pezpallets/xcmp-queue/src/migration.rs

@@ -0,0 +1,427 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! A module that is responsible for migration of storage.
+
+pub mod v5;
+
+use crate::{Config, OverweightIndex, Pallet, QueueConfig, QueueConfigData, DEFAULT_POV_SIZE};
+use alloc::vec::Vec;
+use pezcumulus_primitives_core::XcmpMessageFormat;
+use pezframe_support::{
+	pezpallet_prelude::*,
+	traits::{EnqueueMessage, StorageVersion, UncheckedOnRuntimeUpgrade},
+	weights::{constants::WEIGHT_REF_TIME_PER_MILLIS, Weight},
+};
+
+/// The in-code storage version.
+pub const STORAGE_VERSION: StorageVersion = StorageVersion::new(5);
+
+pub const LOG: &str = "runtime::xcmp-queue-migration";
+
+mod v1 {
+	use super::*;
+	use codec::{Decode, Encode};
+
+	#[pezframe_support::storage_alias]
+	pub(crate) type QueueConfig<T: Config> = StorageValue<Pallet<T>, QueueConfigData, ValueQuery>;
+
+	#[derive(Encode, Decode, Debug)]
+	pub struct QueueConfigData {
+		pub suspend_threshold: u32,
+		pub drop_threshold: u32,
+		pub resume_threshold: u32,
+		pub threshold_weight: u64,
+		pub weight_restrict_decay: u64,
+		pub xcmp_max_individual_weight: u64,
+	}
+
+	impl Default for QueueConfigData {
+		fn default() -> Self {
+			QueueConfigData {
+				suspend_threshold: 2,
+				drop_threshold: 5,
+				resume_threshold: 1,
+				threshold_weight: 100_000,
+				weight_restrict_decay: 2,
+				xcmp_max_individual_weight: 20u64 * WEIGHT_REF_TIME_PER_MILLIS,
+			}
+		}
+	}
+}
+
+pub mod v2 {
+	use super::*;
+
+	#[pezframe_support::storage_alias]
+	pub(crate) type QueueConfig<T: Config> = StorageValue<Pallet<T>, QueueConfigData, ValueQuery>;
+
+	#[derive(Copy, Clone, Eq, PartialEq, Encode, Decode, RuntimeDebug, TypeInfo)]
+	pub struct QueueConfigData {
+		pub suspend_threshold: u32,
+		pub drop_threshold: u32,
+		pub resume_threshold: u32,
+		pub threshold_weight: Weight,
+		pub weight_restrict_decay: Weight,
+		pub xcmp_max_individual_weight: Weight,
+	}
+
+	impl Default for QueueConfigData {
+		fn default() -> Self {
+			Self {
+				suspend_threshold: 2,
+				drop_threshold: 5,
+				resume_threshold: 1,
+				threshold_weight: Weight::from_parts(100_000, 0),
+				weight_restrict_decay: Weight::from_parts(2, 0),
+				xcmp_max_individual_weight: Weight::from_parts(
+					20u64 * WEIGHT_REF_TIME_PER_MILLIS,
+					DEFAULT_POV_SIZE,
+				),
+			}
+		}
+	}
+
+	/// Migrates `QueueConfigData` from v1 (using only reference time weights) to v2 (with
+	/// 2D weights).
+	pub struct UncheckedMigrationToV2<T: Config>(PhantomData<T>);
+
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV2<T> {
+		#[allow(deprecated)]
+		fn on_runtime_upgrade() -> Weight {
+			let translate = |pre: v1::QueueConfigData| -> v2::QueueConfigData {
+				v2::QueueConfigData {
+					suspend_threshold: pre.suspend_threshold,
+					drop_threshold: pre.drop_threshold,
+					resume_threshold: pre.resume_threshold,
+					threshold_weight: Weight::from_parts(pre.threshold_weight, 0),
+					weight_restrict_decay: Weight::from_parts(pre.weight_restrict_decay, 0),
+					xcmp_max_individual_weight: Weight::from_parts(
+						pre.xcmp_max_individual_weight,
+						DEFAULT_POV_SIZE,
+					),
+				}
+			};
+
+			if v2::QueueConfig::<T>::translate(|pre| pre.map(translate)).is_err() {
+				tracing::error!(
+					target: crate::LOG_TARGET,
+					"unexpected error when performing translation of the QueueConfig type \
+					during storage upgrade to v2"
+				);
+			}
+
+			T::DbWeight::get().reads_writes(1, 1)
+		}
+	}
+
+	/// [`UncheckedMigrationToV2`] wrapped in a
+	/// [`VersionedMigration`](pezframe_support::migrations::VersionedMigration), ensuring the
+	/// migration is only performed when on-chain version is 1.
+	#[allow(dead_code)]
+	pub type MigrationToV2<T> = pezframe_support::migrations::VersionedMigration<
+		1,
+		2,
+		UncheckedMigrationToV2<T>,
+		Pallet<T>,
+		<T as pezframe_system::Config>::DbWeight,
+	>;
+}
+
+pub mod v3 {
+	use super::*;
+	use crate::*;
+
+	/// Status of the inbound XCMP channels.
+	#[pezframe_support::storage_alias]
+	pub(crate) type InboundXcmpStatus<T: Config> =
+		StorageValue<Pallet<T>, Vec<InboundChannelDetails>, OptionQuery>;
+
+	/// Inbound aggregate XCMP messages. It can only be one per ParaId/block.
+	#[pezframe_support::storage_alias]
+	pub(crate) type InboundXcmpMessages<T: Config> = StorageDoubleMap<
+		Pallet<T>,
+		Blake2_128Concat,
+		ParaId,
+		Twox64Concat,
+		RelayBlockNumber,
+		Vec<u8>,
+		OptionQuery,
+	>;
+
+	#[pezframe_support::storage_alias]
+	pub(crate) type QueueConfig<T: Config> =
+		StorageValue<Pallet<T>, v2::QueueConfigData, ValueQuery>;
+
+	#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, TypeInfo)]
+	pub struct InboundChannelDetails {
+		/// The `ParaId` of the teyrchain that this channel is connected with.
+		pub sender: ParaId,
+		/// The state of the channel.
+		pub state: InboundState,
+		/// The ordered metadata of each inbound message.
+		///
+		/// Contains info about the relay block number that the message was sent at, and the format
+		/// of the incoming message.
+		pub message_metadata: Vec<(RelayBlockNumber, XcmpMessageFormat)>,
+	}
+
+	#[derive(
+		Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, RuntimeDebug, TypeInfo,
+	)]
+	pub enum InboundState {
+		Ok,
+		Suspended,
+	}
+
+	/// Migrates the pallet storage to v3.
+	pub struct UncheckedMigrationToV3<T: Config>(PhantomData<T>);
+
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV3<T> {
+		fn on_runtime_upgrade() -> Weight {
+			#[pezframe_support::storage_alias]
+			type Overweight<T: Config> =
+				CountedStorageMap<Pallet<T>, Twox64Concat, OverweightIndex, ParaId>;
+			let overweight_messages = Overweight::<T>::initialize_counter() as u64;
+
+			T::DbWeight::get().reads_writes(overweight_messages, 1)
+		}
+	}
+
+	/// [`UncheckedMigrationToV3`] wrapped in a
+	/// [`VersionedMigration`](pezframe_support::migrations::VersionedMigration), ensuring the
+	/// migration is only performed when on-chain version is 2.
+	pub type MigrationToV3<T> = pezframe_support::migrations::VersionedMigration<
+		2,
+		3,
+		UncheckedMigrationToV3<T>,
+		Pallet<T>,
+		<T as pezframe_system::Config>::DbWeight,
+	>;
+
+	pub fn lazy_migrate_inbound_queue<T: Config>() {
+		let Some(mut states) = v3::InboundXcmpStatus::<T>::get() else {
+			tracing::debug!(target: LOG, "Lazy migration finished: item gone");
+			return;
+		};
+		let Some(ref mut next) = states.first_mut() else {
+			tracing::debug!(target: LOG, "Lazy migration finished: item empty");
+			v3::InboundXcmpStatus::<T>::kill();
+			return;
+		};
+		tracing::debug!(
+			target: LOG,
+			sibling=?next.sender,
+			msgs_left=%next.message_metadata.len(),
+			"Migrating inbound HRMP channel."
+		);
+		// We take the last element since the MQ is a FIFO and we want to keep the order.
+		let Some((block_number, format)) = next.message_metadata.pop() else {
+			states.remove(0);
+			v3::InboundXcmpStatus::<T>::put(states);
+			return;
+		};
+		if format != XcmpMessageFormat::ConcatenatedVersionedXcm {
+			tracing::warn!(
+				target: LOG,
+				?format,
+				"Dropping message (not ConcatenatedVersionedXcm)"
+			);
+			v3::InboundXcmpMessages::<T>::remove(&next.sender, &block_number);
+			v3::InboundXcmpStatus::<T>::put(states);
+			return;
+		}
+
+		let Some(msg) = v3::InboundXcmpMessages::<T>::take(&next.sender, &block_number) else {
+			defensive!("Storage corrupted: HRMP message missing:", (next.sender, block_number));
+			v3::InboundXcmpStatus::<T>::put(states);
+			return;
+		};
+
+		let Ok(msg): Result<BoundedVec<_, _>, _> = msg.try_into() else {
+			tracing::error!(target: LOG, "Message dropped: too big");
+			v3::InboundXcmpStatus::<T>::put(states);
+			return;
+		};
+
+		// Finally! We have a proper message.
+		T::XcmpQueue::enqueue_message(msg.as_bounded_slice(), next.sender);
+		tracing::debug!(target: LOG, next_sender=?next.sender, ?block_number, "Migrated HRMP message to MQ");
+		v3::InboundXcmpStatus::<T>::put(states);
+	}
+}
+
+pub mod v4 {
+	use super::*;
+
+	/// Migrates `QueueConfigData` to v4, removing deprecated fields and bumping page
+	/// thresholds to at least the default values.
+	pub struct UncheckedMigrationToV4<T: Config>(PhantomData<T>);
+
+	impl<T: Config> UncheckedOnRuntimeUpgrade for UncheckedMigrationToV4<T> {
+		fn on_runtime_upgrade() -> Weight {
+			let translate = |pre: v2::QueueConfigData| -> QueueConfigData {
+				let pre_default = v2::QueueConfigData::default();
+				// If the previous values are the default ones, let's replace them with the new
+				// default.
+				if pre.suspend_threshold == pre_default.suspend_threshold &&
+					pre.drop_threshold == pre_default.drop_threshold &&
+					pre.resume_threshold == pre_default.resume_threshold
+				{
+					return QueueConfigData::default();
+				}
+
+				// If the previous values are not the default ones, let's leave them as they are.
+				QueueConfigData {
+					suspend_threshold: pre.suspend_threshold,
+					drop_threshold: pre.drop_threshold,
+					resume_threshold: pre.resume_threshold,
+				}
+			};
+
+			if QueueConfig::<T>::translate(|pre| pre.map(translate)).is_err() {
+				tracing::error!(
+					target: crate::LOG_TARGET,
+					"unexpected error when performing translation of the QueueConfig type \
+					during storage upgrade to v4"
+				);
+			}
+
+			T::DbWeight::get().reads_writes(1, 1)
+		}
+	}
+
+	/// [`UncheckedMigrationToV4`] wrapped in a
+	/// [`VersionedMigration`](pezframe_support::migrations::VersionedMigration), ensuring the
+	/// migration is only performed when on-chain version is 3.
+	pub type MigrationToV4<T> = pezframe_support::migrations::VersionedMigration<
+		3,
+		4,
+		UncheckedMigrationToV4<T>,
+		Pallet<T>,
+		<T as pezframe_system::Config>::DbWeight,
+	>;
+}
+
+#[cfg(all(feature = "try-runtime", test))]
+mod tests {
+	use super::*;
+	use crate::mock::{new_test_ext, Test};
+	use pezframe_support::traits::OnRuntimeUpgrade;
+
+	#[test]
+	#[allow(deprecated)]
+	fn test_migration_to_v2() {
+		let v1 = v1::QueueConfigData {
+			suspend_threshold: 5,
+			drop_threshold: 12,
+			resume_threshold: 3,
+			threshold_weight: 333_333,
+			weight_restrict_decay: 1,
+			xcmp_max_individual_weight: 10_000_000_000,
+		};
+
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(1);
+			storage_version.put::<Pallet<Test>>();
+
+			pezframe_support::storage::unhashed::put_raw(
+				&crate::QueueConfig::<Test>::hashed_key(),
+				&v1.encode(),
+			);
+
+			let bytes = v2::MigrationToV2::<Test>::pre_upgrade();
+			assert!(bytes.is_ok());
+			v2::MigrationToV2::<Test>::on_runtime_upgrade();
+			assert!(v2::MigrationToV2::<Test>::post_upgrade(bytes.unwrap()).is_ok());
+
+			let v2 = v2::QueueConfig::<Test>::get();
+
+			assert_eq!(v1.suspend_threshold, v2.suspend_threshold);
+			assert_eq!(v1.drop_threshold, v2.drop_threshold);
+			assert_eq!(v1.resume_threshold, v2.resume_threshold);
+			assert_eq!(v1.threshold_weight, v2.threshold_weight.ref_time());
+			assert_eq!(v1.weight_restrict_decay, v2.weight_restrict_decay.ref_time());
+			assert_eq!(v1.xcmp_max_individual_weight, v2.xcmp_max_individual_weight.ref_time());
+		});
+	}
+
+	#[test]
+	#[allow(deprecated)]
+	fn test_migration_to_v4() {
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(3);
+			storage_version.put::<Pallet<Test>>();
+
+			let v2 = v2::QueueConfigData {
+				drop_threshold: 5,
+				suspend_threshold: 2,
+				resume_threshold: 1,
+				..Default::default()
+			};
+
+			pezframe_support::storage::unhashed::put_raw(
+				&crate::QueueConfig::<Test>::hashed_key(),
+				&v2.encode(),
+			);
+
+			let bytes = v4::MigrationToV4::<Test>::pre_upgrade();
+			assert!(bytes.is_ok());
+			v4::MigrationToV4::<Test>::on_runtime_upgrade();
+			assert!(v4::MigrationToV4::<Test>::post_upgrade(bytes.unwrap()).is_ok());
+
+			let v4 = QueueConfig::<Test>::get();
+
+			assert_eq!(
+				v4,
+				QueueConfigData { suspend_threshold: 32, drop_threshold: 48, resume_threshold: 8 }
+			);
+		});
+
+		new_test_ext().execute_with(|| {
+			let storage_version = StorageVersion::new(3);
+			storage_version.put::<Pallet<Test>>();
+
+			let v2 = v2::QueueConfigData {
+				drop_threshold: 100,
+				suspend_threshold: 50,
+				resume_threshold: 40,
+				..Default::default()
+			};
+
+			pezframe_support::storage::unhashed::put_raw(
+				&crate::QueueConfig::<Test>::hashed_key(),
+				&v2.encode(),
+			);
+
+			let bytes = v4::MigrationToV4::<Test>::pre_upgrade();
+			assert!(bytes.is_ok());
+			v4::MigrationToV4::<Test>::on_runtime_upgrade();
+			assert!(v4::MigrationToV4::<Test>::post_upgrade(bytes.unwrap()).is_ok());
+
+			let v4 = QueueConfig::<Test>::get();
+
+			assert_eq!(
+				v4,
+				QueueConfigData {
+					suspend_threshold: 50,
+					drop_threshold: 100,
+					resume_threshold: 40
+				}
+			);
+		});
+	}
+}

pezcumulus/pezpallets/xcmp-queue/src/migration/v5.rs

@@ -0,0 +1,109 @@
+// Copyright (C) Parity Technologies (UK) Ltd.
+// This file is part of Pezcumulus.
+// 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.
+
+//! Migrates the storage to version 5.
+
+use crate::*;
+use alloc::vec::Vec;
+use pezcumulus_primitives_core::ListChannelInfos;
+use pezframe_support::{pezpallet_prelude::*, traits::UncheckedOnRuntimeUpgrade};
+
+/// Configs needed to run the V5 migration.
+pub trait V5Config: Config {
+	/// List all outbound channels with their target `ParaId` and maximum message size.
+	type ChannelList: ListChannelInfos;
+}
+
+/// Ensures that the storage migrates cleanly to V5.
+///
+/// The migration itself is a no-op, but it checks that none of the `BoundedVec`s would truncate on
+/// the next decode after the upgrade was applied.
+pub type MigrateV4ToV5<T> = pezframe_support::migrations::VersionedMigration<
+	4,
+	5,
+	unversioned::UncheckedMigrateV4ToV5<T>,
+	Pallet<T>,
+	<T as pezframe_system::Config>::DbWeight,
+>;
+
+// V4 storage aliases
+mod v4 {
+	use super::*;
+
+	#[pezframe_support::storage_alias]
+	pub(super) type OutboundXcmpStatus<T: Config> =
+		StorageValue<Pallet<T>, Vec<OutboundChannelDetails>, ValueQuery>;
+
+	#[pezframe_support::storage_alias]
+	pub(super) type OutboundXcmpMessages<T: Config> = StorageDoubleMap<
+		Pallet<T>,
+		Blake2_128Concat,
+		ParaId,
+		Twox64Concat,
+		u16,
+		Vec<u8>,
+		ValueQuery,
+	>;
+
+	#[pezframe_support::storage_alias]
+	pub(super) type SignalMessages<T: Config> =
+		StorageMap<Pallet<T>, Blake2_128Concat, ParaId, Vec<u8>, ValueQuery>;
+}
+
+// Private module to hide the migration.
+mod unversioned {
+	/// Please use [`MigrateV4ToV5`] instead.
+	pub struct UncheckedMigrateV4ToV5<T: super::V5Config>(core::marker::PhantomData<T>);
+}
+
+impl<T: V5Config> UncheckedOnRuntimeUpgrade for unversioned::UncheckedMigrateV4ToV5<T> {
+	fn on_runtime_upgrade() -> pezframe_support::weights::Weight {
+		Default::default()
+	}
+
+	#[cfg(feature = "try-runtime")]
+	fn post_upgrade(_: Vec<u8>) -> Result<(), pezsp_runtime::DispatchError> {
+		// We dont need any front-run protection for this since channels are opened by governance.
+		ensure!(
+			v4::OutboundXcmpStatus::<T>::get().len() as u32 <= T::MaxActiveOutboundChannels::get(),
+			"Too many outbound channels. Close some channels or increase `MaxActiveOutboundChannels`."
+		);
+
+		ensure!(T::MaxPageSize::get() >= 16, "Sanity check failed: MaxPageSize too small");
+
+		// Check if any channels have a too large message max sizes.
+		let max_msg_len = T::MaxPageSize::get() - XcmpMessageFormat::max_encoded_len() as u32;
+		for channel in T::ChannelList::outgoing_channels() {
+			let info = T::ChannelInfo::get_channel_info(channel)
+				.expect("All listed channels must provide info");
+
+			if info.max_message_size > max_msg_len {
+				tracing::error!(
+					target: "runtime::xcmp-queue-migration::v5",
+					channel_max=%info.max_message_size,
+					max_page_size=%max_msg_len,
+					"Max message size for channel is too large. This means that the V5 \
+					migration can be front-run and an attacker could place a large message just right \
+					before the migration to make other messages un-decodable. Please either increase \
+					`MaxPageSize` or decrease the `max_message_size` for this channel."
+				);
+				return Err("Migration can be front-run".into());
+			}
+		}
+
+		Ok(())
+	}
+}

pezcumulus/pezpallets/xcmp-queue/src/mock.rs

@@ -0,0 +1,303 @@
+// 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 crate as xcmp_queue;
+use alloc::collections::BTreeMap;
+use core::marker::PhantomData;
+use pezcumulus_pezpallet_teyrchain_system::AnyRelayNumber;
+use pezcumulus_primitives_core::{ChannelInfo, IsSystem, ParaId};
+use pezframe_support::{
+	derive_impl, parameter_types,
+	traits::{BatchesFootprints, ConstU32, Everything, OriginTrait},
+	BoundedSlice,
+};
+use pezframe_system::EnsureRoot;
+use pezsp_core::H256;
+use pezsp_runtime::{
+	traits::{BlakeTwo256, IdentityLookup},
+	BuildStorage,
+};
+use xcm::prelude::*;
+use xcm_executor::traits::ConvertOrigin;
+
+type Block = pezframe_system::mocking::MockBlock<Test>;
+
+// Configure a mock runtime to test the pallet.
+pezframe_support::construct_runtime!(
+	pub enum Test
+	{
+		System: pezframe_system::{Pallet, Call, Config<T>, Storage, Event<T>},
+		Balances: pezpallet_balances::{Pallet, Call, Storage, Config<T>, Event<T>},
+		TeyrchainSystem: pezcumulus_pezpallet_teyrchain_system::{
+			Pallet, Call, Config<T>, Storage, Inherent, Event<T>,
+		},
+		XcmpQueue: xcmp_queue::{Pallet, Call, Storage, Event<T>},
+	}
+);
+
+parameter_types! {
+	pub const SS58Prefix: u8 = 42;
+}
+
+type AccountId = u64;
+
+#[derive_impl(pezframe_system::config_preludes::TestDefaultConfig)]
+impl pezframe_system::Config for Test {
+	type BaseCallFilter = Everything;
+	type BlockWeights = ();
+	type BlockLength = ();
+	type DbWeight = ();
+	type RuntimeOrigin = RuntimeOrigin;
+	type RuntimeCall = RuntimeCall;
+	type Nonce = u64;
+	type Hash = H256;
+	type Hashing = BlakeTwo256;
+	type AccountId = AccountId;
+	type Lookup = IdentityLookup<Self::AccountId>;
+	type Block = Block;
+	type RuntimeEvent = RuntimeEvent;
+	type Version = ();
+	type PalletInfo = PalletInfo;
+	type AccountData = pezpallet_balances::AccountData<u64>;
+	type OnNewAccount = ();
+	type OnKilledAccount = ();
+	type SystemWeightInfo = ();
+	type SS58Prefix = SS58Prefix;
+	type OnSetCode = pezcumulus_pezpallet_teyrchain_system::TeyrchainSetCode<Test>;
+	type MaxConsumers = pezframe_support::traits::ConstU32<16>;
+}
+
+parameter_types! {
+	pub const ExistentialDeposit: u64 = 5;
+}
+
+pub type Balance = u64;
+
+#[derive_impl(pezpallet_balances::config_preludes::TestDefaultConfig)]
+impl pezpallet_balances::Config for Test {
+	type ExistentialDeposit = ExistentialDeposit;
+	type AccountStore = System;
+}
+
+impl pezcumulus_pezpallet_teyrchain_system::Config for Test {
+	type WeightInfo = ();
+	type RuntimeEvent = RuntimeEvent;
+	type OnSystemEvent = ();
+	type SelfParaId = ();
+	type OutboundXcmpMessageSource = XcmpQueue;
+	// Ignore all DMP messages by enqueueing them into `()`:
+	type DmpQueue = pezframe_support::traits::EnqueueWithOrigin<(), pezsp_core::ConstU8<0>>;
+	type ReservedDmpWeight = ();
+	type XcmpMessageHandler = XcmpQueue;
+	type ReservedXcmpWeight = ();
+	type CheckAssociatedRelayNumber = AnyRelayNumber;
+	type ConsensusHook = pezcumulus_pezpallet_teyrchain_system::consensus_hook::ExpectParentIncluded;
+	type RelayParentOffset = ConstU32<0>;
+}
+
+parameter_types! {
+	pub const RelayChain: Location = Location::parent();
+	pub UniversalLocation: InteriorLocation = [Teyrchain(1u32)].into();
+	pub UnitWeightCost: Weight = Weight::from_parts(1_000_000, 1024);
+	pub const MaxInstructions: u32 = 100;
+	pub const MaxAssetsIntoHolding: u32 = 64;
+}
+
+pub struct SystemTeyrchainAsSuperuser<RuntimeOrigin>(PhantomData<RuntimeOrigin>);
+impl<RuntimeOrigin: OriginTrait> ConvertOrigin<RuntimeOrigin>
+	for SystemTeyrchainAsSuperuser<RuntimeOrigin>
+{
+	fn convert_origin(
+		origin: impl Into<Location>,
+		kind: OriginKind,
+	) -> Result<RuntimeOrigin, Location> {
+		let origin = origin.into();
+		if kind == OriginKind::Superuser &&
+			matches!(
+				origin.unpack(),
+				(1,	[Teyrchain(id)]) if ParaId::from(*id).is_system(),
+			) {
+			Ok(RuntimeOrigin::root())
+		} else {
+			Err(origin)
+		}
+	}
+}
+
+parameter_types! {
+	pub static EnqueuedMessages: Vec<(ParaId, Vec<u8>)> = Default::default();
+	pub static FirstPagePos: BTreeMap<ParaId, usize> = Default::default();
+}
+
+/// An `EnqueueMessage` implementation that puts all messages in thread-local storage.
+pub struct EnqueueToLocalStorage<T>(PhantomData<T>);
+
+impl<T: OnQueueChanged<ParaId>> EnqueueMessage<ParaId> for EnqueueToLocalStorage<T> {
+	type MaxMessageLen = pezsp_core::ConstU32<256>;
+
+	fn enqueue_message(message: BoundedSlice<u8, Self::MaxMessageLen>, origin: ParaId) {
+		let mut msgs = EnqueuedMessages::get();
+		msgs.push((origin, message.to_vec()));
+		EnqueuedMessages::set(msgs);
+		T::on_queue_changed(origin, Self::footprint(origin));
+	}
+
+	fn enqueue_messages<'a>(
+		iter: impl Iterator<Item = BoundedSlice<'a, u8, Self::MaxMessageLen>>,
+		origin: ParaId,
+	) {
+		let mut msgs = EnqueuedMessages::get();
+		msgs.extend(iter.map(|m| (origin, m.to_vec())));
+		EnqueuedMessages::set(msgs);
+		T::on_queue_changed(origin, Self::footprint(origin));
+	}
+
+	fn sweep_queue(origin: ParaId) {
+		let mut msgs = EnqueuedMessages::get();
+		msgs.retain(|(o, _)| o != &origin);
+		EnqueuedMessages::set(msgs);
+		T::on_queue_changed(origin, Self::footprint(origin));
+	}
+}
+
+impl<T: OnQueueChanged<ParaId>> QueueFootprintQuery<ParaId> for EnqueueToLocalStorage<T> {
+	type MaxMessageLen = pezsp_core::ConstU32<256>;
+
+	fn footprint(origin: ParaId) -> QueueFootprint {
+		let msgs = EnqueuedMessages::get();
+		let mut footprint = QueueFootprint::default();
+		for (o, m) in msgs {
+			if o == origin {
+				footprint.storage.count += 1;
+				footprint.storage.size += m.len() as u64;
+			}
+		}
+		// Let's consider that we add one message per page
+		footprint.pages = footprint.storage.count as u32;
+		footprint.ready_pages = footprint.pages;
+		footprint
+	}
+
+	fn get_batches_footprints<'a>(
+		origin: ParaId,
+		msgs: impl Iterator<Item = BoundedSlice<'a, u8, Self::MaxMessageLen>>,
+		total_pages_limit: u32,
+	) -> BatchesFootprints {
+		// Let's consider that we add one message per page
+		let footprint = Self::footprint(origin);
+		let mut batches_footprints = BatchesFootprints {
+			first_page_pos: *FirstPagePos::get().entry(origin).or_default(),
+			footprints: vec![],
+		};
+		for (idx, msg) in msgs.enumerate() {
+			if footprint.pages + idx as u32 + 1 > total_pages_limit {
+				break;
+			}
+
+			batches_footprints.push(msg.into(), true);
+		}
+		batches_footprints
+	}
+}
+
+parameter_types! {
+	/// The asset ID for the asset that we use to pay for message delivery fees.
+	pub FeeAssetId: AssetId = AssetId(RelayChain::get());
+	/// The base fee for the message delivery fees.
+	pub const BaseDeliveryFee: Balance = 300_000_000;
+	/// The fee per byte
+	pub const ByteFee: Balance = 1_000_000;
+}
+
+pub type PriceForSiblingTeyrchainDelivery = pezkuwi_runtime_common::xcm_sender::ExponentialPrice<
+	FeeAssetId,
+	BaseDeliveryFee,
+	ByteFee,
+	XcmpQueue,
+>;
+
+impl Config for Test {
+	type RuntimeEvent = RuntimeEvent;
+	type ChannelInfo = MockedChannelInfo;
+	type VersionWrapper = ();
+	type XcmpQueue = EnqueueToLocalStorage<Pallet<Test>>;
+	type MaxInboundSuspended = ConstU32<1_000>;
+	type MaxActiveOutboundChannels = ConstU32<128>;
+	// Most on-chain HRMP channels are configured to use 102400 bytes of max message size, so we
+	// need to set the page size larger than that until we reduce the channel size on-chain.
+	type MaxPageSize = ConstU32<{ 103 * 1024 }>;
+	type ControllerOrigin = EnsureRoot<AccountId>;
+	type ControllerOriginConverter = SystemTeyrchainAsSuperuser<RuntimeOrigin>;
+	type WeightInfo = ();
+	type PriceForSiblingDelivery = PriceForSiblingTeyrchainDelivery;
+}
+
+pub fn new_test_ext() -> pezsp_io::TestExternalities {
+	pezframe_system::GenesisConfig::<Test>::default().build_storage().unwrap().into()
+}
+
+/// A para that we have an HRMP channel with.
+pub const HRMP_PARA_ID: u32 = 7777;
+
+pub struct MockedChannelInfo;
+impl GetChannelInfo for MockedChannelInfo {
+	fn get_channel_status(id: ParaId) -> ChannelStatus {
+		if id == HRMP_PARA_ID.into() {
+			return ChannelStatus::Ready(usize::MAX, usize::MAX);
+		}
+
+		TeyrchainSystem::get_channel_status(id)
+	}
+
+	fn get_channel_info(id: ParaId) -> Option<ChannelInfo> {
+		if id == HRMP_PARA_ID.into() {
+			return Some(ChannelInfo {
+				max_capacity: u32::MAX,
+				max_total_size: u32::MAX,
+				max_message_size: u32::MAX,
+				msg_count: 0,
+				total_size: 0,
+			});
+		}
+
+		TeyrchainSystem::get_channel_info(id)
+	}
+}
+
+pub(crate) fn mk_page() -> Vec<u8> {
+	let mut page = Vec::<u8>::new();
+
+	let newer_xcm_version = xcm::prelude::XCM_VERSION;
+	let older_xcm_version = newer_xcm_version - 1;
+
+	for i in 0..100 {
+		page.extend(match i % 2 {
+			0 => versioned_xcm(older_xcm_version).encode(),
+			1 => versioned_xcm(newer_xcm_version).encode(),
+			// We cannot push an undecodable XCM here since it would break the decode stream.
+			_ => unreachable!(),
+		});
+	}
+
+	page
+}
+
+pub(crate) fn versioned_xcm(version: XcmVersion) -> VersionedXcm<()> {
+	let instr = Instruction::<()>::Trap(1);
+	VersionedXcm::from(Xcm::<()>(vec![instr; 3]))
+		.into_version(version)
+		.expect("Version conversion should work")
+}

pezcumulus/pezpallets/xcmp-queue/src/weights.rs

@@ -0,0 +1,458 @@
+// 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.
+
+//! Autogenerated weights for `pezcumulus_pezpallet_xcmp_queue`
+//!
+//! THIS FILE WAS AUTO-GENERATED USING THE BIZINIKIWI BENCHMARK CLI VERSION 32.0.0
+//! DATE: 2025-09-04, STEPS: `50`, REPEAT: `20`, LOW RANGE: `[]`, HIGH RANGE: `[]`
+//! WORST CASE MAP SIZE: `1000000`
+//! HOSTNAME: `Serbans-MacBook-Pro.local`, CPU: `<UNKNOWN>`
+//! WASM-EXECUTION: `Compiled`, CHAIN: `None`, DB CACHE: `1024`
+
+// Executed Command:
+// ./target/release/pezkuwi-teyrchain
+// benchmark
+// pallet
+// --pallet
+// pezcumulus-pezpallet-xcmp-queue
+// --chain
+// asset-hub-zagros-dev
+// --output
+// pezcumulus/pezpallets/xcmp-queue/src/weights.rs
+// --template
+// bizinikiwi/.maintain/frame-weight-template.hbs
+// --extrinsic
+// 
+
+#![cfg_attr(rustfmt, rustfmt_skip)]
+#![allow(unused_parens)]
+#![allow(unused_imports)]
+#![allow(missing_docs)]
+#![allow(dead_code)]
+
+use pezframe_support::{traits::Get, weights::{Weight, constants::RocksDbWeight}};
+use core::marker::PhantomData;
+
+/// Weight functions needed for `pezcumulus_pezpallet_xcmp_queue`.
+pub trait WeightInfo {
+	fn set_config_with_u32() -> Weight;
+	fn enqueue_n_bytes_xcmp_message(n: u32, ) -> Weight;
+	fn enqueue_n_empty_xcmp_messages(n: u32, ) -> Weight;
+	fn enqueue_empty_xcmp_message_at(n: u32, ) -> Weight;
+	fn enqueue_n_full_pages(n: u32, ) -> Weight;
+	fn enqueue_1000_small_xcmp_messages() -> Weight;
+	fn suspend_channel() -> Weight;
+	fn resume_channel() -> Weight;
+	fn take_first_concatenated_xcm(n: u32, ) -> Weight;
+	fn on_idle_good_msg() -> Weight;
+	fn on_idle_large_msg() -> Weight;
+}
+
+/// Weights for `pezcumulus_pezpallet_xcmp_queue` using the Bizinikiwi node and recommended hardware.
+pub struct BizinikiwiWeight<T>(PhantomData<T>);
+impl<T: pezframe_system::Config> WeightInfo for BizinikiwiWeight<T> {
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:1)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	fn set_config_with_u32() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `175`
+		//  Estimated: `1497`
+		// Minimum execution time: 3_000_000 picoseconds.
+		Weight::from_parts(4_000_000, 1497)
+			.saturating_add(T::DbWeight::get().reads(1_u64))
+			.saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 105467]`.
+	fn enqueue_n_bytes_xcmp_message(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `251`
+		//  Estimated: `5487`
+		// Minimum execution time: 11_000_000 picoseconds.
+		Weight::from_parts(12_262_916, 5487)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(187, 0).saturating_mul(n.into()))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(3_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 1000]`.
+	fn enqueue_n_empty_xcmp_messages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `251`
+		//  Estimated: `5487`
+		// Minimum execution time: 9_000_000 picoseconds.
+		Weight::from_parts(13_664_993, 5487)
+			// Standard Error: 220
+			.saturating_add(Weight::from_parts(131_195, 0).saturating_mul(n.into()))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(3_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `Measured`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `Measured`)
+	/// Storage: `MessageQueue::Pages` (r:1 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `Measured`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `Measured`)
+	/// The range of component `n` is `[0, 105457]`.
+	fn enqueue_empty_xcmp_message_at(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `434 + n * (1 ±0)`
+		//  Estimated: `3897 + n * (1 ±0)`
+		// Minimum execution time: 16_000_000 picoseconds.
+		Weight::from_parts(16_490_667, 3897)
+			// Standard Error: 2
+			.saturating_add(Weight::from_parts(720, 0).saturating_mul(n.into()))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+			.saturating_add(Weight::from_parts(0, 1).saturating_mul(n.into()))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:100)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 100]`.
+	fn enqueue_n_full_pages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `286`
+		//  Estimated: `5487`
+		// Minimum execution time: 9_000_000 picoseconds.
+		Weight::from_parts(10_000_000, 5487)
+			// Standard Error: 26_047
+			.saturating_add(Weight::from_parts(23_938_725, 0).saturating_mul(n.into()))
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+			.saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(n.into())))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `Measured`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `Measured`)
+	/// Storage: `MessageQueue::Pages` (r:1 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `Measured`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `Measured`)
+	fn enqueue_1000_small_xcmp_messages() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `53167`
+		//  Estimated: `56632`
+		// Minimum execution time: 175_000_000 picoseconds.
+		Weight::from_parts(181_000_000, 56632)
+			.saturating_add(T::DbWeight::get().reads(4_u64))
+			.saturating_add(T::DbWeight::get().writes(2_u64))
+	}
+	/// Storage: `XcmpQueue::OutboundXcmpStatus` (r:1 w:1)
+	/// Proof: `XcmpQueue::OutboundXcmpStatus` (`max_values`: Some(1), `max_size`: Some(1282), added: 1777, mode: `MaxEncodedLen`)
+	fn suspend_channel() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `175`
+		//  Estimated: `2767`
+		// Minimum execution time: 2_000_000 picoseconds.
+		Weight::from_parts(2_000_000, 2767)
+			.saturating_add(T::DbWeight::get().reads(1_u64))
+			.saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	/// Storage: `XcmpQueue::OutboundXcmpStatus` (r:1 w:1)
+	/// Proof: `XcmpQueue::OutboundXcmpStatus` (`max_values`: Some(1), `max_size`: Some(1282), added: 1777, mode: `MaxEncodedLen`)
+	fn resume_channel() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `210`
+		//  Estimated: `2767`
+		// Minimum execution time: 3_000_000 picoseconds.
+		Weight::from_parts(3_000_000, 2767)
+			.saturating_add(T::DbWeight::get().reads(1_u64))
+			.saturating_add(T::DbWeight::get().writes(1_u64))
+	}
+	/// The range of component `n` is `[0, 92]`.
+	fn take_first_concatenated_xcm(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `0`
+		//  Estimated: `0`
+		// Minimum execution time: 1_000_000 picoseconds.
+		Weight::from_parts(1_878_097, 0)
+			// Standard Error: 558
+			.saturating_add(Weight::from_parts(44_593, 0).saturating_mul(n.into()))
+	}
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	fn on_idle_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `105816`
+		//  Estimated: `109281`
+		// Minimum execution time: 63_000_000 picoseconds.
+		Weight::from_parts(67_000_000, 109281)
+			.saturating_add(T::DbWeight::get().reads(6_u64))
+			.saturating_add(T::DbWeight::get().writes(5_u64))
+	}
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	fn on_idle_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65885`
+		//  Estimated: `69350`
+		// Minimum execution time: 42_000_000 picoseconds.
+		Weight::from_parts(47_000_000, 69350)
+			.saturating_add(T::DbWeight::get().reads(6_u64))
+			.saturating_add(T::DbWeight::get().writes(5_u64))
+	}
+}
+
+// For backwards compatibility and tests.
+impl WeightInfo for () {
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:1)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	fn set_config_with_u32() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `175`
+		//  Estimated: `1497`
+		// Minimum execution time: 3_000_000 picoseconds.
+		Weight::from_parts(4_000_000, 1497)
+			.saturating_add(RocksDbWeight::get().reads(1_u64))
+			.saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 105467]`.
+	fn enqueue_n_bytes_xcmp_message(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `251`
+		//  Estimated: `5487`
+		// Minimum execution time: 11_000_000 picoseconds.
+		Weight::from_parts(12_262_916, 5487)
+			// Standard Error: 0
+			.saturating_add(Weight::from_parts(187, 0).saturating_mul(n.into()))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(3_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 1000]`.
+	fn enqueue_n_empty_xcmp_messages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `251`
+		//  Estimated: `5487`
+		// Minimum execution time: 9_000_000 picoseconds.
+		Weight::from_parts(13_664_993, 5487)
+			// Standard Error: 220
+			.saturating_add(Weight::from_parts(131_195, 0).saturating_mul(n.into()))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(3_u64))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `Measured`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `Measured`)
+	/// Storage: `MessageQueue::Pages` (r:1 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `Measured`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `Measured`)
+	/// The range of component `n` is `[0, 105457]`.
+	fn enqueue_empty_xcmp_message_at(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `434 + n * (1 ±0)`
+		//  Estimated: `3897 + n * (1 ±0)`
+		// Minimum execution time: 16_000_000 picoseconds.
+		Weight::from_parts(16_490_667, 3897)
+			// Standard Error: 2
+			.saturating_add(Weight::from_parts(720, 0).saturating_mul(n.into()))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+			.saturating_add(Weight::from_parts(0, 1).saturating_mul(n.into()))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:100)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	/// The range of component `n` is `[0, 100]`.
+	fn enqueue_n_full_pages(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `286`
+		//  Estimated: `5487`
+		// Minimum execution time: 9_000_000 picoseconds.
+		Weight::from_parts(10_000_000, 5487)
+			// Standard Error: 26_047
+			.saturating_add(Weight::from_parts(23_938_725, 0).saturating_mul(n.into()))
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+			.saturating_add(RocksDbWeight::get().writes((1_u64).saturating_mul(n.into())))
+	}
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `Measured`)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `Measured`)
+	/// Storage: `MessageQueue::Pages` (r:1 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `Measured`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `Measured`)
+	fn enqueue_1000_small_xcmp_messages() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `53167`
+		//  Estimated: `56632`
+		// Minimum execution time: 175_000_000 picoseconds.
+		Weight::from_parts(181_000_000, 56632)
+			.saturating_add(RocksDbWeight::get().reads(4_u64))
+			.saturating_add(RocksDbWeight::get().writes(2_u64))
+	}
+	/// Storage: `XcmpQueue::OutboundXcmpStatus` (r:1 w:1)
+	/// Proof: `XcmpQueue::OutboundXcmpStatus` (`max_values`: Some(1), `max_size`: Some(1282), added: 1777, mode: `MaxEncodedLen`)
+	fn suspend_channel() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `175`
+		//  Estimated: `2767`
+		// Minimum execution time: 2_000_000 picoseconds.
+		Weight::from_parts(2_000_000, 2767)
+			.saturating_add(RocksDbWeight::get().reads(1_u64))
+			.saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	/// Storage: `XcmpQueue::OutboundXcmpStatus` (r:1 w:1)
+	/// Proof: `XcmpQueue::OutboundXcmpStatus` (`max_values`: Some(1), `max_size`: Some(1282), added: 1777, mode: `MaxEncodedLen`)
+	fn resume_channel() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `210`
+		//  Estimated: `2767`
+		// Minimum execution time: 3_000_000 picoseconds.
+		Weight::from_parts(3_000_000, 2767)
+			.saturating_add(RocksDbWeight::get().reads(1_u64))
+			.saturating_add(RocksDbWeight::get().writes(1_u64))
+	}
+	/// The range of component `n` is `[0, 92]`.
+	fn take_first_concatenated_xcm(n: u32, ) -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `0`
+		//  Estimated: `0`
+		// Minimum execution time: 1_000_000 picoseconds.
+		Weight::from_parts(1_878_097, 0)
+			// Standard Error: 558
+			.saturating_add(Weight::from_parts(44_593, 0).saturating_mul(n.into()))
+	}
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	fn on_idle_good_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `105816`
+		//  Estimated: `109281`
+		// Minimum execution time: 63_000_000 picoseconds.
+		Weight::from_parts(67_000_000, 109281)
+			.saturating_add(RocksDbWeight::get().reads(6_u64))
+			.saturating_add(RocksDbWeight::get().writes(5_u64))
+	}
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6b345d8e88afa015075c945637c07e8f20` (r:1 w:1)
+	/// Storage: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Proof: UNKNOWN KEY `0x7b3237373ffdfeb1cab4222e3b520d6bedc49980ba3aa32b0a189290fd036649` (r:1 w:1)
+	/// Storage: `MessageQueue::BookStateFor` (r:1 w:1)
+	/// Proof: `MessageQueue::BookStateFor` (`max_values`: None, `max_size`: Some(52), added: 2527, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::ServiceHead` (r:1 w:1)
+	/// Proof: `MessageQueue::ServiceHead` (`max_values`: Some(1), `max_size`: Some(5), added: 500, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::QueueConfig` (r:1 w:0)
+	/// Proof: `XcmpQueue::QueueConfig` (`max_values`: Some(1), `max_size`: Some(12), added: 507, mode: `MaxEncodedLen`)
+	/// Storage: `XcmpQueue::InboundXcmpSuspended` (r:1 w:0)
+	/// Proof: `XcmpQueue::InboundXcmpSuspended` (`max_values`: Some(1), `max_size`: Some(4002), added: 4497, mode: `MaxEncodedLen`)
+	/// Storage: `MessageQueue::Pages` (r:0 w:1)
+	/// Proof: `MessageQueue::Pages` (`max_values`: None, `max_size`: Some(105521), added: 107996, mode: `MaxEncodedLen`)
+	fn on_idle_large_msg() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `65885`
+		//  Estimated: `69350`
+		// Minimum execution time: 42_000_000 picoseconds.
+		Weight::from_parts(47_000_000, 69350)
+			.saturating_add(RocksDbWeight::get().reads(6_u64))
+			.saturating_add(RocksDbWeight::get().writes(5_u64))
+	}
+}

pezcumulus/pezpallets/xcmp-queue/src/weights_ext.rs

@@ -0,0 +1,104 @@
+// 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.
+
+//! Weight-related utilities.
+
+use crate::weights::WeightInfo;
+
+use pezframe_support::{traits::BatchFootprint, weights::Weight};
+use pezsp_runtime::SaturatedConversion;
+
+pub(crate) fn get_average_page_pos(max_message_len: u32) -> u32 {
+	max_message_len / 2
+}
+
+/// Extended weight info.
+pub trait WeightInfoExt: WeightInfo {
+	fn uncached_enqueue_xcmp_messages() -> Weight {
+		Self::enqueue_n_full_pages(0)
+	}
+
+	fn enqueue_xcmp_messages(
+		first_page_pos: u32,
+		batch_footprint: &BatchFootprint,
+		is_first_sender_batch: bool,
+	) -> Weight {
+		let message_count = batch_footprint.msgs_count.saturated_into();
+		let size_in_bytes = batch_footprint.size_in_bytes.saturated_into();
+
+		// The cost of adding `n` empty pages on the message queue.
+		let pages_overhead = {
+			let full_message_overhead = Self::enqueue_n_full_pages(1)
+				.saturating_sub(Self::enqueue_n_empty_xcmp_messages(1));
+			let n_full_messages_overhead =
+				full_message_overhead.saturating_mul(batch_footprint.new_pages_count as u64);
+
+			Self::enqueue_n_full_pages(batch_footprint.new_pages_count)
+				.saturating_sub(Self::enqueue_n_full_pages(0))
+				.saturating_sub(n_full_messages_overhead)
+		};
+
+		// The overhead of enqueueing `n` empty messages on the message queue.
+		let messages_overhead = {
+			Self::enqueue_n_empty_xcmp_messages(message_count)
+				.saturating_sub(Self::enqueue_n_empty_xcmp_messages(0))
+		};
+
+		// The overhead of enqueueing `n` bytes on the message queue.
+		let bytes_overhead = {
+			Self::enqueue_n_bytes_xcmp_message(size_in_bytes)
+				.saturating_sub(Self::enqueue_n_bytes_xcmp_message(0))
+		};
+
+		// If the messages are not added to the beginning of the first page, the page will be
+		// decoded and re-encoded once. Let's account for this.
+		let pos_overhead = {
+			let mut pos_overhead = Self::enqueue_empty_xcmp_message_at(first_page_pos)
+				.saturating_sub(Self::enqueue_empty_xcmp_message_at(0));
+			// We need to account for the PoV size of the first page in the message queue only the
+			// first time when we access it.
+			if !is_first_sender_batch {
+				pos_overhead = pos_overhead.set_proof_size(0);
+			}
+			pos_overhead
+		};
+
+		pages_overhead
+			.saturating_add(messages_overhead)
+			.saturating_add(bytes_overhead)
+			.saturating_add(pos_overhead)
+	}
+
+	fn check_accuracy<MaxMessageLen: bounded_collections::Get<u32>>(err_margin: f64) {
+		assert!(err_margin < 1f64);
+
+		let estimated_weight =
+			Self::uncached_enqueue_xcmp_messages().saturating_add(Self::enqueue_xcmp_messages(
+				get_average_page_pos(MaxMessageLen::get()),
+				&BatchFootprint { msgs_count: 1000, size_in_bytes: 3000, new_pages_count: 0 },
+				true,
+			));
+		let actual_weight = Self::enqueue_1000_small_xcmp_messages();
+
+		// Check that the ref_time diff is less than err_margin
+		approx::assert_relative_eq!(
+			estimated_weight.ref_time() as f64,
+			actual_weight.ref_time() as f64,
+			max_relative = err_margin
+		);
+	}
+}
+
+impl<T: WeightInfo> WeightInfoExt for T {}