Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support customization of fields on derive #129

Merged
merged 13 commits into from
Oct 20, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
Expand Up @@ -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`
Expand Down
2 changes: 0 additions & 2 deletions Cargo.toml
Expand Up @@ -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"]
Expand Down
33 changes: 33 additions & 0 deletions README.md
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion derive/Cargo.toml
Expand Up @@ -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"
Expand Down
117 changes: 117 additions & 0 deletions 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))
}
}