feat: Rebrand Polkadot/Substrate references to PezkuwiChain

This commit systematically rebrands various references from Parity Technologies'
Polkadot/Substrate ecosystem to PezkuwiChain within the kurdistan-sdk.

Key changes include:
- Updated external repository URLs (zombienet-sdk, parity-db, parity-scale-codec, wasm-instrument) to point to pezkuwichain forks.
- Modified internal documentation and code comments to reflect PezkuwiChain naming and structure.
- Replaced direct references to  with  or specific paths within the  for XCM, Pezkuwi, and other modules.
- Cleaned up deprecated  issue and PR references in various  and  files, particularly in  and  modules.
- Adjusted image and logo URLs in documentation to point to PezkuwiChain assets.
- Removed or rephrased comments related to external Polkadot/Substrate PRs and issues.

This is a significant step towards fully customizing the SDK for the PezkuwiChain ecosystem.

pezcumulus/primitives/aura/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "pezcumulus-primitives-aura"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "Apache-2.0"
+description = "Core primitives for Aura in Pezcumulus"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+# Bizinikiwi
+pezsp-api = { workspace = true }
+pezsp-consensus-aura = { workspace = true }
+
+[features]
+default = ["std"]
+std = ["pezsp-api/std", "pezsp-consensus-aura/std"]
+runtime-benchmarks = [
+	"pezsp-api/runtime-benchmarks",
+	"pezsp-consensus-aura/runtime-benchmarks",
+]

pezcumulus/primitives/aura/src/lib.rs

@@ -0,0 +1,54 @@
+// 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.
+
+//! Core primitives for Aura in Pezcumulus.
+//!
+//! In particular, this exposes the [`AuraUnincludedSegmentApi`] which is used to regulate
+//! the behavior of Aura within a teyrchain context.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+pub use pezsp_consensus_aura::Slot;
+
+pezsp_api::decl_runtime_apis! {
+	/// This runtime API is used to inform potential block authors whether they will
+	/// have the right to author at a slot, assuming they have claimed the slot.
+	///
+	/// In particular, this API allows Aura-based teyrchains to regulate their "unincluded segment",
+	/// which is the section of the head of the chain which has not yet been made available in the
+	/// relay chain.
+	///
+	/// When the unincluded segment is short, Aura chains will allow authors to create multiple
+	/// blocks per slot in order to build a backlog. When it is saturated, this API will limit
+	/// the amount of blocks that can be created.
+	///
+	/// Changes:
+	/// - Version 2: Update to `can_build_upon` to take a relay chain `Slot` instead of a teyrchain `Slot`.
+	#[api_version(2)]
+	pub trait AuraUnincludedSegmentApi {
+		/// 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 relay chain 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.
+		fn can_build_upon(included_hash: Block::Hash, slot: Slot) -> bool;
+	}
+}

pezcumulus/primitives/core/Cargo.toml

@@ -0,0 +1,52 @@
+[package]
+name = "pezcumulus-primitives-core"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "Apache-2.0"
+description = "Pezcumulus related core primitive types and traits"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+tracing = { workspace = true }
+
+# Bizinikiwi
+pezsp-api = { workspace = true }
+pezsp-runtime = { workspace = true }
+pezsp-trie = { workspace = true }
+
+# Pezkuwi
+pezkuwi-core-primitives = { workspace = true }
+pezkuwi-primitives = { workspace = true }
+pezkuwi-teyrchain-primitives = { workspace = true }
+xcm = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+	"codec/std",
+	"pezkuwi-core-primitives/std",
+	"pezkuwi-primitives/std",
+	"pezkuwi-teyrchain-primitives/std",
+	"scale-info/std",
+	"pezsp-api/std",
+	"pezsp-runtime/std",
+	"pezsp-trie/std",
+	"tracing/std",
+	"xcm/std",
+]
+runtime-benchmarks = [
+	"pezkuwi-core-primitives/runtime-benchmarks",
+	"pezkuwi-primitives/runtime-benchmarks",
+	"pezkuwi-teyrchain-primitives/runtime-benchmarks",
+	"pezsp-api/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+	"pezsp-trie/runtime-benchmarks",
+	"xcm/runtime-benchmarks",
+]

pezcumulus/primitives/core/src/lib.rs

@@ -0,0 +1,516 @@
+// 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 related core primitive types and traits.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+use alloc::vec::Vec;
+use codec::{Compact, Decode, DecodeAll, DecodeWithMemTracking, Encode, MaxEncodedLen};
+use pezkuwi_teyrchain_primitives::primitives::HeadData;
+use scale_info::TypeInfo;
+use pezsp_runtime::RuntimeDebug;
+
+/// The ref time per core in seconds.
+///
+/// This is the execution time each PoV gets on a core on the relay chain.
+pub const REF_TIME_PER_CORE_IN_SECS: u64 = 2;
+
+pub mod teyrchain_block_data;
+
+pub use pezkuwi_core_primitives::InboundDownwardMessage;
+pub use pezkuwi_primitives::{
+	AbridgedHostConfiguration, AbridgedHrmpChannel, ClaimQueueOffset, CoreSelector,
+	PersistedValidationData,
+};
+pub use pezkuwi_teyrchain_primitives::primitives::{
+	DmpMessageHandler, Id as ParaId, IsSystem, UpwardMessage, ValidationParams, XcmpMessageFormat,
+	XcmpMessageHandler,
+};
+pub use pezsp_runtime::{
+	generic::{Digest, DigestItem},
+	traits::Block as BlockT,
+	ConsensusEngineId,
+};
+pub use teyrchain_block_data::TeyrchainBlockData;
+pub use xcm::latest::prelude::*;
+
+/// A module that re-exports relevant relay chain definitions.
+pub mod relay_chain {
+	pub use pezkuwi_core_primitives::*;
+	pub use pezkuwi_primitives::*;
+}
+
+/// An inbound HRMP message.
+pub type InboundHrmpMessage = pezkuwi_primitives::InboundHrmpMessage<relay_chain::BlockNumber>;
+
+/// And outbound HRMP message
+pub type OutboundHrmpMessage = pezkuwi_primitives::OutboundHrmpMessage<ParaId>;
+
+/// Error description of a message send failure.
+#[derive(Eq, PartialEq, Copy, Clone, RuntimeDebug, Encode, Decode)]
+pub enum MessageSendError {
+	/// The dispatch queue is full.
+	QueueFull,
+	/// There does not exist a channel for sending the message.
+	NoChannel,
+	/// The message is too big to ever fit in a channel.
+	TooBig,
+	/// Some other error.
+	Other,
+	/// There are too many channels open at once.
+	TooManyChannels,
+}
+
+impl From<MessageSendError> for &'static str {
+	fn from(e: MessageSendError) -> Self {
+		use MessageSendError::*;
+		match e {
+			QueueFull => "QueueFull",
+			NoChannel => "NoChannel",
+			TooBig => "TooBig",
+			Other => "Other",
+			TooManyChannels => "TooManyChannels",
+		}
+	}
+}
+
+/// The origin of an inbound message.
+#[derive(
+	Encode, Decode, DecodeWithMemTracking, MaxEncodedLen, Clone, Eq, PartialEq, TypeInfo, Debug,
+)]
+pub enum AggregateMessageOrigin {
+	/// The message came from the para-chain itself.
+	Here,
+	/// The message came from the relay-chain.
+	///
+	/// This is used by the DMP queue.
+	Parent,
+	/// The message came from a sibling para-chain.
+	///
+	/// This is used by the HRMP queue.
+	Sibling(ParaId),
+}
+
+impl From<AggregateMessageOrigin> for Location {
+	fn from(origin: AggregateMessageOrigin) -> Self {
+		match origin {
+			AggregateMessageOrigin::Here => Location::here(),
+			AggregateMessageOrigin::Parent => Location::parent(),
+			AggregateMessageOrigin::Sibling(id) => Location::new(1, Junction::Teyrchain(id.into())),
+		}
+	}
+}
+
+#[cfg(feature = "runtime-benchmarks")]
+impl From<u32> for AggregateMessageOrigin {
+	fn from(x: u32) -> Self {
+		match x {
+			0 => Self::Here,
+			1 => Self::Parent,
+			p => Self::Sibling(ParaId::from(p)),
+		}
+	}
+}
+
+/// Information about an XCMP channel.
+pub struct ChannelInfo {
+	/// The maximum number of messages that can be pending in the channel at once.
+	pub max_capacity: u32,
+	/// The maximum total size of the messages that can be pending in the channel at once.
+	pub max_total_size: u32,
+	/// The maximum message size that could be put into the channel.
+	pub max_message_size: u32,
+	/// The current number of messages pending in the channel.
+	/// Invariant: should be less or equal to `max_capacity`.s`.
+	pub msg_count: u32,
+	/// The total size in bytes of all message payloads in the channel.
+	/// Invariant: should be less or equal to `max_total_size`.
+	pub total_size: u32,
+}
+
+pub trait GetChannelInfo {
+	fn get_channel_status(id: ParaId) -> ChannelStatus;
+	fn get_channel_info(id: ParaId) -> Option<ChannelInfo>;
+}
+
+/// List all open outgoing channels.
+pub trait ListChannelInfos {
+	fn outgoing_channels() -> Vec<ParaId>;
+}
+
+/// Something that should be called when sending an upward message.
+pub trait UpwardMessageSender {
+	/// Send the given UMP message; return the expected number of blocks before the message will
+	/// be dispatched or an error if the message cannot be sent.
+	/// return the hash of the message sent
+	fn send_upward_message(message: UpwardMessage) -> Result<(u32, XcmHash), MessageSendError>;
+
+	/// Pre-check the given UMP message.
+	fn can_send_upward_message(message: &UpwardMessage) -> Result<(), MessageSendError>;
+
+	/// Ensure `[Self::send_upward_message]` is successful when called in benchmarks/tests.
+	#[cfg(any(feature = "std", feature = "runtime-benchmarks", test))]
+	fn ensure_successful_delivery() {}
+}
+
+impl UpwardMessageSender for () {
+	fn send_upward_message(_message: UpwardMessage) -> Result<(u32, XcmHash), MessageSendError> {
+		Err(MessageSendError::NoChannel)
+	}
+
+	fn can_send_upward_message(_message: &UpwardMessage) -> Result<(), MessageSendError> {
+		Err(MessageSendError::Other)
+	}
+}
+
+/// The status of a channel.
+pub enum ChannelStatus {
+	/// Channel doesn't exist/has been closed.
+	Closed,
+	/// Channel is completely full right now.
+	Full,
+	/// Channel is ready for sending; the two parameters are the maximum size a valid message may
+	/// have right now, and the maximum size a message may ever have (this will generally have been
+	/// available during message construction, but it's possible the channel parameters changed in
+	/// the meantime).
+	Ready(usize, usize),
+}
+
+/// A means of figuring out what outbound XCMP messages should be being sent.
+pub trait XcmpMessageSource {
+	/// Take a single XCMP message from the queue for the given `dest`, if one exists.
+	fn take_outbound_messages(maximum_channels: usize) -> Vec<(ParaId, Vec<u8>)>;
+}
+
+impl XcmpMessageSource for () {
+	fn take_outbound_messages(_maximum_channels: usize) -> Vec<(ParaId, Vec<u8>)> {
+		Vec::new()
+	}
+}
+
+/// The "quality of service" considerations for message sending.
+#[derive(Eq, PartialEq, Clone, Copy, Encode, Decode, RuntimeDebug)]
+pub enum ServiceQuality {
+	/// Ensure that this message is dispatched in the same relative order as any other messages
+	/// that were also sent with `Ordered`. This only guarantees message ordering on the dispatch
+	/// side, and not necessarily on the execution side.
+	Ordered,
+	/// Ensure that the message is dispatched as soon as possible, which could result in it being
+	/// dispatched before other messages which are larger and/or rely on relative ordering.
+	Fast,
+}
+
+/// A consensus engine ID indicating that this is a Pezcumulus Teyrchain.
+pub const CUMULUS_CONSENSUS_ID: ConsensusEngineId = *b"CMLS";
+
+/// Information about the core on the relay chain this block will be validated on.
+#[derive(Clone, Debug, Decode, Encode, PartialEq, Eq)]
+pub struct CoreInfo {
+	/// The selector that determines the actual core at `claim_queue_offset`.
+	pub selector: CoreSelector,
+	/// The claim queue offset that determines how far "into the future" the core is selected.
+	pub claim_queue_offset: ClaimQueueOffset,
+	/// The number of cores assigned to the teyrchain at `claim_queue_offset`.
+	pub number_of_cores: Compact<u16>,
+}
+
+/// Return value of [`CumulusDigestItem::core_info_exists_at_max_once`]
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum CoreInfoExistsAtMaxOnce {
+	/// Exists exactly once.
+	Once(CoreInfo),
+	/// Not found.
+	NotFound,
+	/// Found more than once.
+	MoreThanOnce,
+}
+
+/// Identifier for a relay chain block used by [`CumulusDigestItem`].
+#[derive(Clone, Debug, PartialEq, Hash, Eq)]
+pub enum RelayBlockIdentifier {
+	/// The block is identified using its block hash.
+	ByHash(relay_chain::Hash),
+	/// The block is identified using its storage root and block number.
+	ByStorageRoot { storage_root: relay_chain::Hash, block_number: relay_chain::BlockNumber },
+}
+
+/// Consensus header digests for Pezcumulus teyrchains.
+#[derive(Clone, Debug, Decode, Encode, PartialEq)]
+pub enum CumulusDigestItem {
+	/// A digest item indicating the relay-parent a teyrchain block was built against.
+	#[codec(index = 0)]
+	RelayParent(relay_chain::Hash),
+	/// A digest item providing information about the core selected on the relay chain for this
+	/// block.
+	#[codec(index = 1)]
+	CoreInfo(CoreInfo),
+}
+
+impl CumulusDigestItem {
+	/// Encode this as a Bizinikiwi [`DigestItem`].
+	pub fn to_digest_item(&self) -> DigestItem {
+		match self {
+			Self::RelayParent(_) => DigestItem::Consensus(CUMULUS_CONSENSUS_ID, self.encode()),
+			Self::CoreInfo(_) => DigestItem::PreRuntime(CUMULUS_CONSENSUS_ID, self.encode()),
+		}
+	}
+
+	/// Find [`CumulusDigestItem::CoreInfo`] in the given `digest`.
+	///
+	/// If there are multiple valid digests, this returns the value of the first one.
+	pub fn find_core_info(digest: &Digest) -> Option<CoreInfo> {
+		digest.convert_first(|d| match d {
+			DigestItem::PreRuntime(id, val) if id == &CUMULUS_CONSENSUS_ID => {
+				let Ok(CumulusDigestItem::CoreInfo(core_info)) =
+					CumulusDigestItem::decode_all(&mut &val[..])
+				else {
+					return None;
+				};
+
+				Some(core_info)
+			},
+			_ => None,
+		})
+	}
+
+	/// Returns the found [`CoreInfo`] and iff [`Self::CoreInfo`] exists at max once in the given
+	/// `digest`.
+	pub fn core_info_exists_at_max_once(digest: &Digest) -> CoreInfoExistsAtMaxOnce {
+		let mut core_info = None;
+		if digest
+			.logs()
+			.iter()
+			.filter(|l| match l {
+				DigestItem::PreRuntime(CUMULUS_CONSENSUS_ID, d) => {
+					if let Ok(Self::CoreInfo(ci)) = Self::decode_all(&mut &d[..]) {
+						core_info = Some(ci);
+						true
+					} else {
+						false
+					}
+				},
+				_ => false,
+			})
+			.count() <= 1
+		{
+			core_info
+				.map(CoreInfoExistsAtMaxOnce::Once)
+				.unwrap_or(CoreInfoExistsAtMaxOnce::NotFound)
+		} else {
+			CoreInfoExistsAtMaxOnce::MoreThanOnce
+		}
+	}
+
+	/// Returns the [`RelayBlockIdentifier`] from the given `digest`.
+	///
+	/// The identifier corresponds to the relay parent used to build the teyrchain block.
+	pub fn find_relay_block_identifier(digest: &Digest) -> Option<RelayBlockIdentifier> {
+		digest.convert_first(|d| match d {
+			DigestItem::Consensus(id, val) if id == &CUMULUS_CONSENSUS_ID => {
+				let Ok(CumulusDigestItem::RelayParent(hash)) =
+					CumulusDigestItem::decode_all(&mut &val[..])
+				else {
+					return None;
+				};
+
+				Some(RelayBlockIdentifier::ByHash(hash))
+			},
+			DigestItem::Consensus(id, val) if id == &rpsr_digest::RPSR_CONSENSUS_ID => {
+				let Ok((storage_root, block_number)) =
+					rpsr_digest::RpsrType::decode_all(&mut &val[..])
+				else {
+					return None;
+				};
+
+				Some(RelayBlockIdentifier::ByStorageRoot {
+					storage_root,
+					block_number: block_number.into(),
+				})
+			},
+			_ => None,
+		})
+	}
+}
+
+///
+/// If there are multiple valid digests, this returns the value of the first one, although
+/// well-behaving runtimes should not produce headers with more than one.
+pub fn extract_relay_parent(digest: &Digest) -> Option<relay_chain::Hash> {
+	digest.convert_first(|d| match d {
+		DigestItem::Consensus(id, val) if id == &CUMULUS_CONSENSUS_ID =>
+			match CumulusDigestItem::decode(&mut &val[..]) {
+				Ok(CumulusDigestItem::RelayParent(hash)) => Some(hash),
+				_ => None,
+			},
+		_ => None,
+	})
+}
+
+/// Utilities for handling the relay-parent storage root as a digest item.
+///
+/// This is not intended to be part of the public API, as it is a workaround for
+/// <https://github.com/pezkuwichain/kurdistan-sdk/issues/92> via
+/// <https://github.com/pezkuwichain/kurdistan-sdk/issues/169>.
+///
+/// Runtimes using the teyrchain-system pallet are expected to produce this digest item,
+/// but will stop as soon as they are able to provide the relay-parent hash directly.
+///
+/// The relay-chain storage root is, in practice, a unique identifier of a block
+/// in the absence of equivocations (which are slashable). This assumes that the relay chain
+/// uses BABE or SASSAFRAS, because the slot and the author's VRF randomness are both included
+/// in the relay-chain storage root in both cases.
+///
+/// Therefore, the relay-parent storage root is a suitable identifier of unique relay chain
+/// blocks in low-value scenarios such as performance optimizations.
+#[doc(hidden)]
+pub mod rpsr_digest {
+	use super::{relay_chain, ConsensusEngineId, DecodeAll, Digest, DigestItem, Encode};
+	use codec::Compact;
+
+	/// The type used to store the relay-parent storage root and number.
+	pub type RpsrType = (relay_chain::Hash, Compact<relay_chain::BlockNumber>);
+
+	/// A consensus engine ID for relay-parent storage root digests.
+	pub const RPSR_CONSENSUS_ID: ConsensusEngineId = *b"RPSR";
+
+	/// Construct a digest item for relay-parent storage roots.
+	pub fn relay_parent_storage_root_item(
+		storage_root: relay_chain::Hash,
+		number: impl Into<Compact<relay_chain::BlockNumber>>,
+	) -> DigestItem {
+		DigestItem::Consensus(
+			RPSR_CONSENSUS_ID,
+			RpsrType::from((storage_root, number.into())).encode(),
+		)
+	}
+
+	/// Extract the relay-parent storage root and number from the provided header digest. Returns
+	/// `None` if none were found.
+	pub fn extract_relay_parent_storage_root(
+		digest: &Digest,
+	) -> Option<(relay_chain::Hash, relay_chain::BlockNumber)> {
+		digest.convert_first(|d| match d {
+			DigestItem::Consensus(id, val) if id == &RPSR_CONSENSUS_ID => {
+				let (h, n) = RpsrType::decode_all(&mut &val[..]).ok()?;
+
+				Some((h, n.0))
+			},
+			_ => None,
+		})
+	}
+}
+
+/// Information about a collation.
+///
+/// This was used in version 1 of the [`CollectCollationInfo`] runtime api.
+#[derive(Clone, Debug, codec::Decode, codec::Encode, PartialEq)]
+pub struct CollationInfoV1 {
+	/// Messages destined to be interpreted by the Relay chain itself.
+	pub upward_messages: Vec<UpwardMessage>,
+	/// The horizontal messages sent by the teyrchain.
+	pub horizontal_messages: Vec<OutboundHrmpMessage>,
+	/// New validation code.
+	pub new_validation_code: Option<relay_chain::ValidationCode>,
+	/// The number of messages processed from the DMQ.
+	pub processed_downward_messages: u32,
+	/// The mark which specifies the block number up to which all inbound HRMP messages are
+	/// processed.
+	pub hrmp_watermark: relay_chain::BlockNumber,
+}
+
+impl CollationInfoV1 {
+	/// Convert into the latest version of the [`CollationInfo`] struct.
+	pub fn into_latest(self, head_data: HeadData) -> CollationInfo {
+		CollationInfo {
+			upward_messages: self.upward_messages,
+			horizontal_messages: self.horizontal_messages,
+			new_validation_code: self.new_validation_code,
+			processed_downward_messages: self.processed_downward_messages,
+			hrmp_watermark: self.hrmp_watermark,
+			head_data,
+		}
+	}
+}
+
+/// Information about a collation.
+#[derive(Clone, Debug, codec::Decode, codec::Encode, PartialEq, TypeInfo)]
+pub struct CollationInfo {
+	/// Messages destined to be interpreted by the Relay chain itself.
+	pub upward_messages: Vec<UpwardMessage>,
+	/// The horizontal messages sent by the teyrchain.
+	pub horizontal_messages: Vec<OutboundHrmpMessage>,
+	/// New validation code.
+	pub new_validation_code: Option<relay_chain::ValidationCode>,
+	/// The number of messages processed from the DMQ.
+	pub processed_downward_messages: u32,
+	/// The mark which specifies the block number up to which all inbound HRMP messages are
+	/// processed.
+	pub hrmp_watermark: relay_chain::BlockNumber,
+	/// The head data, aka encoded header, of the block that corresponds to the collation.
+	pub head_data: HeadData,
+}
+
+pezsp_api::decl_runtime_apis! {
+	/// Runtime api to collect information about a collation.
+	///
+	/// Version history:
+	/// - Version 2: Changed [`Self::collect_collation_info`] signature
+	/// - Version 3: Signals to the node to use version 1 of [`TeyrchainBlockData`].
+	#[api_version(3)]
+	pub trait CollectCollationInfo {
+		/// Collect information about a collation.
+		#[changed_in(2)]
+		fn collect_collation_info() -> CollationInfoV1;
+		/// Collect information about a collation.
+		///
+		/// The given `header` is the header of the built block for that
+		/// we are collecting the collation info for.
+		fn collect_collation_info(header: &Block::Header) -> CollationInfo;
+	}
+
+	/// Runtime api used to access general info about a teyrchain runtime.
+	pub trait GetTeyrchainInfo {
+		/// Retrieve the teyrchain id used for runtime.
+		fn teyrchain_id() -> ParaId;
+  }
+
+	/// API to tell the node side how the relay parent should be chosen.
+	///
+	/// A larger offset indicates that the relay parent should not be the tip of the relay chain,
+	/// but `N` blocks behind the tip. This offset is then enforced by the runtime.
+	pub trait RelayParentOffsetApi {
+		/// Fetch the slot offset that is expected from the relay chain.
+		fn relay_parent_offset() -> u32;
+	}
+
+	/// API for teyrchain target block rate.
+	///
+	/// This runtime API allows the teyrchain runtime to communicate the target block rate
+	/// to the node side. The target block rate is always valid for the next relay chain slot.
+	///
+	/// The runtime can not enforce this target block rate. It only acts as a maximum, but not more.
+	/// In the end it depends on the collator how many blocks will be produced. If there are no cores
+	/// available or the collator is offline, no blocks at all will be produced.
+	pub trait TargetBlockRate {
+		/// Get the target block rate for this teyrchain.
+		///
+		/// Returns the target number of blocks per relay chain slot.
+		fn target_block_rate() -> u32;
+	}
+}

pezcumulus/primitives/core/src/teyrchain_block_data.rs

@@ -0,0 +1,260 @@
+// 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.
+
+//! Provides [`TeyrchainBlockData`] and its historical versions.
+
+use alloc::vec::Vec;
+use codec::{Decode, Encode};
+use pezsp_runtime::traits::Block as BlockT;
+use pezsp_trie::CompactProof;
+
+/// Special prefix used by [`TeyrchainBlockData`] from version 1 and upwards to distinguish from the
+/// unversioned legacy/v0 version.
+const VERSIONED_TEYRCHAIN_BLOCK_DATA_PREFIX: &[u8] = b"VERSIONEDPBD";
+
+// Struct which allows prepending bytes after reading from an input.
+pub(crate) struct PrependBytesInput<'a, I> {
+	prepend: &'a [u8],
+	read: usize,
+	inner: &'a mut I,
+}
+
+impl<'a, I: codec::Input> codec::Input for PrependBytesInput<'a, I> {
+	fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
+		let remaining_compact = self.prepend.len().saturating_sub(self.read);
+		Ok(self.inner.remaining_len()?.map(|len| len.saturating_add(remaining_compact)))
+	}
+
+	fn read(&mut self, into: &mut [u8]) -> Result<(), codec::Error> {
+		if into.is_empty() {
+			return Ok(());
+		}
+
+		let remaining_compact = self.prepend.len().saturating_sub(self.read);
+		if remaining_compact > 0 {
+			let to_read = into.len().min(remaining_compact);
+			into[..to_read].copy_from_slice(&self.prepend[self.read..][..to_read]);
+			self.read += to_read;
+
+			if to_read < into.len() {
+				// Buffer not full, keep reading the inner.
+				self.inner.read(&mut into[to_read..])
+			} else {
+				// Buffer was filled by the bytes.
+				Ok(())
+			}
+		} else {
+			// Prepended bytes has been read, just read from inner.
+			self.inner.read(into)
+		}
+	}
+}
+
+/// The teyrchain block that is created by a collator.
+///
+/// This is send as PoV (proof of validity block) to the relay-chain validators. There it will be
+/// passed to the teyrchain validation Wasm blob to be validated.
+#[derive(Clone)]
+pub enum TeyrchainBlockData<Block> {
+	V0 { block: [Block; 1], proof: CompactProof },
+	V1 { blocks: Vec<Block>, proof: CompactProof },
+}
+
+impl<Block: Encode> Encode for TeyrchainBlockData<Block> {
+	fn encode(&self) -> Vec<u8> {
+		match self {
+			Self::V0 { block, proof } => (&block[0], &proof).encode(),
+			Self::V1 { blocks, proof } => {
+				let mut res = VERSIONED_TEYRCHAIN_BLOCK_DATA_PREFIX.to_vec();
+				1u8.encode_to(&mut res);
+				blocks.encode_to(&mut res);
+				proof.encode_to(&mut res);
+				res
+			},
+		}
+	}
+}
+
+impl<Block: Decode> Decode for TeyrchainBlockData<Block> {
+	fn decode<I: codec::Input>(input: &mut I) -> Result<Self, codec::Error> {
+		let mut prefix = [0u8; VERSIONED_TEYRCHAIN_BLOCK_DATA_PREFIX.len()];
+		input.read(&mut prefix)?;
+
+		if prefix == VERSIONED_TEYRCHAIN_BLOCK_DATA_PREFIX {
+			match input.read_byte()? {
+				1 => {
+					let blocks = Vec::<Block>::decode(input)?;
+					let proof = CompactProof::decode(input)?;
+
+					Ok(Self::V1 { blocks, proof })
+				},
+				_ => Err("Unknown `TeyrchainBlockData` version".into()),
+			}
+		} else {
+			let mut input = PrependBytesInput { prepend: &prefix, read: 0, inner: input };
+			let block = Block::decode(&mut input)?;
+			let proof = CompactProof::decode(&mut input)?;
+
+			Ok(Self::V0 { block: [block], proof })
+		}
+	}
+}
+
+impl<Block> TeyrchainBlockData<Block> {
+	/// Creates a new instance of `Self`.
+	pub fn new(blocks: Vec<Block>, proof: CompactProof) -> Self {
+		Self::V1 { blocks, proof }
+	}
+
+	/// Returns references to the stored blocks.
+	pub fn blocks(&self) -> &[Block] {
+		match self {
+			Self::V0 { block, .. } => &block[..],
+			Self::V1 { blocks, .. } => &blocks,
+		}
+	}
+
+	/// Returns mutable references to the stored blocks.
+	pub fn blocks_mut(&mut self) -> &mut [Block] {
+		match self {
+			Self::V0 { ref mut block, .. } => block,
+			Self::V1 { ref mut blocks, .. } => blocks,
+		}
+	}
+
+	/// Returns the stored blocks.
+	pub fn into_blocks(self) -> Vec<Block> {
+		match self {
+			Self::V0 { block, .. } => block.into_iter().collect(),
+			Self::V1 { blocks, .. } => blocks,
+		}
+	}
+
+	/// Returns a reference to the stored proof.
+	pub fn proof(&self) -> &CompactProof {
+		match self {
+			Self::V0 { proof, .. } => &proof,
+			Self::V1 { proof, .. } => proof,
+		}
+	}
+
+	/// Deconstruct into the inner parts.
+	pub fn into_inner(self) -> (Vec<Block>, CompactProof) {
+		match self {
+			Self::V0 { block, proof } => (block.into_iter().collect(), proof),
+			Self::V1 { blocks, proof } => (blocks, proof),
+		}
+	}
+}
+
+impl<Block: BlockT> TeyrchainBlockData<Block> {
+	/// Log the size of the individual components (header, extrinsics, storage proof) as info.
+	pub fn log_size_info(&self) {
+		tracing::info!(
+			target: "pezcumulus",
+			header_kb = %self.blocks().iter().map(|b| b.header().encoded_size()).sum::<usize>() as f64 / 1024f64,
+			extrinsics_kb = %self.blocks().iter().map(|b| b.extrinsics().encoded_size()).sum::<usize>() as f64 / 1024f64,
+			storage_proof_kb = %self.proof().encoded_size() as f64 / 1024f64,
+			"PoV size",
+		);
+	}
+
+	/// Converts into [`TeyrchainBlockData::V0`].
+	///
+	/// Returns `None` if there is not exactly one block.
+	pub fn as_v0(&self) -> Option<Self> {
+		match self {
+			Self::V0 { .. } => Some(self.clone()),
+			Self::V1 { blocks, proof } => {
+				if blocks.len() != 1 {
+					return None;
+				}
+
+				blocks
+					.first()
+					.map(|block| Self::V0 { block: [block.clone()], proof: proof.clone() })
+			},
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use pezsp_runtime::testing::*;
+
+	#[derive(codec::Encode, codec::Decode, Clone, PartialEq, Debug)]
+	struct TeyrchainBlockDataV0<B: BlockT> {
+		/// The header of the teyrchain block.
+		pub header: B::Header,
+		/// The extrinsics of the teyrchain block.
+		pub extrinsics: alloc::vec::Vec<B::Extrinsic>,
+		/// The data that is required to emulate the storage accesses executed by all extrinsics.
+		pub storage_proof: pezsp_trie::CompactProof,
+	}
+
+	type TestExtrinsic = TestXt<MockCallU64, ()>;
+	type TestBlock = Block<TestExtrinsic>;
+
+	#[test]
+	fn decoding_encoding_v0_works() {
+		let v0 = TeyrchainBlockDataV0::<TestBlock> {
+			header: Header::new_from_number(10),
+			extrinsics: vec![
+				TestExtrinsic::new_bare(MockCallU64(10)),
+				TestExtrinsic::new_bare(MockCallU64(100)),
+			],
+			storage_proof: CompactProof { encoded_nodes: vec![vec![10u8; 200], vec![20u8; 30]] },
+		};
+
+		let encoded = v0.encode();
+		let decoded = TeyrchainBlockData::<TestBlock>::decode(&mut &encoded[..]).unwrap();
+
+		match &decoded {
+			TeyrchainBlockData::V0 { block, proof } => {
+				assert_eq!(v0.header, block[0].header);
+				assert_eq!(v0.extrinsics, block[0].extrinsics);
+				assert_eq!(&v0.storage_proof, proof);
+			},
+			_ => panic!("Invalid decoding"),
+		}
+
+		let encoded = decoded.as_v0().unwrap().encode();
+
+		let decoded = TeyrchainBlockDataV0::<TestBlock>::decode(&mut &encoded[..]).unwrap();
+		assert_eq!(decoded, v0);
+	}
+
+	#[test]
+	fn decoding_encoding_v1_works() {
+		let v1 = TeyrchainBlockData::<TestBlock>::V1 {
+			blocks: vec![TestBlock::new(
+				Header::new_from_number(10),
+				vec![
+					TestExtrinsic::new_bare(MockCallU64(10)),
+					TestExtrinsic::new_bare(MockCallU64(100)),
+				],
+			)],
+			proof: CompactProof { encoded_nodes: vec![vec![10u8; 200], vec![20u8; 30]] },
+		};
+
+		let encoded = v1.encode();
+		let decoded = TeyrchainBlockData::<TestBlock>::decode(&mut &encoded[..]).unwrap();
+
+		assert_eq!(v1.blocks(), decoded.blocks());
+		assert_eq!(v1.proof(), decoded.proof());
+	}
+}

pezcumulus/primitives/proof-size-hostfunction/Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "pezcumulus-primitives-proof-size-hostfunction"
+version = "0.2.0"
+authors.workspace = true
+edition.workspace = true
+description = "Hostfunction exposing storage proof size to the runtime."
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+pezsp-externalities = { workspace = true }
+pezsp-runtime-interface = { workspace = true }
+pezsp-trie = { workspace = true }
+
+[dev-dependencies]
+pezsp-core = { workspace = true, default-features = true }
+pezsp-io = { workspace = true, default-features = true }
+pezsp-state-machine = { workspace = true, default-features = true }
+
+[features]
+default = ["std"]
+std = ["pezsp-externalities/std", "pezsp-runtime-interface/std", "pezsp-trie/std"]
+runtime-benchmarks = [
+	"pezsp-io/runtime-benchmarks",
+	"pezsp-runtime-interface/runtime-benchmarks",
+	"pezsp-state-machine/runtime-benchmarks",
+	"pezsp-trie/runtime-benchmarks",
+]

pezcumulus/primitives/proof-size-hostfunction/src/lib.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.
+
+//! Tools for reclaiming PoV weight in teyrchain runtimes.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+#[cfg(feature = "std")]
+use pezsp_externalities::ExternalitiesExt;
+
+use pezsp_runtime_interface::runtime_interface;
+
+#[cfg(feature = "std")]
+use pezsp_trie::proof_size_extension::ProofSizeExt;
+
+pub const PROOF_RECORDING_DISABLED: u64 = u64::MAX;
+
+/// Interface that provides access to the current storage proof size.
+///
+/// Should return the current storage proof size if [`ProofSizeExt`] is registered. Otherwise, needs
+/// to return u64::MAX.
+#[runtime_interface]
+pub trait StorageProofSize {
+	/// Returns the current storage proof size.
+	fn storage_proof_size(&mut self) -> u64 {
+		self.extension::<ProofSizeExt>()
+			.map_or(PROOF_RECORDING_DISABLED, |e| e.storage_proof_size())
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use pezsp_core::Blake2Hasher;
+	use pezsp_state_machine::TestExternalities;
+	use pezsp_trie::{
+		proof_size_extension::ProofSizeExt, recorder::Recorder, LayoutV1, PrefixedMemoryDB,
+		TrieDBMutBuilder, TrieMut,
+	};
+
+	use crate::{storage_proof_size, PROOF_RECORDING_DISABLED};
+
+	const TEST_DATA: &[(&[u8], &[u8])] = &[(b"key1", &[1; 64]), (b"key2", &[2; 64])];
+
+	type TestLayout = LayoutV1<pezsp_core::Blake2Hasher>;
+
+	fn get_prepared_test_externalities() -> (TestExternalities<Blake2Hasher>, Recorder<Blake2Hasher>)
+	{
+		let mut db = PrefixedMemoryDB::default();
+		let mut root = Default::default();
+
+		{
+			let mut trie = TrieDBMutBuilder::<TestLayout>::new(&mut db, &mut root).build();
+			for (k, v) in TEST_DATA {
+				trie.insert(k, v).expect("Inserts data");
+			}
+		}
+
+		let recorder: pezsp_trie::recorder::Recorder<Blake2Hasher> = Default::default();
+		let trie_backend = pezsp_state_machine::TrieBackendBuilder::new(db, root)
+			.with_recorder(recorder.clone())
+			.build();
+
+		let mut ext: TestExternalities<Blake2Hasher> = TestExternalities::default();
+		ext.backend = trie_backend;
+		(ext, recorder)
+	}
+
+	#[test]
+	fn host_function_returns_size_from_recorder() {
+		let (mut ext, recorder) = get_prepared_test_externalities();
+		ext.register_extension(ProofSizeExt::new(recorder));
+
+		ext.execute_with(|| {
+			assert_eq!(storage_proof_size::storage_proof_size(), 0);
+			pezsp_io::storage::get(b"key1");
+			assert_eq!(storage_proof_size::storage_proof_size(), 175);
+			pezsp_io::storage::get(b"key2");
+			assert_eq!(storage_proof_size::storage_proof_size(), 275);
+			pezsp_io::storage::get(b"key2");
+			assert_eq!(storage_proof_size::storage_proof_size(), 275);
+		});
+	}
+
+	#[test]
+	fn host_function_returns_max_without_extension() {
+		let (mut ext, _) = get_prepared_test_externalities();
+
+		ext.execute_with(|| {
+			assert_eq!(storage_proof_size::storage_proof_size(), PROOF_RECORDING_DISABLED);
+			pezsp_io::storage::get(b"key1");
+			assert_eq!(storage_proof_size::storage_proof_size(), PROOF_RECORDING_DISABLED);
+			pezsp_io::storage::get(b"key2");
+			assert_eq!(storage_proof_size::storage_proof_size(), PROOF_RECORDING_DISABLED);
+		});
+	}
+}

pezcumulus/primitives/storage-weight-reclaim/Cargo.toml

@@ -0,0 +1,59 @@
+[package]
+name = "pezcumulus-primitives-storage-weight-reclaim"
+version = "1.0.0"
+authors.workspace = true
+edition.workspace = true
+description = "Utilities to reclaim storage weight."
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[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-runtime = { workspace = true }
+
+pezcumulus-primitives-core = { workspace = true }
+pezcumulus-primitives-proof-size-hostfunction = { workspace = true }
+docify = { workspace = true }
+
+[dev-dependencies]
+pezcumulus-test-runtime = { workspace = true }
+pezsp-io = { workspace = true }
+pezsp-trie = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+	"codec/std",
+	"pezcumulus-primitives-core/std",
+	"pezcumulus-primitives-proof-size-hostfunction/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-core/runtime-benchmarks",
+	"pezcumulus-primitives-proof-size-hostfunction/runtime-benchmarks",
+	"pezcumulus-test-runtime/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",
+]

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

@@ -0,0 +1,229 @@
+// 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.
+
+//! Mechanism to reclaim PoV proof size weight after an extrinsic has been applied.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use codec::{Decode, DecodeWithMemTracking, Encode};
+use core::marker::PhantomData;
+use cumulus_primitives_core::Weight;
+use cumulus_primitives_proof_size_hostfunction::{
+	storage_proof_size::storage_proof_size, PROOF_RECORDING_DISABLED,
+};
+use pezframe_support::{
+	dispatch::{DispatchInfo, PostDispatchInfo},
+	weights::WeightMeter,
+};
+use pezframe_system::Config;
+use scale_info::TypeInfo;
+use pezsp_runtime::{
+	impl_tx_ext_default,
+	traits::{DispatchInfoOf, Dispatchable, PostDispatchInfoOf, TransactionExtension},
+	transaction_validity::TransactionValidityError,
+	DispatchResult,
+};
+
+#[cfg(test)]
+mod tests;
+
+const LOG_TARGET: &'static str = "runtime::storage_reclaim";
+
+/// `StorageWeightReclaimer` is a mechanism for manually reclaiming storage weight.
+///
+/// It internally keeps track of the proof size and storage weight at initialization time. At
+/// reclaim  it computes the real consumed storage weight and refunds excess weight.
+///
+/// # Example
+#[doc = docify::embed!("src/tests.rs", simple_reclaimer_example)]
+pub struct StorageWeightReclaimer {
+	previous_remaining_proof_size: u64,
+	previous_reported_proof_size: Option<u64>,
+}
+
+impl StorageWeightReclaimer {
+	/// Creates a new `StorageWeightReclaimer` instance and initializes it with the storage
+	/// size provided by `weight_meter` and reported proof size from the node.
+	#[must_use = "Must call `reclaim_with_meter` to reclaim the weight"]
+	pub fn new(weight_meter: &WeightMeter) -> StorageWeightReclaimer {
+		let previous_remaining_proof_size = weight_meter.remaining().proof_size();
+		let previous_reported_proof_size = get_proof_size();
+		Self { previous_remaining_proof_size, previous_reported_proof_size }
+	}
+
+	/// Check the consumed storage weight and calculate the consumed excess weight.
+	fn reclaim(&mut self, remaining_weight: Weight) -> Option<Weight> {
+		let current_remaining_weight = remaining_weight.proof_size();
+		let current_storage_proof_size = get_proof_size()?;
+		let previous_storage_proof_size = self.previous_reported_proof_size?;
+		let used_weight =
+			self.previous_remaining_proof_size.saturating_sub(current_remaining_weight);
+		let reported_used_size =
+			current_storage_proof_size.saturating_sub(previous_storage_proof_size);
+		let reclaimable = used_weight.saturating_sub(reported_used_size);
+		log::trace!(
+			target: LOG_TARGET,
+			"Found reclaimable storage weight. benchmarked: {used_weight}, consumed: {reported_used_size}"
+		);
+
+		self.previous_remaining_proof_size = current_remaining_weight.saturating_add(reclaimable);
+		self.previous_reported_proof_size = Some(current_storage_proof_size);
+		Some(Weight::from_parts(0, reclaimable))
+	}
+
+	/// Check the consumed storage weight and add the reclaimed
+	/// weight budget back to `weight_meter`.
+	pub fn reclaim_with_meter(&mut self, weight_meter: &mut WeightMeter) -> Option<Weight> {
+		let reclaimed = self.reclaim(weight_meter.remaining())?;
+		weight_meter.reclaim_proof_size(reclaimed.proof_size());
+		Some(reclaimed)
+	}
+}
+
+/// Returns the current storage proof size from the host side.
+///
+/// Returns `None` if proof recording is disabled on the host.
+pub fn get_proof_size() -> Option<u64> {
+	let proof_size = storage_proof_size();
+	(proof_size != PROOF_RECORDING_DISABLED).then_some(proof_size)
+}
+
+// Encapsulate into a mod so that macro generated code doesn't trigger a warning about deprecated
+// usage.
+#[allow(deprecated)]
+mod allow_deprecated {
+	use super::*;
+
+	/// Storage weight reclaim mechanism.
+	///
+	/// This extension checks the size of the node-side storage proof
+	/// before and after executing a given extrinsic. The difference between
+	/// benchmarked and spent weight can be reclaimed.
+	#[deprecated(note = "This extension doesn't provide accurate reclaim for storage intensive \
+		transaction extension pipeline; it ignores the validation and preparation of extensions prior \
+		to itself and ignores the post dispatch logic for extensions subsequent to itself, it also
+		doesn't provide weight information. \
+		Use `StorageWeightReclaim` in the `pezcumulus-pezpallet-weight-reclaim` crate")]
+	#[derive(Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, Default, TypeInfo)]
+	#[scale_info(skip_type_params(T))]
+	pub struct StorageWeightReclaim<T: Config + Send + Sync>(pub(super) PhantomData<T>);
+}
+#[allow(deprecated)]
+pub use allow_deprecated::StorageWeightReclaim;
+
+#[allow(deprecated)]
+impl<T: Config + Send + Sync> StorageWeightReclaim<T> {
+	/// Create a new `StorageWeightReclaim` instance.
+	pub fn new() -> Self {
+		Self(Default::default())
+	}
+}
+
+#[allow(deprecated)]
+impl<T: Config + Send + Sync> core::fmt::Debug for StorageWeightReclaim<T> {
+	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
+		let _ = write!(f, "StorageWeightReclaim");
+		Ok(())
+	}
+}
+
+#[allow(deprecated)]
+impl<T: Config + Send + Sync> TransactionExtension<T::RuntimeCall> for StorageWeightReclaim<T>
+where
+	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
+{
+	const IDENTIFIER: &'static str = "StorageWeightReclaim";
+	type Implicit = ();
+	type Val = ();
+	type Pre = Option<u64>;
+
+	fn prepare(
+		self,
+		_val: Self::Val,
+		_origin: &T::RuntimeOrigin,
+		_call: &T::RuntimeCall,
+		_info: &DispatchInfoOf<T::RuntimeCall>,
+		_len: usize,
+	) -> Result<Self::Pre, TransactionValidityError> {
+		Ok(get_proof_size())
+	}
+
+	fn post_dispatch_details(
+		pre: Self::Pre,
+		info: &DispatchInfoOf<T::RuntimeCall>,
+		post_info: &PostDispatchInfoOf<T::RuntimeCall>,
+		_len: usize,
+		_result: &DispatchResult,
+	) -> Result<Weight, TransactionValidityError> {
+		let Some(pre_dispatch_proof_size) = pre else {
+			return Ok(Weight::zero());
+		};
+
+		let Some(post_dispatch_proof_size) = get_proof_size() else {
+			log::debug!(
+				target: LOG_TARGET,
+				"Proof recording enabled during pre-dispatch, now disabled. This should not happen."
+			);
+			return Ok(Weight::zero());
+		};
+		// Unspent weight according to the `actual_weight` from `PostDispatchInfo`
+		// This unspent weight will be refunded by the `CheckWeight` extension, so we need to
+		// account for that.
+		let unspent = post_info.calc_unspent(info).proof_size();
+		let benchmarked_weight = info.total_weight().proof_size().saturating_sub(unspent);
+		let consumed_weight = post_dispatch_proof_size.saturating_sub(pre_dispatch_proof_size);
+
+		let storage_size_diff = benchmarked_weight.abs_diff(consumed_weight as u64);
+
+		let extrinsic_len = pezframe_system::AllExtrinsicsLen::<T>::get().unwrap_or(0);
+		let node_side_pov_size = post_dispatch_proof_size.saturating_add(extrinsic_len.into());
+
+		// This value will be reclaimed by [`pezframe_system::CheckWeight`], so we need to calculate
+		// that in.
+		pezframe_system::BlockWeight::<T>::mutate(|current| {
+			if consumed_weight > benchmarked_weight {
+				log::error!(
+					target: LOG_TARGET,
+					"Benchmarked storage weight smaller than consumed storage weight. extrinsic: {} benchmarked: {benchmarked_weight} consumed: {consumed_weight} unspent: {unspent}",
+					pezframe_system::Pallet::<T>::extrinsic_index().unwrap_or(0)
+				);
+				current.accrue(Weight::from_parts(0, storage_size_diff), info.class)
+			} else {
+				log::trace!(
+					target: LOG_TARGET,
+					"Reclaiming storage weight. extrinsic: {} benchmarked: {benchmarked_weight} consumed: {consumed_weight} unspent: {unspent}",
+					pezframe_system::Pallet::<T>::extrinsic_index().unwrap_or(0)
+				);
+				current.reduce(Weight::from_parts(0, storage_size_diff), 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 block_weight_proof_size = current.total().proof_size();
+			let missing_from_node = node_side_pov_size.saturating_sub(block_weight_proof_size);
+			if missing_from_node > 0 {
+				log::debug!(
+					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: {missing_from_node}. Setting to node-side proof size."
+				);
+				current.accrue(Weight::from_parts(0, missing_from_node), info.class);
+			}
+		});
+		Ok(Weight::zero())
+	}
+
+	impl_tx_ext_default!(T::RuntimeCall; weight validate);
+}

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

@@ -0,0 +1,721 @@
+// 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 core::marker::PhantomData;
+use pezframe_support::{
+	assert_ok,
+	dispatch::{DispatchClass, PerDispatchClass},
+	weights::{Weight, WeightMeter},
+};
+use pezframe_system::{BlockWeight, CheckWeight};
+use pezsp_runtime::{traits::DispatchTransaction, AccountId32, BuildStorage};
+use pezsp_trie::proof_size_extension::ProofSizeExt;
+
+type Test = cumulus_test_runtime::Runtime;
+const CALL: &<Test as Config>::RuntimeCall =
+	&cumulus_test_runtime::RuntimeCall::System(pezframe_system::Call::set_heap_pages { pages: 0u64 });
+const ALICE: AccountId32 = AccountId32::new([1u8; 32]);
+const LEN: usize = 150;
+
+fn new_test_ext() -> pezsp_io::TestExternalities {
+	let ext: pezsp_io::TestExternalities = cumulus_test_runtime::RuntimeGenesisConfig::default()
+		.build_storage()
+		.unwrap()
+		.into();
+	ext
+}
+
+struct TestRecorder {
+	return_values: Box<[usize]>,
+	counter: core::sync::atomic::AtomicUsize,
+}
+
+impl TestRecorder {
+	fn new(values: &[usize]) -> Self {
+		TestRecorder { return_values: values.into(), counter: Default::default() }
+	}
+}
+
+impl pezsp_trie::ProofSizeProvider for TestRecorder {
+	fn estimate_encoded_size(&self) -> usize {
+		let counter = self.counter.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
+		self.return_values[counter]
+	}
+}
+
+fn setup_test_externalities(proof_values: &[usize]) -> pezsp_io::TestExternalities {
+	let mut test_ext = new_test_ext();
+	let test_recorder = TestRecorder::new(proof_values);
+	test_ext.register_extension(ProofSizeExt::new(test_recorder));
+	test_ext
+}
+
+fn set_current_storage_weight(new_weight: u64) {
+	BlockWeight::<Test>::mutate(|current_weight| {
+		current_weight.set(Weight::from_parts(0, new_weight), DispatchClass::Normal);
+	});
+}
+
+fn get_storage_weight() -> PerDispatchClass<Weight> {
+	BlockWeight::<Test>::get()
+}
+
+#[test]
+#[allow(deprecated)]
+fn basic_refund() {
+	// The real cost will be 100 bytes of storage size
+	let mut test_ext = setup_test_externalities(&[0, 100]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+
+		// Benchmarked storage weight: 500
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 500), ..Default::default() };
+		let post_info = PostDispatchInfo::default();
+
+		// Should add 500 + 150 (len) to weight.
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(0));
+
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		// We expect a refund of 400
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1250);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn underestimating_refund() {
+	// We fixed a bug where `pre dispatch info weight > consumed weight > post info weight`
+	// resulted in error.
+
+	// The real cost will be 100 bytes of storage size
+	let mut test_ext = setup_test_externalities(&[0, 100]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+
+		// Benchmarked storage weight: 500
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 101), ..Default::default() };
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(0, 99)),
+			pays_fee: Default::default(),
+		};
+
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(0));
+
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(())));
+		// We expect an accrue of 1
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(())
+		));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1250);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn sets_to_node_storage_proof_if_higher() {
+	// The storage proof reported by the proof recorder is higher than what is stored on
+	// the runtime side.
+	{
+		let mut test_ext = setup_test_externalities(&[1000, 1005]);
+
+		test_ext.execute_with(|| {
+			// Stored in BlockWeight is 5
+			set_current_storage_weight(5);
+
+			// Benchmarked storage weight: 10
+			let info =
+				DispatchInfo { call_weight: Weight::from_parts(0, 10), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+			assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+			let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+				.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+				.unwrap();
+			assert_eq!(pre, Some(1000));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch_details(
+				(),
+				&info,
+				&post_info,
+				0,
+				&Ok(())
+			));
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+				pre,
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			// We expect that the storage weight was set to the node-side proof size (1005) +
+			// extrinsics length (150)
+			assert_eq!(get_storage_weight().total().proof_size(), 1155);
+		})
+	}
+
+	// In this second scenario the proof size on the node side is only lower
+	// after reclaim happened.
+	{
+		let mut test_ext = setup_test_externalities(&[175, 180]);
+		test_ext.execute_with(|| {
+			set_current_storage_weight(85);
+
+			// Benchmarked storage weight: 100
+			let info =
+				DispatchInfo { call_weight: Weight::from_parts(0, 100), ..Default::default() };
+			let post_info = PostDispatchInfo::default();
+
+			// After this pre_dispatch, the BlockWeight proof size will be
+			// 85 (initial) + 100 (benched) + 150 (tx length) = 335
+			let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+			assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+			let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+				.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+				.unwrap();
+			assert_eq!(pre, Some(175));
+
+			assert_ok!(CheckWeight::<Test>::post_dispatch_details(
+				(),
+				&info,
+				&post_info,
+				0,
+				&Ok(())
+			));
+
+			// First we will reclaim 95, which leaves us with 240 BlockWeight. This is lower
+			// than 180 (proof size hf) + 150 (length), so we expect it to be set to 330.
+			assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+				pre,
+				&info,
+				&post_info,
+				LEN,
+				&Ok(())
+			));
+
+			// We expect that the storage weight was set to the node-side proof weight
+			assert_eq!(get_storage_weight().total().proof_size(), 330);
+		})
+	}
+}
+
+#[test]
+#[allow(deprecated)]
+fn does_nothing_without_extension() {
+	let mut test_ext = new_test_ext();
+
+	// Proof size extension not registered
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+
+		// Benchmarked storage weight: 500
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 500), ..Default::default() };
+		let post_info = PostDispatchInfo::default();
+
+		// Adds 500 + 150 (len) weight
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, None);
+
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1650);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn negative_refund_is_added_to_weight() {
+	let mut test_ext = setup_test_externalities(&[100, 300]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+		// Benchmarked storage weight: 100
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 100), ..Default::default() };
+		let post_info = PostDispatchInfo::default();
+
+		// Weight added should be 100 + 150 (len)
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(100));
+
+		// We expect no refund
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		assert_eq!(
+			get_storage_weight().total().proof_size(),
+			1100 + LEN as u64 + info.total_weight().proof_size()
+		);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_zero_proof_size() {
+	let mut test_ext = setup_test_externalities(&[0, 0]);
+
+	test_ext.execute_with(|| {
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 500), ..Default::default() };
+		let post_info = PostDispatchInfo::default();
+
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(0));
+
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		// Proof size should be exactly equal to extrinsic length
+		assert_eq!(get_storage_weight().total().proof_size(), LEN as u64);
+	});
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_larger_pre_dispatch_proof_size() {
+	let mut test_ext = setup_test_externalities(&[300, 100]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1300);
+
+		let info = DispatchInfo { call_weight: Weight::from_parts(0, 500), ..Default::default() };
+		let post_info = PostDispatchInfo::default();
+
+		// Adds 500 + 150 (len) weight, total weight is 1950
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(300));
+
+		// Refund 500 unspent weight according to `post_info`, total weight is now 1650
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		// Recorded proof size is negative -200, total weight is now 1450
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1450);
+	});
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_incorporates_check_weight_unspent_weight() {
+	let mut test_ext = setup_test_externalities(&[100, 300]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+
+		// Benchmarked storage weight: 300
+		let info = DispatchInfo { call_weight: Weight::from_parts(100, 300), ..Default::default() };
+
+		// Actual weight is 50
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(50, 250)),
+			pays_fee: Default::default(),
+		};
+
+		// Should add 300 + 150 (len) of weight
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(100));
+
+		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+		// we always need to call `post_dispatch` to verify that they interoperate correctly.
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		// Reclaimed 100
+		assert_eq!(get_storage_weight().total().proof_size(), 1350);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_incorporates_check_weight_unspent_weight_on_negative() {
+	let mut test_ext = setup_test_externalities(&[100, 300]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+		// Benchmarked storage weight: 50
+		let info = DispatchInfo { call_weight: Weight::from_parts(100, 50), ..Default::default() };
+
+		// Actual weight is 25
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(50, 25)),
+			pays_fee: Default::default(),
+		};
+
+		// Adds 50 + 150 (len) weight, total weight 1200
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(100));
+
+		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+		// we always need to call `post_dispatch` to verify that they interoperate correctly.
+		// Refunds unspent 25 weight according to `post_info`, 1175
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+		// Adds 200 - 25 (unspent) == 175 weight, total weight 1350
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1350);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_nothing_relcaimed() {
+	let mut test_ext = setup_test_externalities(&[0, 100]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(0);
+		// Benchmarked storage weight: 100
+		let info = DispatchInfo { call_weight: Weight::from_parts(100, 100), ..Default::default() };
+
+		// Actual proof size is 100
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(50, 100)),
+			pays_fee: Default::default(),
+		};
+
+		// Adds benchmarked weight 100 + 150 (len), total weight is now 250
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		// Weight should go up by 150 len + 100 proof size weight, total weight 250
+		assert_eq!(get_storage_weight().total().proof_size(), 250);
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		// Should return `setup_test_externalities` proof recorder value: 100.
+		assert_eq!(pre, Some(0));
+
+		// The `CheckWeight` extension will refund `actual_weight` from `PostDispatchInfo`
+		// we always need to call `post_dispatch` to verify that they interoperate correctly.
+		// Nothing to refund, unspent is 0, total weight 250
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, LEN, &Ok(())));
+		// `setup_test_externalities` proof recorder value: 200, so this means the extrinsic
+		// actually used 100 proof size.
+		// Nothing to refund or add, weight matches proof recorder
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(())
+		));
+
+		// Check block len weight was not reclaimed:
+		// 100 weight + 150 extrinsic len == 250 proof size
+		assert_eq!(get_storage_weight().total().proof_size(), 250);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_incorporates_check_weight_unspent_weight_reverse_order() {
+	let mut test_ext = setup_test_externalities(&[100, 300]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+
+		// Benchmarked storage weight: 300
+		let info = DispatchInfo { call_weight: Weight::from_parts(100, 300), ..Default::default() };
+
+		// Actual weight is 50
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(50, 250)),
+			pays_fee: Default::default(),
+		};
+
+		// Adds 300 + 150 (len) weight, total weight 1450
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(100));
+
+		// This refunds 100 - 50(unspent), total weight is now 1400
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+		// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+		// we always need to call `post_dispatch` to verify that they interoperate correctly.
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+
+		// Above call refunds 50 (unspent), total weight is 1350 now
+		assert_eq!(get_storage_weight().total().proof_size(), 1350);
+	})
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_incorporates_check_weight_unspent_weight_on_negative_reverse_order() {
+	let mut test_ext = setup_test_externalities(&[100, 300]);
+
+	test_ext.execute_with(|| {
+		set_current_storage_weight(1000);
+		// Benchmarked storage weight: 50
+		let info = DispatchInfo { call_weight: Weight::from_parts(100, 50), ..Default::default() };
+
+		// Actual weight is 25
+		let post_info = PostDispatchInfo {
+			actual_weight: Some(Weight::from_parts(50, 25)),
+			pays_fee: Default::default(),
+		};
+
+		// Adds 50 + 150 (len) weight, total weight is 1200
+		let (_, next_len) = CheckWeight::<Test>::do_validate(&info, LEN).unwrap();
+		assert_ok!(CheckWeight::<Test>::do_prepare(&info, LEN, next_len));
+
+		let (pre, _) = StorageWeightReclaim::<Test>(PhantomData)
+			.validate_and_prepare(Some(ALICE.clone()).into(), CALL, &info, LEN, 0)
+			.unwrap();
+		assert_eq!(pre, Some(100));
+
+		// Adds additional 150 weight recorded
+		assert_ok!(StorageWeightReclaim::<Test>::post_dispatch_details(
+			pre,
+			&info,
+			&post_info,
+			LEN,
+			&Ok(()),
+		));
+		// `CheckWeight` gets called after `StorageWeightReclaim` this time.
+		// The `CheckWeight` extension will refunt `actual_weight` from `PostDispatchInfo`
+		// we always need to call `post_dispatch` to verify that they interoperate correctly.
+		assert_ok!(CheckWeight::<Test>::post_dispatch_details((), &info, &post_info, 0, &Ok(()),));
+
+		assert_eq!(get_storage_weight().total().proof_size(), 1350);
+	})
+}
+
+#[test]
+fn storage_size_reported_correctly() {
+	let mut test_ext = setup_test_externalities(&[1000]);
+	test_ext.execute_with(|| {
+		assert_eq!(get_proof_size(), Some(1000));
+	});
+
+	let mut test_ext = new_test_ext();
+
+	let test_recorder = TestRecorder::new(&[0]);
+
+	test_ext.register_extension(ProofSizeExt::new(test_recorder));
+
+	test_ext.execute_with(|| {
+		assert_eq!(get_proof_size(), Some(0));
+	});
+}
+
+#[test]
+fn storage_size_disabled_reported_correctly() {
+	let mut test_ext = setup_test_externalities(&[PROOF_RECORDING_DISABLED as usize]);
+
+	test_ext.execute_with(|| {
+		assert_eq!(get_proof_size(), None);
+	});
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_reclaim_helper() {
+	let mut test_ext = setup_test_externalities(&[1000, 1300, 1800]);
+
+	test_ext.execute_with(|| {
+		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 2000));
+		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+		remaining_weight_meter.consume(Weight::from_parts(0, 500));
+		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+		assert_eq!(reclaimed, Some(Weight::from_parts(0, 200)));
+
+		remaining_weight_meter.consume(Weight::from_parts(0, 800));
+		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+		assert_eq!(reclaimed, Some(Weight::from_parts(0, 300)));
+		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1200));
+	});
+}
+
+#[test]
+#[allow(deprecated)]
+fn test_reclaim_helper_does_not_reclaim_negative() {
+	// Benchmarked weight does not change at all
+	let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+	test_ext.execute_with(|| {
+		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+		assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(0, 1000));
+	});
+
+	// Benchmarked weight increases less than storage proof consumes
+	let mut test_ext = setup_test_externalities(&[1000, 1300]);
+
+	test_ext.execute_with(|| {
+		let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(0, 1000));
+		let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+		remaining_weight_meter.consume(Weight::from_parts(0, 0));
+		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+		assert_eq!(reclaimed, Some(Weight::from_parts(0, 0)));
+	});
+}
+
+/// Just here for doc purposes
+fn get_benched_weight() -> Weight {
+	Weight::from_parts(0, 5)
+}
+
+/// Just here for doc purposes
+fn do_work() {}
+
+#[allow(deprecated)]
+#[docify::export_content(simple_reclaimer_example)]
+fn reclaim_with_weight_meter() {
+	let mut remaining_weight_meter = WeightMeter::with_limit(Weight::from_parts(10, 10));
+
+	let benched_weight = get_benched_weight();
+
+	// It is important to instantiate the `StorageWeightReclaimer` before we consume the weight
+	// for a piece of work from the weight meter.
+	let mut reclaim_helper = StorageWeightReclaimer::new(&remaining_weight_meter);
+
+	if remaining_weight_meter.try_consume(benched_weight).is_ok() {
+		// Perform some work that takes has `benched_weight` storage weight.
+		do_work();
+
+		// Reclaimer will detect that we only consumed 2 bytes, so 3 bytes are reclaimed.
+		let reclaimed = reclaim_helper.reclaim_with_meter(&mut remaining_weight_meter);
+
+		// We reclaimed 3 bytes of storage size!
+		assert_eq!(reclaimed, Some(Weight::from_parts(0, 3)));
+		assert_eq!(get_storage_weight().total().proof_size(), 10);
+		assert_eq!(remaining_weight_meter.remaining(), Weight::from_parts(10, 8));
+	}
+}
+
+#[test]
+fn test_reclaim_helper_works_with_meter() {
+	// The node will report 12 - 10 = 2 consumed storage size between the calls.
+	let mut test_ext = setup_test_externalities(&[10, 12]);
+
+	test_ext.execute_with(|| {
+		// Initial storage size is 10.
+		set_current_storage_weight(10);
+		reclaim_with_weight_meter();
+	});
+}

pezcumulus/primitives/teyrchain-inherent/Cargo.toml

@@ -0,0 +1,42 @@
+[package]
+name = "pezcumulus-primitives-teyrchain-inherent"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+description = "Inherent that needs to be present in every teyrchain block. Contains messages and a relay chain storage-proof."
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+async-trait = { optional = true, workspace = true }
+codec = { features = ["derive"], workspace = true }
+scale-info = { features = ["derive"], workspace = true }
+
+# Bizinikiwi
+pezsp-core = { workspace = true }
+pezsp-inherents = { workspace = true }
+pezsp-trie = { workspace = true }
+
+# Pezcumulus
+pezcumulus-primitives-core = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+	"async-trait",
+	"codec/std",
+	"pezcumulus-primitives-core/std",
+	"scale-info/std",
+	"pezsp-core/std",
+	"pezsp-inherents/std",
+	"pezsp-trie/std",
+]
+runtime-benchmarks = [
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezsp-inherents/runtime-benchmarks",
+	"pezsp-trie/runtime-benchmarks",
+]

pezcumulus/primitives/teyrchain-inherent/src/lib.rs

@@ -0,0 +1,250 @@
+// 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
+//!
+//! The [`TeyrchainInherentData`] is the data that is passed by the collator to the teyrchain
+//! runtime. The runtime will use this data to execute messages from other teyrchains/the relay
+//! chain or to read data from the relay chain state. When the teyrchain is validated by a teyrchain
+//! validator on the relay chain, this data is checked for correctness. If the data passed by the
+//! collator to the runtime isn't correct, the teyrchain candidate is considered invalid.
+//!
+//! To create a [`TeyrchainInherentData`] for a specific relay chain block, there exists the
+//! `TeyrchainInherentDataExt` trait in `pezcumulus-client-teyrchain-inherent` that helps with this.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+use alloc::{collections::btree_map::BTreeMap, vec::Vec};
+use cumulus_primitives_core::{
+	relay_chain::{
+		ApprovedPeerId, BlakeTwo256, BlockNumber as RelayChainBlockNumber, Hash as RelayHash,
+		HashT as _, Header as RelayHeader,
+	},
+	InboundDownwardMessage, InboundHrmpMessage, ParaId, PersistedValidationData,
+};
+use scale_info::TypeInfo;
+use pezsp_inherents::InherentIdentifier;
+
+/// The identifier for the teyrchain inherent.
+pub const TEYRCHAIN_INHERENT_IDENTIFIER_V0: InherentIdentifier = *b"sysi1337";
+pub const INHERENT_IDENTIFIER: InherentIdentifier = *b"sysi1338";
+
+/// Legacy TeyrchainInherentData that is kept around for backward compatibility.
+/// Can be removed once we can safely assume that teyrchain nodes provide the
+/// `relay_parent_descendants` and `collator_peer_id` fields.
+pub mod v0 {
+	use alloc::{collections::BTreeMap, vec::Vec};
+	use cumulus_primitives_core::{
+		InboundDownwardMessage, InboundHrmpMessage, ParaId, PersistedValidationData,
+	};
+	use scale_info::TypeInfo;
+
+	/// The inherent data that is passed by the collator to the teyrchain runtime.
+	#[derive(
+		codec::Encode,
+		codec::Decode,
+		codec::DecodeWithMemTracking,
+		pezsp_core::RuntimeDebug,
+		Clone,
+		PartialEq,
+		TypeInfo,
+	)]
+	pub struct TeyrchainInherentData {
+		pub validation_data: PersistedValidationData,
+		/// A storage proof of a predefined set of keys from the relay-chain.
+		///
+		/// Specifically this witness contains the data for:
+		///
+		/// - the current slot number at the given relay parent
+		/// - active host configuration as per the relay parent,
+		/// - the relay dispatch queue sizes
+		/// - the list of egress HRMP channels (in the list of recipients form)
+		/// - the metadata for the egress HRMP channels
+		pub relay_chain_state: pezsp_trie::StorageProof,
+		/// Downward messages in the order they were sent.
+		pub downward_messages: Vec<InboundDownwardMessage>,
+		/// HRMP messages grouped by channels. The messages in the inner vec must be in order they
+		/// were sent. In combination with the rule of no more than one message in a channel per
+		/// block, this means `sent_at` is **strictly** greater than the previous one (if any).
+		pub horizontal_messages: BTreeMap<ParaId, Vec<InboundHrmpMessage>>,
+	}
+}
+
+/// The inherent data that is passed by the collator to the teyrchain runtime.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct TeyrchainInherentData {
+	pub validation_data: PersistedValidationData,
+	/// A storage proof of a predefined set of keys from the relay-chain.
+	///
+	/// Specifically this witness contains the data for:
+	///
+	/// - the current slot number at the given relay parent
+	/// - active host configuration as per the relay parent,
+	/// - the relay dispatch queue sizes
+	/// - the list of egress HRMP channels (in the list of recipients form)
+	/// - the metadata for the egress HRMP channels
+	pub relay_chain_state: pezsp_trie::StorageProof,
+	/// Downward messages in the order they were sent.
+	pub downward_messages: Vec<InboundDownwardMessage>,
+	/// HRMP messages grouped by channels. The messages in the inner vec must be in order they
+	/// were sent. In combination with the rule of no more than one message in a channel per block,
+	/// this means `sent_at` is **strictly** greater than the previous one (if any).
+	pub horizontal_messages: BTreeMap<ParaId, Vec<InboundHrmpMessage>>,
+	/// Contains the relay parent header and its descendants.
+	/// This information is used to ensure that a teyrchain node builds blocks
+	/// at a specified offset from the chain tip rather than directly at the tip.
+	pub relay_parent_descendants: Vec<RelayHeader>,
+	/// Contains the collator peer ID, which is later sent by the teyrchain to the
+	/// relay chain via a UMP signal to promote the reputation of the given peer ID.
+	pub collator_peer_id: Option<ApprovedPeerId>,
+}
+
+// Upgrades the TeyrchainInherentData v0 to the newest format.
+impl Into<TeyrchainInherentData> for v0::TeyrchainInherentData {
+	fn into(self) -> TeyrchainInherentData {
+		TeyrchainInherentData {
+			validation_data: self.validation_data,
+			relay_chain_state: self.relay_chain_state,
+			downward_messages: self.downward_messages,
+			horizontal_messages: self.horizontal_messages,
+			relay_parent_descendants: Vec::new(),
+			collator_peer_id: None,
+		}
+	}
+}
+
+#[cfg(feature = "std")]
+impl TeyrchainInherentData {
+	/// Transforms [`TeyrchainInherentData`] into [`v0::TeyrchainInherentData`]. Can be used
+	/// to create inherent data compatible with old runtimes.
+	fn as_v0(&self) -> v0::TeyrchainInherentData {
+		v0::TeyrchainInherentData {
+			validation_data: self.validation_data.clone(),
+			relay_chain_state: self.relay_chain_state.clone(),
+			downward_messages: self.downward_messages.clone(),
+			horizontal_messages: self.horizontal_messages.clone(),
+		}
+	}
+}
+
+#[cfg(feature = "std")]
+#[async_trait::async_trait]
+impl pezsp_inherents::InherentDataProvider for TeyrchainInherentData {
+	async fn provide_inherent_data(
+		&self,
+		inherent_data: &mut pezsp_inherents::InherentData,
+	) -> Result<(), pezsp_inherents::Error> {
+		inherent_data.put_data(TEYRCHAIN_INHERENT_IDENTIFIER_V0, &self.as_v0())?;
+		inherent_data.put_data(INHERENT_IDENTIFIER, &self)
+	}
+
+	async fn try_handle_error(
+		&self,
+		_: &pezsp_inherents::InherentIdentifier,
+		_: &[u8],
+	) -> Option<Result<(), pezsp_inherents::Error>> {
+		None
+	}
+}
+
+/// An inbound message whose content was hashed.
+#[derive(
+	codec::Encode,
+	codec::Decode,
+	codec::DecodeWithMemTracking,
+	pezsp_core::RuntimeDebug,
+	Clone,
+	PartialEq,
+	TypeInfo,
+)]
+pub struct HashedMessage {
+	pub sent_at: RelayChainBlockNumber,
+	pub msg_hash: pezsp_core::H256,
+}
+
+impl From<&InboundDownwardMessage<RelayChainBlockNumber>> for HashedMessage {
+	fn from(msg: &InboundDownwardMessage<RelayChainBlockNumber>) -> Self {
+		Self { sent_at: msg.sent_at, msg_hash: MessageQueueChain::hash_msg_data(&msg.msg) }
+	}
+}
+
+impl From<&InboundHrmpMessage> for HashedMessage {
+	fn from(msg: &InboundHrmpMessage) -> Self {
+		Self { sent_at: msg.sent_at, msg_hash: MessageQueueChain::hash_msg_data(&msg.data) }
+	}
+}
+
+/// This struct provides ability to extend a message queue chain (MQC) and compute a new head.
+///
+/// MQC is an instance of a [hash chain] applied to a message queue. Using a hash chain it's
+/// possible to represent a sequence of messages using only a single hash.
+///
+/// A head for an empty chain is agreed to be a zero hash.
+///
+/// An instance is used to track either DMP from the relay chain or HRMP across a channel.
+/// But a given instance is never used to track both. Therefore, you should call either
+/// `extend_downward` or `extend_hrmp`, but not both methods on a single instance.
+///
+/// [hash chain]: https://en.wikipedia.org/wiki/Hash_chain
+#[derive(Default, Clone, codec::Encode, codec::Decode, scale_info::TypeInfo)]
+pub struct MessageQueueChain(RelayHash);
+
+impl MessageQueueChain {
+	/// Create a new instance initialized to `hash`.
+	pub fn new(hash: RelayHash) -> Self {
+		Self(hash)
+	}
+
+	/// Hash the provided message data.
+	fn hash_msg_data(msg: &Vec<u8>) -> pezsp_core::H256 {
+		BlakeTwo256::hash_of(msg)
+	}
+
+	/// Extend the hash chain with a `HashedMessage`.
+	pub fn extend_with_hashed_msg(&mut self, hashed_msg: &HashedMessage) -> &mut Self {
+		let prev_head = self.0;
+		self.0 = BlakeTwo256::hash_of(&(prev_head, hashed_msg.sent_at, &hashed_msg.msg_hash));
+		self
+	}
+
+	/// Extend the hash chain with an HRMP message. This method should be used only when
+	/// this chain is tracking HRMP.
+	pub fn extend_hrmp(&mut self, horizontal_message: &InboundHrmpMessage) -> &mut Self {
+		self.extend_with_hashed_msg(&horizontal_message.into())
+	}
+
+	/// Extend the hash chain with a downward message. This method should be used only when
+	/// this chain is tracking DMP.
+	pub fn extend_downward(&mut self, downward_message: &InboundDownwardMessage) -> &mut Self {
+		self.extend_with_hashed_msg(&downward_message.into())
+	}
+
+	/// Return the current head of the message queue chain.
+	/// This is agreed to be the zero hash for an empty chain.
+	pub fn head(&self) -> RelayHash {
+		self.0
+	}
+}

pezcumulus/primitives/timestamp/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "pezcumulus-primitives-timestamp"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+description = "Provides timestamp related functionality for teyrchains."
+license = "Apache-2.0"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+# Bizinikiwi
+pezsp-inherents = { workspace = true }
+pezsp-timestamp = { workspace = true }
+
+# Pezcumulus
+pezcumulus-primitives-core = { workspace = true }
+
+[features]
+default = ["std"]
+std = ["pezcumulus-primitives-core/std", "pezsp-inherents/std", "pezsp-timestamp/std"]
+runtime-benchmarks = [
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezsp-inherents/runtime-benchmarks",
+	"pezsp-timestamp/runtime-benchmarks",
+]

pezcumulus/primitives/timestamp/src/lib.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.
+
+//! Pezcumulus timestamp related primitives.
+//!
+//! Provides a [`InherentDataProvider`] that should be used in the validation phase of the
+//! teyrchain. It will be used to create the inherent data and that will be used to check the
+//! inherents inside the teyrchain block (in this case the timestamp inherent). As we don't have
+//! access to any clock from the runtime the timestamp is always passed as an inherent into the
+//! runtime. To check this inherent when validating the block, we will use the relay chain slot. As
+//! the relay chain slot is derived from a timestamp, we can easily convert it back to a timestamp
+//! by multiplying it with the slot duration. By comparing the relay chain slot derived timestamp
+//! with the timestamp we can ensure that the teyrchain timestamp is reasonable.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+use core::time::Duration;
+use cumulus_primitives_core::relay_chain::Slot;
+use pezsp_inherents::{Error, InherentData};
+
+pub use pezsp_timestamp::{InherentType, INHERENT_IDENTIFIER};
+
+/// The inherent data provider for the timestamp.
+///
+/// This should be used in the runtime when checking the inherents in the validation phase of the
+/// teyrchain.
+pub struct InherentDataProvider {
+	relay_chain_slot: Slot,
+	relay_chain_slot_duration: Duration,
+}
+
+impl InherentDataProvider {
+	/// Create `Self` from the given relay chain slot and slot duration.
+	pub fn from_relay_chain_slot_and_duration(
+		relay_chain_slot: Slot,
+		relay_chain_slot_duration: Duration,
+	) -> Self {
+		Self { relay_chain_slot, relay_chain_slot_duration }
+	}
+
+	/// Create the inherent data.
+	pub fn create_inherent_data(&self) -> Result<InherentData, Error> {
+		let mut inherent_data = InherentData::new();
+		self.provide_inherent_data(&mut inherent_data).map(|_| inherent_data)
+	}
+
+	/// Provide the inherent data into the given `inherent_data`.
+	pub fn provide_inherent_data(&self, inherent_data: &mut InherentData) -> Result<(), Error> {
+		// As the teyrchain starts building at around `relay_chain_slot + 1` we use that slot to
+		// calculate the timestamp.
+		let data: InherentType = ((*self.relay_chain_slot + 1) *
+			self.relay_chain_slot_duration.as_millis() as u64)
+			.into();
+
+		inherent_data.put_data(INHERENT_IDENTIFIER, &data)
+	}
+}

pezcumulus/primitives/utility/Cargo.toml

@@ -0,0 +1,55 @@
+[package]
+name = "pezcumulus-primitives-utility"
+version = "0.7.0"
+authors.workspace = true
+edition.workspace = true
+license = "Apache-2.0"
+description = "Helper datatypes for Pezcumulus"
+homepage.workspace = true
+repository.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+codec = { features = ["derive"], workspace = true }
+log = { workspace = true }
+
+# Bizinikiwi
+pezframe-support = { workspace = true }
+pezpallet-asset-conversion = { workspace = true }
+pezsp-runtime = { workspace = true }
+
+# Pezkuwi
+pezkuwi-runtime-common = { workspace = true }
+xcm = { workspace = true }
+xcm-builder = { workspace = true }
+xcm-executor = { workspace = true }
+
+# Pezcumulus
+pezcumulus-primitives-core = { workspace = true }
+
+[features]
+default = ["std"]
+std = [
+	"codec/std",
+	"pezcumulus-primitives-core/std",
+	"pezframe-support/std",
+	"log/std",
+	"pezpallet-asset-conversion/std",
+	"pezkuwi-runtime-common/std",
+	"pezsp-runtime/std",
+	"xcm-builder/std",
+	"xcm-executor/std",
+	"xcm/std",
+]
+runtime-benchmarks = [
+	"pezcumulus-primitives-core/runtime-benchmarks",
+	"pezframe-support/runtime-benchmarks",
+	"pezpallet-asset-conversion/runtime-benchmarks",
+	"pezkuwi-runtime-common/runtime-benchmarks",
+	"pezsp-runtime/runtime-benchmarks",
+	"xcm-builder/runtime-benchmarks",
+	"xcm-executor/runtime-benchmarks",
+	"xcm/runtime-benchmarks",
+]

pezcumulus/primitives/utility/src/lib.rs

@@ -0,0 +1,906 @@
+// 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.
+
+//! Helper datatypes for pezcumulus. This includes the [`ParentAsUmp`] routing type which will route
+//! messages into an [`UpwardMessageSender`] if the destination is `Parent`.
+
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+use alloc::{vec, vec::Vec};
+use codec::Encode;
+use core::marker::PhantomData;
+use cumulus_primitives_core::{MessageSendError, UpwardMessageSender};
+use pezframe_support::{
+	defensive,
+	traits::{tokens::fungibles, Get, OnUnbalanced as OnUnbalancedT},
+	weights::{Weight, WeightToFee as WeightToFeeT},
+	CloneNoBound,
+};
+use pezpallet_asset_conversion::SwapCredit as SwapCreditT;
+use pezkuwi_runtime_common::xcm_sender::PriceForMessageDelivery;
+use pezsp_runtime::{
+	traits::{Saturating, Zero},
+	SaturatedConversion,
+};
+use xcm::{latest::prelude::*, VersionedLocation, VersionedXcm, WrapVersion};
+use xcm_builder::{InspectMessageQueues, TakeRevenue};
+use xcm_executor::{
+	traits::{MatchesFungibles, TransactAsset, WeightTrader},
+	AssetsInHolding,
+};
+
+#[cfg(test)]
+mod tests;
+
+/// Xcm router which recognises the `Parent` destination and handles it by sending the message into
+/// the given UMP `UpwardMessageSender` implementation. Thus this essentially adapts an
+/// `UpwardMessageSender` trait impl into a `SendXcm` trait impl.
+///
+/// NOTE: This is a pretty dumb "just send it" router; we will probably want to introduce queuing
+/// to UMP eventually and when we do, the pallet which implements the queuing will be responsible
+/// for the `SendXcm` implementation.
+pub struct ParentAsUmp<T, W, P>(PhantomData<(T, W, P)>);
+impl<T, W, P> SendXcm for ParentAsUmp<T, W, P>
+where
+	T: UpwardMessageSender,
+	W: WrapVersion,
+	P: PriceForMessageDelivery<Id = ()>,
+{
+	type Ticket = Vec<u8>;
+
+	fn validate(dest: &mut Option<Location>, msg: &mut Option<Xcm<()>>) -> SendResult<Vec<u8>> {
+		let d = dest.take().ok_or(SendError::MissingArgument)?;
+
+		if d.contains_parents_only(1) {
+			// An upward message for the relay chain.
+			let xcm = msg.take().ok_or(SendError::MissingArgument)?;
+			let price = P::price_for_delivery((), &xcm);
+			let versioned_xcm =
+				W::wrap_version(&d, xcm).map_err(|()| SendError::DestinationUnsupported)?;
+			versioned_xcm
+				.check_is_decodable()
+				.map_err(|()| SendError::ExceedsMaxMessageSize)?;
+			let data = versioned_xcm.encode();
+
+			// Pre-check with our message sender if everything else is okay.
+			T::can_send_upward_message(&data).map_err(Self::map_upward_sender_err)?;
+
+			Ok((data, price))
+		} else {
+			// Anything else is unhandled. This includes a message that is not meant for us.
+			// We need to make sure that dest/msg is not consumed here.
+			*dest = Some(d);
+			Err(SendError::NotApplicable)
+		}
+	}
+
+	fn deliver(data: Vec<u8>) -> Result<XcmHash, SendError> {
+		let (_, hash) = T::send_upward_message(data).map_err(Self::map_upward_sender_err)?;
+		Ok(hash)
+	}
+
+	#[cfg(feature = "runtime-benchmarks")]
+	fn ensure_successful_delivery(location: Option<Location>) {
+		if location.as_ref().map_or(false, |l| l.contains_parents_only(1)) {
+			T::ensure_successful_delivery();
+		}
+	}
+}
+
+impl<T, W, P> ParentAsUmp<T, W, P> {
+	fn map_upward_sender_err(message_send_error: MessageSendError) -> SendError {
+		match message_send_error {
+			MessageSendError::TooBig => SendError::ExceedsMaxMessageSize,
+			e => SendError::Transport(e.into()),
+		}
+	}
+}
+
+impl<T: UpwardMessageSender + InspectMessageQueues, W, P> InspectMessageQueues
+	for ParentAsUmp<T, W, P>
+{
+	fn clear_messages() {
+		T::clear_messages();
+	}
+
+	fn get_messages() -> Vec<(VersionedLocation, Vec<VersionedXcm<()>>)> {
+		T::get_messages()
+	}
+}
+
+/// Contains information to handle refund/payment for xcm-execution
+#[derive(Clone, Eq, PartialEq, Debug)]
+struct AssetTraderRefunder {
+	// The amount of weight bought minus the weigh already refunded
+	weight_outstanding: Weight,
+	// The concrete asset containing the asset location and outstanding balance
+	outstanding_concrete_asset: Asset,
+}
+
+/// Charges for execution in the first asset of those selected for fee payment
+/// Only succeeds for Concrete Fungible Assets
+/// First tries to convert the this Asset into a local assetId
+/// Then charges for this assetId as described by FeeCharger
+/// Weight, paid balance, local asset Id and the location is stored for
+/// later refund purposes
+/// Important: Errors if the Trader is being called twice by 2 BuyExecution instructions
+/// Alternatively we could just return payment in the aforementioned case
+#[derive(CloneNoBound)]
+pub struct TakeFirstAssetTrader<
+	AccountId: Eq,
+	FeeCharger: ChargeWeightInFungibles<AccountId, ConcreteAssets>,
+	Matcher: MatchesFungibles<ConcreteAssets::AssetId, ConcreteAssets::Balance>,
+	ConcreteAssets: fungibles::Mutate<AccountId> + fungibles::Balanced<AccountId>,
+	HandleRefund: TakeRevenue,
+>(
+	Option<AssetTraderRefunder>,
+	PhantomData<(AccountId, FeeCharger, Matcher, ConcreteAssets, HandleRefund)>,
+);
+impl<
+		AccountId: Eq,
+		FeeCharger: ChargeWeightInFungibles<AccountId, ConcreteAssets>,
+		Matcher: MatchesFungibles<ConcreteAssets::AssetId, ConcreteAssets::Balance>,
+		ConcreteAssets: fungibles::Mutate<AccountId> + fungibles::Balanced<AccountId>,
+		HandleRefund: TakeRevenue,
+	> WeightTrader
+	for TakeFirstAssetTrader<AccountId, FeeCharger, Matcher, ConcreteAssets, HandleRefund>
+{
+	fn new() -> Self {
+		Self(None, PhantomData)
+	}
+	// We take first asset
+	// Check whether we can convert fee to asset_fee (is_sufficient, min_deposit)
+	// If everything goes well, we charge.
+	fn buy_weight(
+		&mut self,
+		weight: Weight,
+		payment: xcm_executor::AssetsInHolding,
+		context: &XcmContext,
+	) -> Result<xcm_executor::AssetsInHolding, XcmError> {
+		log::trace!(target: "xcm::weight", "TakeFirstAssetTrader::buy_weight weight: {:?}, payment: {:?}, context: {:?}", weight, payment, context);
+
+		// Make sure we don't enter twice
+		if self.0.is_some() {
+			return Err(XcmError::NotWithdrawable);
+		}
+
+		// We take the very first asset from payment
+		// (assets are sorted by fungibility/amount after this conversion)
+		let assets: Assets = payment.clone().into();
+
+		// Take the first asset from the selected Assets
+		let first = assets.get(0).ok_or(XcmError::AssetNotFound)?;
+
+		// Get the local asset id in which we can pay for fees
+		let (local_asset_id, _) =
+			Matcher::matches_fungibles(first).map_err(|_| XcmError::AssetNotFound)?;
+
+		// Calculate how much we should charge in the asset_id for such amount of weight
+		// Require at least a payment of minimum_balance
+		// Necessary for fully collateral-backed assets
+		let asset_balance: u128 =
+			FeeCharger::charge_weight_in_fungibles(local_asset_id.clone(), weight)
+				.map(|amount| {
+					let minimum_balance = ConcreteAssets::minimum_balance(local_asset_id);
+					if amount < minimum_balance {
+						minimum_balance
+					} else {
+						amount
+					}
+				})?
+				.try_into()
+				.map_err(|_| XcmError::Overflow)?;
+
+		// Convert to the same kind of asset, with the required fungible balance
+		let required = first.id.clone().into_asset(asset_balance.into());
+
+		// Subtract payment
+		let unused = payment.checked_sub(required.clone()).map_err(|_| XcmError::TooExpensive)?;
+
+		// record weight and asset
+		self.0 = Some(AssetTraderRefunder {
+			weight_outstanding: weight,
+			outstanding_concrete_asset: required,
+		});
+
+		Ok(unused)
+	}
+
+	fn refund_weight(&mut self, weight: Weight, context: &XcmContext) -> Option<Asset> {
+		log::trace!(target: "xcm::weight", "TakeFirstAssetTrader::refund_weight weight: {:?}, context: {:?}", weight, context);
+		if let Some(AssetTraderRefunder {
+			mut weight_outstanding,
+			outstanding_concrete_asset: Asset { id, fun },
+		}) = self.0.clone()
+		{
+			// Get the local asset id in which we can refund fees
+			let (local_asset_id, outstanding_balance) =
+				Matcher::matches_fungibles(&(id.clone(), fun).into()).ok()?;
+
+			let minimum_balance = ConcreteAssets::minimum_balance(local_asset_id.clone());
+
+			// Calculate asset_balance
+			// This read should have already be cached in buy_weight
+			let (asset_balance, outstanding_minus_subtracted) =
+				FeeCharger::charge_weight_in_fungibles(local_asset_id, weight).ok().map(
+					|asset_balance| {
+						// Require at least a drop of minimum_balance
+						// Necessary for fully collateral-backed assets
+						if outstanding_balance.saturating_sub(asset_balance) > minimum_balance {
+							(asset_balance, outstanding_balance.saturating_sub(asset_balance))
+						}
+						// If the amount to be refunded leaves the remaining balance below ED,
+						// we just refund the exact amount that guarantees at least ED will be
+						// dropped
+						else {
+							(outstanding_balance.saturating_sub(minimum_balance), minimum_balance)
+						}
+					},
+				)?;
+
+			// Convert balances into u128
+			let outstanding_minus_subtracted: u128 = outstanding_minus_subtracted.saturated_into();
+			let asset_balance: u128 = asset_balance.saturated_into();
+
+			// Construct outstanding_concrete_asset with the same location id and subtracted
+			// balance
+			let outstanding_concrete_asset: Asset =
+				(id.clone(), outstanding_minus_subtracted).into();
+
+			// Subtract from existing weight and balance
+			weight_outstanding = weight_outstanding.saturating_sub(weight);
+
+			// Override AssetTraderRefunder
+			self.0 = Some(AssetTraderRefunder { weight_outstanding, outstanding_concrete_asset });
+
+			// Only refund if positive
+			if asset_balance > 0 {
+				Some((id, asset_balance).into())
+			} else {
+				None
+			}
+		} else {
+			None
+		}
+	}
+}
+
+impl<
+		AccountId: Eq,
+		FeeCharger: ChargeWeightInFungibles<AccountId, ConcreteAssets>,
+		Matcher: MatchesFungibles<ConcreteAssets::AssetId, ConcreteAssets::Balance>,
+		ConcreteAssets: fungibles::Mutate<AccountId> + fungibles::Balanced<AccountId>,
+		HandleRefund: TakeRevenue,
+	> Drop for TakeFirstAssetTrader<AccountId, FeeCharger, Matcher, ConcreteAssets, HandleRefund>
+{
+	fn drop(&mut self) {
+		if let Some(asset_trader) = self.0.clone() {
+			HandleRefund::take_revenue(asset_trader.outstanding_concrete_asset);
+		}
+	}
+}
+
+/// XCM fee depositor to which we implement the `TakeRevenue` trait.
+/// It receives a `Transact` implemented argument and a 32 byte convertible `AccountId`, and the fee
+/// receiver account's `FungiblesMutateAdapter` should be identical to that implemented by
+/// `WithdrawAsset`.
+pub struct XcmFeesTo32ByteAccount<FungiblesMutateAdapter, AccountId, ReceiverAccount>(
+	PhantomData<(FungiblesMutateAdapter, AccountId, ReceiverAccount)>,
+);
+impl<
+		FungiblesMutateAdapter: TransactAsset,
+		AccountId: Clone + Into<[u8; 32]>,
+		ReceiverAccount: Get<Option<AccountId>>,
+	> TakeRevenue for XcmFeesTo32ByteAccount<FungiblesMutateAdapter, AccountId, ReceiverAccount>
+{
+	fn take_revenue(revenue: Asset) {
+		if let Some(receiver) = ReceiverAccount::get() {
+			let ok = FungiblesMutateAdapter::deposit_asset(
+				&revenue,
+				&([AccountId32 { network: None, id: receiver.into() }].into()),
+				None,
+			)
+			.is_ok();
+
+			debug_assert!(ok, "`deposit_asset` cannot generally fail; qed");
+		}
+	}
+}
+
+/// ChargeWeightInFungibles trait, which converts a given amount of weight
+/// and an assetId, and it returns the balance amount that should be charged
+/// in such assetId for that amount of weight
+pub trait ChargeWeightInFungibles<AccountId, Assets: fungibles::Inspect<AccountId>> {
+	fn charge_weight_in_fungibles(
+		asset_id: <Assets as fungibles::Inspect<AccountId>>::AssetId,
+		weight: Weight,
+	) -> Result<<Assets as fungibles::Inspect<AccountId>>::Balance, XcmError>;
+}
+
+/// Provides an implementation of [`WeightTrader`] to charge for weight using the first asset
+/// specified in the `payment` argument.
+///
+/// The asset used to pay for the weight must differ from the `Target` asset and be exchangeable for
+/// the same `Target` asset through `SwapCredit`.
+///
+/// ### Parameters:
+/// - `Target`: the asset into which the user's payment will be exchanged using `SwapCredit`.
+/// - `SwapCredit`: mechanism used for the exchange of the user's payment asset into the `Target`.
+/// - `WeightToFee`: weight to the `Target` asset fee calculator.
+/// - `Fungibles`: registry of fungible assets.
+/// - `FungiblesAssetMatcher`: utility for mapping [`Asset`] to `Fungibles::AssetId` and
+///   `Fungibles::Balance`.
+/// - `OnUnbalanced`: handler for the fee payment.
+/// - `AccountId`: the account identifier type.
+pub struct SwapFirstAssetTrader<
+	Target: Get<Fungibles::AssetId>,
+	SwapCredit: SwapCreditT<
+		AccountId,
+		Balance = Fungibles::Balance,
+		AssetKind = Fungibles::AssetId,
+		Credit = fungibles::Credit<AccountId, Fungibles>,
+	>,
+	WeightToFee: WeightToFeeT<Balance = Fungibles::Balance>,
+	Fungibles: fungibles::Balanced<AccountId>,
+	FungiblesAssetMatcher: MatchesFungibles<Fungibles::AssetId, Fungibles::Balance>,
+	OnUnbalanced: OnUnbalancedT<fungibles::Credit<AccountId, Fungibles>>,
+	AccountId,
+> where
+	Fungibles::Balance: Into<u128>,
+{
+	/// Accumulated fee paid for XCM execution.
+	total_fee: fungibles::Credit<AccountId, Fungibles>,
+	/// Last asset utilized by a client to settle a fee.
+	last_fee_asset: Option<AssetId>,
+	_phantom_data: PhantomData<(
+		Target,
+		SwapCredit,
+		WeightToFee,
+		Fungibles,
+		FungiblesAssetMatcher,
+		OnUnbalanced,
+		AccountId,
+	)>,
+}
+
+impl<
+		Target: Get<Fungibles::AssetId>,
+		SwapCredit: SwapCreditT<
+			AccountId,
+			Balance = Fungibles::Balance,
+			AssetKind = Fungibles::AssetId,
+			Credit = fungibles::Credit<AccountId, Fungibles>,
+		>,
+		WeightToFee: WeightToFeeT<Balance = Fungibles::Balance>,
+		Fungibles: fungibles::Balanced<AccountId>,
+		FungiblesAssetMatcher: MatchesFungibles<Fungibles::AssetId, Fungibles::Balance>,
+		OnUnbalanced: OnUnbalancedT<fungibles::Credit<AccountId, Fungibles>>,
+		AccountId,
+	> WeightTrader
+	for SwapFirstAssetTrader<
+		Target,
+		SwapCredit,
+		WeightToFee,
+		Fungibles,
+		FungiblesAssetMatcher,
+		OnUnbalanced,
+		AccountId,
+	>
+where
+	Fungibles::Balance: Into<u128>,
+{
+	fn new() -> Self {
+		Self {
+			total_fee: fungibles::Credit::<AccountId, Fungibles>::zero(Target::get()),
+			last_fee_asset: None,
+			_phantom_data: PhantomData,
+		}
+	}
+
+	fn buy_weight(
+		&mut self,
+		weight: Weight,
+		mut payment: AssetsInHolding,
+		_context: &XcmContext,
+	) -> Result<AssetsInHolding, XcmError> {
+		log::trace!(
+			target: "xcm::weight",
+			"SwapFirstAssetTrader::buy_weight weight: {:?}, payment: {:?}",
+			weight,
+			payment,
+		);
+		let first_asset: Asset =
+			payment.fungible.pop_first().ok_or(XcmError::AssetNotFound)?.into();
+		let (fungibles_asset, balance) = FungiblesAssetMatcher::matches_fungibles(&first_asset)
+			.map_err(|error| {
+				log::trace!(
+					target: "xcm::weight",
+					"SwapFirstAssetTrader::buy_weight asset {:?} didn't match. Error: {:?}",
+					first_asset,
+					error,
+				);
+				XcmError::AssetNotFound
+			})?;
+
+		let swap_asset = fungibles_asset.clone().into();
+		if Target::get().eq(&swap_asset) {
+			log::trace!(
+				target: "xcm::weight",
+				"SwapFirstAssetTrader::buy_weight Asset was same as Target, swap not needed.",
+			);
+			// current trader is not applicable.
+			return Err(XcmError::FeesNotMet);
+		}
+
+		let credit_in = Fungibles::issue(fungibles_asset, balance);
+		let fee = WeightToFee::weight_to_fee(&weight);
+
+		// swap the user's asset for the `Target` asset.
+		let (credit_out, credit_change) = SwapCredit::swap_tokens_for_exact_tokens(
+			vec![swap_asset, Target::get()],
+			credit_in,
+			fee,
+		)
+		.map_err(|(credit_in, error)| {
+			log::trace!(
+				target: "xcm::weight",
+				"SwapFirstAssetTrader::buy_weight swap couldn't be done. Error was: {:?}",
+				error,
+			);
+			drop(credit_in);
+			XcmError::FeesNotMet
+		})?;
+
+		match self.total_fee.subsume(credit_out) {
+			Err(credit_out) => {
+				// error may occur if `total_fee.asset` differs from `credit_out.asset`, which does
+				// not apply in this context.
+				defensive!(
+					"`total_fee.asset` must be equal to `credit_out.asset`",
+					(self.total_fee.asset(), credit_out.asset())
+				);
+				return Err(XcmError::FeesNotMet);
+			},
+			_ => (),
+		};
+		self.last_fee_asset = Some(first_asset.id.clone());
+
+		payment.fungible.insert(first_asset.id, credit_change.peek().into());
+		drop(credit_change);
+		Ok(payment)
+	}
+
+	fn refund_weight(&mut self, weight: Weight, _context: &XcmContext) -> Option<Asset> {
+		log::trace!(
+			target: "xcm::weight",
+			"SwapFirstAssetTrader::refund_weight weight: {:?}, self.total_fee: {:?}",
+			weight,
+			self.total_fee,
+		);
+		if self.total_fee.peek().is_zero() {
+			// noting yet paid to refund.
+			return None;
+		}
+		let mut refund_asset = if let Some(asset) = &self.last_fee_asset {
+			// create an initial zero refund in the asset used in the last `buy_weight`.
+			(asset.clone(), Fungible(0)).into()
+		} else {
+			return None;
+		};
+		let refund_amount = WeightToFee::weight_to_fee(&weight);
+		if refund_amount >= self.total_fee.peek() {
+			// not enough was paid to refund the `weight`.
+			return None;
+		}
+
+		let refund_swap_asset = FungiblesAssetMatcher::matches_fungibles(&refund_asset)
+			.map(|(a, _)| a.into())
+			.ok()?;
+
+		let refund = self.total_fee.extract(refund_amount);
+		let refund = match SwapCredit::swap_exact_tokens_for_tokens(
+			vec![Target::get(), refund_swap_asset],
+			refund,
+			None,
+		) {
+			Ok(refund_in_target) => refund_in_target,
+			Err((refund, _)) => {
+				// return an attempted refund back to the `total_fee`.
+				let _ = self.total_fee.subsume(refund).map_err(|refund| {
+					// error may occur if `total_fee.asset` differs from `refund.asset`, which does
+					// not apply in this context.
+					defensive!(
+						"`total_fee.asset` must be equal to `refund.asset`",
+						(self.total_fee.asset(), refund.asset())
+					);
+				});
+				return None;
+			},
+		};
+
+		refund_asset.fun = refund.peek().into().into();
+		drop(refund);
+		Some(refund_asset)
+	}
+}
+
+impl<
+		Target: Get<Fungibles::AssetId>,
+		SwapCredit: SwapCreditT<
+			AccountId,
+			Balance = Fungibles::Balance,
+			AssetKind = Fungibles::AssetId,
+			Credit = fungibles::Credit<AccountId, Fungibles>,
+		>,
+		WeightToFee: WeightToFeeT<Balance = Fungibles::Balance>,
+		Fungibles: fungibles::Balanced<AccountId>,
+		FungiblesAssetMatcher: MatchesFungibles<Fungibles::AssetId, Fungibles::Balance>,
+		OnUnbalanced: OnUnbalancedT<fungibles::Credit<AccountId, Fungibles>>,
+		AccountId,
+	> Drop
+	for SwapFirstAssetTrader<
+		Target,
+		SwapCredit,
+		WeightToFee,
+		Fungibles,
+		FungiblesAssetMatcher,
+		OnUnbalanced,
+		AccountId,
+	>
+where
+	Fungibles::Balance: Into<u128>,
+{
+	fn drop(&mut self) {
+		if self.total_fee.peek().is_zero() {
+			return;
+		}
+		let total_fee = self.total_fee.extract(self.total_fee.peek());
+		OnUnbalanced::on_unbalanced(total_fee);
+	}
+}
+
+#[cfg(test)]
+mod test_xcm_router {
+	use super::*;
+	use cumulus_primitives_core::UpwardMessage;
+	use pezframe_support::assert_ok;
+	use xcm::MAX_XCM_DECODE_DEPTH;
+
+	/// Validates [`validate`] for required Some(destination) and Some(message)
+	struct OkFixedXcmHashWithAssertingRequiredInputsSender;
+	impl OkFixedXcmHashWithAssertingRequiredInputsSender {
+		const FIXED_XCM_HASH: [u8; 32] = [9; 32];
+
+		fn fixed_delivery_asset() -> Assets {
+			Assets::new()
+		}
+
+		fn expected_delivery_result() -> Result<(XcmHash, Assets), SendError> {
+			Ok((Self::FIXED_XCM_HASH, Self::fixed_delivery_asset()))
+		}
+	}
+	impl SendXcm for OkFixedXcmHashWithAssertingRequiredInputsSender {
+		type Ticket = ();
+
+		fn validate(
+			destination: &mut Option<Location>,
+			message: &mut Option<Xcm<()>>,
+		) -> SendResult<Self::Ticket> {
+			assert!(destination.is_some());
+			assert!(message.is_some());
+			Ok(((), OkFixedXcmHashWithAssertingRequiredInputsSender::fixed_delivery_asset()))
+		}
+
+		fn deliver(_: Self::Ticket) -> Result<XcmHash, SendError> {
+			Ok(Self::FIXED_XCM_HASH)
+		}
+	}
+
+	/// Impl [`UpwardMessageSender`] that return `Ok` for `can_send_upward_message`.
+	struct CanSendUpwardMessageSender;
+	impl UpwardMessageSender for CanSendUpwardMessageSender {
+		fn send_upward_message(_: UpwardMessage) -> Result<(u32, XcmHash), MessageSendError> {
+			Err(MessageSendError::Other)
+		}
+
+		fn can_send_upward_message(_: &UpwardMessage) -> Result<(), MessageSendError> {
+			Ok(())
+		}
+	}
+
+	#[test]
+	fn parent_as_ump_does_not_consume_dest_or_msg_on_not_applicable() {
+		// dummy message
+		let message = Xcm(vec![Trap(5)]);
+
+		// ParentAsUmp - check dest is really not applicable
+		let dest = (Parent, Parent, Parent);
+		let mut dest_wrapper = Some(dest.into());
+		let mut msg_wrapper = Some(message.clone());
+		assert_eq!(
+			Err(SendError::NotApplicable),
+			<ParentAsUmp<(), (), ()> as SendXcm>::validate(&mut dest_wrapper, &mut msg_wrapper)
+		);
+
+		// check wrapper were not consumed
+		assert_eq!(Some(dest.into()), dest_wrapper.take());
+		assert_eq!(Some(message.clone()), msg_wrapper.take());
+
+		// another try with router chain with asserting sender
+		assert_eq!(
+			OkFixedXcmHashWithAssertingRequiredInputsSender::expected_delivery_result(),
+			send_xcm::<(ParentAsUmp<(), (), ()>, OkFixedXcmHashWithAssertingRequiredInputsSender)>(
+				dest.into(),
+				message
+			)
+		);
+	}
+
+	#[test]
+	fn parent_as_ump_consumes_dest_and_msg_on_ok_validate() {
+		// dummy message
+		let message = Xcm(vec![Trap(5)]);
+
+		// ParentAsUmp - check dest/msg is valid
+		let dest = (Parent, Here);
+		let mut dest_wrapper = Some(dest.clone().into());
+		let mut msg_wrapper = Some(message.clone());
+		assert!(<ParentAsUmp<CanSendUpwardMessageSender, (), ()> as SendXcm>::validate(
+			&mut dest_wrapper,
+			&mut msg_wrapper
+		)
+		.is_ok());
+
+		// check wrapper were consumed
+		assert_eq!(None, dest_wrapper.take());
+		assert_eq!(None, msg_wrapper.take());
+
+		// another try with router chain with asserting sender
+		assert_eq!(
+			Err(SendError::Transport("Other")),
+			send_xcm::<(
+				ParentAsUmp<CanSendUpwardMessageSender, (), ()>,
+				OkFixedXcmHashWithAssertingRequiredInputsSender
+			)>(dest.into(), message)
+		);
+	}
+
+	#[test]
+	fn parent_as_ump_validate_nested_xcm_works() {
+		let dest = Parent;
+
+		type Router = ParentAsUmp<CanSendUpwardMessageSender, (), ()>;
+
+		// Message that is not too deeply nested:
+		let mut good = Xcm(vec![ClearOrigin]);
+		for _ in 0..MAX_XCM_DECODE_DEPTH - 1 {
+			good = Xcm(vec![SetAppendix(good)]);
+		}
+
+		// Check that the good message is validated:
+		assert_ok!(<Router as SendXcm>::validate(&mut Some(dest.into()), &mut Some(good.clone())));
+
+		// Nesting the message one more time should reject it:
+		let bad = Xcm(vec![SetAppendix(good)]);
+		assert_eq!(
+			Err(SendError::ExceedsMaxMessageSize),
+			<Router as SendXcm>::validate(&mut Some(dest.into()), &mut Some(bad))
+		);
+	}
+}
+#[cfg(test)]
+mod test_trader {
+	use super::*;
+	use pezframe_support::{
+		assert_ok,
+		traits::tokens::{
+			DepositConsequence, Fortitude, Preservation, Provenance, WithdrawConsequence,
+		},
+	};
+	use pezsp_runtime::DispatchError;
+	use xcm_executor::{traits::Error, AssetsInHolding};
+
+	#[test]
+	fn take_first_asset_trader_buy_weight_called_twice_throws_error() {
+		const AMOUNT: u128 = 100;
+
+		// prepare prerequisites to instantiate `TakeFirstAssetTrader`
+		type TestAccountId = u32;
+		type TestAssetId = u32;
+		type TestBalance = u128;
+		struct TestAssets;
+		impl MatchesFungibles<TestAssetId, TestBalance> for TestAssets {
+			fn matches_fungibles(a: &Asset) -> Result<(TestAssetId, TestBalance), Error> {
+				match a {
+					Asset { fun: Fungible(amount), id: AssetId(_id) } => Ok((1, *amount)),
+					_ => Err(Error::AssetNotHandled),
+				}
+			}
+		}
+		impl fungibles::Inspect<TestAccountId> for TestAssets {
+			type AssetId = TestAssetId;
+			type Balance = TestBalance;
+
+			fn total_issuance(_: Self::AssetId) -> Self::Balance {
+				todo!()
+			}
+
+			fn minimum_balance(_: Self::AssetId) -> Self::Balance {
+				0
+			}
+
+			fn balance(_: Self::AssetId, _: &TestAccountId) -> Self::Balance {
+				todo!()
+			}
+
+			fn total_balance(_: Self::AssetId, _: &TestAccountId) -> Self::Balance {
+				todo!()
+			}
+
+			fn reducible_balance(
+				_: Self::AssetId,
+				_: &TestAccountId,
+				_: Preservation,
+				_: Fortitude,
+			) -> Self::Balance {
+				todo!()
+			}
+
+			fn can_deposit(
+				_: Self::AssetId,
+				_: &TestAccountId,
+				_: Self::Balance,
+				_: Provenance,
+			) -> DepositConsequence {
+				todo!()
+			}
+
+			fn can_withdraw(
+				_: Self::AssetId,
+				_: &TestAccountId,
+				_: Self::Balance,
+			) -> WithdrawConsequence<Self::Balance> {
+				todo!()
+			}
+
+			fn asset_exists(_: Self::AssetId) -> bool {
+				todo!()
+			}
+		}
+		impl fungibles::Mutate<TestAccountId> for TestAssets {}
+		impl fungibles::Balanced<TestAccountId> for TestAssets {
+			type OnDropCredit = fungibles::DecreaseIssuance<TestAccountId, Self>;
+			type OnDropDebt = fungibles::IncreaseIssuance<TestAccountId, Self>;
+		}
+		impl fungibles::Unbalanced<TestAccountId> for TestAssets {
+			fn handle_dust(_: fungibles::Dust<TestAccountId, Self>) {
+				todo!()
+			}
+			fn write_balance(
+				_: Self::AssetId,
+				_: &TestAccountId,
+				_: Self::Balance,
+			) -> Result<Option<Self::Balance>, DispatchError> {
+				todo!()
+			}
+
+			fn set_total_issuance(_: Self::AssetId, _: Self::Balance) {
+				todo!()
+			}
+		}
+
+		struct FeeChargerAssetsHandleRefund;
+		impl ChargeWeightInFungibles<TestAccountId, TestAssets> for FeeChargerAssetsHandleRefund {
+			fn charge_weight_in_fungibles(
+				_: <TestAssets as fungibles::Inspect<TestAccountId>>::AssetId,
+				_: Weight,
+			) -> Result<<TestAssets as fungibles::Inspect<TestAccountId>>::Balance, XcmError> {
+				Ok(AMOUNT)
+			}
+		}
+		impl TakeRevenue for FeeChargerAssetsHandleRefund {
+			fn take_revenue(_: Asset) {}
+		}
+
+		// create new instance
+		type Trader = TakeFirstAssetTrader<
+			TestAccountId,
+			FeeChargerAssetsHandleRefund,
+			TestAssets,
+			TestAssets,
+			FeeChargerAssetsHandleRefund,
+		>;
+		let mut trader = <Trader as WeightTrader>::new();
+		let ctx = XcmContext { origin: None, message_id: XcmHash::default(), topic: None };
+
+		// prepare test data
+		let asset: Asset = (Here, AMOUNT).into();
+		let payment = AssetsInHolding::from(asset);
+		let weight_to_buy = Weight::from_parts(1_000, 1_000);
+
+		// lets do first call (success)
+		assert_ok!(trader.buy_weight(weight_to_buy, payment.clone(), &ctx));
+
+		// lets do second call (error)
+		assert_eq!(trader.buy_weight(weight_to_buy, payment, &ctx), Err(XcmError::NotWithdrawable));
+	}
+}
+
+/// Implementation of `xcm_builder::EnsureDelivery` which helps to ensure delivery to the
+/// parent relay chain. Deposits existential deposit for origin (if needed).
+/// Deposits estimated fee to the origin account (if needed).
+/// Allows triggering of additional logic for a specific `ParaId` (e.g. to open an HRMP channel) if
+/// needed.
+#[cfg(feature = "runtime-benchmarks")]
+pub struct ToParentDeliveryHelper<XcmConfig, ExistentialDeposit, PriceForDelivery>(
+	core::marker::PhantomData<(XcmConfig, ExistentialDeposit, PriceForDelivery)>,
+);
+
+#[cfg(feature = "runtime-benchmarks")]
+impl<
+		XcmConfig: xcm_executor::Config,
+		ExistentialDeposit: Get<Option<Asset>>,
+		PriceForDelivery: PriceForMessageDelivery<Id = ()>,
+	> xcm_builder::EnsureDelivery
+	for ToParentDeliveryHelper<XcmConfig, ExistentialDeposit, PriceForDelivery>
+{
+	fn ensure_successful_delivery(
+		origin_ref: &Location,
+		dest: &Location,
+		fee_reason: xcm_executor::traits::FeeReason,
+	) -> (Option<xcm_executor::FeesMode>, Option<Assets>) {
+		use xcm::{latest::MAX_ITEMS_IN_ASSETS, MAX_INSTRUCTIONS_TO_DECODE};
+		use xcm_executor::{traits::FeeManager, FeesMode};
+
+		// check if the destination is relay/parent
+		if dest.ne(&Location::parent()) {
+			return (None, None);
+		}
+
+		// Ensure routers
+		XcmConfig::XcmSender::ensure_successful_delivery(Some(Location::parent()));
+
+		let mut fees_mode = None;
+		if !XcmConfig::FeeManager::is_waived(Some(origin_ref), fee_reason) {
+			// if not waived, we need to set up accounts for paying and receiving fees
+
+			// mint ED to origin if needed
+			if let Some(ed) = ExistentialDeposit::get() {
+				XcmConfig::AssetTransactor::deposit_asset(&ed, &origin_ref, None).unwrap();
+			}
+
+			// overestimate delivery fee
+			let mut max_assets: Vec<Asset> = Vec::new();
+			for i in 0..MAX_ITEMS_IN_ASSETS {
+				max_assets.push((GeneralIndex(i as u128), 100u128).into());
+			}
+			let overestimated_xcm =
+				vec![WithdrawAsset(max_assets.into()); MAX_INSTRUCTIONS_TO_DECODE as usize].into();
+			let overestimated_fees = PriceForDelivery::price_for_delivery((), &overestimated_xcm);
+
+			// mint overestimated fee to origin
+			for fee in overestimated_fees.inner() {
+				XcmConfig::AssetTransactor::deposit_asset(&fee, &origin_ref, None).unwrap();
+			}
+
+			// expected worst case - direct withdraw
+			fees_mode = Some(FeesMode { jit_withdraw: true });
+		}
+		(fees_mode, None)
+	}
+}

pezcumulus/primitives/utility/src/tests/mod.rs

@@ -0,0 +1,17 @@
+// 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.
+
+mod swap_first;

pezcumulus/primitives/utility/src/tests/swap_first.rs

@@ -0,0 +1,551 @@
+// 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_support::{parameter_types, traits::fungibles::Inspect};
+use mock::{setup_pool, AccountId, AssetId, Balance, Fungibles};
+use xcm::latest::AssetId as XcmAssetId;
+use xcm_executor::AssetsInHolding;
+
+fn create_holding_asset(asset_id: AssetId, amount: Balance) -> AssetsInHolding {
+	create_asset(asset_id, amount).into()
+}
+
+fn create_asset(asset_id: AssetId, amount: Balance) -> Asset {
+	Asset { id: create_asset_id(asset_id), fun: Fungible(amount) }
+}
+
+fn create_asset_id(asset_id: AssetId) -> XcmAssetId {
+	AssetId(Location::new(0, [GeneralIndex(asset_id.into())]))
+}
+
+fn xcm_context() -> XcmContext {
+	XcmContext { origin: None, message_id: [0u8; 32], topic: None }
+}
+
+fn weight_worth_of(fee: Balance) -> Weight {
+	Weight::from_parts(fee.try_into().unwrap(), 0)
+}
+
+const TARGET_ASSET: AssetId = 1;
+const CLIENT_ASSET: AssetId = 2;
+const CLIENT_ASSET_2: AssetId = 3;
+
+parameter_types! {
+	pub const TargetAsset: AssetId = TARGET_ASSET;
+}
+
+pub type Trader = SwapFirstAssetTrader<
+	TargetAsset,
+	mock::Swap,
+	mock::WeightToFee,
+	mock::Fungibles,
+	mock::FungiblesMatcher,
+	(),
+	AccountId,
+>;
+
+#[test]
+fn holding_asset_swap_for_target() {
+	let client_asset_total = 15;
+	let fee = 5;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_change = create_holding_asset(CLIENT_ASSET, client_asset_total - fee);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee), holding_asset, &xcm_context()).unwrap(),
+		holding_change
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee);
+}
+
+#[test]
+fn holding_asset_swap_for_target_twice() {
+	let client_asset_total = 20;
+	let fee1 = 5;
+	let fee2 = 6;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_change1 = create_holding_asset(CLIENT_ASSET, client_asset_total - fee1);
+	let holding_change2 = create_holding_asset(CLIENT_ASSET, client_asset_total - fee1 - fee2);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee1), holding_asset, &xcm_context()).unwrap(),
+		holding_change1
+	);
+	assert_eq!(
+		trader
+			.buy_weight(weight_worth_of(fee2), holding_change1, &xcm_context())
+			.unwrap(),
+		holding_change2
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee1 + fee2);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee1 + fee2);
+}
+
+#[test]
+fn buy_and_refund_twice_for_target() {
+	let client_asset_total = 15;
+	let fee = 5;
+	let refund1 = 4;
+	let refund2 = 2;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+	// create pool for refund swap.
+	setup_pool(TARGET_ASSET, 1000, CLIENT_ASSET, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_change = create_holding_asset(CLIENT_ASSET, client_asset_total - fee);
+	let refund_asset = create_asset(CLIENT_ASSET, refund1);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee), holding_asset, &xcm_context()).unwrap(),
+		holding_change
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(trader.refund_weight(weight_worth_of(refund1), &xcm_context()), Some(refund_asset));
+
+	assert_eq!(trader.total_fee.peek(), fee - refund1);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(trader.refund_weight(weight_worth_of(refund2), &xcm_context()), None);
+
+	assert_eq!(trader.total_fee.peek(), fee - refund1);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee - refund1);
+}
+
+#[test]
+fn buy_with_various_assets_and_refund_for_target() {
+	let client_asset_total = 10;
+	let client_asset_2_total = 15;
+	let fee1 = 5;
+	let fee2 = 6;
+	let refund1 = 6;
+	let refund2 = 4;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+	setup_pool(CLIENT_ASSET_2, 1000, TARGET_ASSET, 1000);
+	// create pool for refund swap.
+	setup_pool(TARGET_ASSET, 1000, CLIENT_ASSET_2, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_asset_2 = create_holding_asset(CLIENT_ASSET_2, client_asset_2_total);
+	let holding_change = create_holding_asset(CLIENT_ASSET, client_asset_total - fee1);
+	let holding_change_2 = create_holding_asset(CLIENT_ASSET_2, client_asset_2_total - fee2);
+	// both refunds in the latest buy asset (`CLIENT_ASSET_2`).
+	let refund_asset = create_asset(CLIENT_ASSET_2, refund1);
+	let refund_asset_2 = create_asset(CLIENT_ASSET_2, refund2);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+	let client_total_2 = Fungibles::total_issuance(CLIENT_ASSET_2);
+
+	let mut trader = Trader::new();
+	// first purchase with `CLIENT_ASSET`.
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee1), holding_asset, &xcm_context()).unwrap(),
+		holding_change
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee1);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	// second purchase with `CLIENT_ASSET_2`.
+	assert_eq!(
+		trader
+			.buy_weight(weight_worth_of(fee2), holding_asset_2, &xcm_context())
+			.unwrap(),
+		holding_change_2
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee1 + fee2);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET_2)));
+
+	// first refund in the last asset used with `buy_weight`.
+	assert_eq!(trader.refund_weight(weight_worth_of(refund1), &xcm_context()), Some(refund_asset));
+
+	assert_eq!(trader.total_fee.peek(), fee1 + fee2 - refund1);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET_2)));
+
+	// second refund in the last asset used with `buy_weight`.
+	assert_eq!(
+		trader.refund_weight(weight_worth_of(refund2), &xcm_context()),
+		Some(refund_asset_2)
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee1 + fee2 - refund1 - refund2);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET_2)));
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee1);
+	assert_eq!(
+		Fungibles::total_issuance(CLIENT_ASSET_2),
+		client_total_2 + fee2 - refund1 - refund2
+	);
+}
+
+#[test]
+fn not_enough_to_refund() {
+	let client_asset_total = 15;
+	let fee = 5;
+	let refund = 6;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_change = create_holding_asset(CLIENT_ASSET, client_asset_total - fee);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee), holding_asset, &xcm_context()).unwrap(),
+		holding_change
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(trader.refund_weight(weight_worth_of(refund), &xcm_context()), None);
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee);
+}
+
+#[test]
+fn not_exchangeable_to_refund() {
+	let client_asset_total = 15;
+	let fee = 5;
+	let refund = 1;
+
+	setup_pool(CLIENT_ASSET, 1000, TARGET_ASSET, 1000);
+
+	let holding_asset = create_holding_asset(CLIENT_ASSET, client_asset_total);
+	let holding_change = create_holding_asset(CLIENT_ASSET, client_asset_total - fee);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader.buy_weight(weight_worth_of(fee), holding_asset, &xcm_context()).unwrap(),
+		holding_change
+	);
+
+	assert_eq!(trader.total_fee.peek(), fee);
+	assert_eq!(trader.last_fee_asset, Some(create_asset_id(CLIENT_ASSET)));
+
+	assert_eq!(trader.refund_weight(weight_worth_of(refund), &xcm_context()), None);
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total + fee);
+}
+
+#[test]
+fn nothing_to_refund() {
+	let fee = 5;
+
+	let mut trader = Trader::new();
+	assert_eq!(trader.refund_weight(weight_worth_of(fee), &xcm_context()), None);
+}
+
+#[test]
+fn holding_asset_not_exchangeable_for_target() {
+	let holding_asset = create_holding_asset(CLIENT_ASSET, 10);
+
+	let target_total = Fungibles::total_issuance(TARGET_ASSET);
+	let client_total = Fungibles::total_issuance(CLIENT_ASSET);
+
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader
+			.buy_weight(Weight::from_all(10), holding_asset, &xcm_context())
+			.unwrap_err(),
+		XcmError::FeesNotMet
+	);
+
+	assert_eq!(Fungibles::total_issuance(TARGET_ASSET), target_total);
+	assert_eq!(Fungibles::total_issuance(CLIENT_ASSET), client_total);
+}
+
+#[test]
+fn empty_holding_asset() {
+	let mut trader = Trader::new();
+	assert_eq!(
+		trader
+			.buy_weight(Weight::from_all(10), AssetsInHolding::new(), &xcm_context())
+			.unwrap_err(),
+		XcmError::AssetNotFound
+	);
+}
+
+#[test]
+fn fails_to_match_holding_asset() {
+	let mut trader = Trader::new();
+	let holding_asset = Asset { id: AssetId(Location::new(1, [Teyrchain(1)])), fun: Fungible(10) };
+	assert_eq!(
+		trader
+			.buy_weight(Weight::from_all(10), holding_asset.into(), &xcm_context())
+			.unwrap_err(),
+		XcmError::AssetNotFound
+	);
+}
+
+#[test]
+fn holding_asset_equal_to_target_asset() {
+	let mut trader = Trader::new();
+	let holding_asset = create_holding_asset(TargetAsset::get(), 10);
+	assert_eq!(
+		trader
+			.buy_weight(Weight::from_all(10), holding_asset, &xcm_context())
+			.unwrap_err(),
+		XcmError::FeesNotMet
+	);
+}
+
+pub mod mock {
+	use crate::*;
+	use core::cell::RefCell;
+	use pezframe_support::{
+		ensure,
+		traits::{
+			fungibles::{Balanced, DecreaseIssuance, Dust, IncreaseIssuance, Inspect, Unbalanced},
+			tokens::{
+				DepositConsequence, Fortitude, Fortitude::Polite, Precision::Exact, Preservation,
+				Preservation::Preserve, Provenance, WithdrawConsequence,
+			},
+		},
+	};
+	use pezsp_runtime::{traits::One, DispatchError};
+	use std::collections::HashMap;
+	use xcm::latest::Junction;
+
+	pub type AccountId = u64;
+	pub type AssetId = u32;
+	pub type Balance = u128;
+	pub type Credit = fungibles::Credit<AccountId, Fungibles>;
+
+	thread_local! {
+	   pub static TOTAL_ISSUANCE: RefCell<HashMap<AssetId, Balance>> = RefCell::new(HashMap::new());
+	   pub static ACCOUNT: RefCell<HashMap<(AssetId, AccountId), Balance>> = RefCell::new(HashMap::new());
+	   pub static SWAP: RefCell<HashMap<(AssetId, AssetId), AccountId>> = RefCell::new(HashMap::new());
+	}
+
+	pub struct Swap {}
+	impl SwapCreditT<AccountId> for Swap {
+		type Balance = Balance;
+		type AssetKind = AssetId;
+		type Credit = Credit;
+		fn max_path_len() -> u32 {
+			2
+		}
+		fn swap_exact_tokens_for_tokens(
+			path: Vec<Self::AssetKind>,
+			credit_in: Self::Credit,
+			amount_out_min: Option<Self::Balance>,
+		) -> Result<Self::Credit, (Self::Credit, DispatchError)> {
+			ensure!(2 == path.len(), (credit_in, DispatchError::Unavailable));
+			ensure!(
+				credit_in.peek() >= amount_out_min.unwrap_or(Self::Balance::zero()),
+				(credit_in, DispatchError::Unavailable)
+			);
+			let swap_res = SWAP.with(|b| b.borrow().get(&(path[0], path[1])).map(|v| *v));
+			let pool_account = match swap_res {
+				Some(a) => a,
+				None => return Err((credit_in, DispatchError::Unavailable)),
+			};
+			let credit_out = match Fungibles::withdraw(
+				path[1],
+				&pool_account,
+				credit_in.peek(),
+				Exact,
+				Preserve,
+				Polite,
+			) {
+				Ok(c) => c,
+				Err(_) => return Err((credit_in, DispatchError::Unavailable)),
+			};
+			Fungibles::resolve(&pool_account, credit_in)
+				.map_err(|c| (c, DispatchError::Unavailable))?;
+			Ok(credit_out)
+		}
+		fn swap_tokens_for_exact_tokens(
+			path: Vec<Self::AssetKind>,
+			credit_in: Self::Credit,
+			amount_out: Self::Balance,
+		) -> Result<(Self::Credit, Self::Credit), (Self::Credit, DispatchError)> {
+			ensure!(2 == path.len(), (credit_in, DispatchError::Unavailable));
+			ensure!(credit_in.peek() >= amount_out, (credit_in, DispatchError::Unavailable));
+			let swap_res = SWAP.with(|b| b.borrow().get(&(path[0], path[1])).map(|v| *v));
+			let pool_account = match swap_res {
+				Some(a) => a,
+				None => return Err((credit_in, DispatchError::Unavailable)),
+			};
+			let credit_out = match Fungibles::withdraw(
+				path[1],
+				&pool_account,
+				amount_out,
+				Exact,
+				Preserve,
+				Polite,
+			) {
+				Ok(c) => c,
+				Err(_) => return Err((credit_in, DispatchError::Unavailable)),
+			};
+			let (credit_in, change) = credit_in.split(amount_out);
+			Fungibles::resolve(&pool_account, credit_in)
+				.map_err(|c| (c, DispatchError::Unavailable))?;
+			Ok((credit_out, change))
+		}
+	}
+
+	pub fn pool_account(asset1: AssetId, asset2: AssetId) -> AccountId {
+		(1000 + asset1 * 10 + asset2 * 100).into()
+	}
+
+	pub fn setup_pool(asset1: AssetId, liquidity1: Balance, asset2: AssetId, liquidity2: Balance) {
+		let account = pool_account(asset1, asset2);
+		SWAP.with(|b| b.borrow_mut().insert((asset1, asset2), account));
+		let debt1 = Fungibles::deposit(asset1, &account, liquidity1, Exact);
+		let debt2 = Fungibles::deposit(asset2, &account, liquidity2, Exact);
+		drop(debt1);
+		drop(debt2);
+	}
+
+	pub struct WeightToFee;
+	impl WeightToFeeT for WeightToFee {
+		type Balance = Balance;
+		fn weight_to_fee(weight: &Weight) -> Self::Balance {
+			(weight.ref_time() + weight.proof_size()).into()
+		}
+	}
+
+	pub struct Fungibles {}
+	impl Inspect<AccountId> for Fungibles {
+		type AssetId = AssetId;
+		type Balance = Balance;
+		fn total_issuance(asset: Self::AssetId) -> Self::Balance {
+			TOTAL_ISSUANCE.with(|b| b.borrow().get(&asset).map_or(Self::Balance::zero(), |b| *b))
+		}
+		fn minimum_balance(_: Self::AssetId) -> Self::Balance {
+			Self::Balance::one()
+		}
+		fn total_balance(asset: Self::AssetId, who: &AccountId) -> Self::Balance {
+			ACCOUNT.with(|b| b.borrow().get(&(asset, *who)).map_or(Self::Balance::zero(), |b| *b))
+		}
+		fn balance(asset: Self::AssetId, who: &AccountId) -> Self::Balance {
+			ACCOUNT.with(|b| b.borrow().get(&(asset, *who)).map_or(Self::Balance::zero(), |b| *b))
+		}
+		fn reducible_balance(
+			asset: Self::AssetId,
+			who: &AccountId,
+			_: Preservation,
+			_: Fortitude,
+		) -> Self::Balance {
+			ACCOUNT.with(|b| b.borrow().get(&(asset, *who)).map_or(Self::Balance::zero(), |b| *b))
+		}
+		fn can_deposit(
+			_: Self::AssetId,
+			_: &AccountId,
+			_: Self::Balance,
+			_: Provenance,
+		) -> DepositConsequence {
+			unimplemented!()
+		}
+		fn can_withdraw(
+			_: Self::AssetId,
+			_: &AccountId,
+			_: Self::Balance,
+		) -> WithdrawConsequence<Self::Balance> {
+			unimplemented!()
+		}
+		fn asset_exists(_: Self::AssetId) -> bool {
+			unimplemented!()
+		}
+	}
+
+	impl Unbalanced<AccountId> for Fungibles {
+		fn set_total_issuance(asset: Self::AssetId, amount: Self::Balance) {
+			TOTAL_ISSUANCE.with(|b| b.borrow_mut().insert(asset, amount));
+		}
+		fn handle_dust(_: Dust<AccountId, Self>) {
+			unimplemented!()
+		}
+		fn write_balance(
+			asset: Self::AssetId,
+			who: &AccountId,
+			amount: Self::Balance,
+		) -> Result<Option<Self::Balance>, DispatchError> {
+			let _ = ACCOUNT.with(|b| b.borrow_mut().insert((asset, *who), amount));
+			Ok(None)
+		}
+	}
+
+	impl Balanced<AccountId> for Fungibles {
+		type OnDropCredit = DecreaseIssuance<AccountId, Self>;
+		type OnDropDebt = IncreaseIssuance<AccountId, Self>;
+	}
+
+	pub struct FungiblesMatcher;
+	impl MatchesFungibles<AssetId, Balance> for FungiblesMatcher {
+		fn matches_fungibles(
+			a: &Asset,
+		) -> core::result::Result<(AssetId, Balance), xcm_executor::traits::Error> {
+			match a {
+				Asset { fun: Fungible(amount), id: AssetId(inner_location) } =>
+					match inner_location.unpack() {
+						(0, [Junction::GeneralIndex(id)]) =>
+							Ok(((*id).try_into().unwrap(), *amount)),
+						_ => Err(xcm_executor::traits::Error::AssetNotHandled),
+					},
+				_ => Err(xcm_executor::traits::Error::AssetNotHandled),
+			}
+		}
+	}
+}