Add support for defaulted methods

trait-variant/Cargo.toml

@@ -25,7 +25,7 @@ proc-macro = true
 [dependencies]
 proc-macro2 = "1.0"
 quote = "1.0"
-syn = { version = "2.0", features = ["full"] }
+syn = { version = "2.0", features = ["full", "visit-mut"] }
 
 [dev-dependencies]
 tokio = { version = "1", features = ["rt"] }

trait-variant/examples/variant.rs

@@ -17,9 +17,25 @@ pub trait LocalIntFactory {
         Self: 'a;
 
     async fn make(&self, x: u32, y: &str) -> i32;
+    async fn make_mut(&mut self);
     fn stream(&self) -> impl Iterator<Item = i32>;
     fn call(&self) -> u32;
     fn another_async(&self, input: Result<(), &str>) -> Self::MyFut<'_>;
+    async fn defaulted(&self, x: u32) -> i32 {
+        self.make(x, "10").await
+    }
+    async fn defaulted_mut(&mut self) -> i32 {
+        self.make(10, "10").await
+    }
+    async fn defaulted_mut_2(&mut self) {
+        self.make_mut().await
+    }
+    async fn defaulted_move(self) -> i32
+    where
+        Self: Sized,
+    {
+        self.make(10, "10").await
+    }
 }
 
 #[allow(dead_code)]

trait-variant/src/variant.rs

@@ -8,16 +8,10 @@
 
 use std::iter;
 
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::{
-    parse::{Parse, ParseStream},
-    parse_macro_input, parse_quote,
-    punctuated::Punctuated,
-    token::Plus,
-    Error, FnArg, GenericParam, Ident, ItemTrait, Pat, PatType, Result, ReturnType, Signature,
-    Token, TraitBound, TraitItem, TraitItemConst, TraitItemFn, TraitItemType, Type, TypeGenerics,
-    TypeImplTrait, TypeParam, TypeParamBound,
+    parse::{Parse, ParseStream}, parse_macro_input, parse_quote, punctuated::Punctuated, token::Plus, Error, FnArg, GenericParam, Ident, ItemTrait, Pat, PatIdent, PatType, Receiver, Result, ReturnType, Signature, Token, TraitBound, TraitItem, TraitItemConst, TraitItemFn, TraitItemType, Type, TypeGenerics, TypeImplTrait, TypeParam, TypeParamBound, TypeReference, WhereClause
 };
 
 struct Attrs {
@@ -127,10 +121,10 @@ fn mk_variant(
 
 // Transforms one item declaration within the definition if it has `async fn` and/or `-> impl Trait` return types by adding new bounds.
 fn transform_item(item: &TraitItem, bounds: &Vec<TypeParamBound>) -> TraitItem {
-    let TraitItem::Fn(fn_item @ TraitItemFn { sig, .. }) = item else {
+    let TraitItem::Fn(fn_item @ TraitItemFn { sig, default, .. }) = item else {
         return item.clone();
     };
-    let (arrow, output) = if sig.asyncness.is_some() {
+    let (sig, default) = if sig.asyncness.is_some() {
         let orig = match &sig.output {
             ReturnType::Default => quote! { () },
             ReturnType::Type(_, ty) => quote! { #ty },
@@ -142,7 +136,44 @@ fn transform_item(item: &TraitItem, bounds: &Vec<TypeParamBound>) -> TraitItem {
                 .chain(bounds.iter().cloned())
                 .collect(),
         });
-        (syn::parse2(quote! { -> }).unwrap(), ty)
+        let mut sig = sig.clone();
+        if default.is_some() {
+            add_receiver_bounds(&mut sig);
+        }
+
+        (
+            Signature {
+                asyncness: None,
+                output: ReturnType::Type(syn::parse2(quote! { -> }).unwrap(), Box::new(ty)),
+                ..sig.clone()
+            },
+            fn_item.default.as_ref().map(|b| {
+                let items = sig.inputs.iter().map(|i| match i {
+                    FnArg::Receiver(Receiver { self_token, .. }) => {
+                        quote! { let __self = #self_token; }
+                    }
+                    FnArg::Typed(PatType { pat, .. }) => match pat.as_ref() {
+                        Pat::Ident(PatIdent { ident, .. }) => quote! { let #ident = #ident; },
+                        _ => todo!(),
+                    },
+                });
+
+                struct ReplaceSelfVisitor;
+                impl syn::visit_mut::VisitMut for ReplaceSelfVisitor {
+                    fn visit_ident_mut(&mut self, ident: &mut syn::Ident) {
+                        if ident == "self" {
+                            *ident = syn::Ident::new("__self", ident.span());
+                        }
+                        syn::visit_mut::visit_ident_mut(self, ident);
+                    }
+                }
+
+                let mut block = b.clone();
+                syn::visit_mut::visit_block_mut(&mut ReplaceSelfVisitor, &mut block);
+
+                parse_quote! { { async move { #(#items)* #block} } }
+            }),
+        )
     } else {
         match &sig.output {
             ReturnType::Type(arrow, ty) => match &**ty {
@@ -151,19 +182,22 @@ fn transform_item(item: &TraitItem, bounds: &Vec<TypeParamBound>) -> TraitItem {
                         impl_token: it.impl_token,
                         bounds: it.bounds.iter().chain(bounds).cloned().collect(),
                     });
-                    (*arrow, ty)
+                    (
+                        Signature {
+                            output: ReturnType::Type(*arrow, Box::new(ty)),
+                            ..sig.clone()
+                        },
+                        fn_item.default.clone(),
+                    )
                 }
                 _ => return item.clone(),
             },
             ReturnType::Default => return item.clone(),
         }
     };
     TraitItem::Fn(TraitItemFn {
-        sig: Signature {
-            asyncness: None,
-            output: ReturnType::Type(arrow, Box::new(output)),
-            ..sig.clone()
-        },
+        sig,
+        default,
         ..fn_item.clone()
     })
 }
@@ -182,9 +216,29 @@ fn mk_blanket_impl(variant: &Ident, tr: &ItemTrait) -> TokenStream {
     blanket_generics
         .params
         .push(GenericParam::Type(blanket_bound));
-    let (blanket_impl_generics, _ty, blanket_where_clause) = &blanket_generics.split_for_impl();
+    let (blanket_impl_generics, _ty, blanket_where_clause) = &mut blanket_generics.split_for_impl();
+    let self_is_sync = tr.items.iter().any(|item| {
+        matches!(
+            item,
+            TraitItem::Fn(TraitItemFn {
+                default: Some(_),
+                ..
+            })
+        )
+    });
+
+    let mut blanket_where_clause = blanket_where_clause
+        .map(|w| w.predicates.clone())
+        .unwrap_or_default();
+
+    if self_is_sync {
+        blanket_where_clause.push(parse_quote! { for<'s> &'s Self: Send });
+    }
+
     quote! {
-        impl #blanket_impl_generics #orig #orig_ty_generics for #blanket #blanket_where_clause
+        impl #blanket_impl_generics #orig #orig_ty_generics for #blanket
+            where
+                #blanket_where_clause
         {
             #(#items)*
         }
@@ -229,6 +283,7 @@ fn blanket_impl_item(
             } else {
                 quote! {}
             };
+
             quote! {
                 #sig {
                     <Self as #variant #trait_ty_generics>::#ident(#(#args),*)#maybe_await
@@ -246,3 +301,47 @@ fn blanket_impl_item(
         _ => Error::new_spanned(item, "unsupported item type").into_compile_error(),
     }
 }
+
+fn add_receiver_bounds(sig: &mut Signature) {
+    let Some(FnArg::Receiver(Receiver { ty, reference, .. })) = sig.inputs.first_mut() else {
+        return;
+    };
+    let Type::Reference(
+        recv_ty @ TypeReference {
+            mutability: None, ..
+        },
+    ) = &mut **ty
+    else {
+        return;
+    };
+    let Some((_and, lt)) = reference else {
+        return;
+    };
+
+    let lifetime = syn::Lifetime {
+        apostrophe: Span::mixed_site(),
+        ident: Ident::new("the_self_lt", Span::mixed_site()),
+    };
+    sig.generics.params.insert(
+        0,
+        syn::GenericParam::Lifetime(syn::LifetimeParam {
+            lifetime: lifetime.clone(),
+            colon_token: None,
+            bounds: Default::default(),
+            attrs: Default::default(),
+        }),
+    );
+    recv_ty.lifetime = Some(lifetime.clone());
+    *lt = Some(lifetime);
+    let predicate = parse_quote! { #recv_ty: Send };
+
+    if let Some(wh) = &mut sig.generics.where_clause {
+        wh.predicates.push(predicate);
+    } else {
+        let where_clause = WhereClause {
+            where_token: Token![where](Span::mixed_site()),
+            predicates: Punctuated::from_iter([predicate]),
+        };
+        sig.generics.where_clause = Some(where_clause);
+    }
+}