Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Coin maps helpers #8

Merged
merged 6 commits into from
Jun 13, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
354 changes: 354 additions & 0 deletions src/balance/helpers.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,354 @@
use cosmwasm_std::{Coin, StdError, StdResult, Uint128};
use std::collections::btree_map::Entry;
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(self) -> BTreeMap<String, Uint128>;
fn into_map_unique(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(self) -> BTreeMap<String, Uint128> {
pbukva marked this conversation as resolved.
Show resolved Hide resolved
let mut denom_map = BTreeMap::new();

for coin in self {
if let Some(counter) = denom_map.get_mut(&coin.denom) {
*counter += coin.amount;
} else {
denom_map.insert(coin.denom, coin.amount);
}
}

denom_map
}
pbukva marked this conversation as resolved.
Show resolved Hide resolved

fn into_map_unique(self) -> StdResult<BTreeMap<String, Uint128>> {
pbukva marked this conversation as resolved.
Show resolved Hide resolved
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) => {
return Err(StdError::generic_err(format!(
"Duplicate denom found: {}",
e.key()
)));
}
}
}

Ok(denom_map)
}

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() {
let vec = vec![coin(100, "atom"), coin(50, "btc"), coin(50, "btc")];
let map = vec.into_map();

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();

// 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_unique().unwrap(),
coins.clone().into_map()
);

let result = coins.into_map_unique();
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_unique();
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_unique();
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
Loading