From 98ff7593dc0b2bdd7116df7b9cf5a08997bc2d78 Mon Sep 17 00:00:00 2001
From: Stanislav Eliseev <s.eliseev@dexpa.io>
Date: Thu, 14 Sep 2023 17:30:55 +0200
Subject: [PATCH] Update calculations

---
 src/core/token_wallet/mod.rs      | 49 +++++++++++--------------------
 src/core/transactions_tree/mod.rs |  4 +++
 2 files changed, 21 insertions(+), 32 deletions(-)

diff --git a/src/core/token_wallet/mod.rs b/src/core/token_wallet/mod.rs
index c20d7c917..4bbb0e3d4 100644
--- a/src/core/token_wallet/mod.rs
+++ b/src/core/token_wallet/mod.rs
@@ -116,12 +116,9 @@ impl TokenWallet {
         tokens: BigUint,
         notify_receiver: bool,
         payload: ton_types::Cell,
-        max_simulated_transactions: u32,
     ) -> Result<u64> {
         const FEE_MULTIPLIER: u128 = 2;
 
-        let mut simulated = 0;
-
         // Prepare internal message
         let internal_message =
             self.prepare_transfer(destination, tokens, notify_receiver, payload, 0)?;
@@ -171,40 +168,33 @@ impl TokenWallet {
         let source_tx = tree.next().await?.ok_or(TokenWalletError::NoSourceTx)?;
         check_exit_code(&source_tx, TokenWalletError::SourceTxFailed)?;
         attached_amount += source_tx.total_fees.grams.as_u128();
-        simulated += 1;
 
         if source_tx.outmsg_cnt == 0 {
             return Err(TokenWalletError::NoDestTx.into());
         }
+        //
 
-        if max_simulated_transactions == 2 {
-            tree.retain_message_queue(|message| {
-                message.state_init().is_none() && (message.src_ref() == Some(self.address()))
-            });
-
-            if tree.message_queue().len() != 1  {
-                return Err(TokenWalletError::NoDestTx.into());
+        if let Some(message) = tree.peek() {
+            if message.state_init().is_some() && message.src_ref() == Some(self.address()) {
+                //simulate first deploy transaction (we don't need to count attached amount here because of)
+                let dest_tx = tree.next().await?.ok_or(TokenWalletError::NoDestTx)?;
+                check_exit_code(&dest_tx, TokenWalletError::DestinationTxFailed)?;
             }
+        }
 
+        tree.retain_message_queue(|message| {
+            message.state_init().is_none() && (message.src_ref() == Some(self.address()))
+        });
 
-            // Simulate destination transaction
-            let dest_tx = tree.next().await?.ok_or(TokenWalletError::NoDestTx)?;
-            check_exit_code(&dest_tx, TokenWalletError::DestinationTxFailed)?;
-            attached_amount += dest_tx.total_fees.grams.as_u128();
-        } else {
-            'main: while simulated < max_simulated_transactions   {
-                if let Some(tx) = tree.next().await? {
-                    check_exit_code(&tx, TokenWalletError::DestinationTxFailed)?;
-                    if simulated != 1 {
-                        attached_amount += tx.total_fees.grams.as_u128();
-                    }
-                    simulated += 1;
-                } else {
-                    break 'main;
-                }
-            }
+        if tree.message_queue().len() != 1  {
+            return Err(TokenWalletError::NoDestTx.into());
         }
 
+        // Simulate destination transaction
+        let dest_tx = tree.next().await?.ok_or(TokenWalletError::NoDestTx)?;
+        check_exit_code(&dest_tx, TokenWalletError::DestinationTxFailed)?;
+        attached_amount += dest_tx.total_fees.grams.as_u128();
+
         Ok((attached_amount * FEE_MULTIPLIER) as u64)
     }
 
@@ -841,9 +831,4 @@ mod tests {
         root_state.guess_details(&SimpleClock).unwrap();
     }
 
-    #[test]
-    fn estimate_min_amount() {
-
-    }
-
 }
diff --git a/src/core/transactions_tree/mod.rs b/src/core/transactions_tree/mod.rs
index 30775219e..09874e43c 100644
--- a/src/core/transactions_tree/mod.rs
+++ b/src/core/transactions_tree/mod.rs
@@ -70,6 +70,10 @@ impl TransactionsTreeStream {
         }
     }
 
+    pub fn peek(&self) -> Option<&Message> {
+        self.messages.get(0)
+    }
+
     async fn step(&mut self, mut message: Message) -> TransactionTreeResult<Transaction> {
         const A_LOT: u64 = 1_000_000_000_000_000; // 1'000'000 ever