chore(net): Add proptest roundtrip to rlp types (#829)

crates/storage/codecs/derive/src/arbitrary.rs

@@ -0,0 +1,73 @@
+use proc_macro::{self, TokenStream};
+use proc_macro2::TokenStream as TokenStream2;
+use quote::{format_ident, quote};
+use syn::DeriveInput;
+
+/// If `compact` or `rlp` is passed to `derive_arbitrary`, this function will generate the
+/// corresponding proptest roundtrip tests.
+///
+/// It accepts an optional integer number for the number of proptest cases. Otherwise, it will set
+/// it at 1000.
+pub fn maybe_generate_tests(args: TokenStream, ast: &DeriveInput) -> TokenStream2 {
+    let type_ident = ast.ident.clone();
+
+    // Same as proptest
+    let mut default_cases = 256;
+
+    let mut traits = vec![];
+    let mut roundtrips = vec![];
+
+    for arg in args {
+        if arg.to_string() == "compact" {
+            traits.push(quote! { use super::Compact; });
+            roundtrips.push(quote! {
+                {
+                    let mut buf = vec![];
+
+                    let len = field.clone().to_compact(&mut buf);
+                    let (decoded, _) = super::#type_ident::from_compact(&buf, len);
+
+                    assert!(field == decoded);
+                }
+            });
+        } else if arg.to_string() == "rlp" {
+            traits.push(quote! { use reth_rlp::{Encodable, Decodable}; });
+            roundtrips.push(quote! {
+                {
+                    let mut buf = vec![];
+
+                    let len = field.clone().encode(&mut buf);
+                    let decoded = super::#type_ident::decode(&mut buf.as_slice()).unwrap();
+
+                    assert!(field == decoded);
+                }
+            });
+        } else if let Ok(num) = arg.to_string().parse() {
+            default_cases = num;
+        }
+    }
+
+    let mut tests = TokenStream2::default();
+    if !roundtrips.is_empty() {
+        let mod_tests = format_ident!("{}Tests", ast.ident);
+
+        tests = quote! {
+            #[allow(non_snake_case)]
+            #[cfg(test)]
+            mod #mod_tests {
+                #(#traits)*
+
+                #[test]
+                fn proptest() {
+                    let mut config = proptest::prelude::ProptestConfig::with_cases(#default_cases as u32);
+
+                    proptest::proptest!(config, |(field: super::#type_ident)| {
+                        #(#roundtrips)*
+                    });
+                }
+            }
+        }
+    }
+
+    tests
+}

crates/storage/codecs/derive/src/lib.rs

@@ -1,7 +1,8 @@
-use proc_macro::{self, TokenStream};
+use proc_macro::{self, TokenStream, TokenTree};
 use quote::{format_ident, quote};
 use syn::{parse_macro_input, DeriveInput};
 
+mod arbitrary;
 mod compact;
 
 #[proc_macro_derive(Compact, attributes(maybe_zero))]
@@ -9,6 +10,11 @@ pub fn derive(input: TokenStream) -> TokenStream {
     compact::derive(input)
 }
 
+/// Implements the main codec. If the codec supports it, it will call `derive_arbitrary(..)`.
+/// 
+/// Example usage:
+/// * `#[main_codec(rlp)]`: will implement `derive_arbitrary(rlp)` or `derive_arbitrary(compact, rlp)`, if `compact` is the `main_codec`.
+/// * `#[main_codec(no_arbitrary)]`: will skip `derive_arbitrary`
 #[proc_macro_attribute]
 #[rustfmt::skip]
 #[allow(unreachable_code)]
@@ -70,28 +76,48 @@ pub fn use_postcard(_args: TokenStream, input: TokenStream) -> TokenStream {
 }
 
 #[proc_macro_attribute]
-pub fn use_compact(_args: TokenStream, input: TokenStream) -> TokenStream {
+pub fn use_compact(args: TokenStream, input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as DeriveInput);
+    let compact = quote! {
+        #[derive(Compact, serde::Serialize, serde::Deserialize)]
+        #ast
+    }
+    .into();
 
-    derive_compact_arbitrary(
-        _args,
-        quote! {
-            #[derive(Compact, serde::Serialize, serde::Deserialize)]
-            #ast
+    if let Some(first_arg) = args.clone().into_iter().next() {
+        if first_arg.to_string() == "no_arbitrary" {
+            return compact
         }
-        .into(),
-    )
+    }
+
+    let mut args = args.into_iter().collect::<Vec<_>>();
+    args.push(TokenTree::Ident(proc_macro::Ident::new("compact", proc_macro::Span::call_site())));
+
+    derive_arbitrary(TokenStream::from_iter(args.into_iter()), compact)
 }
 
+/// Adds `Arbitrary` and `proptest::Arbitrary` imports into scope and derives the struct/enum.
+///
+/// If `compact` or `rlp` is passed to `derive_arbitrary`, there will be proptest roundtrip tests
+/// generated. An integer value passed will limit the number of proptest cases generated (default:
+/// 256).
+///
+/// Examples:
+/// * `#[derive_arbitrary]`: will derive arbitrary with no tests.
+/// * `#[derive_arbitrary(rlp)]`: will derive arbitrary and generate rlp roundtrip proptests.
+/// * `#[derive_arbitrary(rlp, 10)]`: will derive arbitrary and generate rlp roundtrip proptests.
+///   Limited to 10 cases.
+/// * `#[derive_arbitrary(compact, rlp)]`. will derive arbitrary and generate rlp and compact
+///   roundtrip proptests.
 #[proc_macro_attribute]
-pub fn derive_compact_arbitrary(_args: TokenStream, input: TokenStream) -> TokenStream {
+pub fn derive_arbitrary(args: TokenStream, input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as DeriveInput);
 
+    let tests = arbitrary::maybe_generate_tests(args, &ast);
+
     // Avoid duplicate names
     let prop_import = format_ident!("{}PropTestArbitratry", ast.ident);
     let arb_import = format_ident!("{}Arbitratry", ast.ident);
-    let mod_tests = format_ident!("{}Tests", ast.ident);
-    let type_ident = ast.ident.clone();
 
     quote! {
         #[cfg(any(test, feature = "arbitrary"))]
@@ -103,23 +129,19 @@ pub fn derive_compact_arbitrary(_args: TokenStream, input: TokenStream) -> Token
         #[cfg_attr(any(test, feature = "arbitrary"), derive(#prop_import, #arb_import))]
         #ast
 
-        #[allow(non_snake_case)]
-        #[cfg(test)]
-        mod #mod_tests {
-            use super::Compact;
+        #tests
+    }
+    .into()
+}
 
-            #[test]
-            fn proptest() {
-                proptest::proptest!(|(field: super::#type_ident)| {
-                    let mut buf = vec![];
-
-                    let len = field.clone().to_compact(&mut buf);
-                    let (decoded, _) = super::#type_ident::from_compact(&buf, len);
-
-                    assert!(field == decoded);
-                });
-            }
-        }
+/// To be used for types that implement `Arbitrary` manually. See [`derive_arbitrary`] for more.
+#[proc_macro_attribute]
+pub fn add_arbitrary_tests(args: TokenStream, input: TokenStream) -> TokenStream {
+    let ast = parse_macro_input!(input as DeriveInput);
+    let tests = arbitrary::maybe_generate_tests(args, &ast);
+    quote! {
+        #ast
+        #tests
     }
     .into()
 }