From b8225b1e604f00f06ea117a74c71b48ca62e8ed2 Mon Sep 17 00:00:00 2001 From: Jiri Vestfal <45360788+MissingNO57@users.noreply.github.com> Date: Thu, 13 Jun 2024 13:54:43 +0200 Subject: [PATCH] feat: Coin maps helpers (#8) --- src/balance/helpers.rs | 363 +++++++++++++++++++++++++++++++++++++++++ src/balance/mod.rs | 3 + src/lib.rs | 1 + 3 files changed, 367 insertions(+) create mode 100644 src/balance/helpers.rs create mode 100644 src/balance/mod.rs diff --git a/src/balance/helpers.rs b/src/balance/helpers.rs new file mode 100644 index 0000000..731de87 --- /dev/null +++ b/src/balance/helpers.rs @@ -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; + fn inplace_sub<'a, I>(&mut self, balance: I) -> StdResult<()> + where + I: IntoIterator; + + fn inplace_add<'a, I>(&mut self, balance: I) -> StdResult<()> + where + I: IntoIterator; +} +impl BTreeMapCoinHelpers for BTreeMap { + fn into_vec(self) -> Vec { + self.into_iter() + .map(|(denom, amount)| Coin { denom, amount }) + .collect() + } + + fn inplace_sub<'a, I>(&mut self, balance: I) -> StdResult<()> + where + I: IntoIterator, + { + // 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, + { + 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 + 'a>; + fn into_map_with_duplicities( + self, + handle_duplicate: F, + ) -> StdResult> + where + F: Fn(OccupiedEntry, Uint128) -> StdResult<()>; + fn into_map(self) -> StdResult>; + fn to_formatted_string(&self) -> String; +} + +impl VecCoinConversions for Vec { + fn to_tuple_iterator<'a>(&'a self) -> Box + 'a> { + Box::new(self.iter().map(|coin| (&coin.denom, &coin.amount))) + } + + fn into_map_with_duplicities( + self, + handle_duplicate: F, + ) -> StdResult> + where + F: Fn(OccupiedEntry, 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> { + 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::>() + .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 = 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 = vec![]; + + let result = coins.into_map(); + assert!(result.is_ok()); + + let map = result.unwrap(); + assert!(map.is_empty()); + } +} diff --git a/src/balance/mod.rs b/src/balance/mod.rs new file mode 100644 index 0000000..927efd7 --- /dev/null +++ b/src/balance/mod.rs @@ -0,0 +1,3 @@ +mod helpers; + +pub use helpers::{BTreeMapCoinHelpers, VecCoinConversions}; diff --git a/src/lib.rs b/src/lib.rs index 84a30c6..a0843a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ +pub mod balance; pub mod crypto; pub mod events; pub mod helpers;