From 90486451fcc8fc385240a2e88225eaf4496d2655 Mon Sep 17 00:00:00 2001
From: Mohammad Nassar <mohammad@starkware.co>
Date: Tue, 13 Aug 2024 18:22:19 +0300
Subject: [PATCH] feat(mempool): add pending transaction type

---
 crates/mempool/src/mempool.rs                 | 18 +++++++--
 crates/mempool/src/transaction_queue.rs       | 39 ++++++++++++++++---
 .../src/executable_transaction.rs             | 20 ++++++++++
 3 files changed, 67 insertions(+), 10 deletions(-)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 922f71fd54..38ed52e26f 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use starknet_api::core::{ContractAddress, Nonce};
 use starknet_api::executable_transaction::Transaction;
-use starknet_api::transaction::{ResourceBoundsMapping, Tip, TransactionHash};
+use starknet_api::transaction::{Resource, ResourceBoundsMapping, Tip, TransactionHash};
 use starknet_mempool_types::errors::MempoolError;
 use starknet_mempool_types::mempool_types::{Account, AccountState, MempoolInput, MempoolResult};
 
@@ -209,9 +209,19 @@ impl TransactionReference {
             sender_address: tx.contract_address(),
             nonce: tx.nonce(),
             tx_hash: tx.tx_hash(),
-            tip: tx.tip().expect("Expected a valid tip value, but received None."),
-            // TODO(Mohammad): add resource bounds to the transaction.
-            resource_bounds: ResourceBoundsMapping::default(),
+            tip: tx.tip().expect("Expected a valid tip value."),
+            resource_bounds: tx
+                .resource_bounds()
+                .expect("Expected a valid resource bounds value.")
+                .clone(),
         }
     }
+
+    pub fn get_l2_gas_price(&self) -> u128 {
+        self.resource_bounds
+            .0
+            .get(&Resource::L2Gas)
+            .map(|bounds| bounds.max_price_per_unit)
+            .expect("Expected a valid L2 gas resource bounds.")
+    }
 }
diff --git a/crates/mempool/src/transaction_queue.rs b/crates/mempool/src/transaction_queue.rs
index c4199166ae..099770e341 100644
--- a/crates/mempool/src/transaction_queue.rs
+++ b/crates/mempool/src/transaction_queue.rs
@@ -68,21 +68,48 @@ impl TransactionQueue {
     }
 }
 
-/// Encapsulates a transaction reference to assess its order (i.e., priority).
+/// Encapsulates a transaction reference to assess its order (i.e., gas price).
+#[derive(Clone, Debug, derive_more::Deref, derive_more::From)]
+struct PendingTransaction(pub TransactionReference);
+
+/// Compare transactions based only on their gas price, using the Eq trait. It ensures that
+/// two gas price are either exactly equal or not.
+impl PartialEq for PendingTransaction {
+    fn eq(&self, other: &PendingTransaction) -> bool {
+        self.get_l2_gas_price() == other.get_l2_gas_price() && self.tx_hash == other.tx_hash
+    }
+}
+
+/// Marks this struct as capable of strict equality comparisons, signaling to the compiler it
+/// adheres to equality semantics.
+// Note: this depends on the implementation of `PartialEq`, see its docstring.
+impl Eq for PendingTransaction {}
+
+impl Ord for PendingTransaction {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.get_l2_gas_price()
+            .cmp(&other.get_l2_gas_price())
+            .then_with(|| self.tx_hash.cmp(&other.tx_hash))
+    }
+}
+
+impl PartialOrd for PendingTransaction {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+/// This struct behaves similarly to `PendingTransaction`, encapsulating a transaction reference
+/// to assess its order (i.e., tip); see its documentation for more details.
 #[derive(Clone, Debug, derive_more::Deref, derive_more::From)]
 struct PriorityTransaction(pub TransactionReference);
 
-/// Compare transactions based only on their tip, a uint, using the Eq trait. It ensures that two
-/// tips are either exactly equal or not.
 impl PartialEq for PriorityTransaction {
     fn eq(&self, other: &PriorityTransaction) -> bool {
         self.tip == other.tip && self.tx_hash == other.tx_hash
     }
 }
 
-/// Marks this struct as capable of strict equality comparisons, signaling to the compiler it
-/// adheres to equality semantics.
-// Note: this depends on the implementation of `PartialEq`, see its docstring.
 impl Eq for PriorityTransaction {}
 
 impl Ord for PriorityTransaction {
diff --git a/crates/starknet_api/src/executable_transaction.rs b/crates/starknet_api/src/executable_transaction.rs
index ed3e3d288c..a9380f9db2 100644
--- a/crates/starknet_api/src/executable_transaction.rs
+++ b/crates/starknet_api/src/executable_transaction.rs
@@ -65,6 +65,7 @@ impl Transaction {
         }
     }
 
+    // TODO(Mohammad): add a getter macro.
     pub fn tip(&self) -> Option<Tip> {
         match self {
             Transaction::Declare(declare_tx) => match &declare_tx.tx {
@@ -82,6 +83,25 @@ impl Transaction {
         }
     }
 
+    pub fn resource_bounds(&self) -> Option<&ResourceBoundsMapping> {
+        match self {
+            Transaction::Declare(declare_tx) => match &declare_tx.tx {
+                crate::transaction::DeclareTransaction::V3(tx_v3) => Some(&tx_v3.resource_bounds),
+                _ => None,
+            },
+            Transaction::DeployAccount(deploy_account_tx) => match &deploy_account_tx.tx {
+                crate::transaction::DeployAccountTransaction::V3(tx_v3) => {
+                    Some(&tx_v3.resource_bounds)
+                }
+                _ => None,
+            },
+            Transaction::Invoke(invoke_tx) => match &invoke_tx.tx {
+                crate::transaction::InvokeTransaction::V3(tx_v3) => Some(&tx_v3.resource_bounds),
+                _ => None,
+            },
+        }
+    }
+
     // TODO(Arni): Update the function to support all transaction types.
     pub fn new_from_rpc_tx(
         rpc_tx: RpcTransaction,