multiversion_macros/
util.rs

1use proc_macro2::{Span, TokenStream};
2use quote::quote;
3use syn::{
4    parse_quote, spanned::Spanned, visit::Visit, visit_mut::VisitMut, BareFnArg, Error, Expr,
5    FnArg, GenericParam, Ident, Lifetime, Pat, PatIdent, PatType, Result, Signature, TypeBareFn,
6    TypeImplTrait,
7};
8
9pub(crate) fn arg_exprs(sig: &Signature) -> Vec<Expr> {
10    sig.inputs
11        .iter()
12        .map(|x| match x {
13            FnArg::Receiver(rec) => {
14                let self_token = rec.self_token;
15                parse_quote! { #self_token }
16            }
17            FnArg::Typed(arg) => {
18                if let Pat::Ident(ident) = &*arg.pat {
19                    let ident = &ident.ident;
20                    parse_quote! { #ident }
21                } else {
22                    panic!("pattern should have been ident")
23                }
24            }
25        })
26        .collect()
27}
28
29pub(crate) fn normalize_signature(sig: &Signature) -> (Signature, Vec<Expr>) {
30    let args = sig
31        .inputs
32        .iter()
33        .enumerate()
34        .map(|(i, x)| match x {
35            FnArg::Receiver(_) => x.clone(),
36            FnArg::Typed(arg) => FnArg::Typed(PatType {
37                pat: Box::new(Pat::Ident(PatIdent {
38                    attrs: Vec::new(),
39                    by_ref: None,
40                    mutability: None,
41                    ident: match arg.pat.as_ref() {
42                        Pat::Ident(pat) => pat.ident.clone(),
43                        _ => Ident::new(&format!("__multiversion_arg_{i}"), x.span()),
44                    },
45                    subpat: None,
46                })),
47                ..arg.clone()
48            }),
49        })
50        .collect::<Vec<_>>();
51    let sig = Signature {
52        inputs: parse_quote! { #(#args),* },
53        ..sig.clone()
54    };
55    let callable_args = arg_exprs(&sig);
56    (sig, callable_args)
57}
58
59pub(crate) fn impl_trait_present(sig: &Signature) -> bool {
60    struct ImplTraitPresent(bool);
61    impl Visit<'_> for ImplTraitPresent {
62        fn visit_type_impl_trait(&mut self, _: &TypeImplTrait) {
63            self.0 = true;
64        }
65    }
66
67    let mut visitor = ImplTraitPresent(false);
68    visitor.visit_signature(sig);
69    visitor.0
70}
71
72struct LifetimeRenamer;
73
74impl VisitMut for LifetimeRenamer {
75    fn visit_lifetime_mut(&mut self, i: &mut Lifetime) {
76        i.ident = Ident::new(&format!("__mv_inner_{}", i.ident), i.ident.span());
77    }
78}
79
80pub(crate) fn fn_type_from_signature(sig: &Signature) -> Result<TypeBareFn> {
81    let lifetimes = sig.generics.lifetimes().collect::<Vec<_>>();
82    let args = sig
83        .inputs
84        .iter()
85        .map(|x| {
86            Ok(BareFnArg {
87                attrs: Vec::new(),
88                name: None,
89                ty: match x {
90                    FnArg::Receiver(rec) => Err(Error::new(
91                        rec.self_token.span,
92                        "cannot determine type of associated fn",
93                    )),
94                    FnArg::Typed(arg) => Ok(arg.ty.as_ref().clone()),
95                }?,
96            })
97        })
98        .collect::<Result<Vec<_>>>()?;
99    assert!(
100        sig.variadic.is_none(),
101        "cannot multiversion function with variadic arguments"
102    );
103    let mut fn_ty = TypeBareFn {
104        lifetimes: if lifetimes.is_empty() {
105            None
106        } else {
107            Some(parse_quote! { for<#(#lifetimes),*> })
108        },
109        unsafety: sig.unsafety,
110        abi: sig.abi.clone(),
111        fn_token: sig.fn_token,
112        paren_token: sig.paren_token,
113        inputs: parse_quote! { #(#args),* },
114        variadic: None,
115        output: sig.output.clone(),
116    };
117    LifetimeRenamer {}.visit_type_bare_fn_mut(&mut fn_ty);
118    Ok(fn_ty)
119}
120
121pub(crate) fn fn_params(sig: &Signature) -> Vec<Ident> {
122    sig.generics
123        .params
124        .iter()
125        .filter_map(|x| match x {
126            GenericParam::Type(ty) => Some(ty.ident.clone()),
127            GenericParam::Const(c) => Some(c.ident.clone()),
128            _ => None,
129        })
130        .collect()
131}
132
133pub(crate) fn await_tokens() -> TokenStream {
134    let kw = Ident::new("await", Span::call_site());
135    quote! { .#kw }
136}