Sui Address Balance support and coin-model refactor (#1163)

Author: 0xh3rman

crates/gem_sui/src/lib.rs

@@ -25,8 +25,8 @@ pub mod tx_builder;
 pub mod signer;
 
 pub use error::SuiError;
-use models::Coin;
 pub use models::ObjectId;
+use models::{Coin, OwnedCoins};
 use std::error::Error;
 use sui_transaction_builder::ObjectInput;
 use sui_types::Address;
@@ -71,14 +71,13 @@ pub fn sui_clock_object_input() -> ObjectInput {
     ObjectInput::shared(sui_clock_object_id(), 1, false)
 }
 
-pub fn validate_enough_balance(coins: &[Coin], amount: u64) -> Option<Box<dyn Error + Send + Sync>> {
-    if coins.is_empty() {
-        return Some("coins list is empty".into());
+pub fn validate_enough_balance(coins: &OwnedCoins<Coin>, amount: u64) -> Option<Box<dyn Error + Send + Sync>> {
+    let total = coins.total();
+    if total == 0 {
+        return Some("no spendable coin objects or address balance".into());
     }
-
-    let total_amount: u64 = coins.iter().map(|x| x.balance).sum();
-    if total_amount < amount {
-        return Some(format!("total amount ({}) is less than amount to send ({})", total_amount, amount).into());
+    if total < amount {
+        return Some(format!("total amount ({}) is less than amount to send ({})", total, amount).into());
     }
     None
 }

crates/gem_sui/src/models/coin.rs

@@ -7,6 +7,7 @@ use serde_serializers::deserialize_bigint_from_str;
 
 #[cfg(feature = "rpc")]
 use super::account::Owner;
+use super::core::Coin;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -24,16 +25,57 @@ pub struct SuiObject {
     pub version: String,
 }
 
-#[cfg(feature = "rpc")]
-#[derive(Debug, Clone, Serialize, Deserialize)]
-#[serde(rename_all = "camelCase")]
-pub struct SuiCoin {
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct OwnedCoins<T> {
     pub coin_type: String,
-    pub coin_object_id: String,
-    #[serde(deserialize_with = "deserialize_bigint_from_str")]
-    pub balance: BigInt,
-    pub version: String,
-    pub digest: String,
+    pub coins: Vec<T>,
+    pub address_balance: u64,
+}
+
+impl<T> Default for OwnedCoins<T> {
+    fn default() -> Self {
+        Self {
+            coin_type: String::new(),
+            coins: Vec::new(),
+            address_balance: 0,
+        }
+    }
+}
+
+impl<T> OwnedCoins<T> {
+    pub fn new(coin_type: String, coins: Vec<T>, address_balance: u64) -> Self {
+        Self {
+            coin_type,
+            coins,
+            address_balance,
+        }
+    }
+
+    pub fn map<U>(self, f: impl FnMut(T) -> U) -> OwnedCoins<U> {
+        OwnedCoins {
+            coin_type: self.coin_type,
+            coins: self.coins.into_iter().map(f).collect(),
+            address_balance: self.address_balance,
+        }
+    }
+
+    pub fn try_map<U, E>(self, f: impl FnMut(T) -> Result<U, E>) -> Result<OwnedCoins<U>, E> {
+        Ok(OwnedCoins {
+            coin_type: self.coin_type,
+            coins: self.coins.into_iter().map(f).collect::<Result<_, _>>()?,
+            address_balance: self.address_balance,
+        })
+    }
+}
+
+impl OwnedCoins<Coin> {
+    pub fn coin_total(&self) -> u64 {
+        self.coins.iter().map(|coin| coin.balance).fold(0, u64::saturating_add)
+    }
+
+    pub fn total(&self) -> u64 {
+        self.coin_total().saturating_add(self.address_balance)
+    }
 }
 
 #[cfg(feature = "rpc")]
@@ -43,6 +85,9 @@ pub struct Balance {
     pub coin_type: String,
     #[serde(deserialize_with = "deserialize_bigint_from_str")]
     pub total_balance: BigInt,
+    #[serde(default)]
+    /// Amount in the per-address balance accumulator.
+    pub address_balance: u64,
 }
 
 #[cfg(feature = "rpc")]

crates/gem_sui/src/models/coin_asset.rs

@@ -1,26 +1,33 @@
+use super::OwnedCoins;
 use bcs;
 use gem_encoding::encode_base64;
 use std::error::Error;
 use sui_transaction_builder::ObjectInput;
-use sui_types::Transaction;
+use sui_types::{Address, Digest, Transaction};
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct Coin {
     pub coin_type: String,
     pub balance: u64,
     pub object: Object,
 }
 
-#[derive(Debug, PartialEq, Clone)]
+impl Coin {
+    pub fn to_input(&self) -> ObjectInput {
+        self.object.to_input()
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct Object {
-    pub object_id: String,
-    pub digest: String,
+    pub object_id: Address,
+    pub digest: Digest,
     pub version: u64,
 }
 
 impl Object {
     pub fn to_input(&self) -> ObjectInput {
-        ObjectInput::owned(self.object_id.parse().unwrap(), self.version, self.digest.parse().unwrap())
+        ObjectInput::owned(self.object_id, self.version, self.digest)
     }
 }
 
@@ -36,7 +43,7 @@ pub struct StakeInput {
     pub validator: String,
     pub stake_amount: u64,
     pub gas: Gas,
-    pub coins: Vec<Coin>,
+    pub coins: OwnedCoins<Coin>,
 }
 
 #[derive(Debug, PartialEq, Clone)]
@@ -52,7 +59,7 @@ pub struct TransferInput {
     pub sender: String,
     pub recipient: String,
     pub amount: u64,
-    pub coins: Vec<Coin>,
+    pub coins: OwnedCoins<Coin>,
     pub send_max: bool,
     pub gas: Gas,
 }
@@ -62,7 +69,7 @@ pub struct TokenTransferInput {
     pub sender: String,
     pub recipient: String,
     pub amount: u64,
-    pub tokens: Vec<Coin>,
+    pub tokens: OwnedCoins<Coin>,
     pub gas: Gas,
     pub gas_coin: Coin,
 }

crates/gem_sui/src/models/core.rs

@@ -1,14 +1,14 @@
 pub mod account;
 pub mod coin;
-pub mod coin_asset;
 pub mod core;
 pub mod inspect;
 pub mod object_id;
 pub mod staking;
+#[cfg(test)]
+pub mod testkit;
 pub mod transaction;
 
 pub use coin::*;
-pub use coin_asset::{CoinAsset, CoinResponse};
 pub use core::*;
 pub use inspect::{InspectCommandResult, InspectEffects, InspectEvent, InspectGasUsed, InspectResult, InspectReturnValue};
 pub use object_id::ObjectId;

crates/gem_sui/src/models/mod.rs

@@ -0,0 +1,28 @@
+use crate::SUI_COIN_TYPE;
+use crate::models::{Coin, Object, OwnedCoins};
+
+impl Object {
+    pub fn mock() -> Self {
+        Self {
+            object_id: "0xabcdef1234567890abcdef1234567890abcdef12".parse().unwrap(),
+            digest: "HdfF7hswRuvbXbEXjGjmUCt7gLybhvbPvvK8zZbCqyD8".parse().unwrap(),
+            version: 100,
+        }
+    }
+}
+
+impl Coin {
+    pub fn mock_sui() -> Self {
+        Self {
+            coin_type: SUI_COIN_TYPE.to_string(),
+            balance: 5_000_000_000,
+            object: Object::mock(),
+        }
+    }
+}
+
+impl OwnedCoins<Coin> {
+    pub fn mock_sui() -> Self {
+        Self::new(SUI_COIN_TYPE.to_string(), vec![Coin::mock_sui()], 0)
+    }
+}

crates/gem_sui/src/models/testkit.rs

@@ -12,7 +12,7 @@ use primitives::{
 use crate::{
     ESTIMATION_GAS_BUDGET, SUI_COIN_TYPE,
     gas_budget::GAS_BUDGET_MULTIPLIER,
-    models::{SuiCoin, SuiObject},
+    models::{Coin, OwnedCoins, SuiObject},
 };
 use crate::{
     provider::preload_mapper::{map_transaction_data, map_transaction_rate_rates},
@@ -27,13 +27,13 @@ impl ChainTransactionLoad for SuiClient {
     }
 
     async fn get_transaction_load(&self, input: TransactionLoadInput) -> Result<TransactionLoadData, Box<dyn Error + Sync + Send>> {
-        let (gas_coins, coins, objects) = self.get_coins_for_input_type(input.sender_address.as_str(), input.input_type.clone()).await?;
+        let (sui_coins, token_coins, objects) = self.get_coins_for_input_type(input.sender_address.as_str(), input.input_type.clone()).await?;
 
-        let estimate_bytes = map_transaction_data(input.clone(), gas_coins.clone(), coins.clone(), objects.clone(), ESTIMATION_GAS_BUDGET)?;
+        let estimate_bytes = map_transaction_data(input.clone(), sui_coins.clone(), token_coins.clone(), objects.clone(), ESTIMATION_GAS_BUDGET)?;
         let fee = self.estimate_fee(&estimate_bytes, &input.gas_price, input.is_max_value).await?;
 
         let message_bytes = match estimated_gas_budget(&input.input_type, &fee)? {
-            Some(budget) => map_transaction_data(input, gas_coins, coins, objects, budget)?,
+            Some(budget) => map_transaction_data(input, sui_coins, token_coins, objects, budget)?,
             None => estimate_bytes,
         };
 
@@ -75,27 +75,27 @@ impl SuiClient {
         &self,
         address: &str,
         input_type: TransactionInputType,
-    ) -> Result<(Vec<SuiCoin>, Vec<SuiCoin>, Vec<SuiObject>), Box<dyn Error + Send + Sync>> {
+    ) -> Result<(OwnedCoins<Coin>, Option<OwnedCoins<Coin>>, Vec<SuiObject>), Box<dyn Error + Send + Sync>> {
         match input_type {
             TransactionInputType::Transfer(asset) => match asset.id.token_id {
-                None => Ok((self.get_coins(address, SUI_COIN_TYPE).await?, Vec::new(), Vec::new())),
+                None => Ok((self.get_coins(address, SUI_COIN_TYPE).await?, None, Vec::new())),
                 Some(token_id) => {
-                    let (gas_coins, coins) = futures::try_join!(self.get_coins(address, SUI_COIN_TYPE), self.get_coins(address, &token_id))?;
-                    Ok((gas_coins, coins, Vec::new()))
+                    let (gas_coins, token_coins) = futures::try_join!(self.get_coins(address, SUI_COIN_TYPE), self.get_coins(address, &token_id))?;
+                    Ok((gas_coins, Some(token_coins), Vec::new()))
                 }
             },
             TransactionInputType::Stake(_, stake_type) => match stake_type {
-                StakeType::Stake(_) => Ok((self.get_coins(address, SUI_COIN_TYPE).await?, Vec::new(), Vec::new())),
+                StakeType::Stake(_) => Ok((self.get_coins(address, SUI_COIN_TYPE).await?, None, Vec::new())),
                 StakeType::Unstake(delegation) => {
                     let (gas_coins, staked_object) = futures::try_join!(self.get_coins(address, SUI_COIN_TYPE), self.get_object(delegation.base.delegation_id.clone()))?;
-                    Ok((gas_coins, Vec::new(), vec![staked_object]))
+                    Ok((gas_coins, None, vec![staked_object]))
                 }
                 StakeType::Redelegate(_) | StakeType::Rewards(_) | StakeType::Withdraw(_) | StakeType::Freeze(_) | StakeType::Unfreeze(_) => {
                     Err("Unsupported stake type for Sui".into())
                 }
             },
-            TransactionInputType::Swap(_, _, _) => Ok((Vec::new(), Vec::new(), Vec::new())),
-            TransactionInputType::Generic(_, _, _) => Ok((Vec::new(), Vec::new(), Vec::new())),
+            TransactionInputType::Swap(_, _, _) => Ok((OwnedCoins::default(), None, Vec::new())),
+            TransactionInputType::Generic(_, _, _) => Ok((OwnedCoins::default(), None, Vec::new())),
             TransactionInputType::TransferNft(_, _) | TransactionInputType::Account(_, _) => Err("Unsupported transaction type for Sui".into()),
             _ => Err("Unsupported transaction type for Sui".into()),
         }