From 01bfd386375f69e20bc157c8bc26217077b8e97b Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sat, 16 Mar 2019 14:30:21 +1100
Subject: [PATCH] Add `default` attribute to `TestRandom` derive.

Allows for generating the item from default instead of randomizing it.
---
 eth2/utils/test_random_derive/src/lib.rs | 48 +++++++++++++++++-------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/eth2/utils/test_random_derive/src/lib.rs b/eth2/utils/test_random_derive/src/lib.rs
index 9a456606c..7920ea695 100644
--- a/eth2/utils/test_random_derive/src/lib.rs
+++ b/eth2/utils/test_random_derive/src/lib.rs
@@ -4,7 +4,20 @@ use crate::proc_macro::TokenStream;
 use quote::quote;
 use syn::{parse_macro_input, DeriveInput};
 
-#[proc_macro_derive(TestRandom)]
+/// Returns true if some field has an attribute declaring it should be generated from default (not
+/// randomized).
+///
+/// The field attribute is: `#[test_random(default)]`
+fn should_use_default(field: &syn::Field) -> bool {
+    for attr in &field.attrs {
+        if attr.tts.to_string() == "( default )" {
+            return true;
+        }
+    }
+    false
+}
+
+#[proc_macro_derive(TestRandom, attributes(test_random))]
 pub fn test_random_derive(input: TokenStream) -> TokenStream {
     let derived_input = parse_macro_input!(input as DeriveInput);
     let name = &derived_input.ident;
@@ -14,14 +27,32 @@ pub fn test_random_derive(input: TokenStream) -> TokenStream {
         _ => panic!("test_random_derive only supports structs."),
     };
 
-    let field_names = get_named_field_idents(&struct_data);
+    // Build quotes for fields that should be generated and those that should be built from
+    // `Default`.
+    let mut quotes = vec![];
+    for field in &struct_data.fields {
+        match &field.ident {
+            Some(ref ident) => {
+                if should_use_default(field) {
+                    quotes.push(quote! {
+                        #ident: <_>::default(),
+                    });
+                } else {
+                    quotes.push(quote! {
+                        #ident: <_>::random_for_test(rng),
+                    });
+                }
+            }
+            _ => panic!("test_random_derive only supports named struct fields."),
+        };
+    }
 
     let output = quote! {
         impl<T: RngCore> TestRandom<T> for #name {
             fn random_for_test(rng: &mut T) -> Self {
                Self {
                     #(
-                        #field_names: <_>::random_for_test(rng),
+                        #quotes
                     )*
                }
             }
@@ -30,14 +61,3 @@ pub fn test_random_derive(input: TokenStream) -> TokenStream {
 
     output.into()
 }
-
-fn get_named_field_idents(struct_data: &syn::DataStruct) -> Vec<(&syn::Ident)> {
-    struct_data
-        .fields
-        .iter()
-        .map(|f| match &f.ident {
-            Some(ref ident) => ident,
-            _ => panic!("test_random_derive only supports named struct fields."),
-        })
-        .collect()
-}