Simplify preallocate tests

zebra-chain/src/block/tests/preallocate.rs

@@ -9,7 +9,10 @@ use crate::{
         header::MIN_COUNTED_HEADER_LEN, CountedHeader, Hash, Header, BLOCK_HASH_SIZE,
         MAX_PROTOCOL_MESSAGE_LEN,
     },
-    serialization::{CompactSizeMessage, TrustedPreallocate, ZcashSerialize},
+    serialization::{
+        arbitrary::max_allocation_is_big_enough, CompactSizeMessage, TrustedPreallocate,
+        ZcashSerialize,
+    },
 };
 
 proptest! {
@@ -25,28 +28,22 @@ proptest! {
     /// 2. The largest allowed vector is small enough to fit in a legal Zcash Wire Protocol message
     #[test]
     fn block_hash_max_allocation(hash in Hash::arbitrary_with(())) {
-        let max_allocation: usize = Hash::max_allocation().try_into().unwrap();
-        let mut smallest_disallowed_vec = Vec::with_capacity(max_allocation + 1);
-        for _ in 0..(Hash::max_allocation()+1) {
-            smallest_disallowed_vec.push(hash);
-        }
+        let (
+            smallest_disallowed_vec_len,
+            smallest_disallowed_serialized_len,
+            largest_allowed_vec_len,
+            largest_allowed_serialized_len,
+        ) = max_allocation_is_big_enough(hash);
 
-        let smallest_disallowed_serialized = smallest_disallowed_vec.zcash_serialize_to_vec().expect("Serialization to vec must succeed");
         // Check that our smallest_disallowed_vec is only one item larger than the limit
-        prop_assert!(((smallest_disallowed_vec.len() - 1) as u64) == Hash::max_allocation());
+        prop_assert!(((smallest_disallowed_vec_len - 1) as u64) == Hash::max_allocation());
         // Check that our smallest_disallowed_vec is too big to send as a protocol message
-        prop_assert!(smallest_disallowed_serialized.len() > MAX_PROTOCOL_MESSAGE_LEN);
-
-        // Create largest_allowed_vec by removing one element from smallest_disallowed_vec without copying (for efficiency)
-        smallest_disallowed_vec.pop();
-        let largest_allowed_vec = smallest_disallowed_vec;
-        let largest_allowed_serialized = largest_allowed_vec.zcash_serialize_to_vec().expect("Serialization to vec must succeed");
+        prop_assert!(smallest_disallowed_serialized_len > MAX_PROTOCOL_MESSAGE_LEN);
 
         // Check that our largest_allowed_vec contains the maximum number of hashes
-        prop_assert!((largest_allowed_vec.len() as u64) == Hash::max_allocation());
+        prop_assert!((largest_allowed_vec_len as u64) == Hash::max_allocation());
         // Check that our largest_allowed_vec is small enough to send as a protocol message
-        prop_assert!(largest_allowed_serialized.len() <= MAX_PROTOCOL_MESSAGE_LEN);
-
+        prop_assert!(largest_allowed_serialized_len <= MAX_PROTOCOL_MESSAGE_LEN);
     }
 
     /// Confirm that each counted header takes at least COUNTED_HEADER_LEN bytes when serialized.
@@ -72,28 +69,24 @@ proptest! {
     fn counted_header_max_allocation(header in any::<Header>()) {
         let header = CountedHeader {
             header,
-            transaction_count: 0.try_into().expect("zero fits in MAX_PROTOCOL_MESSAGE_LEN"),
+            transaction_count: 0.try_into().expect("zero is less than MAX_PROTOCOL_MESSAGE_LEN"),
         };
-        let max_allocation: usize = CountedHeader::max_allocation().try_into().unwrap();
-        let mut smallest_disallowed_vec = Vec::with_capacity(max_allocation + 1);
-        for _ in 0..(CountedHeader::max_allocation()+1) {
-            smallest_disallowed_vec.push(header);
-        }
-        let smallest_disallowed_serialized = smallest_disallowed_vec.zcash_serialize_to_vec().expect("Serialization to vec must succeed");
-        // Check that our smallest_disallowed_vec is only one item larger than the limit
-        prop_assert!(((smallest_disallowed_vec.len() - 1) as u64) == CountedHeader::max_allocation());
-        // Check that our smallest_disallowed_vec is too big to send as a protocol message
-        prop_assert!(smallest_disallowed_serialized.len() > MAX_PROTOCOL_MESSAGE_LEN);
 
+        let (
+            smallest_disallowed_vec_len,
+            smallest_disallowed_serialized_len,
+            largest_allowed_vec_len,
+            largest_allowed_serialized_len,
+        ) = max_allocation_is_big_enough(header);
 
-        // Create largest_allowed_vec by removing one element from smallest_disallowed_vec without copying (for efficiency)
-        smallest_disallowed_vec.pop();
-        let largest_allowed_vec = smallest_disallowed_vec;
-        let largest_allowed_serialized = largest_allowed_vec.zcash_serialize_to_vec().expect("Serialization to vec must succeed");
+        // Check that our smallest_disallowed_vec is only one item larger than the limit
+        prop_assert!(((smallest_disallowed_vec_len - 1) as u64) == CountedHeader::max_allocation());
+        // Check that our smallest_disallowed_vec is too big to send as a protocol message
+        prop_assert!(smallest_disallowed_serialized_len > MAX_PROTOCOL_MESSAGE_LEN);
 
         // Check that our largest_allowed_vec contains the maximum number of CountedHeaders
-        prop_assert!((largest_allowed_vec.len() as u64) == CountedHeader::max_allocation());
+        prop_assert!((largest_allowed_vec_len as u64) == CountedHeader::max_allocation());
         // Check that our largest_allowed_vec is small enough to send as a protocol message
-        prop_assert!(largest_allowed_serialized.len() <= MAX_PROTOCOL_MESSAGE_LEN);
+        prop_assert!(largest_allowed_serialized_len <= MAX_PROTOCOL_MESSAGE_LEN);
     }
 }

zebra-chain/src/orchard/tests/preallocate.rs

@@ -7,7 +7,7 @@ use crate::{
         Action, AuthorizedAction,
     },
     primitives::redpallas::{Signature, SpendAuth},
-    serialization::{tests::max_allocation_is_big_enough, TrustedPreallocate, ZcashSerialize},
+    serialization::{arbitrary::max_allocation_is_big_enough, TrustedPreallocate, ZcashSerialize},
 };
 
 use proptest::{prelude::*, proptest};

zebra-chain/src/sapling/tests/preallocate.rs

@@ -14,7 +14,7 @@ use crate::{
         },
         PerSpendAnchor, SharedAnchor,
     },
-    serialization::{tests::max_allocation_is_big_enough, TrustedPreallocate, ZcashSerialize},
+    serialization::{arbitrary::max_allocation_is_big_enough, TrustedPreallocate, ZcashSerialize},
 };
 
 use proptest::prelude::*;

zebra-chain/src/serialization/arbitrary.rs

@@ -5,7 +5,9 @@ use std::convert::TryInto;
 use chrono::{TimeZone, Utc, MAX_DATETIME, MIN_DATETIME};
 use proptest::{arbitrary::any, prelude::*};
 
-use super::{CompactSizeMessage, DateTime32, MAX_PROTOCOL_MESSAGE_LEN};
+use super::{
+    CompactSizeMessage, DateTime32, TrustedPreallocate, ZcashSerialize, MAX_PROTOCOL_MESSAGE_LEN,
+};
 
 impl Arbitrary for DateTime32 {
     type Parameters = ();
@@ -72,3 +74,40 @@ impl Arbitrary for CompactSizeMessage {
 
     type Strategy = BoxedStrategy<Self>;
 }
+
+/// Allocate a maximum-sized vector of cloned `item`s, and serialize that array.
+///
+/// Returns the following calculations on `item`:
+///   smallest_disallowed_vec_len
+///   smallest_disallowed_serialized_len
+///   largest_allowed_vec_len
+///   largest_allowed_serialized_len
+///
+/// For varible-size types, `largest_allowed_serialized_len` might not fit within
+/// `MAX_PROTOCOL_MESSAGE_LEN` or `MAX_BLOCK_SIZE`.
+pub fn max_allocation_is_big_enough<T>(item: T) -> (usize, usize, usize, usize)
+where
+    T: TrustedPreallocate + ZcashSerialize + Clone,
+{
+    let max_allocation: usize = T::max_allocation().try_into().unwrap();
+    let mut smallest_disallowed_vec = vec![item; max_allocation + 1];
+
+    let smallest_disallowed_serialized = smallest_disallowed_vec
+        .zcash_serialize_to_vec()
+        .expect("Serialization to vec must succeed");
+    let smallest_disallowed_vec_len = smallest_disallowed_vec.len();
+
+    // Create largest_allowed_vec by removing one element from smallest_disallowed_vec without copying (for efficiency)
+    smallest_disallowed_vec.pop();
+    let largest_allowed_vec = smallest_disallowed_vec;
+    let largest_allowed_serialized = largest_allowed_vec
+        .zcash_serialize_to_vec()
+        .expect("Serialization to vec must succeed");
+
+    (
+        smallest_disallowed_vec_len,
+        smallest_disallowed_serialized.len(),
+        largest_allowed_vec.len(),
+        largest_allowed_serialized.len(),
+    )
+}

zebra-chain/src/serialization/tests.rs

@@ -1,4 +1,2 @@
 mod preallocate;
 mod prop;
-
-pub use preallocate::max_allocation_is_big_enough;

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

@@ -5,52 +5,31 @@ use proptest::{collection::size_range, prelude::*};
 use std::{convert::TryInto, matches};
 
 use crate::serialization::{
-    zcash_deserialize::MAX_U8_ALLOCATION, SerializationError, TrustedPreallocate, ZcashDeserialize,
-    ZcashSerialize, MAX_PROTOCOL_MESSAGE_LEN,
+    arbitrary::max_allocation_is_big_enough, zcash_deserialize::MAX_U8_ALLOCATION,
+    SerializationError, TrustedPreallocate, ZcashDeserialize, ZcashSerialize,
+    MAX_PROTOCOL_MESSAGE_LEN,
 };
 
 // Allow direct serialization of Vec<u8> for these tests. We don't usually
 // allow this because some types have specific rules for about serialization
 // of their inner Vec<u8>. This method could be easily misused if it applied
 // more generally.
+//
+// Due to Rust's trait rules, these trait impls apply to all zebra-chain tests,
+// not just the tests in this module. But other crates' tests can't access them.
 impl ZcashSerialize for u8 {
     fn zcash_serialize<W: std::io::Write>(&self, mut writer: W) -> Result<(), std::io::Error> {
         writer.write_all(&[*self])
     }
 }
 
-/// Return the following calculations on `item`:
-///   smallest_disallowed_vec_len
-///   smallest_disallowed_serialized_len
-///   largest_allowed_vec_len
-///   largest_allowed_serialized_len
-pub fn max_allocation_is_big_enough<T>(item: T) -> (usize, usize, usize, usize)
-where
-    T: TrustedPreallocate + ZcashSerialize + Clone,
-{
-    let max_allocation: usize = T::max_allocation().try_into().unwrap();
-    let mut smallest_disallowed_vec = Vec::with_capacity(max_allocation + 1);
-    for _ in 0..(max_allocation + 1) {
-        smallest_disallowed_vec.push(item.clone());
+impl TrustedPreallocate for u8 {
+    fn max_allocation() -> u64 {
+        // MAX_PROTOCOL_MESSAGE_LEN takes up 5 bytes when encoded as a CompactSize.
+        (MAX_PROTOCOL_MESSAGE_LEN - 5)
+            .try_into()
+            .expect("MAX_PROTOCOL_MESSAGE_LEN fits in u64")
     }
-    let smallest_disallowed_serialized = smallest_disallowed_vec
-        .zcash_serialize_to_vec()
-        .expect("Serialization to vec must succeed");
-    let smallest_disallowed_vec_len = smallest_disallowed_vec.len();
-
-    // Create largest_allowed_vec by removing one element from smallest_disallowed_vec without copying (for efficiency)
-    smallest_disallowed_vec.pop();
-    let largest_allowed_vec = smallest_disallowed_vec;
-    let largest_allowed_serialized = largest_allowed_vec
-        .zcash_serialize_to_vec()
-        .expect("Serialization to vec must succeed");
-
-    (
-        smallest_disallowed_vec_len,
-        smallest_disallowed_serialized.len(),
-        largest_allowed_vec.len(),
-        largest_allowed_serialized.len(),
-    )
 }
 
 proptest! {
@@ -111,25 +90,20 @@ fn u8_size_is_correct() {
 /// 1. The smallest disallowed `Vec<u8>` is too big to include in a Zcash Wire Protocol message
 /// 2. The largest allowed `Vec<u8>`is exactly the size of a maximal Zcash Wire Protocol message
 fn u8_max_allocation_is_correct() {
-    let mut shortest_disallowed_vec = vec![0u8; MAX_U8_ALLOCATION + 1];
-    let shortest_disallowed_serialized = shortest_disallowed_vec
-        .zcash_serialize_to_vec()
-        .expect("Serialization to vec must succeed");
+    let (
+        smallest_disallowed_vec_len,
+        smallest_disallowed_serialized_len,
+        largest_allowed_vec_len,
+        largest_allowed_serialized_len,
+    ) = max_allocation_is_big_enough(0u8);
 
     // Confirm that shortest_disallowed_vec is only one item larger than the limit
-    assert_eq!((shortest_disallowed_vec.len() - 1), MAX_U8_ALLOCATION);
+    assert_eq!((smallest_disallowed_vec_len - 1), MAX_U8_ALLOCATION);
     // Confirm that shortest_disallowed_vec is too large to be included in a valid zcash message
-    assert!(shortest_disallowed_serialized.len() > MAX_PROTOCOL_MESSAGE_LEN);
-
-    // Create largest_allowed_vec by removing one element from smallest_disallowed_vec without copying (for efficiency)
-    shortest_disallowed_vec.pop();
-    let longest_allowed_vec = shortest_disallowed_vec;
-    let longest_allowed_serialized = longest_allowed_vec
-        .zcash_serialize_to_vec()
-        .expect("serialization to vec must succed");
+    assert!(smallest_disallowed_serialized_len > MAX_PROTOCOL_MESSAGE_LEN);
 
     // Check that our largest_allowed_vec contains the maximum number of items
-    assert_eq!(longest_allowed_vec.len(), MAX_U8_ALLOCATION);
+    assert_eq!(largest_allowed_vec_len, MAX_U8_ALLOCATION);
     // Check that our largest_allowed_vec is the size of a maximal protocol message
-    assert_eq!(longest_allowed_serialized.len(), MAX_PROTOCOL_MESSAGE_LEN);
+    assert_eq!(largest_allowed_serialized_len, MAX_PROTOCOL_MESSAGE_LEN);
 }

zebra-chain/src/serialization/zcash_deserialize.rs

@@ -2,6 +2,7 @@ use std::{
     convert::{TryFrom, TryInto},
     io,
     net::Ipv6Addr,
+    sync::Arc,
 };
 
 use super::{AtLeastOne, CompactSizeMessage, SerializationError, MAX_PROTOCOL_MESSAGE_LEN};
@@ -171,6 +172,15 @@ pub trait TrustedPreallocate {
     fn max_allocation() -> u64;
 }
 
+impl<T> TrustedPreallocate for Arc<T>
+where
+    T: TrustedPreallocate,
+{
+    fn max_allocation() -> u64 {
+        T::max_allocation()
+    }
+}
+
 /// The length of the longest valid `Vec<u8>` that can be received over the network
 ///
 /// It takes 5 bytes to encode a CompactSize representing any number netween 2^16 and (2^32 - 1)