walletdb: create explicit type for balances

src/bin/darkfid.rs

@@ -135,9 +135,7 @@ impl Darkfid {
             let balances = self.client.lock().await.get_balances().await?;
             let mut symbols: HashMap<String, (String, String)> = HashMap::new();
 
-            for (id, value) in balances.iter() {
-                let id: jubjub::Fr = deserialize(&id)?;
-
+            for balance in balances.list.iter() {
                 // this is hardcoded for SOL
                 // TODO: if id == btc_id:
                 //          network = bitcoin
@@ -146,8 +144,8 @@ impl Darkfid {
 
                 let network = "solana";
 
-                if let Some(symbol) = self.drk_tokenlist.symbol_from_id(id)? {
-                    let amount = encode_base10(*value, 8);
+                if let Some(symbol) = self.drk_tokenlist.symbol_from_id(balance.token_id)? {
+                    let amount = encode_base10(balance.value, 8);
                     symbols.insert(symbol, (amount, network.to_string()));
                 }
             }

src/client.rs

@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::net::SocketAddr;
 
 use async_std::sync::{Arc, Mutex};
@@ -19,7 +18,7 @@ use crate::{
     service::{GatewayClient, GatewaySlabsSubscriber},
     state::{state_transition, ProgramState, StateUpdate},
     tx,
-    wallet::{CashierDbPtr, Keypair, WalletPtr},
+    wallet::{CashierDbPtr, Keypair, WalletPtr, walletdb::Balances},
     Result,
 };
 
@@ -350,7 +349,7 @@ impl Client {
         self.wallet.key_gen()
     }
 
-    pub async fn get_balances(&self) -> Result<HashMap<Vec<u8>, u64>> {
+    pub async fn get_balances(&self) -> Result<Balances> {
         self.wallet.get_balances()
     }
 

src/wallet/walletdb.rs

@@ -23,6 +23,29 @@ pub struct Keypair {
     pub private: jubjub::Fr,
 }
 
+#[derive(Debug, Clone)]
+pub struct Balance {
+    pub token_id: jubjub::Fr,
+    pub value: u64,
+}
+
+#[derive(Debug, Clone)]
+pub struct Balances {
+    pub list: Vec<Balance>,
+}
+impl Balances {
+    pub fn add(&mut self, balance: &Balance) {
+        if let Some(mut saved_balance) = self
+            .list
+            .iter_mut()
+            .find(|b| b.token_id == balance.token_id)
+        {
+            saved_balance.value += balance.value;
+        } else {
+            self.list.push(balance.clone());
+        }
+    }
+}
 
 //#[derive(Clone)]
 pub struct WalletDb {
@@ -297,7 +320,7 @@ impl WalletDb {
         Ok(())
     }
 
-    pub fn get_balances(&self) -> Result<HashMap<Vec<u8>, u64>> {
+    pub fn get_balances(&self) -> Result<Balances> {
         debug!(target: "WALLETDB", "Get token and balances...");
         let conn = Connection::open(&self.path)?;
         conn.pragma_update(None, "key", &self.password)?;
@@ -310,19 +333,15 @@ impl WalletDb {
             Ok((row.get(0)?, row.get(1)?))
         })?;
 
-        let mut balances = HashMap::new();
+        let mut balances = Balances { list: Vec::new() };
 
         for row in rows {
             let row = row?;
             let value: u64 = row.0;
-            let token_id: Vec<u8> = row.1;
-
-            if let Some(val) = balances.get_mut(&token_id) {
-                *val += value;
-            } else {
-                balances.insert(token_id, value);
-            }
+            let token_id: jubjub::Fr = self.get_value_deserialized(&row.1)?;
+            balances.add(&Balance { token_id, value });
         }
+
         Ok(balances)
     }
 
@@ -374,7 +393,6 @@ mod tests {
 
     use super::*;
     use crate::crypto::{coin::Coin, OwnCoin};
-    use crate::serial::serialize;
     use crate::util::join_config_path;
     use ff::PrimeField;
 
@@ -492,9 +510,10 @@ mod tests {
 
         let balances = wallet.get_balances()?;
 
-        let token_id = serialize(&token_id);
 
-        assert_eq!(balances[&token_id], 440);
+        assert_eq!(balances.list.len(), 1);
+        assert_eq!(balances.list[0].value, 440);
+        assert_eq!(balances.list[0].token_id, token_id);
 
         std::fs::remove_file(walletdb_path)?;