From 380e27f1b6c0964c55efbc08eed847818dc7d88b 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      | 18 ++++++++++++++++++
 crates/mempool/src/mempool_test.rs |  2 +-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 8b1b4c84e..85c759678 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,
 };
@@ -21,6 +22,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 {
@@ -55,6 +58,11 @@ impl Mempool {
             eligible_txs.push(tx);
         }
 
+        // Update the mempool state with the new nonces.
+        for tx in eligible_txs.iter() {
+            self.mempool_state.entry(tx.sender_address).or_default().nonce = tx.nonce;
+        }
+
         Ok(eligible_txs)
     }
 
@@ -62,6 +70,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)
     }
 
@@ -108,6 +117,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 8ee7256f8..e1d5ed8ad 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() }
     }
 }