wallet-rpc: Remove code duplication4

wallet/wallet-rpc-lib/src/rpc/interface.rs

@@ -15,11 +15,9 @@
 
 use common::chain::SignedTransaction;
 
-use crate::types::{AmountString, NewAccountInfo, TransactionOptions, UtxoInfo};
-
-use super::types::{
-    AccountIndexArg, AddressInfo, AddressWithUsageInfo, BalanceInfo, BlockInfo, EmptyArgs,
-    HexEncoded, TxOptionsOverrides,
+use crate::types::{
+    AccountIndexArg, AddressInfo, AddressWithUsageInfo, Balances, BlockInfo, DecimalAmount,
+    EmptyArgs, HexEncoded, NewAccountInfo, TransactionOptions, TxOptionsOverrides, UtxoInfo,
 };
 
 #[rpc::rpc(server)]
@@ -43,7 +41,7 @@ trait WalletRpc {
     async fn issue_address(&self, account_index: AccountIndexArg) -> rpc::RpcResult<AddressInfo>;
 
     #[method(name = "get_balance")]
-    async fn get_balance(&self, account_index: AccountIndexArg) -> rpc::RpcResult<BalanceInfo>;
+    async fn get_balance(&self, account_index: AccountIndexArg) -> rpc::RpcResult<Balances>;
 
     #[method(name = "get_utxos")]
     async fn get_utxos(&self, account_index: AccountIndexArg) -> rpc::RpcResult<Vec<UtxoInfo>>;
@@ -60,7 +58,7 @@ trait WalletRpc {
         &self,
         account_index: AccountIndexArg,
         address: String,
-        amount: AmountString,
+        amount: DecimalAmount,
         options: TransactionOptions,
     ) -> rpc::RpcResult<()>;
 }

wallet/wallet-rpc-lib/src/rpc/server_impl.rs

@@ -1,4 +1,4 @@
-// Copyright (c) 2023 RBB S.r.l
+// Copyright (c) 2024 RBB S.r.l
 // [email protected]
 // SPDX-License-Identifier: MIT
 // Licensed under the MIT License;
@@ -15,23 +15,18 @@
 
 use std::collections::BTreeMap;
 
-use common::{address::Address, chain::SignedTransaction, primitives::Amount};
+use common::{address::Address, chain::SignedTransaction};
 use utils::shallow_clone::ShallowClone;
-use wallet::account::Currency;
 use wallet_controller::{ControllerConfig, ControllerError, NodeInterface, UtxoStates, UtxoTypes};
 use wallet_types::with_locked::WithLocked;
 
 use crate::{
-    service::WalletControllerError,
-    types::{AmountString, BalanceInfo, NewAccountInfo, TransactionOptions, UtxoInfo},
-};
-
-use super::{
+    rpc::{WalletRpc, WalletRpcServer},
     types::{
-        AccountIndexArg, AddressInfo, AddressWithUsageInfo, BlockInfo, EmptyArgs, HexEncoded,
-        RpcError, TxOptionsOverrides,
+        AccountIndexArg, AddressInfo, AddressWithUsageInfo, Balances, BlockInfo, DecimalAmount,
+        EmptyArgs, HexEncoded, NewAccountInfo, RpcError, TransactionOptions, TxOptionsOverrides,
+        UtxoInfo,
     },
-    WalletRpc, WalletRpcServer,
 };
 
 #[async_trait::async_trait]
@@ -84,38 +79,16 @@ impl WalletRpcServer for WalletRpc {
         Ok(result)
     }
 
-    async fn get_balance(&self, account_index: AccountIndexArg) -> rpc::RpcResult<BalanceInfo> {
+    async fn get_balance(&self, account_index: AccountIndexArg) -> rpc::RpcResult<Balances> {
         let account_idx = account_index.index()?;
         let with_locked = WithLocked::Unlocked; // TODO make user-defined
-        let coin_decimals = self.chain_config.coin_decimals();
 
-        let balances: BalanceInfo = rpc::handle_result(
+        let balances: Balances = rpc::handle_result(
             self.wallet
                 .call_async(move |w| {
                     Box::pin(async move {
-                        let mut amounts = w
-                            .readonly_controller(account_idx)
-                            .get_balance(UtxoStates::ALL, with_locked)?;
-
-                        let coins = amounts.remove(&Currency::Coin).unwrap_or(Amount::ZERO);
-                        let coins = AmountString::new(coins, coin_decimals);
-
-                        let tokens = {
-                            let mut tokens = BTreeMap::new();
-                            for (curr, amt) in amounts {
-                                match curr {
-                                    Currency::Coin => panic!("Coins removed in the previous step"),
-                                    Currency::Token(tok_id) => {
-                                        let decimals =
-                                            w.get_token_number_of_decimals(tok_id).await?;
-                                        tokens.insert(tok_id, AmountString::new(amt, decimals));
-                                    }
-                                }
-                            }
-                            tokens
-                        };
-
-                        Ok::<_, WalletControllerError>(BalanceInfo::new(coins, tokens))
+                        let c = w.readonly_controller(account_idx);
+                        c.get_decimal_balance(UtxoStates::ALL, with_locked).await
                     })
                 })
                 .await,
@@ -151,10 +124,11 @@ impl WalletRpcServer for WalletRpc {
         &self,
         account_index: AccountIndexArg,
         address: String,
-        amount_str: AmountString,
+        amount_str: DecimalAmount,
         options: TransactionOptions,
     ) -> rpc::RpcResult<()> {
-        let amount = amount_str.amount(self.chain_config.coin_decimals())?;
+        let decimals = self.chain_config.coin_decimals();
+        let amount = amount_str.to_amount(decimals).ok_or(RpcError::InvalidCoinAmount)?;
         let address = Address::from_str(&self.chain_config, &address)
             .map_err(|_| RpcError::InvalidAddress)?;
         let acct = account_index.index()?;

wallet/wallet-rpc-lib/src/rpc/types.rs

@@ -15,30 +15,16 @@
 
 //! Types supporting the RPC interface
 
-use std::collections::BTreeMap;
-
 use common::{
     address::Address,
-    chain::{tokens::TokenId, Destination, GenBlock, TxOutput, UtxoOutPoint},
-    primitives::{Amount, BlockHeight, Id},
+    chain::{Destination, GenBlock, TxOutput, UtxoOutPoint},
+    primitives::{BlockHeight, Id},
 };
 use crypto::key::hdkd::{child_number::ChildNumber, u31::U31};
 
 pub use mempool_types::tx_options::TxOptionsOverrides;
 pub use serialization::hex_encoded::HexEncoded;
-
-#[derive(Debug, Eq, PartialEq, Clone, serde::Serialize, serde::Deserialize)]
-pub struct AmountString(pub String);
-
-impl AmountString {
-    pub fn new(amt: Amount, decimals: u8) -> Self {
-        Self(amt.into_fixedpoint_str(decimals))
-    }
-
-    pub fn amount(&self, decimals: u8) -> Result<Amount, RpcError> {
-        Amount::from_fixedpoint_str(&self.0, decimals).ok_or(RpcError::InvalidCoinAmount)
-    }
-}
+pub use wallet_controller::types::{Balances, DecimalAmount};
 
 #[derive(Debug, thiserror::Error)]
 pub enum RpcError {
@@ -50,9 +36,6 @@ pub enum RpcError {
 
     #[error("Invalid address")]
     InvalidAddress,
-
-    #[error("Malformed amount")]
-    MalformedAmount,
 }
 
 impl From<RpcError> for rpc::Error {
@@ -120,18 +103,6 @@ impl AddressWithUsageInfo {
     }
 }
 
-#[derive(Debug, Eq, PartialEq, Clone, serde::Serialize, serde::Deserialize)]
-pub struct BalanceInfo {
-    pub coins: AmountString,
-    pub tokens: BTreeMap<TokenId, AmountString>,
-}
-
-impl BalanceInfo {
-    pub fn new(coins: AmountString, tokens: BTreeMap<TokenId, AmountString>) -> Self {
-        Self { coins, tokens }
-    }
-}
-
 #[derive(Debug, Clone, serde::Serialize)]
 pub struct UtxoInfo {
     pub outpoint: UtxoOutPoint,

wallet/wallet-rpc-lib/tests/basic.rs

@@ -24,7 +24,7 @@ use common::{
 };
 use utils::{make_seedable_rng, ClientT, JsonValue, Seed, ACCOUNT0_ARG, ACCOUNT1_ARG};
 use wallet_rpc_lib::types::{
-    AddressInfo, BalanceInfo, BlockInfo, EmptyArgs, NewAccountInfo, TransactionOptions,
+    AddressInfo, Balances, BlockInfo, EmptyArgs, NewAccountInfo, TransactionOptions,
 };
 
 #[rstest]
@@ -69,8 +69,8 @@ async fn send_coins_to_acct1(#[case] seed: Seed) {
     log::info!("acct1_addr: {acct1_addr:?}");
 
     // Get balance info
-    let balances: BalanceInfo = wallet_rpc.request("get_balance", [ACCOUNT0_ARG]).await.unwrap();
-    let coins_before = balances.coins.amount(coin_decimals).unwrap();
+    let balances: Balances = wallet_rpc.request("get_balance", [ACCOUNT0_ARG]).await.unwrap();
+    let coins_before = balances.coins().to_amount(coin_decimals).unwrap();
     log::info!("Balances: {balances:?}");
     let utxos: JsonValue = wallet_rpc.request("get_utxos", [ACCOUNT0_ARG]).await.unwrap();
     log::info!("UTXOs: {utxos:#}");
@@ -107,12 +107,12 @@ async fn send_coins_to_acct1(#[case] seed: Seed) {
         wallet_rpc.request("send_coins", params).await.unwrap()
     };
 
-    let balances: BalanceInfo = wallet_rpc.request("get_balance", [ACCOUNT0_ARG]).await.unwrap();
-    let coins_after = balances.coins.amount(coin_decimals).unwrap();
+    let balances: Balances = wallet_rpc.request("get_balance", [ACCOUNT0_ARG]).await.unwrap();
+    let coins_after = balances.coins().to_amount(coin_decimals).unwrap();
     assert!(coins_after <= (coins_before / 2).unwrap());
     assert!(coins_after >= (coins_before / 3).unwrap());
 
-    let balances: BalanceInfo = wallet_rpc.request("get_balance", [ACCOUNT1_ARG]).await.unwrap();
+    let balances: Balances = wallet_rpc.request("get_balance", [ACCOUNT1_ARG]).await.unwrap();
     log::info!("acct1 balances: {balances:?}");
 
     tf.stop().await;