From d43df3481f90b0d5088466f35b2645ea5bc97f0d 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      | 30 ++++++++++++++++++++++++++++++
 crates/mempool/src/mempool_test.rs |  2 +-
 crates/mempool_types/src/errors.rs |  3 +++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index d6ee63827..bf2546d2a 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -21,6 +21,8 @@ pub struct Mempool {
     tx_pool: TransactionPool,
     // Transactions eligible for sequencing.
     tx_queue: TransactionQueue,
+    // Represents the current state of the mempool during block creation.
+    mempool_state: HashMap<ContractAddress, AccountState>,
 }
 
 impl Mempool {
@@ -55,6 +57,11 @@ impl Mempool {
             eligible_txs.push(tx);
         }
 
+        // Update the mempool state with the new nonces.
+        for tx in &eligible_txs {
+            self.mempool_state.entry(tx.sender_address).or_default().nonce = tx.nonce;
+        }
+
         Ok(eligible_txs)
     }
 
@@ -62,6 +69,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.validate_input(&input)?;
         self.insert_tx(input)
     }
 
@@ -111,6 +119,28 @@ impl Mempool {
         Ok(())
     }
 
+    fn validate_input(&self, input: &MempoolInput) -> MempoolResult<()> {
+        let MempoolInput { tx, account } = input;
+        if let Some(AccountState { nonce }) = self.mempool_state.get(&tx.sender_address) {
+            if nonce >= &tx.nonce {
+                return Err(MempoolError::DuplicateNonce {
+                    address: tx.sender_address,
+                    nonce: tx.nonce,
+                });
+            }
+        }
+
+        let Account { state: AccountState { nonce }, .. } = account;
+        if nonce > &tx.nonce {
+            return Err(MempoolError::DuplicateNonce {
+                address: tx.sender_address,
+                nonce: tx.nonce,
+            });
+        }
+
+        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 8f015548a..402f91c48 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() }
     }
 }
 
diff --git a/crates/mempool_types/src/errors.rs b/crates/mempool_types/src/errors.rs
index ebe6ad50d..b7eb248a0 100644
--- a/crates/mempool_types/src/errors.rs
+++ b/crates/mempool_types/src/errors.rs
@@ -1,8 +1,11 @@
+use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::transaction::TransactionHash;
 use thiserror::Error;
 
 #[derive(Clone, Debug, Error, PartialEq, Eq)]
 pub enum MempoolError {
+    #[error("Duplicate transaction, sender address: {address}, nonce: {:?}", nonce)]
+    DuplicateNonce { address: ContractAddress, nonce: Nonce },
     #[error("Duplicate transaction, with hash: {tx_hash}")]
     DuplicateTransaction { tx_hash: TransactionHash },
     #[error("Transaction with hash: {tx_hash} not found")]