feat(mempool): replenishing queue in get_txs

crates/mempool/src/mempool.rs

@@ -54,6 +54,7 @@ impl Mempool {
             let tx = self.tx_pool.remove(tx_hash)?;
             eligible_txs.push(tx);
         }
+        self.enqueue_next_eligible_txs(&eligible_txs)?;
 
         Ok(eligible_txs)
     }
@@ -111,6 +112,21 @@ impl Mempool {
         Ok(())
     }
 
+    fn enqueue_next_eligible_txs(&mut self, txs: &[ThinTransaction]) -> MempoolResult<()> {
+        for tx in txs {
+            let current_account_state = Account {
+                sender_address: tx.sender_address,
+                state: AccountState { nonce: tx.nonce },
+            };
+            if let Some(next_tx_reference) =
+                self.tx_pool.get_next_eligible_tx(current_account_state)?
+            {
+                self.tx_queue.insert(*next_tx_reference);
+            }
+        }
+        Ok(())
+    }
+
     #[cfg(test)]
     pub(crate) fn _tx_pool(&self) -> &TransactionPool {
         &self.tx_pool

crates/mempool/src/mempool_test.rs

@@ -175,8 +175,6 @@ fn test_get_txs(#[case] requested_txs: usize) {
 }
 
 #[rstest]
-// TODO(Ayelet): remove ignore once replenishing is merged.
-#[ignore]
 fn test_get_txs_multi_nonce() {
     // Setup.
     let tx_address_0_nonce_0 =
@@ -192,8 +190,11 @@ fn test_get_txs_multi_nonce() {
     let txs = mempool.get_txs(2).unwrap();
 
     // Assert that the account's next tx was added the queue.
-    assert_eq!(txs, &[tx_address_0_nonce_0, tx_address_0_nonce_1]);
-    let expected_mempool_state = MempoolState::new([], []);
+    // TODO(Ayelet): all transactions should be returned after replenishing.
+    assert_eq!(txs, &[tx_address_0_nonce_0]);
+    let expected_queue_txs = [TransactionReference::new(&tx_address_0_nonce_1)];
+    let expected_pool_txs = [tx_address_0_nonce_1];
+    let expected_mempool_state = MempoolState::new(expected_pool_txs, expected_queue_txs);
     expected_mempool_state.assert_eq_mempool_state(&mempool);
 }
 

crates/mempool/src/transaction_pool.rs

@@ -3,7 +3,9 @@ use std::collections::{hash_map, BTreeMap, HashMap};
 use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::TransactionHash;
 use starknet_mempool_types::errors::MempoolError;
-use starknet_mempool_types::mempool_types::{MempoolResult, ThinTransaction};
+use starknet_mempool_types::mempool_types::{
+    Account, AccountState, MempoolResult, ThinTransaction,
+};
 
 use crate::mempool::TransactionReference;
 
@@ -86,6 +88,16 @@ impl TransactionPool {
         self.txs_by_account.get(address, nonce)
     }
 
+    pub fn get_next_eligible_tx(
+        &self,
+        current_account_state: Account,
+    ) -> MempoolResult<Option<&TransactionReference>> {
+        let Account { sender_address, state: AccountState { nonce } } = current_account_state;
+        // TOOD(Ayelet): Change to StarknetApiError.
+        let next_nonce = nonce.try_increment().map_err(|_| MempoolError::FeltOutOfRange)?;
+        Ok(self.get_by_address_and_nonce(sender_address, next_nonce))
+    }
+
     #[cfg(test)]
     pub(crate) fn _tx_pool(&self) -> &HashToTransaction {
         &self.tx_pool