From 4d00627559ff3800833a7ee2ed34db4a00cf2fe0 Mon Sep 17 00:00:00 2001
From: Yohann Kazoula <yoh.kzl@gmail.com>
Date: Fri, 31 Jan 2025 02:30:38 +0200
Subject: [PATCH] Add helper functions to `NewPooledTransactionHashes68` struct
 (#14118)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/net/eth-wire-types/src/broadcast.rs | 115 ++++++++++++++++++++-
 1 file changed, 114 insertions(+), 1 deletion(-)

diff --git a/crates/net/eth-wire-types/src/broadcast.rs b/crates/net/eth-wire-types/src/broadcast.rs
index 95e9019c9..828924112 100644
--- a/crates/net/eth-wire-types/src/broadcast.rs
+++ b/crates/net/eth-wire-types/src/broadcast.rs
@@ -375,6 +375,35 @@ impl NewPooledTransactionHashes68 {
     pub fn metadata_iter(&self) -> impl Iterator<Item = (&B256, (u8, usize))> {
         self.hashes.iter().zip(self.types.iter().copied().zip(self.sizes.iter().copied()))
     }
+
+    /// Appends a transaction
+    pub fn push<T: SignedTransaction>(&mut self, tx: &T) {
+        self.hashes.push(*tx.tx_hash());
+        self.sizes.push(tx.encode_2718_len());
+        self.types.push(tx.ty());
+    }
+
+    /// Appends the provided transactions
+    pub fn extend<'a, T: SignedTransaction>(&mut self, txs: impl IntoIterator<Item = &'a T>) {
+        for tx in txs {
+            self.push(tx);
+        }
+    }
+
+    /// Consumes and appends a transaction
+    pub fn with_transaction<T: SignedTransaction>(mut self, tx: &T) -> Self {
+        self.push(tx);
+        self
+    }
+
+    /// Consumes and appends the provided transactions
+    pub fn with_transactions<'a, T: SignedTransaction>(
+        mut self,
+        txs: impl IntoIterator<Item = &'a T>,
+    ) -> Self {
+        self.extend(txs);
+        self
+    }
 }
 
 impl Encodable for NewPooledTransactionHashes68 {
@@ -739,7 +768,10 @@ impl FromIterator<(TxHash, Eth68TxMetadata)> for RequestTxHashes {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use alloy_primitives::{b256, hex};
+    use alloy_consensus::Typed2718;
+    use alloy_eips::eip2718::Encodable2718;
+    use alloy_primitives::{b256, hex, PrimitiveSignature as Signature, U256};
+    use reth_ethereum_primitives::{Transaction, TransactionSigned};
     use std::str::FromStr;
 
     /// Takes as input a struct / encoded hex message pair, ensuring that we encode to the exact hex
@@ -948,4 +980,85 @@ mod tests {
         assert_eq!(0, hashes.len());
         assert_eq!(5, rest.len());
     }
+
+    fn signed_transaction() -> impl SignedTransaction {
+        TransactionSigned::new_unhashed(
+            Transaction::Legacy(Default::default()),
+            Signature::new(
+                U256::from_str(
+                    "0x64b1702d9298fee62dfeccc57d322a463ad55ca201256d01f62b45b2e1c21c12",
+                )
+                .unwrap(),
+                U256::from_str(
+                    "0x64b1702d9298fee62dfeccc57d322a463ad55ca201256d01f62b45b2e1c21c10",
+                )
+                .unwrap(),
+                false,
+            ),
+        )
+    }
+
+    #[test]
+    fn test_pooled_tx_hashes_68_push() {
+        let tx = signed_transaction();
+        let mut tx_hashes =
+            NewPooledTransactionHashes68 { types: vec![], sizes: vec![], hashes: vec![] };
+        tx_hashes.push(&tx);
+        assert_eq!(tx_hashes.types.len(), 1);
+        assert_eq!(tx_hashes.sizes.len(), 1);
+        assert_eq!(tx_hashes.hashes.len(), 1);
+        assert_eq!(tx_hashes.types[0], tx.ty());
+        assert_eq!(tx_hashes.sizes[0], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[0], *tx.tx_hash());
+    }
+
+    #[test]
+    fn test_pooled_tx_hashes_68_extend() {
+        let tx = signed_transaction();
+        let txs = vec![tx.clone(), tx.clone()];
+        let mut tx_hashes =
+            NewPooledTransactionHashes68 { types: vec![], sizes: vec![], hashes: vec![] };
+        tx_hashes.extend(&txs);
+        assert_eq!(tx_hashes.types.len(), 2);
+        assert_eq!(tx_hashes.sizes.len(), 2);
+        assert_eq!(tx_hashes.hashes.len(), 2);
+        assert_eq!(tx_hashes.types[0], tx.ty());
+        assert_eq!(tx_hashes.sizes[0], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[0], *tx.tx_hash());
+        assert_eq!(tx_hashes.types[1], tx.ty());
+        assert_eq!(tx_hashes.sizes[1], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[1], *tx.tx_hash());
+    }
+
+    #[test]
+    fn test_pooled_tx_hashes_68_with_transaction() {
+        let tx = signed_transaction();
+        let tx_hashes =
+            NewPooledTransactionHashes68 { types: vec![], sizes: vec![], hashes: vec![] }
+                .with_transaction(&tx);
+        assert_eq!(tx_hashes.types.len(), 1);
+        assert_eq!(tx_hashes.sizes.len(), 1);
+        assert_eq!(tx_hashes.hashes.len(), 1);
+        assert_eq!(tx_hashes.types[0], tx.ty());
+        assert_eq!(tx_hashes.sizes[0], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[0], *tx.tx_hash());
+    }
+
+    #[test]
+    fn test_pooled_tx_hashes_68_with_transactions() {
+        let tx = signed_transaction();
+        let txs = vec![tx.clone(), tx.clone()];
+        let tx_hashes =
+            NewPooledTransactionHashes68 { types: vec![], sizes: vec![], hashes: vec![] }
+                .with_transactions(&txs);
+        assert_eq!(tx_hashes.types.len(), 2);
+        assert_eq!(tx_hashes.sizes.len(), 2);
+        assert_eq!(tx_hashes.hashes.len(), 2);
+        assert_eq!(tx_hashes.types[0], tx.ty());
+        assert_eq!(tx_hashes.sizes[0], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[0], *tx.tx_hash());
+        assert_eq!(tx_hashes.types[1], tx.ty());
+        assert_eq!(tx_hashes.sizes[1], tx.encode_2718_len());
+        assert_eq!(tx_hashes.hashes[1], *tx.tx_hash());
+    }
 }