From 9a7a733a087a01b243a12890e6630b4e0312638d Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 22 Nov 2024 20:41:46 +0100
Subject: [PATCH] feat: impl InMemorySize for PooledTx (#12791)

---
 crates/primitives-traits/src/size.rs         | 19 ++++++++++++++++++-
 crates/primitives/src/transaction/pooled.rs  | 13 +++++++++++++
 crates/primitives/src/transaction/sidecar.rs | 11 +++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/crates/primitives-traits/src/size.rs b/crates/primitives-traits/src/size.rs
index da3b39888f46..4d721dd00b30 100644
--- a/crates/primitives-traits/src/size.rs
+++ b/crates/primitives-traits/src/size.rs
@@ -10,15 +10,32 @@ pub trait InMemorySize {
 
 impl<T: InMemorySize> InMemorySize for alloy_consensus::Signed<T> {
     fn size(&self) -> usize {
-        T::size(self.tx()) + core::mem::size_of::<Signature>() + core::mem::size_of::<TxHash>()
+        T::size(self.tx()) + self.signature().size() + self.hash().size()
     }
 }
 
+/// Implement `InMemorySize` for a type with `size_of`
+macro_rules! impl_in_mem_size_size_of {
+    ($($ty:ty),*) => {
+        $(
+            impl InMemorySize for $ty {
+                #[inline]
+                fn size(&self) -> usize {
+                    core::mem::size_of::<Self>()
+                }
+            }
+        )*
+    };
+}
+
+impl_in_mem_size_size_of!(Signature, TxHash);
+
 /// Implement `InMemorySize` for a type with a native `size` method.
 macro_rules! impl_in_mem_size {
     ($($ty:ty),*) => {
         $(
             impl InMemorySize for $ty {
+                #[inline]
                 fn size(&self) -> usize {
                    Self::size(self)
                 }
diff --git a/crates/primitives/src/transaction/pooled.rs b/crates/primitives/src/transaction/pooled.rs
index bb8406147034..09111c61a170 100644
--- a/crates/primitives/src/transaction/pooled.rs
+++ b/crates/primitives/src/transaction/pooled.rs
@@ -20,6 +20,7 @@ use alloy_primitives::{
 use alloy_rlp::{Decodable, Encodable, Error as RlpError, Header};
 use bytes::Buf;
 use derive_more::{AsRef, Deref};
+use reth_primitives_traits::InMemorySize;
 use serde::{Deserialize, Serialize};
 
 /// A response to `GetPooledTransactions`. This can include either a blob transaction, or a
@@ -559,6 +560,18 @@ impl alloy_consensus::Transaction for PooledTransactionsElement {
     }
 }
 
+impl InMemorySize for PooledTransactionsElement {
+    fn size(&self) -> usize {
+        match self {
+            Self::Legacy(tx) => tx.size(),
+            Self::Eip2930(tx) => tx.size(),
+            Self::Eip1559(tx) => tx.size(),
+            Self::Eip7702(tx) => tx.size(),
+            Self::BlobTransaction(tx) => tx.size(),
+        }
+    }
+}
+
 impl TryFrom<TransactionSigned> for PooledTransactionsElement {
     type Error = TransactionConversionError;
 
diff --git a/crates/primitives/src/transaction/sidecar.rs b/crates/primitives/src/transaction/sidecar.rs
index c1b1b029afcb..2cf04bc8e741 100644
--- a/crates/primitives/src/transaction/sidecar.rs
+++ b/crates/primitives/src/transaction/sidecar.rs
@@ -4,6 +4,7 @@ use crate::{Transaction, TransactionSigned};
 use alloy_consensus::{transaction::RlpEcdsaTx, Signed, TxEip4844WithSidecar};
 use alloy_eips::eip4844::BlobTransactionSidecar;
 use derive_more::Deref;
+use reth_primitives_traits::InMemorySize;
 use serde::{Deserialize, Serialize};
 
 /// A response to `GetPooledTransactions` that includes blob data, their commitments, and their
@@ -73,6 +74,16 @@ impl BlobTransaction {
     }
 }
 
+impl InMemorySize for BlobTransaction {
+    fn size(&self) -> usize {
+        // TODO(mattsse): replace with next alloy bump
+        self.0.hash().size() +
+            self.0.signature().size() +
+            self.0.tx().tx().size() +
+            self.0.tx().sidecar.size()
+    }
+}
+
 #[cfg(all(test, feature = "c-kzg"))]
 mod tests {
     use super::*;