Return transaction iterators from mempool storage

Also rename transaction methods for consistency.

zebrad/src/components/mempool.rs

@@ -181,7 +181,7 @@ impl Mempool {
         txid: UnminedTxId,
     ) -> Result<(), MempoolError> {
         // Check if the transaction is already in the mempool.
-        if storage.contains(&txid) {
+        if storage.contains_transaction_exact(&txid) {
             return Err(MempoolError::InMempool);
         }
         if let Some(error) = storage.rejection_error(&txid) {
@@ -258,12 +258,12 @@ impl Service<Request> for Mempool {
                 tx_downloads,
             } => match req {
                 Request::TransactionIds => {
-                    let res = storage.tx_ids();
+                    let res = storage.tx_ids().collect();
                     async move { Ok(Response::TransactionIds(res)) }.boxed()
                 }
                 Request::TransactionsById(ids) => {
-                    let rsp = Ok(storage.transactions(ids)).map(Response::Transactions);
-                    async move { rsp }.boxed()
+                    let res = storage.transactions_exact(ids).cloned().collect();
+                    async move { Ok(Response::Transactions(res)) }.boxed()
                 }
                 Request::RejectedTransactionIds(ids) => {
                     let rsp = Ok(storage.rejected_transactions(ids))
@@ -314,7 +314,7 @@ fn remove_expired_transactions(
 ) {
     let mut txid_set = HashSet::new();
 
-    for t in storage.transactions_all() {
+    for t in storage.transactions() {
         if let Some(expiry_height) = t.transaction.expiry_height() {
             if tip_height >= expiry_height {
                 txid_set.insert(t.id.mined_id());

zebrad/src/components/mempool/storage.rs

@@ -120,12 +120,6 @@ impl Storage {
         Ok(tx_id)
     }
 
-    /// Returns `true` if a [`UnminedTx`] matching an [`UnminedTxId`] is in
-    /// the mempool.
-    pub fn contains(&self, txid: &UnminedTxId) -> bool {
-        self.verified.iter().any(|tx| &tx.id == txid)
-    }
-
     /// Remove [`UnminedTx`]es from the mempool via exact [`UnminedTxId`].
     ///
     /// For v5 transactions, transactions are matched by WTXID, using both the:
@@ -172,22 +166,39 @@ impl Storage {
     }
 
     /// Returns the set of [`UnminedTxId`]s in the mempool.
-    pub fn tx_ids(&self) -> Vec<UnminedTxId> {
-        self.verified.iter().map(|tx| tx.id).collect()
-    }
-
-    /// Returns the set of [`Transaction`]s matching `tx_ids` in the mempool.
-    pub fn transactions(&self, tx_ids: HashSet<UnminedTxId>) -> Vec<UnminedTx> {
-        self.verified
-            .iter()
-            .filter(|tx| tx_ids.contains(&tx.id))
-            .cloned()
-            .collect()
+    pub fn tx_ids(&self) -> impl Iterator<Item = UnminedTxId> + '_ {
+        self.verified.iter().map(|tx| tx.id)
     }
 
     /// Returns the set of [`Transaction`]s in the mempool.
-    pub fn transactions_all(&self) -> Vec<UnminedTx> {
-        self.verified.iter().cloned().collect()
+    pub fn transactions(&self) -> impl Iterator<Item = &UnminedTx> {
+        self.verified.iter()
+    }
+
+    /// Returns the number of transactions in the mempool.
+    #[allow(dead_code)]
+    pub fn transaction_count(&self) -> usize {
+        self.verified.len()
+    }
+
+    /// Returns the set of [`Transaction`]s with exactly matching `tx_ids` in the mempool.
+    ///
+    /// This matches the exact transaction, with identical blockchain effects, signatures, and proofs.
+    pub fn transactions_exact(
+        &self,
+        tx_ids: HashSet<UnminedTxId>,
+    ) -> impl Iterator<Item = &UnminedTx> {
+        self.verified
+            .iter()
+            .filter(move |tx| tx_ids.contains(&tx.id))
+    }
+
+    /// Returns `true` if a [`UnminedTx`] exactly matching an [`UnminedTxId`] is in
+    /// the mempool.
+    ///
+    /// This matches the exact transaction, with identical blockchain effects, signatures, and proofs.
+    pub fn contains_transaction_exact(&self, txid: &UnminedTxId) -> bool {
+        self.verified.iter().any(|tx| &tx.id == txid)
     }
 
     /// Returns `true` if a [`UnminedTx`] matching the supplied [`UnminedTxId`] is in

zebrad/src/components/mempool/storage/tests/vectors.rs

@@ -24,14 +24,14 @@ fn mempool_storage_crud_exact_mainnet() {
     let _ = storage.insert(unmined_tx.clone());
 
     // Check that it is in the mempool, and not rejected.
-    assert!(storage.contains(&unmined_tx.id));
+    assert!(storage.contains_transaction_exact(&unmined_tx.id));
 
     // Remove tx
     let removal_count = storage.remove_exact(&iter::once(unmined_tx.id).collect());
 
     // Check that it is /not/ in the mempool.
     assert_eq!(removal_count, 1);
-    assert!(!storage.contains(&unmined_tx.id));
+    assert!(!storage.contains_transaction_exact(&unmined_tx.id));
 }
 
 #[test]
@@ -52,7 +52,7 @@ fn mempool_storage_crud_same_effects_mainnet() {
     let _ = storage.insert(unmined_tx.clone());
 
     // Check that it is in the mempool, and not rejected.
-    assert!(storage.contains(&unmined_tx.id));
+    assert!(storage.contains_transaction_exact(&unmined_tx.id));
 
     // Remove tx
     let removal_count =
@@ -60,7 +60,7 @@ fn mempool_storage_crud_same_effects_mainnet() {
 
     // Check that it is /not/ in the mempool.
     assert_eq!(removal_count, 1);
-    assert!(!storage.contains(&unmined_tx.id));
+    assert!(!storage.contains_transaction_exact(&unmined_tx.id));
 }
 
 #[test]
@@ -103,12 +103,12 @@ fn mempool_storage_basic_for_network(network: Network) -> Result<()> {
 
     // Make sure the last MEMPOOL_SIZE transactions we sent are in the verified
     for tx in expected_in_mempool {
-        assert!(storage.contains(&tx.id));
+        assert!(storage.contains_transaction_exact(&tx.id));
     }
 
     // Anything greater should not be in the verified
     for tx in expected_to_be_rejected {
-        assert!(!storage.contains(&tx.id));
+        assert!(!storage.contains_transaction_exact(&tx.id));
     }
 
     // Query all the ids we have for rejected, get back `total - MEMPOOL_SIZE`

zebrad/src/components/mempool/tests/prop.rs

@@ -57,7 +57,7 @@ proptest! {
             // The first call to `poll_ready` shouldn't clear the storage yet.
             mempool.dummy_call().await;
 
-            prop_assert_eq!(mempool.storage().tx_ids().len(), 1);
+            prop_assert_eq!(mempool.storage().transaction_count(), 1);
 
             // Simulate a chain reset.
             chain_tip_sender.set_finalized_tip(chain_tip);
@@ -65,7 +65,7 @@ proptest! {
             // This time a call to `poll_ready` should clear the storage.
             mempool.dummy_call().await;
 
-            prop_assert!(mempool.storage().tx_ids().is_empty());
+            prop_assert_eq!(mempool.storage().transaction_count(), 0);
 
             peer_set.expect_no_requests().await?;
             state_service.expect_no_requests().await?;
@@ -110,7 +110,7 @@ proptest! {
             // The first call to `poll_ready` shouldn't clear the storage yet.
             mempool.dummy_call().await;
 
-            prop_assert_eq!(mempool.storage().tx_ids().len(), 1);
+            prop_assert_eq!(mempool.storage().transaction_count(), 1);
 
             // Simulate the synchronizer catching up to the network chain tip.
             mempool.disable(&mut recent_syncs).await;
@@ -121,7 +121,7 @@ proptest! {
             // Enable the mempool again so the storage can be accessed.
             mempool.enable(&mut recent_syncs).await;
 
-            prop_assert!(mempool.storage().tx_ids().is_empty());
+            prop_assert_eq!(mempool.storage().transaction_count(), 0);
 
             peer_set.expect_no_requests().await?;
             state_service.expect_no_requests().await?;