Added support for non-timestamped assets.

apps/src/lib/client/rpc.rs

@@ -43,9 +43,7 @@ use namada::types::hash::Hash;
 use namada::types::ibc::{is_ibc_denom, IbcTokenHash};
 use namada::types::io::Io;
 use namada::types::key::*;
-use namada::types::masp::{
-    encode_asset_type, BalanceOwner, ExtendedViewingKey, PaymentAddress,
-};
+use namada::types::masp::{BalanceOwner, ExtendedViewingKey, PaymentAddress};
 use namada::types::storage::{BlockHeight, BlockResults, Epoch, Key, KeySeg};
 use namada::types::token::{Change, MaspDenom};
 use namada::types::{storage, token};
@@ -479,7 +477,7 @@ pub async fn query_pinned_balance(
                 for (token_alias, token) in &tokens {
                     let total_balance = balance
                         .0
-                        .get(&(epoch, token.clone()))
+                        .get(&token.clone())
                         .cloned()
                         .unwrap_or_default();
 
@@ -511,11 +509,7 @@ pub async fn query_pinned_balance(
             (Ok((balance, epoch)), None) => {
                 let mut found_any = false;
 
-                for ((_, token_addr), value) in balance
-                    .0
-                    .iter()
-                    .filter(|((token_epoch, _), _)| *token_epoch == epoch)
-                {
+                for (token_addr, value) in balance.0.iter() {
                     if !found_any {
                         display_line!(
                             context.io(),
@@ -780,26 +774,6 @@ pub async fn query_proposal_by_id<C: namada::ledger::queries::Client + Sync>(
     namada_sdk::rpc::query_proposal_by_id(client, proposal_id).await
 }
 
-/// Get given epoched token from the supplied value sum
-pub fn extract_amount(
-    value_sum: I128Sum,
-    token: &Address,
-    epoch: Epoch,
-) -> token::Amount {
-    let mut amount = token::Amount::zero();
-    for denom in MaspDenom::iter() {
-        let asset_type = encode_asset_type(epoch, token, denom).expect(
-            "unable to make asset type given token, epoch, and denomination",
-        );
-        let value: u128 = value_sum[&asset_type]
-            .try_into()
-            .expect("MASP balance is negative");
-        amount += token::Amount::from_masp_denominated_u128(value, denom)
-            .expect("MASP balance is too high");
-    }
-    amount
-}
-
 /// Query token shielded balance(s)
 pub async fn query_shielded_balance(
     context: &impl Namada,
@@ -845,18 +819,15 @@ pub async fn query_shielded_balance(
                 // Query the multi-asset balance at the given spending key
                 let viewing_key =
                     ExtendedFullViewingKey::from(viewing_keys[0]).fvk.vk;
+                let mut shielded = context.shielded_mut().await;
                 let balance = if no_conversions {
-                    context
-                        .shielded_mut()
-                        .await
+                    shielded
                         .compute_shielded_balance(&viewing_key)
                         .await
                         .unwrap()
                         .expect("context should contain viewing key")
                 } else {
-                    context
-                        .shielded_mut()
-                        .await
+                    shielded
                         .compute_exchanged_balance(
                             context.client(),
                             context.io(),
@@ -868,7 +839,14 @@ pub async fn query_shielded_balance(
                         .expect("context should contain viewing key")
                 };
 
-                let total_balance = extract_amount(balance, &token, epoch);
+                let total_balance = shielded
+                    .decode_combine_sum_to_epoch(
+                        context.client(),
+                        balance,
+                        epoch,
+                    )
+                    .await
+                    .get(&token);
                 if total_balance.is_zero() {
                     display_line!(
                         context.io(),
@@ -880,7 +858,9 @@ pub async fn query_shielded_balance(
                         context.io(),
                         "{}: {}",
                         token_alias,
-                        context.format_amount(&token, total_balance,).await
+                        context
+                            .format_amount(&token, total_balance.into())
+                            .await
                     );
                 }
             }
@@ -966,18 +946,15 @@ pub async fn query_shielded_balance(
                 for fvk in &viewing_keys {
                     // Query the multi-asset balance at the given spending key
                     let viewing_key = ExtendedFullViewingKey::from(*fvk).fvk.vk;
+                    let mut shielded = context.shielded_mut().await;
                     let balance = if no_conversions {
-                        context
-                            .shielded_mut()
-                            .await
+                        shielded
                             .compute_shielded_balance(&viewing_key)
                             .await
                             .unwrap()
                             .expect("context should contain viewing key")
                     } else {
-                        context
-                            .shielded_mut()
-                            .await
+                        shielded
                             .compute_exchanged_balance(
                                 context.client(),
                                 context.io(),
@@ -989,12 +966,20 @@ pub async fn query_shielded_balance(
                             .expect("context should contain viewing key")
                     };
 
-                    let total_balance = extract_amount(balance, &token, epoch);
+                    let total_balance = shielded
+                        .decode_combine_sum_to_epoch(
+                            context.client(),
+                            balance,
+                            epoch,
+                        )
+                        .await
+                        .get(&token);
                     if !total_balance.is_zero() {
                         found_any = true;
                     }
-                    let formatted =
-                        context.format_amount(&token, total_balance).await;
+                    let formatted = context
+                        .format_amount(&token, total_balance.into())
+                        .await;
                     display_line!(
                         context.io(),
                         "  {}, owned by {}",
@@ -1092,13 +1077,22 @@ pub async fn print_decoded_balance_with_epoch(
             .get(&token_addr)
             .map(|a| a.to_string())
             .unwrap_or_else(|| token_addr.to_string());
-        display_line!(
-            context.io(),
-            "{} | {} : {}",
-            alias,
-            epoch,
-            context.format_amount(&token_addr, asset_value).await,
-        );
+        if let Some(epoch) = epoch {
+            display_line!(
+                context.io(),
+                "{} | {} : {}",
+                alias,
+                epoch,
+                context.format_amount(&token_addr, asset_value).await,
+            );
+        } else {
+            display_line!(
+                context.io(),
+                "{} : {}",
+                alias,
+                context.format_amount(&token_addr, asset_value).await,
+            );
+        }
     }
 }
 

core/src/ledger/masp_conversions.rs

@@ -191,6 +191,12 @@ where
     Ok((noterized_inflation, precision))
 }
 
+// Safely sum a pair of amounts
+#[cfg(any(feature = "wasm-runtime", test))]
+fn sum_pair(pair: (token::Amount, token::Amount)) -> Option<token::Amount> {
+    pair.0.checked_add(pair.1)
+}
+
 // This is only enabled when "wasm-runtime" is on, because we're using rayon
 #[cfg(any(feature = "wasm-runtime", test))]
 /// Update the MASP's allowed conversions
@@ -246,13 +252,13 @@ where
     // notes clients have to use. This trick works under the assumption that
     // reward tokens will then be reinflated back to the current epoch.
     let reward_assets = [
-        encode_asset_type(Epoch(0), &native_token, MaspDenom::Zero)
+        encode_asset_type(Some(Epoch(0)), &native_token, MaspDenom::Zero)
             .into_storage_result()?,
-        encode_asset_type(Epoch(0), &native_token, MaspDenom::One)
+        encode_asset_type(Some(Epoch(0)), &native_token, MaspDenom::One)
             .into_storage_result()?,
-        encode_asset_type(Epoch(0), &native_token, MaspDenom::Two)
+        encode_asset_type(Some(Epoch(0)), &native_token, MaspDenom::Two)
             .into_storage_result()?,
-        encode_asset_type(Epoch(0), &native_token, MaspDenom::Three)
+        encode_asset_type(Some(Epoch(0)), &native_token, MaspDenom::Three)
             .into_storage_result()?,
     ];
     // Conversions from the previous to current asset for each address
@@ -276,12 +282,18 @@ where
             // Provide an allowed conversion from previous timestamp. The
             // negative sign allows each instance of the old asset to be
             // cancelled out/replaced with the new asset
-            let old_asset =
-                encode_asset_type(wl_storage.storage.last_epoch, addr, denom)
-                    .into_storage_result()?;
-            let new_asset =
-                encode_asset_type(wl_storage.storage.block.epoch, addr, denom)
-                    .into_storage_result()?;
+            let old_asset = encode_asset_type(
+                Some(wl_storage.storage.last_epoch),
+                addr,
+                denom,
+            )
+            .into_storage_result()?;
+            let new_asset = encode_asset_type(
+                Some(wl_storage.storage.block.epoch),
+                addr,
+                denom,
+            )
+            .into_storage_result()?;
             // Get the last rewarded amount of the native token
             let normed_inflation = wl_storage
                 .storage
@@ -329,11 +341,12 @@ where
                 if denom == MaspDenom::Three {
                     // The reward for each reward.1 units of the current asset
                     // is reward.0 units of the reward token
-                    total_reward += (addr_bal
-                        * (new_normed_inflation, *normed_inflation))
-                        .0
-                        .checked_sub(addr_bal)
-                        .unwrap_or_default();
+                    total_reward += sum_pair(
+                        addr_bal * (new_normed_inflation, *normed_inflation),
+                    )
+                    .unwrap_or(token::Amount::max())
+                    .checked_sub(addr_bal)
+                    .unwrap_or_default();
                     // Save the new normed inflation
                     *normed_inflation = new_normed_inflation;
                 }
@@ -374,9 +387,7 @@ where
                 if denom == MaspDenom::Three {
                     // The reward for each reward.1 units of the current asset
                     // is reward.0 units of the reward token
-                    total_reward += ((addr_bal * (real_reward, reward.1)).0
-                        * (*normed_inflation, ref_inflation))
-                        .0;
+                    total_reward += (addr_bal * (reward.0, reward.1)).0;
                 }
             }
             // Add a conversion from the previous asset type
@@ -467,9 +478,12 @@ where
         for denom in token::MaspDenom::iter() {
             // Add the decoding entry for the new asset type. An uncommitted
             // node position is used since this is not a conversion.
-            let new_asset =
-                encode_asset_type(wl_storage.storage.block.epoch, &addr, denom)
-                    .into_storage_result()?;
+            let new_asset = encode_asset_type(
+                Some(wl_storage.storage.block.epoch),
+                &addr,
+                denom,
+            )
+            .into_storage_result()?;
             wl_storage.storage.conversion_state.assets.insert(
                 new_asset,
                 (

core/src/proto/types.rs

@@ -767,7 +767,7 @@ pub struct MaspBuilder {
     pub target: crate::types::hash::Hash,
     /// The decoded set of asset types used by the transaction. Useful for
     /// offline wallets trying to display AssetTypes.
-    pub asset_types: HashSet<(Address, MaspDenom, Epoch)>,
+    pub asset_types: HashSet<(Address, MaspDenom, Option<Epoch>)>,
     /// Track how Info objects map to descriptors and outputs
     #[serde(
         serialize_with = "borsh_serde::<SaplingMetadataSerde, _>",

core/src/types/masp.rs

@@ -19,12 +19,12 @@ use crate::types::token::MaspDenom;
 
 /// Make asset type corresponding to given address and epoch
 pub fn encode_asset_type(
-    epoch: Epoch,
+    epoch: Option<Epoch>,
     token: &Address,
     denom: MaspDenom,
 ) -> Result<AssetType, std::io::Error> {
     // Timestamp the chosen token with the current epoch
-    let token_bytes = (token, denom, epoch.0).serialize_to_vec();
+    let token_bytes = (token, denom, epoch).serialize_to_vec();
     // Generate the unique asset identifier from the unique token address
     AssetType::new(token_bytes.as_ref()).map_err(|_| {
         std::io::Error::new(

core/src/types/token.rs

@@ -67,6 +67,11 @@ impl Amount {
         }
     }
 
+    /// Convert a [`u128`] to an [`Amount`].
+    pub fn from_u128(x: u128) -> Self {
+        Self { raw: Uint::from(x) }
+    }
+
     /// Get the amount as a [`Change`]
     pub fn change(&self) -> Change {
         self.raw.try_into().unwrap()
@@ -227,7 +232,7 @@ impl Amount {
         if hi != 0 && hi_pos >= 4 {
             return None;
         } else if hi != 0 {
-            raw[denom as usize + 1] = hi;
+            raw[hi_pos] = hi;
         }
         Some(Self { raw: Uint(raw) })
     }