chrometracer_attributes/
lib.rs

1use std::collections::HashSet;
2
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use quote::ToTokens;
6use syn::ext::IdentExt;
7use syn::{
8    parse::{Parse, ParseStream, Parser},
9    parse_quote,
10    punctuated::Punctuated,
11    Expr, Ident, LitInt, LitStr, Path, Token,
12};
13
14#[derive(Default)]
15struct ChromeEventArgs {
16    level: Option<Level>,
17    target: Option<LitStr>,
18    event: Option<Event>,
19    fields: Fields,
20    skips: HashSet<Ident>,
21}
22
23#[derive(Default)]
24struct Fields(Punctuated<Field, Token![,]>);
25
26impl ToTokens for Fields {
27    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
28        self.0.to_tokens(tokens)
29    }
30}
31
32impl Parse for Fields {
33    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
34        let _ = input.parse::<kw::fields>();
35        let content;
36        let _ = syn::parenthesized!(content in input);
37        let fields: Punctuated<_, Token![,]> = content.parse_terminated(Field::parse)?;
38        Ok(Self(fields))
39    }
40}
41
42#[derive(Clone)]
43enum Level {
44    Str(LitStr),
45    Int(LitInt),
46    Path(Path),
47}
48
49impl Parse for Level {
50    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
51        let _ = input.parse::<kw::level>()?;
52        let _ = input.parse::<Token![=]>()?;
53        let lookahead = input.lookahead1();
54        if lookahead.peek(LitStr) {
55            Ok(Self::Str(input.parse()?))
56        } else if lookahead.peek(LitInt) {
57            Ok(Self::Int(input.parse()?))
58        } else if lookahead.peek(Ident) {
59            Ok(Self::Path(input.parse()?))
60        } else {
61            Err(lookahead.error())
62        }
63    }
64}
65
66struct Skips(HashSet<Ident>);
67
68impl Parse for Skips {
69    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
70        let _ = input.parse::<kw::skip>();
71        let content;
72        let _ = syn::parenthesized!(content in input);
73        let names: Punctuated<Ident, Token![,]> = content.parse_terminated(Ident::parse_any)?;
74        let mut skips = HashSet::new();
75        for name in names {
76            if skips.contains(&name) {
77                return Err(syn::Error::new(
78                    name.span(),
79                    "tried to skip the same field twice",
80                ));
81            } else {
82                skips.insert(name);
83            }
84        }
85        Ok(Self(skips))
86    }
87}
88
89impl ToTokens for ChromeEventArgs {
90    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
91        self.fields.to_tokens(tokens);
92    }
93}
94
95impl Parse for ChromeEventArgs {
96    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
97        let mut args = Self::default();
98
99        while !input.is_empty() {
100            let lookahead = input.lookahead1();
101            if lookahead.peek(kw::event) {
102                args.event = Some(Event::parse(input)?);
103            } else if lookahead.peek(kw::level) {
104                args.level = Some(Level::parse(input)?)
105            } else if lookahead.peek(kw::fields) {
106                args.fields = Fields::parse(input)?;
107            } else if lookahead.peek(kw::skip) {
108                let Skips(skips) = input.parse()?;
109                args.skips = skips;
110            } else if lookahead.peek(kw::target) {
111                let target = input.parse::<StrArg<kw::target>>()?.value;
112                args.target = Some(target);
113            } else if lookahead.peek(Token![,]) {
114                let _ = input.parse::<Token![,]>()?;
115            } else {
116                panic!(
117                    "Unknown fields, expected one of \"event\", \"level\", \"fields\", \"skip\", \"target\"",
118                )
119            }
120        }
121        Ok(args)
122    }
123}
124
125struct StrArg<T> {
126    value: LitStr,
127    _p: std::marker::PhantomData<T>,
128}
129
130impl<T: Parse> Parse for StrArg<T> {
131    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
132        let _ = input.parse::<T>()?;
133        let _ = input.parse::<Token![=]>()?;
134        let value = input.parse()?;
135        Ok(Self {
136            value,
137            _p: std::marker::PhantomData,
138        })
139    }
140}
141
142struct Event {
143    keyword: kw::event,
144    colon_token: Token![:],
145    event: LitStr,
146    comma_token: Token![,],
147}
148
149impl Parse for Event {
150    fn parse(input: ParseStream) -> syn::Result<Self> {
151        Ok(Event {
152            keyword: input.parse()?,
153            colon_token: input.parse()?,
154            event: input.parse()?,
155            comma_token: input.parse()?,
156        })
157    }
158}
159
160impl ToTokens for Event {
161    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
162        self.keyword.to_tokens(tokens);
163        self.colon_token.to_tokens(tokens);
164        self.event.to_tokens(tokens);
165        self.comma_token.to_tokens(tokens);
166    }
167}
168
169struct Field {
170    path: Path,
171    eq_token: Token![=],
172    value: Expr,
173}
174
175impl Parse for Field {
176    fn parse(input: ParseStream) -> syn::Result<Self> {
177        Ok(Field {
178            path: input.parse()?,
179            eq_token: input.parse()?,
180            value: input.parse()?,
181        })
182    }
183}
184
185impl ToTokens for Field {
186    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
187        self.path.to_tokens(tokens);
188        self.eq_token.to_tokens(tokens);
189        self.value.to_tokens(tokens);
190    }
191}
192
193#[proc_macro_attribute]
194pub fn instrument(attr: TokenStream, item: TokenStream) -> TokenStream {
195    let args = syn::parse_macro_input!(attr as ChromeEventArgs);
196
197    let event = args
198        .event
199        .map(|e| e.event)
200        .unwrap_or_else(|| LitStr::new("", Span::call_site()));
201
202    let fields = args.fields.0.iter();
203    let fields2 = args.fields.0.iter();
204    let fields3 = args.fields.0.iter();
205
206    let mut input = syn::Item::parse.parse(item).unwrap();
207
208    if let syn::Item::Fn(ref mut item) = input {
209        let original = &item.block;
210        item.block = Box::new(parse_quote! {{
211            let start = chrometracer::current(|tracer| tracer.map(|t| t.start));
212
213            if let Some(start) = start {
214                let event = match #event {
215                    "async" => Some((chrometracer::EventType::AsyncStart, chrometracer::EventType::AsyncEnd)),
216                    "" => None,
217                    _ => panic!("Unknown event, expected one of \"async\"")
218                };
219
220                let ret = if let Some(event) = event {
221                    chrometracer::event!(#(#fields,)* ph = event.0,
222                        ts = ::std::time::SystemTime::now().duration_since(start).unwrap().as_nanos() as f64 / 1000.0);
223                    let ret = #original;
224                    chrometracer::event!(#(#fields2,)* ph = event.1,
225                        ts = ::std::time::SystemTime::now().duration_since(start).unwrap().as_nanos() as f64 / 1000.0);
226                    ret
227                } else {
228                    let now = ::std::time::SystemTime::now();
229                    let ts = now.duration_since(start).unwrap().as_nanos() as f64 / 1000.0;
230                    let ret = #original;
231                    let dur = ::std::time::SystemTime::now().duration_since(now).unwrap().as_nanos() as f64 / 1000.0;
232
233                    chrometracer::event!(#(#fields3,)* ph = chrometracer::EventType::Complete, dur = dur, ts = ts);
234                    ret
235                };
236
237                ret
238            } else {
239                #original
240            }
241        }});
242    } else {
243        unreachable!()
244    }
245
246    input.into_token_stream().into()
247}
248
249mod kw {
250    syn::custom_keyword!(event);
251    syn::custom_keyword!(skip);
252    syn::custom_keyword!(fields);
253    syn::custom_keyword!(level);
254    syn::custom_keyword!(target);
255}
chrometracer_attributes/lib.rs

chrometracer_attributes/
lib.rs
