From bd34515e1afd77ef5ce9bc3db0984c95731a5fd8 Mon Sep 17 00:00:00 2001
From: mohammad-starkware <130282237+MohammadNassar1@users.noreply.github.com>
Date: Tue, 23 Jul 2024 16:16:06 +0300
Subject: [PATCH] feat(mempool): implement remove from nonce for account
 transaction index (#443)

---
 crates/mempool/src/transaction_pool.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/crates/mempool/src/transaction_pool.rs b/crates/mempool/src/transaction_pool.rs
index fa21c30a3..3d151855e 100644
--- a/crates/mempool/src/transaction_pool.rs
+++ b/crates/mempool/src/transaction_pool.rs
@@ -104,4 +104,25 @@ impl AccountTransactionIndex {
     fn get(&self, address: ContractAddress, nonce: Nonce) -> Option<&TransactionReference> {
         self.0.get(&address)?.get(&nonce)
     }
+
+    fn _remove_up_to_nonce(
+        &mut self,
+        address: ContractAddress,
+        nonce: Nonce,
+    ) -> Vec<TransactionReference> {
+        let Some(account_txs) = self.0.get_mut(&address) else {
+            return Vec::default();
+        };
+
+        // Split the transactions at the given nonce.
+        let txs_with_higher_or_equal_nonce = account_txs.split_off(&nonce);
+        let txs_with_lower_nonce = std::mem::replace(account_txs, txs_with_higher_or_equal_nonce);
+
+        if account_txs.is_empty() {
+            self.0.remove(&address);
+        }
+
+        // Collect and return the transactions with lower nonces.
+        txs_with_lower_nonce.into_values().collect()
+    }
 }