Skip to content

Commit

Permalink
feat: Coin maps helpers (#8)
Browse files Browse the repository at this point in the history
  • Loading branch information
MissingNO57 authored Jun 13, 2024
1 parent a9fa7e4 commit b8225b1
Show file tree
Hide file tree
Showing 3 changed files with 367 additions and 0 deletions.
363 changes: 363 additions & 0 deletions src/balance/helpers.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,363 @@
use cosmwasm_std::{Coin, StdError, StdResult, Uint128};
use std::collections::btree_map::Entry;
use std::collections::btree_map::OccupiedEntry;
use std::collections::BTreeMap;

pub trait BTreeMapCoinHelpers {
fn into_vec(self) -> Vec<Coin>;
fn inplace_sub<'a, I>(&mut self, balance: I) -> StdResult<()>
where
I: IntoIterator<Item = (&'a String, &'a Uint128)>;

fn inplace_add<'a, I>(&mut self, balance: I) -> StdResult<()>
where
I: IntoIterator<Item = (&'a String, &'a Uint128)>;
}
impl BTreeMapCoinHelpers for BTreeMap<String, Uint128> {
fn into_vec(self) -> Vec<Coin> {
self.into_iter()
.map(|(denom, amount)| Coin { denom, amount })
.collect()
}

fn inplace_sub<'a, I>(&mut self, balance: I) -> StdResult<()>
where
I: IntoIterator<Item = (&'a String, &'a Uint128)>,
{
// Decrease total remaining supply
for (denom, amount) in balance {
if let Some(r) = self.get_mut(denom) {
if amount > r {
return Err(StdError::generic_err(format!(
"Subtract overflow for denom: {}",
denom
)));
}
*r -= amount;
} else {
return Err(StdError::generic_err(format!("Unknown denom {}", denom)));
}

// Remove denom if balance is now 0
if self.get(denom) == Some(&Uint128::zero()) {
self.remove(denom);
}
}
Ok(())
}

fn inplace_add<'a, I>(&mut self, balance: I) -> StdResult<()>
where
I: IntoIterator<Item = (&'a String, &'a Uint128)>,
{
for (denom, amount) in balance {
if let Some(counter) = self.get_mut(denom) {
*counter = counter.checked_add(*amount).map_err(|_| {
StdError::generic_err(format!("Addition overflow for denom: {}", denom))
})?;
} else {
self.insert(denom.clone(), *amount);
}
}
Ok(())
}
}

pub trait VecCoinConversions {
fn to_tuple_iterator<'a>(&'a self) -> Box<dyn Iterator<Item = (&'a String, &'a Uint128)> + 'a>;
fn into_map_with_duplicities<F>(
self,
handle_duplicate: F,
) -> StdResult<BTreeMap<String, Uint128>>
where
F: Fn(OccupiedEntry<String, Uint128>, Uint128) -> StdResult<()>;
fn into_map(self) -> StdResult<BTreeMap<String, Uint128>>;
fn to_formatted_string(&self) -> String;
}

impl VecCoinConversions for Vec<Coin> {
fn to_tuple_iterator<'a>(&'a self) -> Box<dyn Iterator<Item = (&'a String, &'a Uint128)> + 'a> {
Box::new(self.iter().map(|coin| (&coin.denom, &coin.amount)))
}

fn into_map_with_duplicities<F>(
self,
handle_duplicate: F,
) -> StdResult<BTreeMap<String, Uint128>>
where
F: Fn(OccupiedEntry<String, Uint128>, Uint128) -> StdResult<()>,
{
let mut denom_map = BTreeMap::new();

for coin in self {
match denom_map.entry(coin.denom) {
Entry::Vacant(e) => {
e.insert(coin.amount);
}
Entry::Occupied(e) => {
handle_duplicate(e, coin.amount)?;
}
}
}

Ok(denom_map)
}

fn into_map(self) -> StdResult<BTreeMap<String, Uint128>> {
self.into_map_with_duplicities(|e, _| {
Err(StdError::generic_err(format!(
"Duplicate denom found: {}",
e.key()
)))
})
}

fn to_formatted_string(&self) -> String {
self.iter()
.map(|coin| format!("{}{}", coin.amount, coin.denom))
.collect::<Vec<_>>()
.join(", ")
}
}

#[cfg(test)]
mod tests {
use super::*;
use cosmwasm_std::{coin, Uint128};
use std::collections::BTreeMap;

#[test]
fn test_into_vec() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(100));
map.insert("btc".to_string(), Uint128::new(50));

let vec = map.into_vec();
assert_eq!(vec, vec![coin(100, "atom"), coin(50, "btc")]);
}

#[test]
fn test_inplace_sub() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(100));
map.insert("btc".to_string(), Uint128::new(50));

let balance = vec![
("atom".to_string(), Uint128::new(30)),
("btc".to_string(), Uint128::new(20)),
];
map.inplace_sub(balance.iter().map(|(d, a)| (d, a)))
.unwrap();

assert_eq!(map.get("atom"), Some(&Uint128::new(70)));
assert_eq!(map.get("btc"), Some(&Uint128::new(30)));
}

#[test]
fn test_inplace_sub_overflow() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(100));

let balance = vec![("atom".to_string(), Uint128::new(150))];
let result = map.inplace_sub(balance.iter().map(|(d, a)| (d, a)));

assert_eq!(
result.err(),
Some(StdError::generic_err("Subtract overflow for denom: atom"))
);
}

#[test]
fn test_inplace_add() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(100));
map.insert("btc".to_string(), Uint128::new(50));

let balance = vec![
("atom".to_string(), Uint128::new(30)),
("btc".to_string(), Uint128::new(20)),
("eth".to_string(), Uint128::new(10)),
];
assert!(map.inplace_add(balance.iter().map(|(d, a)| (d, a))).is_ok());

assert_eq!(map.get("atom"), Some(&Uint128::new(130)));
assert_eq!(map.get("btc"), Some(&Uint128::new(70)));
assert_eq!(map.get("eth"), Some(&Uint128::new(10)));
}

#[test]
fn test_to_tuple_iterator() {
let vec = vec![coin(100, "atom"), coin(50, "btc")];
let mut iter = vec.to_tuple_iterator();

assert_eq!(iter.next(), Some((&"atom".to_string(), &Uint128::new(100))));
assert_eq!(iter.next(), Some((&"btc".to_string(), &Uint128::new(50))));
assert_eq!(iter.next(), None);
}

#[test]
fn test_into_map_with_duplicities_sum() {
let vec = vec![coin(100, "atom"), coin(50, "btc"), coin(50, "btc")];
let map = vec
.into_map_with_duplicities(|mut e, new_value| Ok(*e.get_mut() += new_value))
.unwrap();

assert_eq!(map.get("atom"), Some(&Uint128::new(100)));
assert_eq!(map.get("btc"), Some(&Uint128::new(100)));
}

#[test]
fn test_to_formatted_string() {
let vec = vec![coin(100, "atom"), coin(50, "btc")];
let formatted = vec.to_formatted_string();

assert_eq!(formatted, "100atom, 50btc");
}

#[test]
fn test_empty_map_add_subtract() {
let mut map = BTreeMap::new();

let balance = vec![coin(30, "atom"), coin(20, "btc"), coin(10, "eth")];
assert!(map.inplace_add(balance.to_tuple_iterator()).is_ok());

assert_eq!(map.get("atom"), Some(&Uint128::new(30)));
assert_eq!(map.get("btc"), Some(&Uint128::new(20)));
assert_eq!(map.get("eth"), Some(&Uint128::new(10)));

assert!(map.inplace_sub(balance.to_tuple_iterator()).is_ok());

assert!(map.is_empty())
}
#[test]
fn test_subtracting_nonexistent_denom() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(100));

let balance = vec![("btc".to_string(), Uint128::new(50))];
let result = map.inplace_sub(balance.iter().map(|(d, a)| (d, a)));

assert!(result.is_err());
assert_eq!(
result.err(),
Some(StdError::generic_err("Unknown denom btc"))
);
}

#[test]
fn test_addition_with_overflow() {
let mut map = BTreeMap::new();
map.insert("atom".to_string(), Uint128::new(u128::MAX));

let balance = vec![("atom".to_string(), Uint128::new(1))];
let result = map.inplace_add(balance.iter().map(|(d, a)| (d, a)));

assert_eq!(
result.err(),
Some(StdError::generic_err("Addition overflow for denom: atom"))
);

assert_eq!(map.get("atom"), Some(&Uint128::new(u128::MAX)));
}

#[test]
fn test_ordering() {
let coins = vec![
coin(50, "btc"),
coin(100, "atom"),
coin(200, "eth"),
coin(150, "usd"),
coin(75, "eur"),
];

let sorted_coins = vec![
coin(100, "atom"),
coin(50, "btc"),
coin(200, "eth"),
coin(75, "eur"),
coin(150, "usd"),
];
let sorted_keys: Vec<String> = sorted_coins.iter().map(|res| res.denom.clone()).collect();

// Convert vector of coins into map
let map = coins.into_map().unwrap();

// Ensure that keys in map are sorted
let keys: Vec<_> = map.keys().cloned().collect();
assert_eq!(keys, sorted_keys);

// Ensure that vec output is sorted
assert_eq!(map.into_vec(), sorted_coins);
}

#[test]
fn test_vec_of_coins_into_btreemap_unique() {
let coins = vec![
coin(100, "atom"),
coin(50, "btc"),
coin(200, "eth"),
coin(150, "usd"),
coin(75, "eur"),
];

// Resulting maps are equivalent if there is no duplicity
assert_eq!(
coins.clone().into_map().unwrap(),
coins
.clone()
.into_map_with_duplicities(|mut e, b| { Ok(*e.get_mut() = b) })
.unwrap()
);

let result = coins.into_map();
assert!(result.is_ok());

let map = result.unwrap();
let keys: Vec<_> = map.keys().cloned().collect();
assert_eq!(
keys,
vec![
"atom".to_string(),
"btc".to_string(),
"eth".to_string(),
"eur".to_string(),
"usd".to_string()
]
);

assert_eq!(map.get("atom"), Some(&Uint128::new(100)));
assert_eq!(map.get("btc"), Some(&Uint128::new(50)));
assert_eq!(map.get("eth"), Some(&Uint128::new(200)));
assert_eq!(map.get("usd"), Some(&Uint128::new(150)));
assert_eq!(map.get("eur"), Some(&Uint128::new(75)));
}

#[test]
fn test_vec_of_coins_into_btreemap_unique_with_duplicates() {
let coins = vec![
coin(100, "atom"),
coin(50, "btc"),
coin(200, "eth"),
coin(100, "btc"), // Duplicate denom
coin(75, "eth"), // Duplicate denom
];

let result = coins.into_map();
assert!(result.is_err());

if let Err(err) = result {
assert_eq!(err, StdError::generic_err("Duplicate denom found: btc"));
}
}

#[test]
fn test_vec_of_coins_into_btreemap_unique_with_empty_vec() {
let coins: Vec<Coin> = vec![];

let result = coins.into_map();
assert!(result.is_ok());

let map = result.unwrap();
assert!(map.is_empty());
}
}
3 changes: 3 additions & 0 deletions src/balance/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
mod helpers;

pub use helpers::{BTreeMapCoinHelpers, VecCoinConversions};
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod balance;
pub mod crypto;
pub mod events;
pub mod helpers;
Expand Down

0 comments on commit b8225b1

Please sign in to comment.