Redefined MaspAmount using the ValueSum type from the MASP crate.

anoma · Jan 10, 2024 · e046c9f · e046c9f
1 parent 29a8484
commit e046c9f
Show file tree

Hide file tree

Showing 12 changed files with 68 additions and 76 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/apps/src/lib/client/rpc.rs b/apps/src/lib/client/rpc.rs
@@ -188,7 +188,7 @@ pub async fn query_transfers(
                 ) || shielded_accounts
                     .values()
                     .cloned()
-                    .any(|x| x.iter().any(check))
+                    .any(|x| !x.get(token).is_zero())
             }
             None => true,
         };
@@ -229,9 +229,9 @@ pub async fn query_transfers(
         for (account, masp_change) in shielded_accounts {
             if fvk_map.contains_key(&account) {
                 display!(context.io(), "  {}:", fvk_map[&account]);
-                for (token_addr, val) in masp_change {
+                for (token_addr, val) in masp_change.components() {
                     let token_alias =
-                        lookup_token_alias(context, &token_addr, &MASP).await;
+                        lookup_token_alias(context, token_addr, &MASP).await;
                     let sign = match val.cmp(&Change::zero()) {
                         Ordering::Greater => "+",
                         Ordering::Less => "-",
@@ -241,7 +241,7 @@ pub async fn query_transfers(
                         context.io(),
                         " {}{} {}",
                         sign,
-                        context.format_amount(&token_addr, val.into()).await,
+                        context.format_amount(token_addr, (*val).into()).await,
                         token_alias,
                     );
                 }
@@ -912,11 +912,11 @@ pub async fn query_shielded_balance(
                 let balance = shielded
                     .decode_amount(context.client(), balance, epoch)
                     .await;
-                for (key, value) in balance.iter() {
-                    if !balances.contains_key(key) {
-                        balances.insert(key.clone(), Vec::new());
-                    }
-                    balances.get_mut(key).unwrap().push((fvk, *value));
+                for (key, value) in balance.components() {
+                    balances
+                        .entry(key.clone())
+                        .or_insert(Vec::new())
+                        .push((fvk, *value));
                 }
             }
 
@@ -1065,12 +1065,12 @@ pub async fn print_decoded_balance(
             .await
             .decode_amount(context.client(), balance, epoch)
             .await;
-        for (token_addr, amount) in decoded_balance {
+        for (token_addr, amount) in decoded_balance.components() {
             display_line!(
                 context.io(),
                 "{} : {}",
-                lookup_token_alias(context, &token_addr, &MASP).await,
-                context.format_amount(&token_addr, amount.into()).await,
+                lookup_token_alias(context, token_addr, &MASP).await,
+                context.format_amount(token_addr, (*amount).into()).await,
             );
         }
     }

diff --git a/core/src/types/uint.rs b/core/src/types/uint.rs
@@ -8,7 +8,7 @@ use std::ops::{Add, AddAssign, BitAnd, Div, Mul, Neg, Rem, Sub, SubAssign};
 use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
 use impl_num_traits::impl_uint_num_traits;
 use num_integer::Integer;
-use num_traits::CheckedMul;
+use num_traits::{CheckedAdd, CheckedMul, CheckedSub};
 use uint::construct_uint;
 
 use super::dec::{Dec, POS_DECIMAL_PRECISION};
@@ -555,28 +555,6 @@ impl I256 {
         sign
     }
 
-    /// Adds two [`I256`]'s if the absolute value does
-    /// not exceed [`MAX_SIGNED_VALUE`], else returns `None`.
-    pub fn checked_add(&self, other: &Self) -> Option<Self> {
-        if self.non_negative() == other.non_negative() {
-            self.abs().checked_add(other.abs()).and_then(|val| {
-                Self::try_from(val)
-                    .ok()
-                    .map(|val| if !self.non_negative() { -val } else { val })
-            })
-        } else {
-            Some(*self + *other)
-        }
-    }
-
-    /// Subtracts two [`I256`]'s if the absolute value does
-    /// not exceed [`MAX_SIGNED_VALUE`], else returns `None`.
-    pub fn checked_sub(&self, other: &Self) -> Option<Self> {
-        self.checked_add(&other.neg())
-    }
-
-    ///
-
     /// Changed the inner Uint into a canonical representation.
     fn canonical(self) -> Self {
         Self(self.0.canonical())
@@ -738,6 +716,30 @@ impl Mul<Uint> for I256 {
     }
 }
 
+impl CheckedAdd for I256 {
+    /// Adds two [`I256`]'s if the absolute value does
+    /// not exceed [`MAX_SIGNED_VALUE`], else returns `None`.
+    fn checked_add(&self, other: &Self) -> Option<Self> {
+        if self.non_negative() == other.non_negative() {
+            self.abs().checked_add(other.abs()).and_then(|val| {
+                Self::try_from(val)
+                    .ok()
+                    .map(|val| if !self.non_negative() { -val } else { val })
+            })
+        } else {
+            Some(*self + *other)
+        }
+    }
+}
+
+impl CheckedSub for I256 {
+    /// Subtracts two [`I256`]'s if the absolute value does
+    /// not exceed [`MAX_SIGNED_VALUE`], else returns `None`.
+    fn checked_sub(&self, other: &Self) -> Option<Self> {
+        self.checked_add(&other.neg())
+    }
+}
+
 impl CheckedMul for I256 {
     fn checked_mul(&self, v: &Self) -> Option<Self> {
         let is_negative = self.is_negative() != v.is_negative();

diff --git a/proof_of_stake/Cargo.toml b/proof_of_stake/Cargo.toml
@@ -22,6 +22,7 @@ namada_core = {path = "../core", default-features = false}
 borsh.workspace = true
 data-encoding.workspace = true
 derivative.workspace = true
+num-traits.workspace = true
 once_cell.workspace = true
 proptest = {workspace = true, optional = true}
 serde.workspace = true

diff --git a/proof_of_stake/src/storage.rs b/proof_of_stake/src/storage.rs
@@ -16,6 +16,7 @@ use namada_core::types::key::{
 };
 use namada_core::types::storage::Epoch;
 use namada_core::types::token;
+use num_traits::CheckedAdd;
 
 use crate::storage_key::consensus_keys_key;
 use crate::types::{

diff --git a/proof_of_stake/src/validator_set_update.rs b/proof_of_stake/src/validator_set_update.rs
@@ -10,6 +10,7 @@ use namada_core::types::address::Address;
 use namada_core::types::key::PublicKeyTmRawHash;
 use namada_core::types::storage::Epoch;
 use namada_core::types::token;
+use num_traits::ops::checked::CheckedAdd;
 use once_cell::unsync::Lazy;
 
 use crate::storage::{

diff --git a/sdk/Cargo.toml b/sdk/Cargo.toml
@@ -75,6 +75,7 @@ namada_core = {path = "../core", default-features = false, features = ["rand"]}
 namada_ethereum_bridge = {path = "../ethereum_bridge", default-features = false}
 namada_proof_of_stake = {path = "../proof_of_stake", default-features = false}
 num256.workspace = true
+num-traits.workspace = true
 orion.workspace = true
 owo-colors = "3.5.0"
 parse_duration = "2.1.1"

diff --git a/sdk/src/eth_bridge/bridge_pool.rs b/sdk/src/eth_bridge/bridge_pool.rs
@@ -20,6 +20,7 @@ use namada_core::types::ethereum_events::EthAddress;
 use namada_core::types::keccak::KeccakHash;
 use namada_core::types::token::{balance_key, Amount};
 use namada_core::types::voting_power::FractionalVotingPower;
+use num_traits::ops::checked::CheckedSub;
 use owo_colors::OwoColorize;
 use serde::Serialize;
 

diff --git a/sdk/src/masp.rs b/sdk/src/masp.rs
@@ -35,7 +35,7 @@ use masp_primitives::transaction::components::sapling::builder::SaplingMetadata;
 use masp_primitives::transaction::components::transparent::builder::TransparentBuilder;
 use masp_primitives::transaction::components::{
     ConvertDescription, I128Sum, OutputDescription, SpendDescription, TxOut,
-    U64Sum,
+    U64Sum, ValueSum,
 };
 use masp_primitives::transaction::fees::fixed::FeeRule;
 use masp_primitives::transaction::sighash::{signature_hash, SignableInput};
@@ -58,7 +58,7 @@ use namada_core::types::storage::{BlockHeight, Epoch, Key, KeySeg, TxIndex};
 use namada_core::types::time::{DateTimeUtc, DurationSecs};
 use namada_core::types::token;
 use namada_core::types::token::{
-    Change, MaspDenom, Transfer, HEAD_TX_KEY, PIN_KEY_PREFIX, TX_KEY_PREFIX,
+    MaspDenom, Transfer, HEAD_TX_KEY, PIN_KEY_PREFIX, TX_KEY_PREFIX,
 };
 use namada_core::types::transaction::WrapperTx;
 use rand_core::{CryptoRng, OsRng, RngCore};
@@ -75,7 +75,6 @@ use crate::queries::Client;
 use crate::rpc::{query_conversion, query_storage_value};
 use crate::tendermint_rpc::query::Query;
 use crate::tendermint_rpc::Order;
-use crate::tx::decode_component;
 use crate::{display_line, edisplay_line, rpc, MaybeSend, MaybeSync, Namada};
 
 /// Env var to point to a dir with MASP parameters. When not specified,
@@ -484,10 +483,7 @@ pub struct MaspChange {
 }
 
 /// a masp amount
-#[derive(
-    BorshSerialize, BorshDeserialize, Debug, Clone, Default, PartialEq, Eq,
-)]
-pub struct MaspAmount(pub HashMap<(Epoch, Address), token::Change>);
+pub type MaspAmount = ValueSum<(Epoch, Address), token::Change>;
 
 /// An extension of Option's cloned method for pair types
 fn cloned_pair<T: Clone, U: Clone>((a, b): (&T, &U)) -> (T, U) {
@@ -1344,20 +1340,19 @@ impl<U: ShieldedUtils + MaybeSend + MaybeSync> ShieldedContext<U> {
         client: &C,
         amt: I128Sum,
         target_epoch: Epoch,
-    ) -> HashMap<Address, token::Change> {
-        let mut res = HashMap::new();
+    ) -> ValueSum<Address, token::Change> {
+        let mut res = ValueSum::zero();
         for (asset_type, val) in amt.components() {
             // Decode the asset type
             let decoded = self.decode_asset_type(client, *asset_type).await;
             // Only assets with the target timestamp count
             match decoded {
-                Some(asset_type @ (_, _, epoch)) if epoch == target_epoch => {
-                    decode_component(
-                        asset_type,
-                        *val,
-                        &mut res,
-                        |address, _| address,
-                    );
+                Some((address, denom, epoch)) if epoch == target_epoch => {
+                    let decoded_change =
+                        token::Change::from_masp_denominated(*val, denom)
+                            .expect("expected this to fit");
+                    res += ValueSum::from_pair(address, decoded_change)
+                        .expect("expected this to fit");
                 }
                 _ => {}
             }
@@ -1372,18 +1367,20 @@ impl<U: ShieldedUtils + MaybeSend + MaybeSync> ShieldedContext<U> {
         client: &C,
         amt: I128Sum,
     ) -> MaspAmount {
-        let mut res: HashMap<(Epoch, Address), Change> = HashMap::default();
+        let mut res = MaspAmount::zero();
         for (asset_type, val) in amt.components() {
             // Decode the asset type
-            if let Some(decoded) =
+            if let Some((addr, denom, epoch)) =
                 self.decode_asset_type(client, *asset_type).await
             {
-                decode_component(decoded, *val, &mut res, |address, epoch| {
-                    (epoch, address)
-                })
+                let decoded_change =
+                    token::Change::from_masp_denominated(*val, denom)
+                        .expect("expected this to fit");
+                res += MaspAmount::from_pair((epoch, addr), decoded_change)
+                    .expect("unable to construct decoded amount");
             }
         }
-        MaspAmount(res)
+        res
     }
 
     /// Make shielded components to embed within a Transfer object. If no

diff --git a/sdk/src/tx.rs b/sdk/src/tx.rs
@@ -1,6 +1,6 @@
 //! SDK functions to construct different types of transactions
 use std::borrow::Cow;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::collections::{BTreeMap, HashSet};
 use std::fs::File;
 use std::path::{Path, PathBuf};
 use std::time::Duration;
@@ -494,24 +494,6 @@ pub fn display_inner_resp(context: &impl Namada, resp: &TxResponse) {
     );
 }
 
-/// decode components of a masp note
-pub fn decode_component<K, F>(
-    (addr, denom, epoch): (Address, MaspDenom, Epoch),
-    val: i128,
-    res: &mut HashMap<K, token::Change>,
-    mk_key: F,
-) where
-    F: FnOnce(Address, Epoch) -> K,
-    K: Eq + std::hash::Hash,
-{
-    let decoded_change = token::Change::from_masp_denominated(val, denom)
-        .expect("expected this to fit");
-
-    res.entry(mk_key(addr, epoch))
-        .and_modify(|val| *val += decoded_change)
-        .or_insert(decoded_change);
-}
-
 /// Save accounts initialized from a tx into the wallet, if any.
 pub async fn save_initialized_accounts<N: Namada>(
     context: &N,

diff --git a/wasm/Cargo.lock b/wasm/Cargo.lock
diff --git a/wasm_for_tests/wasm_source/Cargo.lock b/wasm_for_tests/wasm_source/Cargo.lock