refactor: get_proofs returns ProofInfo

bindings/cdk-js/src/wallet.rs

@@ -47,8 +47,10 @@ impl JsWallet {
     }
 
     #[wasm_bindgen(js_name = totalBalance)]
-    pub async fn total_balance(&self) -> Result<JsAmount> {
-        Ok(self.inner.total_balance().await.map_err(into_err)?.into())
+    pub async fn total_balance(&self) -> Result<JsValue> {
+        Ok(serde_wasm_bindgen::to_value(
+            &self.inner.total_balance().await.map_err(into_err)?,
+        )?)
     }
 
     #[wasm_bindgen(js_name = totalPendingBalance)]

crates/cdk-redb/src/wallet.rs

@@ -401,13 +401,13 @@ impl WalletDatabase for RedbWalletDatabase {
         mint_url: Option<UncheckedUrl>,
         state: Option<Vec<State>>,
         spending_conditions: Option<Vec<SpendingConditions>>,
-    ) -> Result<Option<Proofs>, Self::Err> {
+    ) -> Result<Option<Vec<ProofInfo>>, Self::Err> {
         let db = self.db.lock().await;
         let read_txn = db.begin_read().map_err(Error::from)?;
 
         let table = read_txn.open_table(PROOFS_TABLE).map_err(Error::from)?;
 
-        let proofs: Proofs = table
+        let proofs: Vec<ProofInfo> = table
             .iter()
             .map_err(Error::from)?
             .flatten()
@@ -416,7 +416,7 @@ impl WalletDatabase for RedbWalletDatabase {
 
                 if let Ok(proof_info) = serde_json::from_str::<ProofInfo>(v.value()) {
                     match proof_info.matches_conditions(&mint_url, &state, &spending_conditions) {
-                        Ok(true) => proof = Some(proof_info.proof),
+                        Ok(true) => proof = Some(proof_info),
                         Ok(false) => (),
                         Err(_) => (),
                     }

crates/cdk-rexie/src/wallet.rs

@@ -450,7 +450,7 @@ impl WalletDatabase for RexieWalletDatabase {
         mint_url: Option<UncheckedUrl>,
         state: Option<Vec<State>>,
         spending_conditions: Option<Vec<SpendingConditions>>,
-    ) -> Result<Option<Proofs>, Self::Err> {
+    ) -> Result<Option<Vec<ProofInfo>>, Self::Err> {
         let rexie = self.db.lock().await;
 
         let transaction = rexie
@@ -464,7 +464,7 @@ impl WalletDatabase for RexieWalletDatabase {
             .await
             .map_err(Error::from)?;
 
-        let proofs: Proofs = proofs
+        let proofs: Vec<ProofInfo> = proofs
             .into_iter()
             .filter_map(|(_k, v)| {
                 let mut proof = None;
@@ -475,7 +475,7 @@ impl WalletDatabase for RexieWalletDatabase {
                         &state,
                         &spending_conditions,
                     ) {
-                        Ok(true) => Some(proof_info.proof),
+                        Ok(true) => Some(proof_info),
                         Ok(false) => None,
                         Err(_) => None,
                     };

crates/cdk/src/cdk_database/mod.rs

@@ -83,7 +83,7 @@ pub trait WalletDatabase {
         mint_url: Option<UncheckedUrl>,
         state: Option<Vec<State>>,
         spending_conditions: Option<Vec<SpendingConditions>>,
-    ) -> Result<Option<Proofs>, Self::Err>;
+    ) -> Result<Option<Vec<ProofInfo>>, Self::Err>;
     async fn remove_proofs(&self, proofs: &Proofs) -> Result<(), Self::Err>;
 
     async fn set_proof_state(&self, y: PublicKey, state: State) -> Result<(), Self::Err>;

crates/cdk/src/cdk_database/wallet_memory.rs

@@ -171,15 +171,15 @@ impl WalletDatabase for WalletMemoryDatabase {
         mint_url: Option<UncheckedUrl>,
         state: Option<Vec<State>>,
         spending_conditions: Option<Vec<SpendingConditions>>,
-    ) -> Result<Option<Proofs>, Error> {
+    ) -> Result<Option<Vec<ProofInfo>>, Error> {
         let proofs = self.proofs.read().await;
 
-        let proofs: Proofs = proofs
+        let proofs: Vec<ProofInfo> = proofs
             .clone()
             .into_values()
             .filter_map(|proof_info| {
                 match proof_info.matches_conditions(&mint_url, &state, &spending_conditions) {
-                    Ok(true) => Some(proof_info.proof),
+                    Ok(true) => Some(proof_info),
                     Ok(false) => None,
                     Err(_) => None,
                 }

crates/cdk/src/types.rs

@@ -100,10 +100,16 @@ pub struct ProofInfo {
     pub mint_url: UncheckedUrl,
     pub state: State,
     pub spending_condition: Option<SpendingConditions>,
+    pub unit: CurrencyUnit,
 }
 
 impl ProofInfo {
-    pub fn new(proof: Proof, mint_url: UncheckedUrl, state: State) -> Result<Self, Error> {
+    pub fn new(
+        proof: Proof,
+        mint_url: UncheckedUrl,
+        state: State,
+        unit: CurrencyUnit,
+    ) -> Result<Self, Error> {
         let y = proof
             .y()
             .map_err(|_| Error::CustomError("Could not find y".to_string()))?;
@@ -116,6 +122,7 @@ impl ProofInfo {
             mint_url,
             state,
             spending_condition,
+            unit,
         })
     }
 

crates/cdk/src/wallet.rs

@@ -138,20 +138,23 @@ impl Wallet {
 
     /// Total Balance of wallet
     #[instrument(skip(self))]
-    pub async fn total_balance(&self) -> Result<Amount, Error> {
-        let mut balance = Amount::ZERO;
+    pub async fn total_balance(&self) -> Result<HashMap<CurrencyUnit, Amount>, Error> {
+        let mut balances = HashMap::new();
 
         if let Some(proofs) = self
             .localstore
             .get_proofs(None, Some(vec![State::Unspent]), None)
             .await?
         {
-            let amount = proofs.iter().map(|p| p.amount).sum();
-
-            balance += amount;
+            for proof in proofs {
+                balances
+                    .entry(proof.unit)
+                    .and_modify(|ps| *ps += proof.proof.amount)
+                    .or_insert(proof.proof.amount);
+            }
         }
 
-        Ok(balance)
+        Ok(balances)
     }
 
     /// Total Balance of wallet
@@ -164,7 +167,7 @@ impl Wallet {
             .get_proofs(None, Some(vec![State::Pending]), None)
             .await?
         {
-            let amount = proofs.iter().map(|p| p.amount).sum();
+            let amount = proofs.iter().map(|p| p.proof.amount).sum();
 
             balance += amount;
         }
@@ -173,10 +176,12 @@ impl Wallet {
     }
 
     #[instrument(skip(self))]
-    pub async fn mint_balances(&self) -> Result<HashMap<UncheckedUrl, Amount>, Error> {
+    pub async fn mint_balances(
+        &self,
+    ) -> Result<HashMap<UncheckedUrl, HashMap<CurrencyUnit, Amount>>, Error> {
         let mints = self.localstore.get_mints().await?;
 
-        let mut balances = HashMap::new();
+        let mut mint_balances = HashMap::new();
 
         for (mint, _) in mints {
             if let Some(proofs) = self
@@ -184,15 +189,22 @@ impl Wallet {
                 .get_proofs(Some(mint.clone()), None, None)
                 .await?
             {
-                let amount = proofs.iter().map(|p| p.amount).sum();
+                let mut balances = HashMap::new();
+
+                for proof in proofs {
+                    balances
+                        .entry(proof.unit)
+                        .and_modify(|ps| *ps += proof.proof.amount)
+                        .or_insert(proof.proof.amount);
+                }
 
-                balances.insert(mint, amount);
+                mint_balances.insert(mint, balances);
             } else {
-                balances.insert(mint, Amount::ZERO);
+                mint_balances.insert(mint, HashMap::new());
             }
         }
 
-        Ok(balances)
+        Ok(mint_balances)
     }
 
     #[instrument(skip(self), fields(mint_url = %mint_url))]
@@ -200,7 +212,8 @@ impl Wallet {
         Ok(self
             .localstore
             .get_proofs(Some(mint_url), Some(vec![State::Unspent]), None)
-            .await?)
+            .await?
+            .map(|p| p.into_iter().map(|p| p.proof).collect()))
     }
 
     #[instrument(skip(self), fields(mint_url = %mint_url))]
@@ -366,7 +379,10 @@ impl Wallet {
                 .await?
             {
                 let states = self
-                    .check_proofs_spent(mint.clone(), proofs.clone())
+                    .check_proofs_spent(
+                        mint.clone(),
+                        proofs.clone().into_iter().map(|p| p.proof).collect(),
+                    )
                     .await?;
 
                 // Both `State::Pending` and `State::Unspent` should be included in the pending table.
@@ -378,13 +394,15 @@ impl Wallet {
                     .map(|s| s.y)
                     .collect();
 
-                let (pending_proofs, non_pending_proofs): (Proofs, Proofs) = proofs
+                let (pending_proofs, non_pending_proofs): (Vec<ProofInfo>, Vec<ProofInfo>) = proofs
                     .into_iter()
-                    .partition(|p| p.y().map(|y| pending_states.contains(&y)).unwrap_or(false));
+                    .partition(|p| pending_states.contains(&p.y));
 
-                let amount = pending_proofs.iter().map(|p| p.amount).sum();
+                let amount = pending_proofs.iter().map(|p| p.proof.amount).sum();
 
-                self.localstore.remove_proofs(&non_pending_proofs).await?;
+                self.localstore
+                    .remove_proofs(&non_pending_proofs.into_iter().map(|p| p.proof).collect())
+                    .await?;
 
                 balance += amount;
             }
@@ -611,7 +629,14 @@ impl Wallet {
 
         let proofs = proofs
             .into_iter()
-            .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Unspent))
+            .flat_map(|proof| {
+                ProofInfo::new(
+                    proof,
+                    mint_url.clone(),
+                    State::Unspent,
+                    quote_info.unit.clone(),
+                )
+            })
             .collect();
 
         // Add new proofs to store
@@ -703,7 +728,9 @@ impl Wallet {
                 let send_proofs_info = send_proofs
                     .clone()
                     .into_iter()
-                    .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Reserved))
+                    .flat_map(|proof| {
+                        ProofInfo::new(proof, mint_url.clone(), State::Reserved, unit.clone())
+                    })
                     .collect();
 
                 self.localstore.add_proofs(send_proofs_info).await?;
@@ -726,7 +753,7 @@ impl Wallet {
 
         let keep_proofs = keep_proofs
             .into_iter()
-            .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Unspent))
+            .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Unspent, unit.clone()))
             .collect();
 
         self.localstore.add_proofs(keep_proofs).await?;
@@ -938,11 +965,14 @@ impl Wallet {
         unit: &CurrencyUnit,
         amount: Amount,
     ) -> Result<Proofs, Error> {
-        let mint_proofs = self
+        let mint_proofs: Proofs = self
             .localstore
             .get_proofs(Some(mint_url.clone()), Some(vec![State::Unspent]), None)
             .await?
-            .ok_or(Error::InsufficientFunds)?;
+            .ok_or(Error::InsufficientFunds)?
+            .into_iter()
+            .map(|p| p.proof)
+            .collect();
 
         let mint_keysets = self
             .localstore
@@ -1077,7 +1107,14 @@ impl Wallet {
 
             let change_proofs_info = change_proofs
                 .into_iter()
-                .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Unspent))
+                .flat_map(|proof| {
+                    ProofInfo::new(
+                        proof,
+                        mint_url.clone(),
+                        State::Unspent,
+                        quote_info.unit.clone(),
+                    )
+                })
                 .collect();
 
             self.localstore.add_proofs(change_proofs_info).await?;
@@ -1239,7 +1276,7 @@ impl Wallet {
             total_amount += proofs.iter().map(|p| p.amount).sum();
             let proofs = proofs
                 .into_iter()
-                .flat_map(|proof| ProofInfo::new(proof, mint.clone(), State::Unspent))
+                .flat_map(|proof| ProofInfo::new(proof, mint.clone(), State::Unspent, unit.clone()))
                 .collect();
             self.localstore.add_proofs(proofs).await?;
         }
@@ -1436,7 +1473,9 @@ impl Wallet {
 
                 let unspent_proofs = unspent_proofs
                     .into_iter()
-                    .flat_map(|proof| ProofInfo::new(proof, mint_url.clone(), State::Unspent))
+                    .flat_map(|proof| {
+                        ProofInfo::new(proof, mint_url.clone(), State::Unspent, keyset.unit.clone())
+                    })
                     .collect();
 
                 self.localstore.add_proofs(unspent_proofs).await?;