use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::ops;
use cosmwasm_std::{Coin, StdError, StdResult};
#[derive(Serialize, Deserialize, Clone, Default, Debug, PartialEq, JsonSchema)]
pub struct Balance(pub Vec<Coin>);
impl Balance {
pub fn into_vec(self) -> Vec<Coin> {
self.0
}
pub fn has(&self, required: &Coin) -> bool {
self.0
.iter()
.find(|c| c.denom == required.denom)
.map(|m| m.amount >= required.amount)
.unwrap_or(false)
}
pub fn normalize(&mut self) {
self.0.retain(|c| c.amount.u128() != 0);
self.0.sort_unstable_by(|a, b| a.denom.cmp(&b.denom));
let mut dups: Vec<usize> = self
.0
.iter()
.enumerate()
.filter_map(|(i, c)| {
if i != 0 && c.denom == self.0[i - 1].denom {
Some(i)
} else {
None
}
})
.collect();
dups.reverse();
for dup in dups {
let add = self.0[dup].amount;
self.0[dup - 1].amount += add;
self.0.remove(dup);
}
}
fn find(&self, denom: &str) -> Option<(usize, &Coin)> {
self.0.iter().enumerate().find(|(_i, c)| c.denom == denom)
}
fn insert_pos(&self, denom: &str) -> Option<usize> {
self.0.iter().position(|c| c.denom.as_str() >= denom)
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn sub_saturating(mut self, other: Coin) -> StdResult<Self> {
match self.find(&other.denom) {
Some((i, c)) => {
if c.amount <= other.amount {
self.0.remove(i);
} else {
self.0[i].amount = (self.0[i].amount - other.amount)?;
}
}
None => return Err(StdError::underflow(0, other.amount.u128())),
};
Ok(self)
}
}
impl ops::AddAssign<Coin> for Balance {
fn add_assign(&mut self, other: Coin) {
match self.find(&other.denom) {
Some((i, c)) => {
self.0[i].amount = c.amount + other.amount;
}
None => match self.insert_pos(&other.denom) {
Some(idx) => self.0.insert(idx, other),
None => self.0.push(other),
},
};
}
}
impl ops::Add<Coin> for Balance {
type Output = Self;
fn add(mut self, other: Coin) -> Self {
self += other;
self
}
}
impl ops::AddAssign<Balance> for Balance {
fn add_assign(&mut self, other: Balance) {
for coin in other.0.into_iter() {
self.add_assign(coin);
}
}
}
impl ops::Add<Balance> for Balance {
type Output = Self;
fn add(mut self, other: Balance) -> Self {
self += other;
self
}
}
impl ops::Sub<Coin> for Balance {
type Output = StdResult<Self>;
fn sub(mut self, other: Coin) -> StdResult<Self> {
match self.find(&other.denom) {
Some((i, c)) => {
let remainder = (c.amount - other.amount)?;
if remainder.u128() == 0 {
self.0.remove(i);
} else {
self.0[i].amount = remainder;
}
}
None => return Err(StdError::underflow(0, other.amount.u128())),
};
Ok(self)
}
}
impl ops::Sub<Vec<Coin>> for Balance {
type Output = StdResult<Self>;
fn sub(self, amount: Vec<Coin>) -> StdResult<Self> {
let mut res = self;
for coin in amount {
res = res.sub(coin.clone())?;
}
Ok(res)
}
}
#[cfg(test)]
mod test {
use super::*;
use cosmwasm_std::coin;
#[test]
fn balance_has_works() {
let balance = Balance(vec![coin(555, "BTC"), coin(12345, "ETH")]);
assert!(balance.has(&coin(777, "ETH")));
assert!(balance.has(&coin(555, "BTC")));
assert!(!balance.has(&coin(12346, "ETH")));
assert!(!balance.has(&coin(456, "ETC")));
}
#[test]
fn balance_add_works() {
let balance = Balance(vec![coin(555, "BTC"), coin(12345, "ETH")]);
let more_eth = balance.clone() + coin(54321, "ETH");
assert_eq!(
more_eth,
Balance(vec![coin(555, "BTC"), coin(66666, "ETH")])
);
let add_atom = balance.clone() + coin(777, "ATOM");
assert_eq!(
add_atom,
Balance(vec![
coin(777, "ATOM"),
coin(555, "BTC"),
coin(12345, "ETH"),
])
);
}
#[test]
fn balance_in_place_addition() {
let mut balance = Balance(vec![coin(555, "BTC")]);
balance += coin(777, "ATOM");
assert_eq!(
&balance,
&Balance(vec![coin(777, "ATOM"), coin(555, "BTC")])
);
balance += Balance(vec![coin(666, "ETH"), coin(123, "ATOM")]);
assert_eq!(
&balance,
&Balance(vec![coin(900, "ATOM"), coin(555, "BTC"), coin(666, "ETH")])
);
let foo = balance + Balance(vec![coin(234, "BTC")]);
assert_eq!(
&foo,
&Balance(vec![coin(900, "ATOM"), coin(789, "BTC"), coin(666, "ETH")])
);
}
#[test]
fn balance_subtract_works() {
let balance = Balance(vec![coin(555, "BTC"), coin(12345, "ETH")]);
let less_eth = (balance.clone() - coin(2345, "ETH")).unwrap();
assert_eq!(
less_eth,
Balance(vec![coin(555, "BTC"), coin(10000, "ETH")])
);
let no_btc = (balance.clone() - coin(555, "BTC")).unwrap();
assert_eq!(no_btc, Balance(vec![coin(12345, "ETH")]));
let underflow = balance.clone() - coin(666, "BTC");
assert!(underflow.is_err());
let missing = balance.clone() - coin(1, "ATOM");
assert!(missing.is_err());
}
#[test]
fn balance_subtract_saturating_works() {
let balance = Balance(vec![coin(555, "BTC"), coin(12345, "ETH")]);
let less_eth = balance.clone().sub_saturating(coin(2345, "ETH")).unwrap();
assert_eq!(
less_eth,
Balance(vec![coin(555, "BTC"), coin(10000, "ETH")])
);
let no_btc = balance.clone().sub_saturating(coin(555, "BTC")).unwrap();
assert_eq!(no_btc, Balance(vec![coin(12345, "ETH")]));
let saturating = balance.clone().sub_saturating(coin(666, "BTC"));
assert!(saturating.is_ok());
assert_eq!(saturating.unwrap(), Balance(vec![coin(12345, "ETH")]));
let missing = balance.clone() - coin(1, "ATOM");
assert!(missing.is_err());
}
#[test]
fn normalize_balance() {
let mut balance = Balance(vec![coin(123, "ETH"), coin(0, "BTC"), coin(8990, "ATOM")]);
balance.normalize();
assert_eq!(balance, Balance(vec![coin(8990, "ATOM"), coin(123, "ETH")]));
let mut balance = Balance(vec![
coin(123, "ETH"),
coin(789, "BTC"),
coin(321, "ETH"),
coin(11, "BTC"),
]);
balance.normalize();
assert_eq!(balance, Balance(vec![coin(800, "BTC"), coin(444, "ETH")]));
}
}