From f17c5222d1938dfb7b615189f0c2277137a74081 Mon Sep 17 00:00:00 2001
From: Elin <elin@starkware.co>
Date: Tue, 5 Nov 2024 16:40:46 +0200
Subject: [PATCH] fix(mempool): fix nonce-rewind in `commit_block` for known
 non-included addresses (#1837)

---
 crates/mempool/src/mempool.rs     |  5 +++--
 crates/mempool/tests/flow_test.rs | 35 +++++++++++++++++++++++--------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 138be0e192..6b691dc326 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -157,14 +157,15 @@ impl Mempool {
 
         // Rewind nonces of addresses that were not included in block.
         let known_addresses_not_included_in_block =
-            self.mempool_state.keys().filter(|&key| !address_to_nonce.contains_key(key));
+            self.mempool_state.keys().filter(|&key| !address_to_nonce.contains_key(key)).copied();
         for address in known_addresses_not_included_in_block {
             // Account nonce is the minimal nonce of this address: it was proposed but not included.
             let tx_reference = self
                 .tx_pool
-                .account_txs_sorted_by_nonce(*address)
+                .account_txs_sorted_by_nonce(address)
                 .next()
                 .expect("Address {address} should appear in transaction pool.");
+            self.tx_queue.remove(address);
             self.tx_queue.insert(*tx_reference);
         }
 
diff --git a/crates/mempool/tests/flow_test.rs b/crates/mempool/tests/flow_test.rs
index 8236a1625b..59492b345d 100644
--- a/crates/mempool/tests/flow_test.rs
+++ b/crates/mempool/tests/flow_test.rs
@@ -219,30 +219,47 @@ fn test_commit_block_fills_nonce_gap(mut mempool: Mempool) {
 }
 
 #[rstest]
-fn test_flow_commit_block_rewinds_queued_nonce(mut mempool: Mempool) {
+fn test_commit_block_rewinds_queued_nonce(mut mempool: Mempool) {
     // Setup.
-    let tx_nonce_2 = add_tx_input!(tx_hash: 1, address: "0x0", tx_nonce: 2, account_nonce: 2);
-    let tx_nonce_3 = add_tx_input!(tx_hash: 2, address: "0x0", tx_nonce: 3, account_nonce: 2);
-    let tx_nonce_4 = add_tx_input!(tx_hash: 3, address: "0x0", tx_nonce: 4, account_nonce: 2);
+    let tx_address_0_nonce_2 =
+        add_tx_input!(tx_hash: 1, address: "0x0", tx_nonce: 2, account_nonce: 2);
+    let tx_address_0_nonce_3 =
+        add_tx_input!(tx_hash: 2, address: "0x0", tx_nonce: 3, account_nonce: 2);
+    let tx_address_1_nonce_2 =
+        add_tx_input!(tx_hash: 3, address: "0x1", tx_nonce: 2, account_nonce: 2);
+    let tx_address_1_nonce_3 =
+        add_tx_input!(tx_hash: 4, address: "0x1", tx_nonce: 3, account_nonce: 2);
 
-    for input in [&tx_nonce_2, &tx_nonce_3, &tx_nonce_4] {
+    for input in
+        [&tx_address_0_nonce_2, &tx_address_0_nonce_3, &tx_address_1_nonce_2, &tx_address_1_nonce_3]
+    {
         add_tx(&mut mempool, input);
     }
 
     get_txs_and_assert_expected(
         &mut mempool,
-        3,
-        &[tx_nonce_2.tx, tx_nonce_3.tx.clone(), tx_nonce_4.tx.clone()],
+        4,
+        &[
+            tx_address_1_nonce_2.tx.clone(),
+            tx_address_0_nonce_2.tx,
+            tx_address_1_nonce_3.tx,
+            tx_address_0_nonce_3.tx.clone(),
+        ],
     );
 
     // Test.
     let nonces = [("0x0", 3)];
     let tx_hashes = [1];
-    // Nonce 2 was accepted, but 3 and 4 were not, so are rewound.
+    // Address 0x0: nonce 2 was accepted, but 3 was not, so is rewound.
+    // Address 0x1: nonce 2 was not accepted, both 2 and 3 were rewound.
     commit_block(&mut mempool, nonces, tx_hashes);
 
     // Nonces 3 and 4 were re-enqueued correctly.
-    get_txs_and_assert_expected(&mut mempool, 2, &[tx_nonce_3.tx, tx_nonce_4.tx]);
+    get_txs_and_assert_expected(
+        &mut mempool,
+        2,
+        &[tx_address_1_nonce_2.tx, tx_address_0_nonce_3.tx],
+    );
 }
 
 #[rstest]