From 9774833079aa3aad86db557dfe6b2f52c85ea80f Mon Sep 17 00:00:00 2001
From: Mohammad Nassar <mohammad@starkware.co>
Date: Thu, 18 Jul 2024 09:48:05 +0300
Subject: [PATCH] feat(mempool): impl mempool state

---
 crates/mempool/src/mempool.rs      | 23 +++++++++++++++++++++++
 crates/mempool/src/mempool_test.rs |  2 +-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index e109800e4..00c459cc8 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::{Tip, TransactionHash};
+use starknet_mempool_types::errors::MempoolError;
 use starknet_mempool_types::mempool_types::{
     Account, AccountState, MempoolInput, MempoolResult, ThinTransaction,
 };
@@ -20,6 +21,8 @@ pub struct Mempool {
     tx_pool: TransactionPool,
     // Transactions eligible for sequencing.
     tx_queue: TransactionQueue,
+    // Represents the current state of the mempool during batch processing.
+    mempool_state: HashMap<ContractAddress, AccountState>,
 }
 
 impl Mempool {
@@ -54,6 +57,16 @@ impl Mempool {
             eligible_txs.push(tx);
         }
 
+        for tx in &eligible_txs {
+            self.mempool_state.entry(tx.sender_address).and_modify(|account| {
+                account.nonce = tx.nonce;
+            });
+        }
+
+        for tx in &eligible_txs {
+            self.mempool_state.entry(tx.sender_address).or_default().nonce = tx.nonce;
+        }
+
         Ok(eligible_txs)
     }
 
@@ -61,6 +74,7 @@ impl Mempool {
     /// TODO: support fee escalation and transactions with future nonces.
     /// TODO: check Account nonce and balance.
     pub fn add_tx(&mut self, input: MempoolInput) -> MempoolResult<()> {
+        self.is_duplicated_tx(&input.tx)?;
         self.insert_tx(input)
     }
 
@@ -100,6 +114,15 @@ impl Mempool {
         Ok(())
     }
 
+    fn is_duplicated_tx(&self, tx: &ThinTransaction) -> MempoolResult<()> {
+        if let Some(AccountState { nonce }) = self.mempool_state.get(&tx.sender_address) {
+            if *nonce >= tx.nonce {
+                return Err(MempoolError::DuplicateTransaction { tx_hash: tx.tx_hash });
+            }
+        }
+        Ok(())
+    }
+
     #[cfg(test)]
     pub(crate) fn _tx_pool(&self) -> &TransactionPool {
         &self.tx_pool
diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs
index 30eae9a28..8bba89c99 100644
--- a/crates/mempool/src/mempool_test.rs
+++ b/crates/mempool/src/mempool_test.rs
@@ -44,7 +44,7 @@ impl MempoolState {
 impl From<MempoolState> for Mempool {
     fn from(mempool_state: MempoolState) -> Mempool {
         let MempoolState { tx_pool, tx_queue } = mempool_state;
-        Mempool { tx_pool, tx_queue }
+        Mempool { tx_pool, tx_queue, ..Default::default() }
     }
 }