rewrite_impl_trait/
lib.rs

1//! Converts any `impl Trait` in traits, impl blocks, or bare functions into method generics
2//!
3//! This can be used to work around language issues with `impl Trait`, such as a lack of support in type aliases.
4//! It also allows trait mocking with [mockall](https://crates.io/crates/mockall), for traits that use impl Trait in arguments.
5
6use proc_macro::TokenStream;
7use quote::quote;
8use quote::quote_spanned;
9use syn::GenericParam;
10use syn::Ident;
11use syn::ImplItem;
12use syn::ItemImpl;
13use syn::ItemTrait;
14use syn::Signature;
15use syn::TraitItem;
16use syn::TypeParam;
17use syn::{parse_macro_input, spanned::Spanned, FnArg, Item, ItemFn, Type};
18
19/// Rewrites `impl Trait` into a method generic on functions, trait definitions, and trait implementations.
20#[proc_macro_attribute]
21pub fn into_generic(_attr: TokenStream, input: TokenStream) -> TokenStream {
22    let item = parse_macro_input!(input as Item);
23
24    match item {
25        Item::Trait(item_trait) => trait_into_generic(item_trait),
26        Item::Impl(item_impl) => trait_impl_into_generic(item_impl),
27        Item::Fn(function) => function_into_generic(function),
28        _ => {
29            quote_spanned! { item.span() => compile_error!("RewriteImplTrait must be used on a Trait, Impl, or Fn definition.") #[item] }.into()
30        }
31    }
32}
33
34fn trait_into_generic(mut item_trait: ItemTrait) -> TokenStream {
35    for item in &mut item_trait.items {
36        if let TraitItem::Method(method) = item {
37            sig_into_generic(&mut method.sig);
38        }
39    }
40
41    quote! {
42        #item_trait
43    }
44    .into()
45}
46
47fn trait_impl_into_generic(mut item_trait: ItemImpl) -> TokenStream {
48    for item in &mut item_trait.items {
49        if let ImplItem::Method(method) = item {
50            sig_into_generic(&mut method.sig);
51        }
52    }
53
54    quote! {
55        #item_trait
56    }
57    .into()
58}
59
60fn function_into_generic(mut function: ItemFn) -> TokenStream {
61    sig_into_generic(&mut function.sig);
62
63    quote! {
64        #function
65    }
66    .into()
67}
68
69fn sig_into_generic(sig: &mut Signature) {
70    let input_span = sig.span();
71    let mut index = 0;
72    let mut new_generics: Vec<GenericParam> = vec![];
73
74    for input in &mut sig.inputs {
75        if let FnArg::Typed(pat_ty) = input {
76            let bounds = if let Type::ImplTrait(impl_trait) = &*pat_ty.ty {
77                impl_trait.bounds.clone()
78            } else {
79                continue;
80            };
81
82            let param = format!("RewriteImplTrait{}", index);
83            let ident = Ident::new(&param, input_span);
84            index += 1;
85
86            new_generics.push(GenericParam::Type(TypeParam {
87                ident: ident.clone(),
88                bounds,
89                attrs: vec![],
90                colon_token: None,
91                eq_token: None,
92                default: None,
93            }));
94
95            pat_ty.ty = Box::new(Type::Verbatim(quote! { #ident }));
96        }
97    }
98
99    for new_param in new_generics {
100        sig.generics.params.push(new_param);
101    }
102}