From f83d33cbedc09d866c4bcc5e987fbe54974c40ac Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <138376632+ayeletstarkware@users.noreply.github.com>
Date: Tue, 20 Aug 2024 11:20:37 +0300
Subject: [PATCH] fix(mempool): in add_tx, delete queue transaction with lower
 nonce that account nonce (#452)

---
 crates/mempool/src/mempool.rs      |  6 +++++-
 crates/mempool/src/mempool_test.rs | 29 +++++++++++++++++++++++++----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 7fd506a49b7..1747620b118 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -120,7 +120,11 @@ impl Mempool {
         let MempoolInput { tx, account: Account { sender_address, state: AccountState { nonce } } } =
             input;
 
-        // Remove transactions with lower nonce than the account nonce.
+        // Note: != is actually equivalent to > here, as lower nonces are rejected in validation.
+        if self.tx_queue.get_nonce(sender_address).is_some_and(|queued_nonce| queued_nonce != nonce)
+        {
+            self.tx_queue.remove(sender_address);
+        }
         self.tx_pool.remove_up_to_nonce(sender_address, nonce);
 
         self.tx_pool.insert(tx)?;
diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs
index 36d9d3026f3..733862ef5dc 100644
--- a/crates/mempool/src/mempool_test.rs
+++ b/crates/mempool/src/mempool_test.rs
@@ -438,6 +438,27 @@ fn test_add_tx_lower_than_queued_nonce() {
     expected_mempool_content.assert_eq_queue_content(&mempool);
 }
 
+#[rstest]
+fn test_add_tx_updates_queue_with_higher_account_nonce() {
+    // Setup.
+    let input =
+        add_tx_input!(tx_hash: 1, sender_address: "0x0", tx_nonce: 0_u8, account_nonce: 0_u8);
+    let higher_account_nonce_input =
+        add_tx_input!(tx_hash: 2, sender_address: "0x0", tx_nonce: 1_u8, account_nonce: 1_u8);
+
+    let queue_txs = [TransactionReference::new_from_thin_tx(&input.tx)];
+    let mut mempool: Mempool = MempoolContent::with_queue(queue_txs).into();
+
+    // Test.
+    add_tx(&mut mempool, &higher_account_nonce_input);
+
+    // Assert: the higher account nonce transaction is in the queue.
+    let expected_queue_txs =
+        [TransactionReference::new_from_thin_tx(&higher_account_nonce_input.tx)];
+    let expected_mempool_content = MempoolContent::with_queue(expected_queue_txs);
+    expected_mempool_content.assert_eq_queue_content(&mempool);
+}
+
 #[rstest]
 fn test_add_tx_with_identical_tip_succeeds(mut mempool: Mempool) {
     // Setup.
@@ -478,11 +499,11 @@ fn test_add_tx_delete_tx_with_lower_nonce_than_account_nonce() {
     add_tx(&mut mempool, &tx_nonce_1_account_nonce_1);
 
     // Assert the transaction with the lower nonce is removed.
-    // TODO(Ayelet): Assert the queue after modifying add_tx to delete lower nonce transactions
-    // in the queue.
+    let expected_queue_txs =
+        [TransactionReference::new_from_thin_tx(&tx_nonce_1_account_nonce_1.tx)];
     let expected_pool_txs = [tx_nonce_1_account_nonce_1.tx];
-    let expected_mempool_content = MempoolContent::with_pool(expected_pool_txs);
-    expected_mempool_content.assert_eq_pool_content(&mempool);
+    let expected_mempool_content = MempoolContent::new(expected_pool_txs, expected_queue_txs);
+    expected_mempool_content.assert_eq_mempool_content(&mempool);
 }
 
 #[rstest]