Make on_unbalanceds work with fungibles imbalances (#4564)

Make `on_unbalanceds` work with `fungibles` `imbalances`.

The `fungibles` `imbalances` cannot be handled by the default
implementation of `on_unbalanceds` from the `OnUnbalanced` trait. This
is because the `fungibles` `imbalances` types do not implement the
`Imbalance` trait (and cannot with its current semantics). The
`on_unbalanceds` function requires only the `merge` function for the
imbalance type. In this PR, we provide the `TryMerge` trait, which can
be implemented by all imbalance types and make `OnUnbalanced` require it
instead `Imbalance`.

### Migration for `OnUnbalanced` trait implementations:
In case if you have a custom implementation of `on_unbalanceds` trait
function, remove it's `<B>` type argument.

### Migration for custom imbalance types:
If you have your own imbalance types implementations, implement the
`TryMerge` trait for it introduced with this update.
        
The applicability of the `on_unbalanceds` function to fungibles
imbalances is useful in cases like -
[link](https://github.com/paritytech/polkadot-sdk/blob/3a8e675e9f6f283514c00c14d3d1d90ed5bf59c0/substrate/frame/transaction-payment/asset-conversion-tx-payment/src/payment.rs#L267)
from #4488.

---------

Co-authored-by: Oliver Tale-Yazdi <[email protected]>

Author: muharem

cumulus/parachains/common/src/impls.rs

@@ -67,7 +67,7 @@ where
 	AccountIdOf<R>: From<polkadot_primitives::AccountId> + Into<polkadot_primitives::AccountId>,
 	<R as frame_system::Config>::RuntimeEvent: From<pallet_balances::Event<R>>,
 {
-	fn on_unbalanceds<B>(
+	fn on_unbalanceds(
 		mut fees_then_tips: impl Iterator<
 			Item = fungible::Credit<R::AccountId, pallet_balances::Pallet<R>>,
 		>,

polkadot/runtime/common/src/impls.rs

@@ -51,7 +51,7 @@ where
 	<R as frame_system::Config>::AccountId: From<polkadot_primitives::AccountId>,
 	<R as frame_system::Config>::AccountId: Into<polkadot_primitives::AccountId>,
 {
-	fn on_unbalanceds<B>(
+	fn on_unbalanceds(
 		mut fees_then_tips: impl Iterator<Item = Credit<R::AccountId, pallet_balances::Pallet<R>>>,
 	) {
 		if let Some(fees) = fees_then_tips.next() {

prdoc/pr_4564.prdoc

@@ -0,0 +1,32 @@
+title: "Make `OnUnbalanced::on_unbalanceds` work with `fungibles` `imbalances`"
+
+doc:
+  - audience: Runtime Dev
+    description: |
+        The `on_unbalanceds` function of `OnUnbalanced` trait accepts the `fungibles` `imbalances` 
+        imbalances. This is done by replacing the `Imbalance` trait bound on imbalance type with 
+        the `TryMerge` trait bound. The `TryMerge` trait is implemented for all imbalance types.
+
+        ### Migration for `OnUnbalanced` trait implementations:
+        In case if you have a custom implementation of `on_unbalanceds` trait function, remove 
+        it's `<B>` type argument.
+
+        ### Migration for custom imbalance types:
+        If you have your own imbalance types implementations, implement the `TryMerge` trait for it 
+        introduced with this update.
+
+crates:
+  - name: frame-support
+    bump: major
+  - name: pallet-balances
+    bump: minor
+  - name: pallet-asset-conversion-tx-payment
+    bump: patch
+  - name: pallet-transaction-payment
+    bump: patch
+  - name: kitchensink-runtime
+    bump: patch
+  - name: polkadot-runtime-common
+    bump: patch
+  - name: parachains-common
+    bump: minor

substrate/bin/node/runtime/src/lib.rs

@@ -189,7 +189,7 @@ type NegativeImbalance = <Balances as Currency<AccountId>>::NegativeImbalance;
 
 pub struct DealWithFees;
 impl OnUnbalanced<NegativeImbalance> for DealWithFees {
-	fn on_unbalanceds<B>(mut fees_then_tips: impl Iterator<Item = NegativeImbalance>) {
+	fn on_unbalanceds(mut fees_then_tips: impl Iterator<Item = NegativeImbalance>) {
 		if let Some(fees) = fees_then_tips.next() {
 			// for fees, 80% to treasury, 20% to author
 			let mut split = fees.ration(80, 20);

substrate/frame/balances/src/impl_currency.rs

@@ -41,7 +41,7 @@ use sp_runtime::traits::Bounded;
 mod imbalances {
 	use super::*;
 	use core::mem;
-	use frame_support::traits::SameOrOther;
+	use frame_support::traits::{tokens::imbalance::TryMerge, SameOrOther};
 
 	/// Opaque, move-only struct with private fields that serves as a token denoting that
 	/// funds have been created without any equal and opposite accounting.
@@ -133,6 +133,12 @@ mod imbalances {
 		}
 	}
 
+	impl<T: Config<I>, I: 'static> TryMerge for PositiveImbalance<T, I> {
+		fn try_merge(self, other: Self) -> Result<Self, (Self, Self)> {
+			Ok(self.merge(other))
+		}
+	}
+
 	impl<T: Config<I>, I: 'static> TryDrop for NegativeImbalance<T, I> {
 		fn try_drop(self) -> result::Result<(), Self> {
 			self.drop_zero()
@@ -197,6 +203,12 @@ mod imbalances {
 		}
 	}
 
+	impl<T: Config<I>, I: 'static> TryMerge for NegativeImbalance<T, I> {
+		fn try_merge(self, other: Self) -> Result<Self, (Self, Self)> {
+			Ok(self.merge(other))
+		}
+	}
+
 	impl<T: Config<I>, I: 'static> Drop for PositiveImbalance<T, I> {
 		/// Basic drop handler will just square up the total issuance.
 		fn drop(&mut self) {

substrate/frame/support/src/traits/tokens/fungible/imbalance.rs

@@ -24,7 +24,7 @@ use super::{super::Imbalance as ImbalanceT, Balanced, *};
 use crate::traits::{
 	fungibles,
 	misc::{SameOrOther, TryDrop},
-	tokens::{AssetId, Balance},
+	tokens::{imbalance::TryMerge, AssetId, Balance},
 };
 use core::marker::PhantomData;
 use frame_support_procedural::{EqNoBound, PartialEqNoBound, RuntimeDebugNoBound};
@@ -157,6 +157,14 @@ impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalance
 	}
 }
 
+impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>> TryMerge
+	for Imbalance<B, OnDrop, OppositeOnDrop>
+{
+	fn try_merge(self, other: Self) -> Result<Self, (Self, Self)> {
+		Ok(self.merge(other))
+	}
+}
+
 /// Converts a `fungibles` `imbalance` instance to an instance of a `fungible` imbalance type.
 ///
 /// This function facilitates imbalance conversions within the implementations of

substrate/frame/support/src/traits/tokens/fungibles/imbalance.rs

@@ -24,7 +24,10 @@ use super::*;
 use crate::traits::{
 	fungible,
 	misc::{SameOrOther, TryDrop},
-	tokens::{imbalance::Imbalance as ImbalanceT, AssetId, Balance},
+	tokens::{
+		imbalance::{Imbalance as ImbalanceT, TryMerge},
+		AssetId, Balance,
+	},
 };
 use core::marker::PhantomData;
 use frame_support_procedural::{EqNoBound, PartialEqNoBound, RuntimeDebugNoBound};
@@ -176,6 +179,18 @@ impl<
 	}
 }
 
+impl<
+		A: AssetId,
+		B: Balance,
+		OnDrop: HandleImbalanceDrop<A, B>,
+		OppositeOnDrop: HandleImbalanceDrop<A, B>,
+	> TryMerge for Imbalance<A, B, OnDrop, OppositeOnDrop>
+{
+	fn try_merge(self, other: Self) -> Result<Self, (Self, Self)> {
+		self.merge(other)
+	}
+}
+
 /// Converts a `fungible` `imbalance` instance to an instance of a `fungibles` imbalance type using
 /// a specified `asset`.
 ///

substrate/frame/support/src/traits/tokens/imbalance.rs

@@ -58,7 +58,7 @@ pub use split_two_ways::SplitTwoWays;
 ///
 /// You can always retrieve the raw balance value using `peek`.
 #[must_use]
-pub trait Imbalance<Balance>: Sized + TryDrop + Default {
+pub trait Imbalance<Balance>: Sized + TryDrop + Default + TryMerge {
 	/// The oppositely imbalanced type. They come in pairs.
 	type Opposite: Imbalance<Balance>;
 
@@ -182,6 +182,13 @@ pub trait Imbalance<Balance>: Sized + TryDrop + Default {
 	fn peek(&self) -> Balance;
 }
 
+/// Try to merge two imbalances.
+pub trait TryMerge: Sized {
+	/// Consume `self` and an `other` to return a new instance that combines both. Errors with
+	/// Err(self, other) if the imbalances cannot be merged (e.g. imbalances of different assets).
+	fn try_merge(self, other: Self) -> Result<Self, (Self, Self)>;
+}
+
 #[cfg(feature = "std")]
 impl<Balance: Default> Imbalance<Balance> for () {
 	type Opposite = ();
@@ -236,3 +243,10 @@ impl<Balance: Default> Imbalance<Balance> for () {
 		Default::default()
 	}
 }
+
+#[cfg(feature = "std")]
+impl TryMerge for () {
+	fn try_merge(self, _: Self) -> Result<Self, (Self, Self)> {
+		Ok(())
+	}
+}

substrate/frame/support/src/traits/tokens/imbalance/on_unbalanced.rs

@@ -35,11 +35,26 @@ pub trait OnUnbalanced<Imbalance: TryDrop> {
 	/// Handler for some imbalances. The different imbalances might have different origins or
 	/// meanings, dependent on the context. Will default to simply calling on_unbalanced for all
 	/// of them. Infallible.
-	fn on_unbalanceds<B>(amounts: impl Iterator<Item = Imbalance>)
+	fn on_unbalanceds(mut amounts: impl Iterator<Item = Imbalance>)
 	where
-		Imbalance: crate::traits::Imbalance<B>,
+		Imbalance: crate::traits::tokens::imbalance::TryMerge,
 	{
-		Self::on_unbalanced(amounts.fold(Imbalance::zero(), |i, x| x.merge(i)))
+		let mut sum: Option<Imbalance> = None;
+		while let Some(next) = amounts.next() {
+			sum = match sum {
+				Some(sum) => match sum.try_merge(next) {
+					Ok(sum) => Some(sum),
+					Err((sum, next)) => {
+						Self::on_unbalanced(next);
+						Some(sum)
+					},
+				},
+				None => Some(next),
+			}
+		}
+		if let Some(sum) = sum {
+			Self::on_unbalanced(sum)
+		}
 	}
 
 	/// Handler for some imbalance. Infallible.

substrate/frame/transaction-payment/asset-conversion-tx-payment/src/mock.rs

@@ -131,7 +131,7 @@ pub struct DealWithFees;
 impl OnUnbalanced<fungible::Credit<<Runtime as frame_system::Config>::AccountId, Balances>>
 	for DealWithFees
 {
-	fn on_unbalanceds<B>(
+	fn on_unbalanceds(
 		mut fees_then_tips: impl Iterator<
 			Item = fungible::Credit<<Runtime as frame_system::Config>::AccountId, Balances>,
 		>,