use super::{super::Imbalance as ImbalanceT, Balanced, *};
use crate::traits::{
fungibles,
misc::{SameOrOther, TryDrop},
tokens::{imbalance::TryMerge, AssetId, Balance},
};
use core::marker::PhantomData;
use frame_support_procedural::{EqNoBound, PartialEqNoBound, RuntimeDebugNoBound};
use sp_runtime::traits::Zero;
pub trait HandleImbalanceDrop<Balance> {
fn handle(amount: Balance);
}
impl<Balance> HandleImbalanceDrop<Balance> for () {
fn handle(_: Balance) {}
}
#[must_use]
#[derive(EqNoBound, PartialEqNoBound, RuntimeDebugNoBound)]
pub struct Imbalance<
B: Balance,
OnDrop: HandleImbalanceDrop<B>,
OppositeOnDrop: HandleImbalanceDrop<B>,
> {
amount: B,
_phantom: PhantomData<(OnDrop, OppositeOnDrop)>,
}
impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>> Drop
for Imbalance<B, OnDrop, OppositeOnDrop>
{
fn drop(&mut self) {
if !self.amount.is_zero() {
OnDrop::handle(self.amount)
}
}
}
impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>> TryDrop
for Imbalance<B, OnDrop, OppositeOnDrop>
{
fn try_drop(self) -> Result<(), Self> {
self.drop_zero()
}
}
impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>> Default
for Imbalance<B, OnDrop, OppositeOnDrop>
{
fn default() -> Self {
Self::zero()
}
}
impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>>
Imbalance<B, OnDrop, OppositeOnDrop>
{
pub(crate) fn new(amount: B) -> Self {
Self { amount, _phantom: PhantomData }
}
pub(crate) fn forget(imbalance: Self) {
core::mem::forget(imbalance);
}
}
impl<B: Balance, OnDrop: HandleImbalanceDrop<B>, OppositeOnDrop: HandleImbalanceDrop<B>>
ImbalanceT<B> for Imbalance<B, OnDrop, OppositeOnDrop>
{
type Opposite = Imbalance<B, OppositeOnDrop, OnDrop>;
fn zero() -> Self {
Self { amount: Zero::zero(), _phantom: PhantomData }
}
fn drop_zero(self) -> Result<(), Self> {
if self.amount.is_zero() {
core::mem::forget(self);
Ok(())
} else {
Err(self)
}
}
fn split(self, amount: B) -> (Self, Self) {
let first = self.amount.min(amount);
let second = self.amount - first;
core::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);
core::mem::forget(other);
self
}
fn subsume(&mut self, other: Self) {
self.amount = self.amount.saturating_add(other.amount);
core::mem::forget(other);
}
fn offset(
self,
other: Imbalance<B, OppositeOnDrop, OnDrop>,
) -> SameOrOther<Self, Imbalance<B, OppositeOnDrop, OnDrop>> {
let (a, b) = (self.amount, other.amount);
core::mem::forget((self, other));
if a == b {
SameOrOther::None
} else if a > b {
SameOrOther::Same(Imbalance::new(a - b))
} else {
SameOrOther::Other(Imbalance::<B, OppositeOnDrop, OnDrop>::new(b - a))
}
}
fn peek(&self) -> B {
self.amount
}
}
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))
}
}
pub(crate) fn from_fungibles<
A: AssetId,
B: Balance,
OnDropIn: fungibles::HandleImbalanceDrop<A, B>,
OppositeIn: fungibles::HandleImbalanceDrop<A, B>,
OnDropOut: HandleImbalanceDrop<B>,
OppositeOut: HandleImbalanceDrop<B>,
>(
imbalance: fungibles::Imbalance<A, B, OnDropIn, OppositeIn>,
) -> Imbalance<B, OnDropOut, OppositeOut> {
let new = Imbalance::new(imbalance.peek());
fungibles::Imbalance::forget(imbalance);
new
}
pub type Debt<AccountId, B> = Imbalance<
<B as Inspect<AccountId>>::Balance,
<B as Balanced<AccountId>>::OnDropDebt,
<B as Balanced<AccountId>>::OnDropCredit,
>;
pub type Credit<AccountId, B> = Imbalance<
<B as Inspect<AccountId>>::Balance,
<B as Balanced<AccountId>>::OnDropCredit,
<B as Balanced<AccountId>>::OnDropDebt,
>;