Merge pull request #129 from greyblake/customize-fields

Support customization of fields on derive

CHANGELOG.md

@@ -32,6 +32,10 @@ Released YYYY-MM-DD.
 
 Released 2022-09-08.
 
+### Added
+
+* Support custom arbitrary implementation for fields on derive. [#129](https://github.com/rust-fuzz/arbitrary/pull/129)
+
 ### Fixed
 
 * Fixed a potential panic due to an off-by-one error in the `Arbitrary`

Cargo.toml

@@ -22,8 +22,6 @@ rust-version = "1.63.0"
 [dependencies]
 derive_arbitrary = { version = "1.1.6", path = "./derive", optional = true }
 
-[dev-dependencies]
-
 [features]
 # Turn this feature on to enable support for `#[derive(Arbitrary)]`.
 derive = ["derive_arbitrary"]

README.md

@@ -61,6 +61,39 @@ pub struct Rgb {
 }
 ```
 
+#### Customizing single fields
+
+This can be particular handy if your structure uses a type that does not implement `Arbitrary` or you want to have more customization for particular fields.
+
+```rust
+#[derive(Arbitrary)]
+pub struct Rgba {
+    // set `r` to Default::default()
+    #[arbitrary(default)]
+    pub r: u8,
+
+    // set `g` to 255
+    #[arbitrary(value = 255)]
+    pub g: u8,
+
+    // Generate `b` with a custom function of type
+    //
+    //    fn(&mut Unstructured) -> arbitrary::Result<T>
+    //
+    // where `T` is the field's type.
+    #[arbitrary(with = arbitrary_b)]
+    pub b: u8,
+
+    // Generate `a` with a custom closure (shortuct to avoid a custom funciton)
+    #[arbitrary(with = |u: &mut Unstructured| u.int_in_range(0..=64))]
+    pub a: u8,
+}
+
+fn arbitrary_b(u: &mut Unstructured) -> arbitrary::Result<u8> {
+    u.int_in_range(64..=128)
+}
+```
+
 ### Implementing `Arbitrary` By Hand
 
 Alternatively, you can write an `Arbitrary` implementation by hand:

derive/Cargo.toml

@@ -9,7 +9,7 @@ authors = [
     "Corey Farwell <coreyf@rwell.org>",
 ]
 categories = ["development-tools::testing"]
-edition = "2018"
+edition = "2021"
 keywords = ["arbitrary", "testing", "derive", "macro"]
 readme = "README.md"
 description = "Derives arbitrary traits"

derive/src/field_attributes.rs

@@ -0,0 +1,117 @@
+use proc_macro2::{Group, Span, TokenStream, TokenTree};
+use quote::quote;
+use syn::{spanned::Spanned, *};
+
+/// Used to filter out necessary field attribute and within error messages.
+static ARBITRARY_ATTRIBUTE_NAME: &str = "arbitrary";
+
+/// Determines how a value for a field should be constructed.
+#[cfg_attr(test, derive(Debug))]
+pub enum FieldConstructor {
+    /// Assume that Arbitrary is defined for the type of this field and use it (default)
+    Arbitrary,
+
+    /// Places `Default::default()` as a field value.
+    Default,
+
+    /// Use custom function or closure to generate a value for a field.
+    With(TokenStream),
+
+    /// Set a field always to the given value.
+    Value(TokenStream),
+}
+
+pub fn determine_field_constructor(field: &Field) -> Result<FieldConstructor> {
+    let opt_attr = fetch_attr_from_field(field)?;
+    let ctor = match opt_attr {
+        Some(attr) => parse_attribute(attr)?,
+        None => FieldConstructor::Arbitrary,
+    };
+    Ok(ctor)
+}
+
+fn fetch_attr_from_field(field: &Field) -> Result<Option<&Attribute>> {
+    let found_attributes: Vec<_> = field
+        .attrs
+        .iter()
+        .filter(|a| {
+            let path = &a.path;
+            let name = quote!(#path).to_string();
+            name == ARBITRARY_ATTRIBUTE_NAME
+        })
+        .collect();
+    if found_attributes.len() > 1 {
+        let name = field.ident.as_ref().unwrap();
+        let msg = format!(
+            "Multiple conflicting #[{ARBITRARY_ATTRIBUTE_NAME}] attributes found on field `{name}`"
+        );
+        return Err(syn::Error::new(field.span(), msg));
+    }
+    Ok(found_attributes.into_iter().next())
+}
+
+fn parse_attribute(attr: &Attribute) -> Result<FieldConstructor> {
+    let group = {
+        let mut tokens_iter = attr.clone().tokens.into_iter();
+        let token = tokens_iter.next().ok_or_else(|| {
+            let msg = format!("#[{ARBITRARY_ATTRIBUTE_NAME}] cannot be empty.");
+            syn::Error::new(attr.span(), msg)
+        })?;
+        match token {
+            TokenTree::Group(g) => g,
+            t => {
+                let msg = format!("#[{ARBITRARY_ATTRIBUTE_NAME}] must contain a group, got: {t})");
+                return Err(syn::Error::new(attr.span(), msg));
+            }
+        }
+    };
+    parse_attribute_internals(group)
+}
+
+fn parse_attribute_internals(group: Group) -> Result<FieldConstructor> {
+    let stream = group.stream();
+    let mut tokens_iter = stream.into_iter();
+    let token = tokens_iter.next().ok_or_else(|| {
+        let msg = format!("#[{ARBITRARY_ATTRIBUTE_NAME}] cannot be empty.");
+        syn::Error::new(group.span(), msg)
+    })?;
+    match token.to_string().as_ref() {
+        "default" => Ok(FieldConstructor::Default),
+        "with" => {
+            let func_path = parse_assigned_value("with", tokens_iter, group.span())?;
+            Ok(FieldConstructor::With(func_path))
+        }
+        "value" => {
+            let value = parse_assigned_value("value", tokens_iter, group.span())?;
+            Ok(FieldConstructor::Value(value))
+        }
+        _ => {
+            let msg = format!("Unknown option for #[{ARBITRARY_ATTRIBUTE_NAME}]: `{token}`");
+            Err(syn::Error::new(token.span(), msg))
+        }
+    }
+}
+
+// Input:
+//     = 2 + 2
+// Output:
+//     2 + 2
+fn parse_assigned_value(
+    opt_name: &str,
+    mut tokens_iter: impl Iterator<Item = TokenTree>,
+    default_span: Span,
+) -> Result<TokenStream> {
+    let eq_sign = tokens_iter.next().ok_or_else(|| {
+        let msg = format!(
+            "Invalid syntax for #[{ARBITRARY_ATTRIBUTE_NAME}], `{opt_name}` is missing assignment."
+        );
+        syn::Error::new(default_span, msg)
+    })?;
+
+    if eq_sign.to_string() == "=" {
+        Ok(tokens_iter.collect())
+    } else {
+        let msg = format!("Invalid syntax for #[{ARBITRARY_ATTRIBUTE_NAME}], expected `=` after `{opt_name}`, got: `{eq_sign}`");
+        Err(syn::Error::new(eq_sign.span(), msg))
+    }
+}