scx_stats_derive/
lib.rs

1use quote::{format_ident, quote, quote_spanned};
2use scx_stats::{StatsData, StatsKind, StatsMetaAux};
3use std::sync::atomic::{AtomicU64, Ordering};
4use syn::parse_macro_input;
5use syn::spanned::Spanned;
6use syn::{Attribute, Data, DeriveInput, Fields, Lit};
7
8static ASSERT_IDX: AtomicU64 = AtomicU64::new(0);
9
10#[proc_macro_derive(Stats, attributes(stat))]
11pub fn stat(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
12    let stats_aux = parse_macro_input!(input as StatsMetaAux);
13    let (meta, ident, paths) = (stats_aux.meta, stats_aux.ident, stats_aux.paths);
14
15    let mut output = proc_macro2::TokenStream::new();
16
17    for (_fname, field) in meta.fields.iter() {
18        match &field.data {
19            StatsData::Datum(datum)
20            | StatsData::Array(datum)
21            | StatsData::Dict { key: _, datum } => {
22                if let StatsKind::Struct(name) = &datum {
23                    let path = &paths[name.as_str()];
24                    let idx = ASSERT_IDX.fetch_add(1, Ordering::Relaxed);
25                    let assert_id = format_ident!("_AssertStatsMeta_{}", idx);
26                    #[rustfmt::skip]
27                    let assert = quote_spanned! {path.span()=>
28                          struct #assert_id where #path: scx_stats::Meta;
29                    };
30                    output.extend(assert.into_iter());
31                }
32            }
33        }
34    }
35
36    let body = serde_json::to_string(&meta).unwrap();
37    let trait_body = quote! {
38    #[rustfmt::skip]
39    impl scx_stats::Meta for #ident {
40        fn meta() -> scx_stats::StatsMeta {
41            let body = #body;
42            scx_stats::serde_json::from_str(body).unwrap()
43        }
44    }
45    };
46    output.extend(trait_body);
47
48    output.into()
49}
50
51#[proc_macro_attribute]
52pub fn stat_doc(
53    _attr: proc_macro::TokenStream,
54    item: proc_macro::TokenStream,
55) -> proc_macro::TokenStream {
56    let input = parse_macro_input!(item as DeriveInput);
57    let ident = input.ident;
58    let vis = input.vis;
59    let attrs = input.attrs;
60    let generics = input.generics;
61    let data = input.data;
62
63    let mut output = proc_macro2::TokenStream::new();
64
65    if let Data::Struct(data_struct) = data {
66        let fields = match data_struct.fields {
67            Fields::Named(fields_named) => fields_named.named,
68            _ => {
69                return syn::Error::new_spanned(
70                    ident,
71                    "stat attribute can only be used on structs with named fields",
72                )
73                .to_compile_error()
74                .into();
75            }
76        };
77
78        let mut new_fields = Vec::new();
79
80        for mut field in fields {
81            let mut doc_string = None;
82            let mut new_attrs = Vec::new();
83
84            for attr in field.attrs.clone() {
85                if attr.path().is_ident("stat") {
86                    // Parse the arguments within #[stat(...)]
87                    attr.parse_nested_meta(|meta| {
88                        if meta.path.is_ident("desc") {
89                            // Extract the literal string value from `desc`
90                            let desc_literal: Lit = meta.value()?.parse()?;
91                            if let Lit::Str(lit_str) = desc_literal {
92                                doc_string = Some(lit_str.value());
93                            }
94                        }
95                        Ok(())
96                    })
97                    .unwrap_or_else(|err| {
98                        panic!("Failed to parse the stat attribute: {}", err);
99                    });
100                }
101                new_attrs.push(attr);
102            }
103
104            // If a description string was found, add a #[doc = "..."] attribute
105            if let Some(description) = doc_string {
106                let doc_attr = Attribute {
107                    pound_token: syn::token::Pound::default(),
108                    style: syn::AttrStyle::Outer,
109                    bracket_token: syn::token::Bracket::default(),
110                    meta: syn::Meta::NameValue(syn::MetaNameValue {
111                        path: syn::Path::from(format_ident!("doc")),
112                        eq_token: syn::token::Eq::default(),
113                        value: syn::Expr::Lit(syn::ExprLit {
114                            lit: Lit::Str(syn::LitStr::new(&description, field.span())),
115                            attrs: vec![],
116                        }),
117                    }),
118                };
119                new_attrs.push(doc_attr);
120            }
121
122            field.attrs = new_attrs;
123            new_fields.push(field);
124        }
125
126        // Rebuild the struct with the modified fields
127        let struct_def = quote! {
128            #(#attrs)*
129            #vis struct #ident #generics {
130                #(#new_fields),*
131            }
132        };
133
134        output.extend(struct_def);
135        return output.into();
136    }
137
138    // If not a struct with named fields, return an error
139    syn::Error::new_spanned(
140        ident,
141        "stat attribute can only be used on structs with named fields",
142    )
143    .to_compile_error()
144    .into()
145}
scx_stats_derive/lib.rs

scx_stats_derive/
lib.rs
