aptos_log_derive_link/
lib.rs

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

aptos_log_derive_link/
lib.rs
