Add `zcash_serialized_size()` to `ZcashSerialize` trait (#2824)

* add a zcash_serialized_size()

* add a size field to `UnminedTx`

* refactor zcash_serialized_size() to don't allocate RAM

* improve performance

Co-authored-by: teor <teor@riseup.net>

* clippy

Co-authored-by: teor <teor@riseup.net>

zebra-chain/src/serialization.rs

@@ -32,7 +32,7 @@ pub use zcash_deserialize::{
 };
 pub use zcash_serialize::{
     zcash_serialize_bytes, zcash_serialize_bytes_external_count, zcash_serialize_external_count,
-    ZcashSerialize, MAX_PROTOCOL_MESSAGE_LEN,
+    FakeWriter, ZcashSerialize, MAX_PROTOCOL_MESSAGE_LEN,
 };
 
 #[cfg(test)]

zebra-chain/src/serialization/tests/prop.rs

@@ -4,7 +4,10 @@ use proptest::prelude::*;
 
 use std::io::Cursor;
 
-use crate::serialization::{ReadZcashExt, WriteZcashExt};
+use crate::{
+    serialization::{ReadZcashExt, WriteZcashExt, ZcashSerialize},
+    transaction::UnminedTx,
+};
 
 proptest! {
     #[test]
@@ -35,4 +38,11 @@ proptest! {
             prop_assert_eq!(bytes, expect_bytes);
         }
     }
+
+    #[test]
+    fn transaction_serialized_size(transaction in any::<UnminedTx>()) {
+        zebra_test::init();
+
+        prop_assert_eq!(transaction.transaction.zcash_serialized_size().unwrap(), transaction.size);
+    }
 }

zebra-chain/src/serialization/zcash_serialize.rs

@@ -27,6 +27,29 @@ pub trait ZcashSerialize: Sized {
         self.zcash_serialize(&mut data)?;
         Ok(data)
     }
+
+    /// Get the size of `self` by using a fake writer.
+    fn zcash_serialized_size(&self) -> Result<usize, io::Error> {
+        let mut writer = FakeWriter(0);
+        self.zcash_serialize(&mut writer)
+            .expect("writer should never fail");
+        Ok(writer.0)
+    }
+}
+
+/// A fake writer helper used to get object lengths without allocating RAM.
+pub struct FakeWriter(pub usize);
+
+impl std::io::Write for FakeWriter {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.0 += buf.len();
+
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
 }
 
 /// Serialize a `Vec` as a compactsize number of items, then the items. This is

zebra-chain/src/transaction/arbitrary.rs

@@ -30,7 +30,7 @@ use crate::{
 
 use itertools::Itertools;
 
-use super::{FieldNotPresent, JoinSplitData, LockTime, Memo, Transaction};
+use super::{FieldNotPresent, JoinSplitData, LockTime, Memo, Transaction, UnminedTx};
 
 /// The maximum number of arbitrary transactions, inputs, or outputs.
 ///
@@ -768,6 +768,16 @@ impl Arbitrary for Transaction {
     type Strategy = BoxedStrategy<Self>;
 }
 
+impl Arbitrary for UnminedTx {
+    type Parameters = ();
+
+    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
+        any::<Transaction>().prop_map_into().boxed()
+    }
+
+    type Strategy = BoxedStrategy<Self>;
+}
+
 // Utility functions
 
 /// Convert `trans` into a fake v5 transaction,

zebra-chain/src/transaction/unmined.rs

@@ -17,6 +17,8 @@ use std::{fmt, sync::Arc};
 #[cfg(any(test, feature = "proptest-impl"))]
 use proptest_derive::Arbitrary;
 
+use crate::serialization::ZcashSerialize;
+
 use super::{
     AuthDigest, Hash,
     Transaction::{self, *},
@@ -144,13 +146,15 @@ impl UnminedTxId {
 
 /// An unmined transaction, and its pre-calculated unique identifying ID.
 #[derive(Clone, Debug, Eq, PartialEq)]
-#[cfg_attr(any(test, feature = "proptest-impl"), derive(Arbitrary))]
 pub struct UnminedTx {
     /// A unique identifier for this unmined transaction.
     pub id: UnminedTxId,
 
     /// The unmined transaction itself.
     pub transaction: Arc<Transaction>,
+
+    /// The size in bytes of the serialized transaction data
+    pub size: usize,
 }
 
 // Each of these conversions is implemented slightly differently,
@@ -160,6 +164,9 @@ impl From<Transaction> for UnminedTx {
     fn from(transaction: Transaction) -> Self {
         Self {
             id: (&transaction).into(),
+            size: transaction
+                .zcash_serialized_size()
+                .expect("all transactions have a size"),
             transaction: Arc::new(transaction),
         }
     }
@@ -170,6 +177,9 @@ impl From<&Transaction> for UnminedTx {
         Self {
             id: transaction.into(),
             transaction: Arc::new(transaction.clone()),
+            size: transaction
+                .zcash_serialized_size()
+                .expect("all transactions have a size"),
         }
     }
 }
@@ -178,6 +188,9 @@ impl From<Arc<Transaction>> for UnminedTx {
     fn from(transaction: Arc<Transaction>) -> Self {
         Self {
             id: transaction.as_ref().into(),
+            size: transaction
+                .zcash_serialized_size()
+                .expect("all transactions have a size"),
             transaction,
         }
     }
@@ -188,6 +201,9 @@ impl From<&Arc<Transaction>> for UnminedTx {
         Self {
             id: transaction.as_ref().into(),
             transaction: transaction.clone(),
+            size: transaction
+                .zcash_serialized_size()
+                .expect("all transactions have a size"),
         }
     }
 }

zebra-network/src/protocol/external/codec.rs

@@ -15,8 +15,9 @@ use zebra_chain::{
     block::{self, Block},
     parameters::Network,
     serialization::{
-        sha256d, zcash_deserialize_bytes_external_count, ReadZcashExt, SerializationError as Error,
-        WriteZcashExt, ZcashDeserialize, ZcashSerialize, MAX_PROTOCOL_MESSAGE_LEN,
+        sha256d, zcash_deserialize_bytes_external_count, FakeWriter, ReadZcashExt,
+        SerializationError as Error, WriteZcashExt, ZcashDeserialize, ZcashSerialize,
+        MAX_PROTOCOL_MESSAGE_LEN,
     },
     transaction::Transaction,
 };
@@ -181,21 +182,8 @@ impl Codec {
     /// for large data structures like lists, blocks, and transactions.
     /// See #1774.
     fn body_length(&self, msg: &Message) -> usize {
-        struct FakeWriter(usize);
+        let mut writer = FakeWriter { 0: 0 };
 
-        impl std::io::Write for FakeWriter {
-            fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
-                self.0 += buf.len();
-
-                Ok(buf.len())
-            }
-
-            fn flush(&mut self) -> std::io::Result<()> {
-                Ok(())
-            }
-        }
-
-        let mut writer = FakeWriter(0);
         self.write_body(msg, &mut writer)
             .expect("writer should never fail");
         writer.0

zebrad/src/components/inbound/tests.rs

@@ -134,11 +134,7 @@ async fn mempool_advertise_transaction_ids() -> Result<(), crate::BoxError> {
         peer_set
             .expect_request(Request::TransactionsById(txs))
             .map(|responder| {
-                let unmined_transaction = UnminedTx {
-                    id: test_transaction_id,
-                    transaction: test_transaction,
-                };
-
+                let unmined_transaction = UnminedTx::from(test_transaction);
                 responder.respond(Response::Transactions(vec![unmined_transaction]))
             });
     // Simulate a successful transaction verification