diem_log_derive/
lib.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8use proc_macro::TokenStream;
9use proc_macro2::Ident;
10use quote::{format_ident, quote};
11use syn::{
12    parse_macro_input, AngleBracketedGenericArguments, Attribute, Data,
13    DataStruct, DeriveInput, Fields, FieldsNamed, GenericArgument, Meta,
14    MetaList, NestedMeta, Path, PathArguments, PathSegment, Type, TypePath,
15};
16
17#[proc_macro_derive(Schema, attributes(schema))]
18pub fn derive(input: TokenStream) -> TokenStream {
19    // Parse the input tokens into a syntax tree.
20    let input = parse_macro_input!(input as DeriveInput);
21    let name = input.ident;
22    let fields = match input.data {
23        Data::Struct(DataStruct {
24            fields: Fields::Named(FieldsNamed { named, .. }),
25            ..
26        }) => named,
27        _ => panic!("derive(Schema) only supports structs with named fields"),
28    };
29    let (impl_generics, ty_generics, where_clause) =
30        input.generics.split_for_impl();
31
32    let fields: Vec<StructField> = fields
33        .iter()
34        .map(|f| {
35            let ty = f.ty.clone();
36            let value_type = extract_attr(&f.attrs);
37
38            let inner_ty = extract_internal_type(&ty).cloned();
39            StructField {
40                value_type,
41                ident: f.ident.clone(),
42                ty: f.ty.clone(),
43                inner_ty,
44            }
45        })
46        .collect();
47
48    let setters = fields.iter().map(|f| {
49        let ident = &f.ident;
50
51        if let Some(ty) = &f.inner_ty {
52            quote! {
53                pub fn #ident(mut self, #ident: #ty) -> Self {
54                    self.#ident = ::std::option::Option::Some(#ident);
55                    self
56                }
57            }
58        } else {
59            let ty = &f.ty;
60            quote! {
61                pub fn #ident(mut self, #ident: #ty) -> Self {
62                    self.#ident = #ident;
63                    self
64                }
65            }
66        }
67    });
68
69    // Calls to visit_pair
70    let visitor = format_ident!("visitor");
71    let from_serde = quote! { ::diem_logger::Value::from_serde };
72    let from_display = quote! { ::diem_logger::Value::from_display };
73    let from_debug = quote! { ::diem_logger::Value::from_debug };
74    let key_new = quote! { ::diem_logger::Key::new };
75    let visits = fields.iter().map(|f| {
76        let ident = f.ident.as_ref().unwrap();
77        let ident_str = ident.to_string();
78
79        let from_fn = match f.value_type {
80            Some(ValueType::Display) => &from_display,
81            Some(ValueType::Debug) => &from_debug,
82            Some(ValueType::Serde) | None => &from_serde,
83        };
84        if f.inner_ty.is_some() {
85            quote! {
86                if let Some(#ident) = &self.#ident {
87                    #visitor.visit_pair(#key_new(#ident_str), #from_fn(#ident));
88                }
89            }
90        } else {
91            quote! {
92                #visitor.visit_pair(#key_new(#ident_str), #from_fn(&self.#ident));
93            }
94        }
95    });
96
97    // Build the output, possibly using quasi-quotation
98    let expanded = quote! {
99        impl #impl_generics #name #ty_generics #where_clause {
100            #(#setters)*
101        }
102
103        impl #impl_generics ::diem_logger::Schema for #name #ty_generics #where_clause {
104            fn visit(&self, visitor: &mut dyn ::diem_logger::Visitor) {
105                #(#visits)*
106            }
107        }
108    };
109
110    // Hand the output tokens back to the compiler
111    TokenStream::from(expanded)
112}
113
114#[derive(Debug)]
115enum ValueType {
116    Debug,
117    Display,
118    Serde,
119}
120
121#[derive(Debug)]
122struct StructField {
123    value_type: Option<ValueType>,
124    ident: Option<Ident>,
125    ty: Type,
126    /// Indicates that the type is wrapped by an Option type
127    inner_ty: Option<Type>,
128}
129
130fn extract_internal_type(ty: &Type) -> Option<&Type> {
131    if let Type::Path(TypePath {
132        qself: None,
133        path: Path { segments, .. },
134    }) = ty
135    {
136        if let Some(PathSegment { ident, arguments }) = segments.first() {
137            // Extract the inner type if it is "Option"
138            if ident == "Option" {
139                if let PathArguments::AngleBracketed(
140                    AngleBracketedGenericArguments { args, .. },
141                ) = arguments
142                {
143                    if let Some(GenericArgument::Type(ty)) = args.first() {
144                        return Some(ty);
145                    }
146                }
147            }
148        }
149    }
150
151    None
152}
153
154fn extract_attr(attrs: &[Attribute]) -> Option<ValueType> {
155    for attr in attrs {
156        if let Meta::List(MetaList { path, nested, .. }) =
157            attr.parse_meta().unwrap()
158        {
159            for segment in path.segments {
160                // Only handle schema attrs
161                if segment.ident != "schema" {
162                    continue;
163                }
164
165                let meta = if let Some(meta) = nested.first() {
166                    meta
167                } else {
168                    continue;
169                };
170
171                let path = if let NestedMeta::Meta(Meta::Path(path)) = meta {
172                    path
173                } else {
174                    panic!("unsupported schema attribute");
175                };
176
177                let answer = match path
178                    .segments
179                    .first()
180                    .unwrap()
181                    .ident
182                    .to_string()
183                    .as_ref()
184                {
185                    "debug" => ValueType::Debug,
186                    "display" => ValueType::Display,
187                    "serde" => ValueType::Serde,
188                    _ => panic!("unsupported schema attribute"),
189                };
190
191                return Some(answer);
192            }
193        }
194    }
195
196    None
197}