proc_macro_assertions/ident_generator/
counting.rs

1use std::collections::BTreeMap;
2
3use proc_macro2::Span;
4use syn::Ident;
5
6use super::IdentGenerator;
7
8#[derive(Debug, Default, Clone)]
9#[allow(clippy::module_name_repetitions)]
10pub struct CountingIdentGenerator {
11    idents: BTreeMap<Option<&'static str>, usize>,
12}
13
14/// A [`IdentGenerator`] that works by counting up from 0 for each prefix
15/// for each invocation.
16impl CountingIdentGenerator {
17    /// Create a new counting ident generator.
18    #[must_use]
19    pub fn new() -> Self {
20        Self::default()
21    }
22}
23
24impl IdentGenerator for CountingIdentGenerator {
25    fn generate(&mut self, prefix: Option<&'static str>, span: Span) -> Ident {
26        let counter = self.idents.entry(prefix).or_default();
27
28        *counter += 1;
29
30        prefix.map_or_else(
31            || Ident::new(format!("C_{counter}").as_str(), span),
32            |ident| Ident::new(format!("C_{ident}_{counter}").as_str(), span),
33        )
34    }
35}
36
37#[cfg(test)]
38mod test {
39    use proc_macro2::Span;
40    use syn::Ident;
41
42    use crate::ident_generator::IdentGenerator;
43
44    use super::CountingIdentGenerator;
45
46    #[test]
47    fn ident() {
48        let mut ident_gen = CountingIdentGenerator::new();
49
50        for i in 1..=100 {
51            assert_eq!(
52                ident_gen.ident(),
53                Ident::new(format!("C_{i}").as_str(), Span::call_site())
54            );
55        }
56    }
57
58    #[test]
59    fn prefixed() {
60        let mut ident_gen = CountingIdentGenerator::new();
61
62        for i in 1..=100 {
63            assert_eq!(
64                ident_gen.prefixed("prefix"),
65                Ident::new(format!("C_prefix_{i}").as_str(), Span::call_site())
66            );
67        }
68    }
69
70    #[test]
71    fn spanned() {
72        let mut ident_gen = CountingIdentGenerator::new();
73
74        assert_eq!(
75            ident_gen.spanned(Span::mixed_site()),
76            Ident::new("C_1", Span::mixed_site())
77        );
78    }
79
80    #[test]
81    fn prefixed_spanned() {
82        let mut ident_gen = CountingIdentGenerator::new();
83
84        for i in 1..=100 {
85            assert_eq!(
86                ident_gen.generate(Some("prefix"), Span::mixed_site()),
87                Ident::new(format!("C_prefix_{i}").as_str(), Span::mixed_site())
88            );
89        }
90    }
91}