From eae2ed3462d684283f0775c3856d142a5e3ffb56 Mon Sep 17 00:00:00 2001 From: muharem Date: Wed, 11 Oct 2023 11:50:18 +0200 Subject: [PATCH] extract for imbalance --- substrate/frame/balances/src/impl_currency.rs | 10 ++++++++++ .../support/src/traits/tokens/fungible/imbalance.rs | 7 +++++++ .../support/src/traits/tokens/fungibles/imbalance.rs | 9 +++++++++ substrate/frame/support/src/traits/tokens/imbalance.rs | 7 +++++++ 4 files changed, 33 insertions(+) diff --git a/substrate/frame/balances/src/impl_currency.rs b/substrate/frame/balances/src/impl_currency.rs index c64a4929dd56..1ac882ade70d 100644 --- a/substrate/frame/balances/src/impl_currency.rs +++ b/substrate/frame/balances/src/impl_currency.rs @@ -100,6 +100,11 @@ mod imbalances { mem::forget(self); (Self(first), Self(second)) } + fn extract(&mut self, amount: T::Balance) -> Self { + let new = self.0.min(amount); + self.0 = self.0 - new; + Self(new) + } fn merge(mut self, other: Self) -> Self { self.0 = self.0.saturating_add(other.0); mem::forget(other); @@ -159,6 +164,11 @@ mod imbalances { mem::forget(self); (Self(first), Self(second)) } + fn extract(&mut self, amount: T::Balance) -> Self { + let new = self.0.min(amount); + self.0 = self.0 - new; + Self(new) + } fn merge(mut self, other: Self) -> Self { self.0 = self.0.saturating_add(other.0); mem::forget(other); diff --git a/substrate/frame/support/src/traits/tokens/fungible/imbalance.rs b/substrate/frame/support/src/traits/tokens/fungible/imbalance.rs index 32a63fd25b29..995797bc8f66 100644 --- a/substrate/frame/support/src/traits/tokens/fungible/imbalance.rs +++ b/substrate/frame/support/src/traits/tokens/fungible/imbalance.rs @@ -113,6 +113,13 @@ impl, OppositeOnDrop: HandleImbalance sp_std::mem::forget(self); (Imbalance::new(first), Imbalance::new(second)) } + + fn extract(&mut self, amount: B) -> Self { + let new = self.amount.min(amount); + self.amount = self.amount - new; + Imbalance::new(new) + } + fn merge(mut self, other: Self) -> Self { self.amount = self.amount.saturating_add(other.amount); sp_std::mem::forget(other); diff --git a/substrate/frame/support/src/traits/tokens/fungibles/imbalance.rs b/substrate/frame/support/src/traits/tokens/fungibles/imbalance.rs index 9f660a5f1689..7c0d7721a2e6 100644 --- a/substrate/frame/support/src/traits/tokens/fungibles/imbalance.rs +++ b/substrate/frame/support/src/traits/tokens/fungibles/imbalance.rs @@ -109,6 +109,15 @@ impl< sp_std::mem::forget(self); (Imbalance::new(asset.clone(), first), Imbalance::new(asset, second)) } + + /// Mutate `self` by extracting a new instance with at most `amount` value, reducing `self` + /// accordingly. + pub fn extract(&mut self, amount: B) -> Self { + let new = self.amount.min(amount); + self.amount = self.amount - new; + Imbalance::new(self.asset.clone(), new) + } + pub fn merge(mut self, other: Self) -> Result { if self.asset == other.asset { self.amount = self.amount.saturating_add(other.amount); diff --git a/substrate/frame/support/src/traits/tokens/imbalance.rs b/substrate/frame/support/src/traits/tokens/imbalance.rs index 403321725042..e487b8122341 100644 --- a/substrate/frame/support/src/traits/tokens/imbalance.rs +++ b/substrate/frame/support/src/traits/tokens/imbalance.rs @@ -72,6 +72,10 @@ pub trait Imbalance: Sized + TryDrop + Default { /// is guaranteed to be at most `amount` and the second will be the remainder. fn split(self, amount: Balance) -> (Self, Self); + /// Mutate `self` by extracting a new instance with at most `amount` value, reducing `self` + /// accordingly. + fn extract(&mut self, amount: Balance) -> Self; + /// Consume `self` and return two independent instances; the amounts returned will be in /// approximately the same ratio as `first`:`second`. /// @@ -190,6 +194,9 @@ impl Imbalance for () { fn split(self, _: Balance) -> (Self, Self) { ((), ()) } + fn extract(&mut self, _: Balance) -> Self { + () + } fn ration(self, _: u32, _: u32) -> (Self, Self) where Balance: From + Saturating + Div,