Skip to main content

iroh_metrics_derive/
lib.rs

1use heck::ToSnakeCase;
2use proc_macro::TokenStream;
3use quote::quote;
4use syn::{
5    Attribute, Data, DeriveInput, Error, Expr, ExprLit, Fields, Ident, Lit, LitStr,
6    meta::ParseNestedMeta, parse_macro_input, spanned::Spanned,
7};
8
9#[proc_macro_derive(MetricsGroup, attributes(metrics, default))]
10pub fn derive_metrics_group(input: TokenStream) -> TokenStream {
11    let input = parse_macro_input!(input as DeriveInput);
12    let mut out = proc_macro2::TokenStream::new();
13    out.extend(expand_metrics(&input).unwrap_or_else(Error::into_compile_error));
14    out.extend(expand_iterable(&input).unwrap_or_else(Error::into_compile_error));
15    out.into()
16}
17
18#[proc_macro_derive(Iterable)]
19pub fn derive_iterable(input: TokenStream) -> TokenStream {
20    let input = parse_macro_input!(input as DeriveInput);
21    let out = expand_iterable(&input).unwrap_or_else(Error::into_compile_error);
22    out.into()
23}
24
25#[proc_macro_derive(MetricsGroupSet, attributes(metrics))]
26pub fn derive_metrics_group_set(input: TokenStream) -> TokenStream {
27    let input = parse_macro_input!(input as DeriveInput);
28    let mut out = proc_macro2::TokenStream::new();
29    out.extend(expand_metrics_group_set(&input).unwrap_or_else(Error::into_compile_error));
30    out.into()
31}
32
33fn expand_iterable(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
34    let (name, fields) = parse_named_struct(input)?;
35
36    let count = fields.len();
37
38    let mut match_arms = quote! {};
39    for (i, field) in fields.iter().enumerate() {
40        let ident = field.ident.as_ref().unwrap();
41        let ident_str = ident.to_string();
42        let attr = parse_metrics_attr(&field.attrs)?;
43        let help = attr
44            .help
45            .or_else(|| parse_doc_first_line(&field.attrs))
46            .unwrap_or_else(|| ident_str.clone());
47        match_arms.extend(quote! {
48            #i => Some(::iroh_metrics::MetricItem::new(#ident_str, #help, &self.#ident as &dyn ::iroh_metrics::Metric)),
49        });
50    }
51
52    Ok(quote! {
53        impl ::iroh_metrics::iterable::Iterable for #name {
54            fn field_count(&self) -> usize {
55                #count
56            }
57
58            fn field_ref(&self, n: usize) -> Option<::iroh_metrics::MetricItem<'_>> {
59                match n {
60                    #match_arms
61                    _ => None,
62                }
63            }
64        }
65    })
66}
67
68fn expand_metrics(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
69    let (name, fields) = parse_named_struct(input)?;
70    let attr = parse_metrics_attr(&input.attrs)?;
71    let name_str = attr
72        .name
73        .unwrap_or_else(|| name.to_string().to_snake_case());
74
75    let default = if attr.default {
76        let mut items = vec![];
77        for field in fields.iter() {
78            let ident = field
79                .ident
80                .as_ref()
81                .ok_or_else(|| Error::new(field.span(), "Only named fields are supported"))?;
82            let attr = parse_default_attr(&field.attrs)?;
83            let item = if let Some(expr) = attr {
84                quote!( #ident: #expr)
85            } else {
86                quote!( #ident: ::std::default::Default::default() )
87            };
88            items.push(item);
89        }
90        Some(quote! {
91            impl ::std::default::Default for #name {
92                fn default() -> Self {
93                    Self {
94                        #(#items),*
95                    }
96                }
97            }
98        })
99    } else {
100        None
101    };
102
103    Ok(quote! {
104        impl ::iroh_metrics::MetricsGroup for #name {
105            fn name(&self) -> &'static str {
106                #name_str
107            }
108        }
109
110        #default
111    })
112}
113
114fn expand_metrics_group_set(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
115    let (name, fields) = parse_named_struct(input)?;
116    let attr = parse_metrics_attr(&input.attrs)?;
117    let name_str = attr
118        .name
119        .unwrap_or_else(|| name.to_string().to_snake_case());
120
121    let mut cloned = quote! {};
122    let mut refs = quote! {};
123    for field in fields {
124        let name = field.ident.as_ref().unwrap();
125        cloned.extend(quote! {
126            self.#name.clone() as ::std::sync::Arc<dyn ::iroh_metrics::MetricsGroup>,
127        });
128        refs.extend(quote! {
129            &*self.#name as &dyn ::iroh_metrics::MetricsGroup,
130        });
131    }
132
133    Ok(quote! {
134        impl ::iroh_metrics::MetricsGroupSet for #name {
135            fn name(&self) -> &'static str {
136                #name_str
137            }
138
139            fn groups_cloned(&self) -> impl ::std::iter::Iterator<Item = ::std::sync::Arc<dyn ::iroh_metrics::MetricsGroup>> {
140                [#cloned].into_iter()
141            }
142
143            fn groups(&self) -> impl ::std::iter::Iterator<Item = &dyn ::iroh_metrics::MetricsGroup> {
144                [#refs].into_iter()
145            }
146        }
147    })
148}
149
150fn parse_doc_first_line(attrs: &[Attribute]) -> Option<String> {
151    attrs
152        .iter()
153        .filter(|attr| attr.path().is_ident("doc"))
154        .flat_map(|attr| attr.meta.require_name_value())
155        .find_map(|name_value| {
156            let Expr::Lit(ExprLit { lit, .. }) = &name_value.value else {
157                return None;
158            };
159            let Lit::Str(str) = lit else { return None };
160            Some(str.value().trim().to_string())
161        })
162}
163
164#[derive(Default)]
165struct MetricsAttr {
166    name: Option<String>,
167    help: Option<String>,
168    default: bool,
169}
170
171fn parse_default_attr(attrs: &[Attribute]) -> Result<Option<syn::Expr>, syn::Error> {
172    let Some(attr) = attrs.iter().find(|attr| attr.path().is_ident("default")) else {
173        return Ok(None);
174    };
175    let expr = attr.parse_args::<Expr>()?;
176    Ok(Some(expr))
177}
178
179fn parse_metrics_attr(attrs: &[Attribute]) -> Result<MetricsAttr, syn::Error> {
180    let mut out = MetricsAttr::default();
181    for attr in attrs.iter().filter(|attr| attr.path().is_ident("metrics")) {
182        attr.parse_nested_meta(|meta| {
183            if meta.path.is_ident("name") {
184                out.name = Some(parse_lit_str(&meta)?);
185                Ok(())
186            } else if meta.path.is_ident("help") {
187                out.help = Some(parse_lit_str(&meta)?);
188                Ok(())
189            } else if meta.path.is_ident("default") {
190                out.default = true;
191                Ok(())
192            } else {
193                Err(meta.error(
194                    "The `metrics` attribute supports only `name`, `help` and `default` fields.",
195                ))
196            }
197        })?;
198    }
199    Ok(out)
200}
201
202fn parse_lit_str(meta: &ParseNestedMeta<'_>) -> Result<String, Error> {
203    let s: LitStr = meta.value()?.parse()?;
204    Ok(s.value().trim().to_string())
205}
206
207fn parse_named_struct(input: &DeriveInput) -> Result<(&Ident, &Fields), Error> {
208    match &input.data {
209        Data::Struct(data) if matches!(data.fields, Fields::Named(_)) => {
210            Ok((&input.ident, &data.fields))
211        }
212        _ => Err(Error::new(
213            input.span(),
214            "The `MetricsGroup` and `Iterable` derives support only structs.",
215        )),
216    }
217}