1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

// SPDX-FileCopyrightText: 2020 Robin Krahl <robin.krahl@ireas.org>
// SPDX-License-Identifier: Apache-2.0 or MIT

//! A derive macro for the [`merge::Merge`][] trait.
//!
//! See the documentation for the [`merge`][] crate for more information.
//!
//! [`merge`]: https://lib.rs/crates/merge
//! [`merge::Merge`]: https://docs.rs/merge/latest/merge/trait.Merge.html

extern crate proc_macro;

use proc_macro2::TokenStream;
use proc_macro_error::{abort, abort_call_site, dummy::set_dummy, proc_macro_error, ResultExt};
use quote::{quote, quote_spanned};
use syn::Token;

struct Field {
    name: syn::Member,
    span: proc_macro2::Span,
    attrs: FieldAttrs,
}

#[derive(Default)]
struct FieldAttrs {
    skip: bool,
    strategy: Option<syn::Path>,
}

enum FieldAttr {
    Skip,
    Strategy(syn::Path),
}

#[proc_macro_derive(Merge, attributes(merge))]
#[proc_macro_error]
pub fn merge_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = syn::parse(input).unwrap();
    impl_merge(&ast).into()
}

fn impl_merge(ast: &syn::DeriveInput) -> TokenStream {
    let name = &ast.ident;

    set_dummy(quote! {
        impl ::merge::Merge for #name {
            fn merge(&mut self, other: Self) {
                unimplemented!()
            }
        }
    });

    if let syn::Data::Struct(syn::DataStruct { ref fields, .. }) = ast.data {
        impl_merge_for_struct(name, fields)
    } else {
        abort_call_site!("merge::Merge can only be derived for structs")
    }
}

fn impl_merge_for_struct(name: &syn::Ident, fields: &syn::Fields) -> TokenStream {
    let assignments = gen_assignments(fields);

    quote! {
        impl ::merge::Merge for #name {
            fn merge(&mut self, other: Self) {
                #assignments
            }
        }
    }
}

fn gen_assignments(fields: &syn::Fields) -> TokenStream {
    let fields = fields.iter().enumerate().map(Field::from);
    let assignments = fields.filter(|f| !f.attrs.skip).map(|f| gen_assignment(&f));
    quote! {
        #( #assignments )*
    }
}

fn gen_assignment(field: &Field) -> TokenStream {
    use syn::spanned::Spanned;

    let name = &field.name;
    if let Some(strategy) = &field.attrs.strategy {
        quote_spanned!(strategy.span()=> #strategy(&mut self.#name, other.#name);)
    } else {
        quote_spanned!(field.span=> ::merge::Merge::merge(&mut self.#name, other.#name);)
    }
}

impl From<(usize, &syn::Field)> for Field {
    fn from(data: (usize, &syn::Field)) -> Self {
        use syn::spanned::Spanned;

        let (index, field) = data;
        Field {
            name: if let Some(ident) = &field.ident {
                syn::Member::Named(ident.clone())
            } else {
                syn::Member::Unnamed(index.into())
            },
            span: field.span(),
            attrs: field.attrs.iter().into(),
        }
    }
}

impl FieldAttrs {
    fn apply(&mut self, attr: FieldAttr) {
        match attr {
            FieldAttr::Skip => self.skip = true,
            FieldAttr::Strategy(path) => self.strategy = Some(path),
        }
    }
}

impl<'a, I: Iterator<Item = &'a syn::Attribute>> From<I> for FieldAttrs {
    fn from(iter: I) -> Self {
        let mut field_attrs = Self::default();

        for attr in iter {
            if !attr.path.is_ident("merge") {
                continue;
            }

            let parser = syn::punctuated::Punctuated::<FieldAttr, Token![,]>::parse_terminated;
            for attr in attr.parse_args_with(parser).unwrap_or_abort() {
                field_attrs.apply(attr);
            }
        }

        field_attrs
    }
}

impl syn::parse::Parse for FieldAttr {
    fn parse(input: syn::parse::ParseStream) -> syn::parse::Result<Self> {
        let name: syn::Ident = input.parse()?;
        if name == "skip" {
            // TODO check remaining stream
            Ok(FieldAttr::Skip)
        } else if name == "strategy" {
            let _: Token![=] = input.parse()?;
            let path: syn::Path = input.parse()?;
            Ok(FieldAttr::Strategy(path))
        } else {
            abort!(name, "Unexpected attribute: {}", name)
        }
    }
}