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

use quote::quote;
use syn::parse::Error;
use syn::parse_quote;
use syn::spanned::Spanned;
use syn::{Data, DeriveInput, Fields, Index, Member};

extern crate proc_macro;

#[allow(non_snake_case)]
#[proc_macro_derive(DropTake, attributes(drop_take))]
pub fn DropTake(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    syn::parse(input)
        .and_then(impl_drop_take_semantics)
        .unwrap_or_else(|err| err.to_compile_error())
        .into()
}

fn impl_drop_take_semantics(input: DeriveInput) -> syn::parse::Result<proc_macro2::TokenStream> {
    let name = input.ident;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    match input.data {
        Data::Enum(data) => Err(Error::new(
            data.enum_token.span(),
            "DropTake only works for structs",
        )),
        Data::Union(data) => Err(Error::new(
            data.union_token.span(),
            "DropTake only works for structs",
        )),
        Data::Struct(data) => {
            let take = extract_drop_take_fields(data.fields);
            Ok(quote! {
                #[automatically_derived]
                impl #impl_generics ::core::ops::Drop for #name #ty_generics #where_clause {
                    #[allow(warnings)]
                    fn drop(&mut self) {
                        extern crate drop_take;
                        <Self as drop_take::DropTake>::drop_take({ #take })
                    }
                }
            })
        }
    }
}

fn extract_drop_take_fields(fields: Fields) -> proc_macro2::TokenStream {
    let fields = match fields {
        Fields::Unit => Default::default(),
        Fields::Named(fields) => fields.named,
        Fields::Unnamed(fields) => fields.unnamed,
    };
    let members = fields
        .into_iter()
        .enumerate()
        .filter(|(_i, field)| {
            field
                .attrs
                .iter()
                .any(|attr| attr.path == parse_quote!(drop_take))
        })
        .map(|(i, field)| {
            if let Some(name) = field.ident {
                Member::Named(name)
            } else {
                Member::Unnamed(Index {
                    index: i as u32,
                    span: field.ty.span(),
                })
            }
        });

    quote!(unsafe { (#(drop_take::Take::take(&mut self. #members)),*) })
}