scx_stats/
stats.rs

1use quote::ToTokens;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4use syn::parse::{Parse, ParseBuffer};
5use syn::spanned::Spanned;
6use syn::{
7    Attribute, Error, Field, Fields, GenericArgument, Ident, ItemStruct, LitStr, Path,
8    PathArguments, Token, Type, TypePath,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize)]
12pub enum StatsKind {
13    #[serde(rename = "i64")]
14    I64,
15    #[serde(rename = "u64")]
16    U64,
17    #[serde(rename = "float")]
18    Float,
19    #[serde(rename = "string")]
20    String,
21    #[serde(rename = "struct")]
22    Struct(String),
23}
24
25impl StatsKind {
26    pub fn new(ty: &Type, paths: &mut BTreeMap<String, Path>) -> syn::Result<Self> {
27        match ty {
28            Type::Reference(reference) => return Self::new(&reference.elem, paths),
29            Type::Path(TypePath { qself: _, path }) => {
30                if let Some(ident) = path.get_ident() {
31                    match ident.to_string().as_str() {
32                        "String" | "str" => return Ok(Self::String),
33                        "i8" | "i16" | "i32" | "i64" | "isize" => return Ok(Self::I64),
34                        "u8" | "u16" | "u32" | "u64" | "usize" => return Ok(Self::U64),
35                        "f32" | "f64" => return Ok(Self::Float),
36                        _ => {}
37                    }
38                }
39                let name = path.to_token_stream().to_string();
40                paths.insert(name.to_string(), path.clone());
41                return Ok(Self::Struct(name));
42            }
43            _ => {}
44        }
45        Err(Error::new(ty.span(), "scx_stats: Unsupported element type"))
46    }
47
48    pub fn can_be_dict_key(&self) -> bool {
49        matches!(self, Self::I64 | Self::U64 | Self::String)
50    }
51}
52
53impl std::fmt::Display for StatsKind {
54    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
55        match self {
56            Self::I64 => write!(f, "i64"),
57            Self::U64 => write!(f, "u64"),
58            Self::Float => write!(f, "float"),
59            Self::String => write!(f, "string"),
60            Self::Struct(name) => write!(f, "{}", name),
61        }
62    }
63}
64
65#[derive(Clone, Debug, Serialize, Deserialize)]
66pub enum StatsData {
67    #[serde(rename = "datum")]
68    Datum(StatsKind),
69    #[serde(rename = "array")]
70    Array(StatsKind),
71    #[serde(rename = "dict")]
72    Dict { key: StatsKind, datum: StatsKind },
73}
74
75impl StatsData {
76    fn new_array(path: &Path, paths: &mut BTreeMap<String, Path>) -> syn::Result<Option<Self>> {
77        if path.leading_colon.is_some() {
78            return Ok(None);
79        }
80
81        let is_vec = match path.segments.len() {
82            1 => path.segments[0].ident == "Vec",
83            3 => {
84                path.segments[0].ident == "std"
85                    && path.segments[1].ident == "vec"
86                    && path.segments[2].ident == "Vec"
87            }
88            _ => false,
89        };
90
91        if !is_vec {
92            return Ok(None);
93        }
94
95        if let PathArguments::AngleBracketed(ab) = &path.segments.last().unwrap().arguments {
96            let args = &ab.args;
97            if args.is_empty() {
98                return Err(Error::new(
99                    args.span(),
100                    "scx_stats: T generic argument missing",
101                ));
102            }
103
104            match &args[0] {
105                GenericArgument::Type(ty) => Ok(Some(Self::Array(StatsKind::new(ty, paths)?))),
106                _ => Ok(None),
107            }
108        } else {
109            Ok(None)
110        }
111    }
112
113    fn new_dict(path: &Path, paths: &mut BTreeMap<String, Path>) -> syn::Result<Option<Self>> {
114        if path.leading_colon.is_some() {
115            return Ok(None);
116        }
117
118        let is_btree_map = match path.segments.len() {
119            1 => path.segments[0].ident == "BTreeMap",
120            3 => {
121                path.segments[0].ident == "std"
122                    && path.segments[1].ident == "collections"
123                    && path.segments[2].ident == "BTreeMap"
124            }
125            _ => false,
126        };
127
128        if !is_btree_map {
129            return Ok(None);
130        }
131
132        if let PathArguments::AngleBracketed(ab) = &path.segments.last().unwrap().arguments {
133            let args = &ab.args;
134            if args.len() < 2 {
135                return Err(Error::new(
136                    args.span(),
137                    "scx_stats: K, V generic arguments missing",
138                ));
139            }
140
141            match (&args[0], &args[1]) {
142                (GenericArgument::Type(ty0), GenericArgument::Type(ty1)) => {
143                    let kind0 = StatsKind::new(ty0, paths)?;
144                    let kind1 = StatsKind::new(ty1, paths)?;
145
146                    if kind0.can_be_dict_key() {
147                        Ok(Some(Self::Dict {
148                            key: kind0,
149                            datum: kind1,
150                        }))
151                    } else {
152                        Err(Error::new(
153                            ty0.span(),
154                            "scx_stats: K must be an integer or String",
155                        ))
156                    }
157                }
158                _ => Ok(None),
159            }
160        } else {
161            Ok(None)
162        }
163    }
164
165    pub fn new(ty: &Type, paths: &mut BTreeMap<String, Path>) -> syn::Result<Self> {
166        let kind = StatsKind::new(ty, paths)?;
167        if let StatsKind::Struct(_) = &kind {
168            if let Type::Path(path) = ty {
169                if let Some(ar) = Self::new_array(&path.path, paths)? {
170                    return Ok(ar);
171                }
172                if let Some(dict) = Self::new_dict(&path.path, paths)? {
173                    return Ok(dict);
174                }
175            }
176        }
177        Ok(Self::Datum(kind))
178    }
179}
180
181impl std::fmt::Display for StatsData {
182    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
183        match self {
184            Self::Datum(kind) => write!(f, "{}", kind),
185            Self::Array(kind) => write!(f, "[{}]", kind),
186            Self::Dict { key, datum } => write!(f, "{{{}:{}}}", key, datum),
187        }
188    }
189}
190
191#[derive(Clone, Debug, Serialize, Deserialize)]
192pub enum StatsAttr {
193    Top,
194    Desc(String),
195    User(String, String),
196}
197
198struct StatsAttrVec {
199    attrs: Vec<StatsAttr>,
200}
201
202impl Parse for StatsAttrVec {
203    fn parse(input: &ParseBuffer) -> syn::Result<Self> {
204        let mut attrs = vec![];
205        loop {
206            let ident = input.parse::<Ident>()?;
207            match ident.to_string().as_str() {
208                "top" => attrs.push(StatsAttr::Top),
209                "desc" => {
210                    input.parse::<Token!(=)>()?;
211                    attrs.push(StatsAttr::Desc(input.parse::<LitStr>()?.value()))
212                }
213                key if key.starts_with("_") => {
214                    let val = match input.peek(Token!(=)) {
215                        true => {
216                            input.parse::<Token!(=)>()?;
217                            input.parse::<LitStr>()?.value()
218                        }
219                        false => "true".to_string(),
220                    };
221                    attrs.push(StatsAttr::User(key.to_string(), val));
222                }
223                _ => Err(Error::new(ident.span(), "scx_stats: Unknown attribute"))?,
224            }
225            if !input.is_empty() {
226                input.parse::<Token!(,)>()?;
227            }
228            if input.is_empty() {
229                break;
230            }
231        }
232        Ok(Self { attrs })
233    }
234}
235
236#[derive(Clone, Debug, Default, Serialize, Deserialize)]
237pub struct StatsFieldAttrs {
238    #[serde(default, skip_serializing_if = "Option::is_none")]
239    pub desc: Option<String>,
240    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
241    pub user: BTreeMap<String, String>,
242}
243
244impl StatsFieldAttrs {
245    pub fn new(attrs: &[Attribute]) -> syn::Result<Self> {
246        let mut fattrs: Self = Self::default();
247
248        for attr in attrs {
249            if attr.path().is_ident("stat") {
250                let vec = attr.parse_args::<StatsAttrVec>()?;
251                for elem in vec.attrs.into_iter() {
252                    match elem {
253                        StatsAttr::Desc(v) => fattrs.desc = Some(v),
254                        StatsAttr::User(k, v) => {
255                            fattrs.user.insert(k, v);
256                        }
257                        v => Err(Error::new(
258                            attr.span(),
259                            format!("Not a field attribute: {:?}", &v),
260                        ))?,
261                    }
262                }
263            }
264        }
265
266        Ok(fattrs)
267    }
268}
269
270#[derive(Clone, Debug, Serialize, Deserialize)]
271pub struct StatsField {
272    #[serde(flatten)]
273    pub data: StatsData,
274    #[serde(flatten)]
275    pub attrs: StatsFieldAttrs,
276}
277
278impl StatsField {
279    pub fn new(field: &Field, paths: &mut BTreeMap<String, Path>) -> syn::Result<(String, Self)> {
280        Ok((
281            field.ident.as_ref().unwrap().to_string(),
282            Self {
283                data: StatsData::new(&field.ty, paths)?,
284                attrs: StatsFieldAttrs::new(&field.attrs)?,
285            },
286        ))
287    }
288}
289
290#[derive(Clone, Debug, Default, Serialize, Deserialize)]
291pub struct StatsStructAttrs {
292    #[serde(default, skip_serializing_if = "Option::is_none")]
293    pub top: Option<String>,
294    #[serde(default, skip_serializing_if = "Option::is_none")]
295    pub desc: Option<String>,
296    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
297    pub user: BTreeMap<String, String>,
298}
299
300impl StatsStructAttrs {
301    pub fn new(attrs: &[Attribute]) -> syn::Result<Self> {
302        let mut sattrs: Self = Self::default();
303
304        for attr in attrs {
305            if attr.path().is_ident("stat") {
306                let vec = attr.parse_args::<StatsAttrVec>()?;
307                for elem in vec.attrs.into_iter() {
308                    match elem {
309                        StatsAttr::Top => sattrs.top = Some("true".into()),
310                        StatsAttr::Desc(v) => sattrs.desc = Some(v),
311                        StatsAttr::User(k, v) => {
312                            sattrs.user.insert(k, v);
313                        }
314                    }
315                }
316            }
317        }
318
319        Ok(sattrs)
320    }
321}
322
323#[derive(Clone, Debug, Serialize, Deserialize)]
324pub struct StatsMeta {
325    pub name: String,
326    #[serde(flatten)]
327    pub attrs: StatsStructAttrs,
328    pub fields: BTreeMap<String, StatsField>,
329}
330
331#[derive(Clone, Debug)]
332pub struct StatsMetaAux {
333    pub meta: StatsMeta,
334    pub ident: Ident,
335    pub paths: BTreeMap<String, Path>,
336}
337
338impl Parse for StatsMetaAux {
339    fn parse(input: &ParseBuffer) -> syn::Result<Self> {
340        let mut paths = BTreeMap::new();
341        let mut fields = BTreeMap::new();
342
343        let item_struct: ItemStruct = input.parse()?;
344        let attrs = StatsStructAttrs::new(&item_struct.attrs)?;
345
346        if let Fields::Named(named_fields) = &item_struct.fields {
347            for field in named_fields.named.iter() {
348                let (name, sf) = StatsField::new(field, &mut paths)?;
349                fields.insert(name, sf);
350            }
351        }
352
353        Ok(Self {
354            meta: StatsMeta {
355                name: item_struct.ident.to_string(),
356                attrs,
357                fields,
358            },
359            ident: item_struct.ident,
360            paths,
361        })
362    }
363}
364
365pub trait Meta {
366    fn meta() -> StatsMeta;
367}
scx_stats/stats.rs

scx_stats/
stats.rs
