semaphore_rs_depth_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use semaphore_rs_depth_config::get_supported_depths;
4use syn::{
5    parse::{Parse, ParseStream},
6    parse_macro_input, parse_quote,
7    visit_mut::VisitMut,
8    Ident, Token,
9};
10
11/// Multi-depth test generator
12///
13/// This macro is used to generate a test for each supported depth.
14/// It expects to annotate a function with a single argument, and will generate
15/// test cases delegating to that function for each supported depth.
16///
17/// For example,
18/// ```
19/// use semaphore_rs_depth_macros::test_all_depths;
20/// #[test_all_depths]
21/// fn test_depth_non_zero(depth: usize) {
22///     assert!(depth > 0);
23/// }
24/// ```
25/// with `depth_16` and `depth_30` features active will generate the following
26/// code:
27/// ```no_run
28/// fn test_depth_non_zero(depth: usize) {
29///     assert!(depth > 0);
30/// }
31///
32/// #[test]
33/// fn test_depth_non_zero_depth_16() {
34///     test_depth_non_zero(16);
35/// }
36///
37/// #[test]
38/// fn test_depth_non_zero_depth_30() {
39///     test_depth_non_zero(30);
40/// }
41/// ```
42#[proc_macro_attribute]
43pub fn test_all_depths(_attr: TokenStream, item: TokenStream) -> TokenStream {
44    let fun = parse_macro_input!(item as syn::ItemFn);
45    let fun_name = &fun.sig.ident;
46
47    let original_fun = quote! { #fun };
48    let mut result = TokenStream::from(original_fun);
49
50    for depth in get_supported_depths() {
51        let fun_name_versioned = format_ident!("{}_depth_{}", fun_name, depth);
52        let tokens = quote! {
53            #[test]
54            fn #fun_name_versioned() {
55                #fun_name(#depth);
56            }
57        };
58        result.extend(TokenStream::from(tokens));
59    }
60    result
61}
62
63#[derive(Debug)]
64struct ArrayForDepthsInput {
65    replaced_ident: Ident,
66    expr: syn::Expr,
67}
68
69#[derive(Debug)]
70struct MacroArgs {
71    args: Vec<syn::Expr>,
72}
73
74impl Parse for MacroArgs {
75    fn parse(input: ParseStream) -> syn::Result<Self> {
76        let mut args = Vec::new();
77        while !input.is_empty() {
78            args.push(input.parse::<syn::Expr>()?);
79            if input.is_empty() {
80                break;
81            }
82            input.parse::<Token![,]>()?;
83        }
84        Ok(MacroArgs { args })
85    }
86}
87
88impl MacroArgs {
89    fn tokens(&self) -> proc_macro2::TokenStream {
90        let args = &self.args;
91        quote! { #(#args),* }
92    }
93}
94
95struct IdentReplacer(Ident, syn::Expr);
96
97impl VisitMut for IdentReplacer {
98    fn visit_expr_mut(&mut self, expr: &mut syn::Expr) {
99        match expr {
100            syn::Expr::Path(ident) => {
101                if ident.path.is_ident(&self.0) {
102                    *expr = self.1.clone();
103                }
104            }
105            syn::Expr::Macro(mcr) => {
106                let Ok(mut args) = mcr.mac.parse_body::<MacroArgs>() else {
107                    return;
108                };
109                for arg in &mut args.args {
110                    self.visit_expr_mut(arg);
111                }
112                mcr.mac.tokens = args.tokens();
113            }
114            _ => syn::visit_mut::visit_expr_mut(self, expr),
115        }
116    }
117}
118
119impl Parse for ArrayForDepthsInput {
120    fn parse(input: ParseStream) -> syn::Result<Self> {
121        input.parse::<Token![|]>()?;
122        let replaced_ident = input.parse::<Ident>()?;
123        input.parse::<Token![|]>()?;
124        let expr = input.parse::<syn::Expr>()?;
125        Ok(ArrayForDepthsInput {
126            replaced_ident,
127            expr,
128        })
129    }
130}
131
132/// Macro to generate code for multiple depths.
133///
134/// Generates an array of expressions, where the given identifier is replaced
135/// with each supported depth. The argument must use closure syntax, but this
136/// is pure syntactic, the closure expression gets unrolled statically.
137///
138/// This macro also descends into other macros, as long as they use standard
139/// Rust syntax for arguments. Any non-standard macros will be ignored in the
140/// expansion of this.
141///
142/// For example, `array_for_depths!(|depth| depth + 5)`, with only `depth_16`
143/// and `depth_30` supported, will generate `[16 + 5, 30 + 5]`, and
144/// `array_for_depths!(|depth| concat!("foo", depth))` will generate
145/// `[concat!("foo", 16), concat!("foo", 30)]`.
146#[proc_macro]
147pub fn array_for_depths(input: TokenStream) -> TokenStream {
148    let input = parse_macro_input!(input as ArrayForDepthsInput);
149    let items = get_supported_depths()
150        .iter()
151        .map(|depth| {
152            let mut replacer = IdentReplacer(input.replaced_ident.clone(), parse_quote!(#depth));
153            let mut expr = input.expr.clone();
154            replacer.visit_expr_mut(&mut expr);
155            expr
156        })
157        .collect::<Vec<_>>();
158    let array = quote! { [#(#items),*] };
159    array.into()
160}