iroh_metrics_derive/
lib.rs

1use heck::{ToKebabCase, ToLowerCamelCase, ToPascalCase, ToShoutySnakeCase, ToSnakeCase};
2use proc_macro::TokenStream;
3use quote::quote;
4use syn::{
5    Attribute, Data, DeriveInput, Error, Expr, ExprLit, Fields, GenericArgument, Ident, Lit,
6    LitStr, PathArguments, Type, 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
33#[proc_macro_derive(EncodeLabelSet, attributes(label))]
34pub fn derive_encode_label_set(input: TokenStream) -> TokenStream {
35    let input = parse_macro_input!(input as DeriveInput);
36    expand_encode_label_set(&input)
37        .unwrap_or_else(Error::into_compile_error)
38        .into()
39}
40
41/// Derives [`EncodeLabelValue`] for an enum with only unit variants.
42///
43/// Maps each variant to its name (snake_case by default). Use
44/// `#[label(rename_all = "...")]` on the enum or `#[label(name = "...")]`
45/// per-variant to customize.
46#[proc_macro_derive(EncodeLabelValue, attributes(label))]
47pub fn derive_encode_label_value(input: TokenStream) -> TokenStream {
48    let input = parse_macro_input!(input as DeriveInput);
49    expand_encode_label_value(&input)
50        .unwrap_or_else(Error::into_compile_error)
51        .into()
52}
53
54fn expand_iterable(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
55    let (name, fields) = parse_named_struct(input)?;
56
57    // Partition into scalars and families. `Family<_, _>` is detected by the
58    // last path segment of the field type, so type aliases require an
59    // explicit `#[metrics(family)]` override.
60    let mut metrics = Vec::new();
61    let mut families = Vec::new();
62    for field in fields.iter() {
63        let attr = parse_metrics_attr(&field.attrs)?;
64        let info = field_info(field, attr)?;
65        if info.is_family {
66            families.push(info);
67        } else {
68            metrics.push(info);
69        }
70    }
71
72    let metric_count = metrics.len();
73    let family_count = families.len();
74
75    let metric_arms = metrics.iter().enumerate().map(|(i, f)| {
76        let (ident, ident_str, help) = (f.ident, &f.ident_str, &f.help);
77        quote! {
78            #i => Some(::iroh_metrics::MetricItem::new(#ident_str, #help, &self.#ident as &dyn ::iroh_metrics::Metric)),
79        }
80    });
81    let family_arms = families.iter().enumerate().map(|(i, f)| {
82        let (ident, ident_str, help) = (f.ident, &f.ident_str, &f.help);
83        quote! {
84            #i => Some(::iroh_metrics::FamilyItem::new(#ident_str, #help, &self.#ident as &dyn ::iroh_metrics::FamilyEncoder)),
85        }
86    });
87
88    let family_impl = (family_count > 0).then(|| {
89        quote! {
90            fn family_field_count(&self) -> usize { #family_count }
91            fn family_field_ref(&self, n: usize) -> Option<::iroh_metrics::FamilyItem<'_>> {
92                match n {
93                    #(#family_arms)*
94                    _ => None,
95                }
96            }
97        }
98    });
99
100    Ok(quote! {
101        impl ::iroh_metrics::iterable::Iterable for #name {
102            fn metric_field_count(&self) -> usize { #metric_count }
103            fn metric_field_ref(&self, n: usize) -> Option<::iroh_metrics::MetricItem<'_>> {
104                match n {
105                    #(#metric_arms)*
106                    _ => None,
107                }
108            }
109            #family_impl
110        }
111    })
112}
113
114/// Per-field info pre-computed once for `expand_iterable`.
115struct FieldInfo<'a> {
116    ident: &'a Ident,
117    ident_str: String,
118    /// Help text: explicit `#[metrics(help = "...")]` > first doc line > field name.
119    help: String,
120    is_family: bool,
121}
122
123fn field_info<'a>(field: &'a syn::Field, attr: MetricsAttr) -> Result<FieldInfo<'a>, Error> {
124    let ident = field
125        .ident
126        .as_ref()
127        .ok_or_else(|| Error::new(field.span(), "Only named fields are supported"))?;
128    let ident_str = ident.to_string();
129    let help = attr
130        .help
131        .or_else(|| parse_doc_first_line(&field.attrs))
132        .unwrap_or_else(|| ident_str.clone());
133    let is_family = attr.family || is_family_type(&field.ty);
134    Ok(FieldInfo {
135        ident,
136        ident_str,
137        help,
138        is_family,
139    })
140}
141
142/// Checks if a type is `Family<_, _>`.
143fn is_family_type(ty: &Type) -> bool {
144    if let Type::Path(type_path) = ty {
145        if let Some(segment) = type_path.path.segments.last() {
146            if segment.ident == "Family" {
147                if let PathArguments::AngleBracketed(args) = &segment.arguments {
148                    // Family should have 2 generic arguments
149                    return args
150                        .args
151                        .iter()
152                        .filter(|arg| matches!(arg, GenericArgument::Type(_)))
153                        .count()
154                        == 2;
155                }
156            }
157        }
158    }
159    false
160}
161
162fn expand_metrics(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
163    let (name, fields) = parse_named_struct(input)?;
164    let attr = parse_metrics_attr(&input.attrs)?;
165    let name_str = attr
166        .name
167        .unwrap_or_else(|| name.to_string().to_snake_case());
168
169    let default = if attr.default {
170        let mut items = vec![];
171        for field in fields.iter() {
172            let ident = field
173                .ident
174                .as_ref()
175                .ok_or_else(|| Error::new(field.span(), "Only named fields are supported"))?;
176            let attr = parse_default_attr(&field.attrs)?;
177            let item = if let Some(expr) = attr {
178                quote!( #ident: #expr)
179            } else {
180                quote!( #ident: ::std::default::Default::default() )
181            };
182            items.push(item);
183        }
184        Some(quote! {
185            impl ::std::default::Default for #name {
186                fn default() -> Self {
187                    Self {
188                        #(#items),*
189                    }
190                }
191            }
192        })
193    } else {
194        None
195    };
196
197    Ok(quote! {
198        impl ::iroh_metrics::MetricsGroup for #name {
199            fn name(&self) -> &'static str {
200                #name_str
201            }
202        }
203
204        #default
205    })
206}
207
208fn expand_metrics_group_set(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
209    let (name, fields) = parse_named_struct(input)?;
210    let attr = parse_metrics_attr(&input.attrs)?;
211    let name_str = attr
212        .name
213        .unwrap_or_else(|| name.to_string().to_snake_case());
214
215    let mut cloned = quote! {};
216    let mut refs = quote! {};
217    for field in fields {
218        let name = field.ident.as_ref().unwrap();
219        cloned.extend(quote! {
220            self.#name.clone() as ::std::sync::Arc<dyn ::iroh_metrics::MetricsGroup>,
221        });
222        refs.extend(quote! {
223            &*self.#name as &dyn ::iroh_metrics::MetricsGroup,
224        });
225    }
226
227    Ok(quote! {
228        impl ::iroh_metrics::MetricsGroupSet for #name {
229            fn name(&self) -> &'static str {
230                #name_str
231            }
232
233            fn groups_cloned(&self) -> impl ::std::iter::Iterator<Item = ::std::sync::Arc<dyn ::iroh_metrics::MetricsGroup>> {
234                [#cloned].into_iter()
235            }
236
237            fn groups(&self) -> impl ::std::iter::Iterator<Item = &dyn ::iroh_metrics::MetricsGroup> {
238                [#refs].into_iter()
239            }
240        }
241    })
242}
243
244fn expand_encode_label_set(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
245    let (name, fields) = parse_named_struct(input)?;
246    let struct_attr = parse_label_struct_attr(&input.attrs)?;
247
248    let mut label_pairs = vec![];
249    for field in fields.iter() {
250        let ident = field
251            .ident
252            .as_ref()
253            .ok_or_else(|| Error::new(field.span(), "Only named fields are supported"))?;
254        let attr = parse_label_attr(&field.attrs)?;
255
256        // Skip fields marked with #[label(skip)]
257        if attr.skip {
258            continue;
259        }
260
261        // Field-level `name = ...` wins; otherwise apply the struct-level
262        // `rename_all` transformation to the field ident.
263        let label_name = match attr.name {
264            Some(n) => n,
265            None => match struct_attr.rename_all {
266                Some(rule) => rule.apply(&ident.to_string()),
267                None => ident.to_string(),
268            },
269        };
270
271        // Borrow the field through `EncodeLabelValue` so string fields
272        // don't allocate on every scrape.
273        label_pairs.push(quote! {
274            (#label_name, ::iroh_metrics::EncodeLabelValue::encode_label_value(&self.#ident))
275        });
276    }
277
278    Ok(quote! {
279        impl ::iroh_metrics::EncodeLabelSet for #name {
280            fn encode_label_pairs(&self) -> ::std::vec::Vec<::iroh_metrics::LabelPair<'_>> {
281                ::std::vec![#(#label_pairs),*]
282            }
283        }
284    })
285}
286
287fn expand_encode_label_value(input: &DeriveInput) -> Result<proc_macro2::TokenStream, Error> {
288    let name = &input.ident;
289    let Data::Enum(data) = &input.data else {
290        return Err(Error::new(
291            input.span(),
292            "EncodeLabelValue can only be derived for enums with unit variants.",
293        ));
294    };
295    let enum_attr = parse_label_struct_attr(&input.attrs)?;
296    let rule = enum_attr.rename_all.unwrap_or(RenameRule::SnakeCase);
297
298    let mut arms = Vec::new();
299    for variant in &data.variants {
300        if !matches!(variant.fields, Fields::Unit) {
301            return Err(Error::new(
302                variant.span(),
303                "EncodeLabelValue only supports unit variants.",
304            ));
305        }
306        let attr = parse_label_attr(&variant.attrs)?;
307        let ident = &variant.ident;
308        let label = attr.name.unwrap_or_else(|| rule.apply(&ident.to_string()));
309        arms.push(quote! {
310            Self::#ident => ::iroh_metrics::LabelValue::Str(::std::borrow::Cow::Borrowed(#label)),
311        });
312    }
313
314    Ok(quote! {
315        impl ::iroh_metrics::EncodeLabelValue for #name {
316            fn encode_label_value(&self) -> ::iroh_metrics::LabelValue<'_> {
317                match self {
318                    #(#arms)*
319                }
320            }
321        }
322    })
323}
324
325#[derive(Clone, Copy)]
326enum RenameRule {
327    SnakeCase,
328    CamelCase,
329    PascalCase,
330    ScreamingSnakeCase,
331    KebabCase,
332    Lowercase,
333    Uppercase,
334}
335
336impl RenameRule {
337    fn parse(s: &str, span: proc_macro2::Span) -> Result<Self, Error> {
338        match s {
339            "snake_case" => Ok(Self::SnakeCase),
340            "camelCase" => Ok(Self::CamelCase),
341            "PascalCase" => Ok(Self::PascalCase),
342            "SCREAMING_SNAKE_CASE" => Ok(Self::ScreamingSnakeCase),
343            "kebab-case" => Ok(Self::KebabCase),
344            "lowercase" => Ok(Self::Lowercase),
345            "UPPERCASE" => Ok(Self::Uppercase),
346            other => Err(Error::new(
347                span,
348                format!(
349                    "unknown rename_all value `{other}`. Supported: snake_case, camelCase, \
350                     PascalCase, SCREAMING_SNAKE_CASE, kebab-case, lowercase, UPPERCASE.",
351                ),
352            )),
353        }
354    }
355
356    fn apply(self, ident: &str) -> String {
357        match self {
358            Self::SnakeCase => ident.to_snake_case(),
359            Self::CamelCase => ident.to_lower_camel_case(),
360            Self::PascalCase => ident.to_pascal_case(),
361            Self::ScreamingSnakeCase => ident.to_shouty_snake_case(),
362            Self::KebabCase => ident.to_kebab_case(),
363            Self::Lowercase => ident.to_lowercase(),
364            Self::Uppercase => ident.to_uppercase(),
365        }
366    }
367}
368
369#[derive(Default)]
370struct LabelStructAttr {
371    rename_all: Option<RenameRule>,
372}
373
374fn parse_label_struct_attr(attrs: &[Attribute]) -> Result<LabelStructAttr, Error> {
375    let mut out = LabelStructAttr::default();
376    for attr in attrs.iter().filter(|attr| attr.path().is_ident("label")) {
377        attr.parse_nested_meta(|meta| {
378            if meta.path.is_ident("rename_all") {
379                let s: LitStr = meta.value()?.parse()?;
380                out.rename_all = Some(RenameRule::parse(&s.value(), s.span())?);
381                Ok(())
382            } else {
383                Err(meta.error("The struct-level `label` attribute supports only `rename_all`."))
384            }
385        })?;
386    }
387    Ok(out)
388}
389
390fn parse_doc_first_line(attrs: &[Attribute]) -> Option<String> {
391    attrs
392        .iter()
393        .filter(|attr| attr.path().is_ident("doc"))
394        .flat_map(|attr| attr.meta.require_name_value())
395        .find_map(|name_value| {
396            let Expr::Lit(ExprLit { lit, .. }) = &name_value.value else {
397                return None;
398            };
399            let Lit::Str(str) = lit else { return None };
400            Some(str.value().trim().to_string())
401        })
402}
403
404#[derive(Default)]
405struct MetricsAttr {
406    name: Option<String>,
407    help: Option<String>,
408    default: bool,
409    /// `#[metrics(family)]` — force-treat the field as a `Family<_, _>`
410    /// even when the type is hidden behind an alias.
411    family: bool,
412}
413
414#[derive(Default)]
415struct LabelAttr {
416    name: Option<String>,
417    skip: bool,
418}
419
420fn parse_default_attr(attrs: &[Attribute]) -> Result<Option<syn::Expr>, syn::Error> {
421    let Some(attr) = attrs.iter().find(|attr| attr.path().is_ident("default")) else {
422        return Ok(None);
423    };
424    let expr = attr.parse_args::<Expr>()?;
425    Ok(Some(expr))
426}
427
428fn parse_metrics_attr(attrs: &[Attribute]) -> Result<MetricsAttr, syn::Error> {
429    let mut out = MetricsAttr::default();
430    for attr in attrs.iter().filter(|attr| attr.path().is_ident("metrics")) {
431        attr.parse_nested_meta(|meta| {
432            if meta.path.is_ident("name") {
433                out.name = Some(parse_lit_str(&meta)?);
434                Ok(())
435            } else if meta.path.is_ident("help") {
436                out.help = Some(parse_lit_str(&meta)?);
437                Ok(())
438            } else if meta.path.is_ident("default") {
439                out.default = true;
440                Ok(())
441            } else if meta.path.is_ident("family") {
442                out.family = true;
443                Ok(())
444            } else {
445                Err(meta.error(
446                    "The `metrics` attribute supports only `name`, `help`, `default`, and `family`.",
447                ))
448            }
449        })?;
450    }
451    Ok(out)
452}
453
454fn parse_label_attr(attrs: &[Attribute]) -> Result<LabelAttr, syn::Error> {
455    let mut out = LabelAttr::default();
456    for attr in attrs.iter().filter(|attr| attr.path().is_ident("label")) {
457        attr.parse_nested_meta(|meta| {
458            if meta.path.is_ident("name") {
459                out.name = Some(parse_lit_str(&meta)?);
460                Ok(())
461            } else if meta.path.is_ident("skip") {
462                out.skip = true;
463                Ok(())
464            } else {
465                Err(meta.error("The `label` attribute supports only `name` and `skip` fields."))
466            }
467        })?;
468    }
469    Ok(out)
470}
471
472fn parse_lit_str(meta: &ParseNestedMeta<'_>) -> Result<String, Error> {
473    let s: LitStr = meta.value()?.parse()?;
474    Ok(s.value().trim().to_string())
475}
476
477fn parse_named_struct(input: &DeriveInput) -> Result<(&Ident, &Fields), Error> {
478    match &input.data {
479        Data::Struct(data) if matches!(data.fields, Fields::Named(_)) => {
480            Ok((&input.ident, &data.fields))
481        }
482        _ => Err(Error::new(
483            input.span(),
484            "The `MetricsGroup` and `Iterable` derives support only structs.",
485        )),
486    }
487}
iroh_metrics_derive/lib.rs

iroh_metrics_derive/
lib.rs
