refactor: update priority queue to use ThinPriorityTransaction

crates/mempool/src/mempool.rs

@@ -34,35 +34,31 @@ impl Mempool {
             state: HashMap::default(),
         };
 
-        mempool.txs_queue = TransactionPriorityQueue::from(
-            inputs
-                .into_iter()
-                .map(|input| {
-                    // Attempts to insert a key-value pair into the mempool's state. Returns `None`
-                    // if the key was not present, otherwise returns the old value while updating
-                    // the new value.
-                    let prev_value =
-                        mempool.state.insert(input.account.address, input.account.state);
-                    assert!(
-                        prev_value.is_none(),
-                        "Sender address: {:?} already exists in the mempool. Can't add {:?} to \
-                         the mempool.",
-                        input.account.address,
-                        input.tx
-                    );
-
-                    // Insert the transaction into the tx_store.
-                    let res = mempool.tx_store.push(input.tx.clone());
-                    assert!(
-                        res.is_ok(),
-                        "Transaction: {:?} already exists in the mempool.",
-                        input.tx.tx_hash
-                    );
-
-                    input.tx
-                })
-                .collect::<Vec<ThinTransaction>>(),
-        );
+        for input in inputs.into_iter() {
+            // Attempts to insert a key-value pair into the mempool's state. Returns `None`
+            // if the key was not present, otherwise returns the old value while updating
+            // the new value.
+            let address = input.account.address;
+            let state = input.account.state;
+            let existing_account_state = mempool.state.insert(address, state);
+            assert!(
+                existing_account_state.is_none(),
+                "Sender address: {:?} already exists in the mempool. Can't add {:?} to the \
+                 mempool.",
+                input.account.address,
+                input.tx
+            );
+
+            // Insert the transaction into the tx_store.
+            let res = mempool.tx_store.push(input.tx.clone());
+            assert!(
+                res.is_ok(),
+                "Transaction: {:?} already exists in the mempool.",
+                input.tx.tx_hash
+            );
+
+            mempool.txs_queue.push(input.tx.clone().into());
+        }
 
         mempool
     }
@@ -77,10 +73,13 @@ impl Mempool {
     // back. TODO: Consider renaming to `pop_txs` to be more consistent with the standard
     // library.
     pub fn get_txs(&mut self, n_txs: usize) -> MempoolResult<Vec<ThinTransaction>> {
-        let txs = self.txs_queue.pop_last_chunk(n_txs);
-        for tx in &txs {
+        let pq_txs = self.txs_queue.pop_last_chunk(n_txs);
+
+        let mut txs: Vec<ThinTransaction> = Vec::default();
+        for pq_tx in &pq_txs {
+            let tx = self.tx_store.remove(&pq_tx.tx_hash)?;
             self.state.remove(&tx.sender_address);
-            self.tx_store.remove(&tx.tx_hash)?;
+            txs.push(tx);
         }
 
         Ok(txs)
@@ -94,7 +93,7 @@ impl Mempool {
             Occupied(_) => Err(MempoolError::DuplicateTransaction { tx_hash: tx.tx_hash }),
             Vacant(entry) => {
                 entry.insert(account.state);
-                self.txs_queue.push(tx.clone());
+                self.txs_queue.push(tx.clone().into());
                 self.tx_store.push(tx)?;
 
                 Ok(())

crates/mempool/src/mempool_test.rs

@@ -10,6 +10,7 @@ use starknet_mempool_types::mempool_types::ThinTransaction;
 use starknet_mempool_types::utils::create_thin_tx_for_testing;
 
 use crate::mempool::{Account, Mempool, MempoolInput};
+use crate::priority_queue::ThinPriorityTransaction;
 
 /// Creates a valid input for mempool's `add_tx` with optional default value for
 /// `sender_address`.
@@ -132,5 +133,8 @@ fn test_add_same_tx(mut mempool: Mempool) {
 fn check_mempool_txs_eq(mempool: &Mempool, expected_txs: &[ThinTransaction]) {
     let mempool_txs = mempool.txs_queue.iter();
     // Deref the inner mempool tx type.
-    expected_txs.iter().zip(mempool_txs).all(|(a, b)| *a == **b);
+    expected_txs
+        .iter()
+        .zip(mempool_txs)
+        .all(|(a, b)| <ThinTransaction as Into<ThinPriorityTransaction>>::into(a.clone()) == *b);
 }

crates/mempool/src/priority_queue.rs

@@ -1,58 +1,61 @@
 use std::cmp::Ordering;
 use std::collections::BTreeSet;
 
+use starknet_api::core::Nonce;
+use starknet_api::transaction::{Tip, TransactionHash};
 use starknet_mempool_types::mempool_types::ThinTransaction;
 // Assumption: for the MVP only one transaction from the same contract class can be in the mempool
 // at a time. When this changes, saving the transactions themselves on the queu might no longer be
 // appropriate, because we'll also need to stores transactions without indexing them. For example,
 // transactions with future nonces will need to be stored, and potentially indexed on block commits.
 #[derive(Clone, Debug, Default, derive_more::Deref, derive_more::DerefMut)]
-pub struct TransactionPriorityQueue(BTreeSet<PrioritizedTransaction>);
+pub struct TransactionPriorityQueue(BTreeSet<ThinPriorityTransaction>);
 
 impl TransactionPriorityQueue {
-    pub fn push(&mut self, tx: ThinTransaction) {
-        let mempool_tx = PrioritizedTransaction(tx);
-        self.insert(mempool_tx);
+    pub fn push(&mut self, tx: ThinPriorityTransaction) {
+        self.insert(tx);
     }
 
     // TODO(gilad): remove collect
-    pub fn pop_last_chunk(&mut self, n_txs: usize) -> Vec<ThinTransaction> {
-        (0..n_txs).filter_map(|_| self.pop_last().map(|tx| tx.0)).collect()
+    pub fn pop_last_chunk(&mut self, n_txs: usize) -> Vec<ThinPriorityTransaction> {
+        (0..n_txs).filter_map(|_| self.pop_last()).collect()
     }
 }
 
-impl From<Vec<ThinTransaction>> for TransactionPriorityQueue {
-    fn from(transactions: Vec<ThinTransaction>) -> Self {
-        TransactionPriorityQueue(BTreeSet::from_iter(
-            transactions.into_iter().map(PrioritizedTransaction),
-        ))
-    }
+#[derive(Clone, Debug, Default)]
+pub struct ThinPriorityTransaction {
+    pub nonce: Nonce,
+    pub tx_hash: TransactionHash,
+    pub tip: Tip,
 }
 
-#[derive(Clone, Debug, derive_more::Deref, derive_more::From)]
-pub struct PrioritizedTransaction(pub ThinTransaction);
-
 /// 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 PrioritizedTransaction {
-    fn eq(&self, other: &PrioritizedTransaction) -> bool {
+impl PartialEq for ThinPriorityTransaction {
+    fn eq(&self, other: &ThinPriorityTransaction) -> bool {
         self.tip == other.tip
     }
 }
 
 /// 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 PrioritizedTransaction {}
+impl Eq for ThinPriorityTransaction {}
 
-impl Ord for PrioritizedTransaction {
+impl Ord for ThinPriorityTransaction {
     fn cmp(&self, other: &Self) -> Ordering {
         self.tip.cmp(&other.tip)
     }
 }
 
-impl PartialOrd for PrioritizedTransaction {
+impl PartialOrd for ThinPriorityTransaction {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
+
+impl From<ThinTransaction> for ThinPriorityTransaction {
+    fn from(tx: ThinTransaction) -> Self {
+        ThinPriorityTransaction { nonce: tx.nonce, tx_hash: tx.tx_hash, tip: tx.tip }
+    }
+}