core: state: Add Wallet::augment_balances_for_orders

fmt

common/src/types/wallet.rs

@@ -202,6 +202,36 @@ impl Wallet {
         self.fees.retain(|fee| !fee.is_default());
     }
 
+    /// Given the current orders in the wallet, augment the balances with all base and quote mints
+    pub fn augment_balances_for_orders(&mut self) {
+        self.remove_default_elements(); // TODO: Do we really need this?
+        for order in self.orders.clone().values() {
+            self.augment_balance(&order.quote_mint);
+            self.augment_balance(&order.base_mint);
+        }
+        // Re-derive the public wallet share, given that the underlying balance data has changed
+        let (new_private_share, new_public_share) = create_wallet_shares_from_private(
+            &self.clone().into(),
+            &self.private_shares.clone(),
+            self.blinder,
+        );
+        self.private_shares = new_private_share;
+        self.blinded_public_shares = new_public_share;
+    }
+
+    /// If the mint does not exist in the wallet, add it with a balance of zero
+    fn augment_balance(&mut self, mint: &BigUint) {
+        if !self.balances.contains_key(mint) {
+            self.balances.insert(
+                mint.clone(),
+                Balance {
+                    mint: mint.clone(),
+                    amount: 0u64,
+                },
+            );
+        }
+    }
+
     /// Get the balance, fee, and fee_balance for an order by specifying the order directly
     ///
     /// We allow orders to be matched when undercapitalized; i.e. the respective balance does

config/src/lib.rs

@@ -267,7 +267,7 @@ impl Clone for RelayerConfig {
             starknet_private_keys: self.starknet_private_keys.clone(),
             eth_websocket_addr: self.eth_websocket_addr.clone(),
             debug: self.debug,
-            demo: self.demo
+            demo: self.demo,
         }
     }
 }
@@ -366,7 +366,7 @@ fn parse_config_from_args(full_args: Vec<String>) -> Result<RelayerConfig, Strin
         starknet_private_keys: cli_args.starknet_private_keys,
         eth_websocket_addr: cli_args.eth_websocket_addr,
         debug: cli_args.debug,
-        demo: cli_args.demo
+        demo: cli_args.demo,
     };
     set_contract_from_file(&mut config, cli_args.deployments_file)?;
 

task-driver/src/helpers.rs

@@ -137,7 +137,7 @@ pub(super) fn construct_wallet_commitment_proof(
     valid_reblind_witness: &SizedValidReblindWitness,
     proof_manager_work_queue: CrossbeamSender<ProofManagerJob>,
     disable_fee_validation: bool,
-    demo: bool
+    demo: bool,
 ) -> Result<(SizedValidCommitmentsWitness, TokioReceiver<ProofBundle>), String> {
     // Choose the first fee. If no fee is found and fee validation is disabled, use a zero fee
     let first_fee = wallet.fees.iter().find(|f| !f.is_default()).cloned();
@@ -304,6 +304,10 @@ pub(super) async fn update_wallet_validity_proofs(
         return Ok(());
     }
 
+    // Before proving `VALID REBLIND` and `VALID COMMITMENTS`, ensure that all balances are augmented
+    let wallet = &mut wallet.clone();
+    wallet.augment_balances_for_orders();
+
     // Dispatch a proof of `VALID REBLIND` for the wallet
     let (reblind_witness, reblind_response_channel) =
         construct_wallet_reblind_proof(wallet.clone(), proof_manager_work_queue.clone())?;
@@ -319,7 +323,7 @@ pub(super) async fn update_wallet_validity_proofs(
             &wallet_reblind_witness,
             proof_manager_work_queue.clone(),
             global_state.disable_fee_validation,
-            global_state.demo
+            global_state.demo,
         )?;
 
         let order_commitment_witness = Box::new(commitments_witness);

task-driver/src/initialize_state.rs

@@ -246,7 +246,7 @@ impl InitializeStateTask {
                     &wallet_reblind_witness,
                     self.proof_manager_work_queue.clone(),
                     self.global_state.disable_fee_validation,
-                    self.global_state.demo
+                    self.global_state.demo,
                 )
                 .map_err(InitializeStateTaskError::ProveValidCommitments)?;
 

workers/handshake-manager/src/manager/match.rs

@@ -165,7 +165,8 @@ impl HandshakeExecutor {
         }
 
         // Prove `VALID MATCH MPC` with the counterparty
-        let (commitment, proof) = Self::prove_valid_match(&witness, &fabric, self.global_state.demo).await?;
+        let (commitment, proof) =
+            Self::prove_valid_match(&witness, &fabric, self.global_state.demo).await?;
 
         // Verify the commitment links between the match proof and the two parties'
         // proofs of VALID COMMITMENTS