local_impl/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::{quote, ToTokens};
4use syn::{parse, parse2, ImplItem, ImplItemFn, ItemImpl, Result};
5
6/// Create and impl an extension trait.
7///
8/// # Examples
9///
10/// ```
11/// #[local_impl::local_impl]
12/// impl<T> VecExt for Vec<T> {
13///     fn not_empty(&self) -> bool {
14///         !self.is_empty()
15///     }
16/// }
17///
18/// # fn main() {
19/// let mut v = Vec::new();
20/// assert!(!v.not_empty());
21/// v.push(1);
22/// assert!(v.not_empty());
23/// # }
24/// ```
25///
26/// Imported across modules/traits and using trait bounds.
27///
28/// ```
29/// mod other_module {
30///     #[local_impl::local_impl]
31///     impl<T: Default> VecExt for Vec<T> {
32///         fn push_default(&mut self) {
33///             self.push(T::default());
34///         }
35///     }
36/// }
37///
38/// # fn main() {
39/// use other_module::VecExt;
40/// let mut v: Vec<()> = Vec::new();
41/// v.push_default();
42/// v.push_default();
43/// assert_eq!(v, vec![(), ()]);
44/// # }
45/// ```
46#[proc_macro_attribute]
47pub fn local_impl(attrs: TokenStream, input: TokenStream) -> TokenStream {
48    let mut orig_input = input.clone();
49
50    local_impl_impl(attrs.into(), input.into()).map_or_else(
51        |err| {
52            orig_input.extend::<proc_macro::TokenStream>(err.to_compile_error().into());
53            orig_input
54        },
55        Into::into,
56    )
57}
58
59fn local_impl_impl(attrs: TokenStream2, input: TokenStream2) -> Result<TokenStream2> {
60    let _ = parse2::<parse::Nothing>(attrs)?;
61    let input = parse2::<ItemImpl>(input)?;
62
63    let (impl_generics, generics, where_clause) = input.generics.split_for_impl();
64    let trait_name = input.trait_.clone().expect("Expected a trait name").1;
65
66    let method_heads = input
67        .items
68        .iter()
69        .map(|item| match item {
70            ImplItem::Fn(ImplItemFn { attrs: _, sig, .. }) => {
71                quote! {
72                    #sig;
73                }
74            }
75            _ => panic!("Only methods are allowed"),
76        })
77        .collect::<Vec<_>>();
78
79    let trait_def = quote!(
80        pub(crate) trait #trait_name #generics #where_clause {
81            #(#method_heads)*
82        }
83    );
84
85    let items = input.items;
86    let self_ty = input.self_ty;
87
88    let impl_block = quote!(
89        impl #impl_generics #trait_name #generics for #self_ty {
90            #(#items)*
91        }
92    );
93
94    Ok(quote!(
95        #trait_def
96        #impl_block
97    )
98    .into_token_stream())
99}