Skip to content

Commit

Permalink
extract amount method for fungible/s Imbalance (#1847)
Browse files Browse the repository at this point in the history
Introduces an `extract` amount method for `fungible/s` `Imbalance`.
  • Loading branch information
muharem authored Oct 16, 2023
1 parent 646ecd0 commit c422e3f
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 0 deletions.
10 changes: 10 additions & 0 deletions substrate/frame/balances/src/impl_currency.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,13 @@ impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, 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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<Self, (Self, Self)> {
if self.asset == other.asset {
self.amount = self.amount.saturating_add(other.amount);
Expand Down
7 changes: 7 additions & 0 deletions substrate/frame/support/src/traits/tokens/imbalance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,10 @@ pub trait Imbalance<Balance>: 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`.
///
Expand Down Expand Up @@ -190,6 +194,9 @@ impl<Balance: Default> Imbalance<Balance> for () {
fn split(self, _: Balance) -> (Self, Self) {
((), ())
}
fn extract(&mut self, _: Balance) -> Self {
()
}
fn ration(self, _: u32, _: u32) -> (Self, Self)
where
Balance: From<u32> + Saturating + Div<Output = Balance>,
Expand Down

0 comments on commit c422e3f

Please sign in to comment.