make on_unbalanceds work fungibles imbalances

cumulus/parachains/common/src/impls.rs

@@ -66,7 +66,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<primitives::AccountId>,
 	<R as frame_system::Config>::AccountId: Into<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() {

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

@@ -185,7 +185,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

@@ -39,7 +39,7 @@ pub use imbalances::{NegativeImbalance, PositiveImbalance};
 // of the inner member.
 mod imbalances {
 	use super::{result, Config, Imbalance, RuntimeDebug, Saturating, TryDrop, Zero};
-	use frame_support::traits::SameOrOther;
+	use frame_support::traits::{tokens::imbalance::TryMerge, SameOrOther};
 	use sp_std::mem;
 
 	/// Opaque, move-only struct with private fields that serves as a token denoting that
@@ -132,6 +132,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()
@@ -196,6 +202,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 frame_support_procedural::{EqNoBound, PartialEqNoBound, RuntimeDebugNoBound};
 use sp_runtime::traits::Zero;
@@ -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 frame_support_procedural::{EqNoBound, PartialEqNoBound, RuntimeDebugNoBound};
 use sp_runtime::traits::Zero;
@@ -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 are not identical (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

@@ -159,7 +159,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>,
 		>,

substrate/frame/transaction-payment/src/mock.rs

@@ -140,7 +140,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>,
 		>,