Add support for defaulted methods

Author: mendess

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) -> i32 {
+        self.make(10, "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,16 @@
 
 use std::iter;
 
-use proc_macro2::TokenStream;
+use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
 use syn::{
     parse::{Parse, ParseStream},
-    parse_macro_input,
+    parse_macro_input, parse_quote,
     punctuated::Punctuated,
     token::{Comma, Plus},
-    Error, FnArg, GenericParam, Generics, Ident, ItemTrait, Lifetime, Pat, PatType, Result,
-    ReturnType, Signature, Token, TraitBound, TraitItem, TraitItemConst, TraitItemFn,
-    TraitItemType, Type, TypeImplTrait, TypeParamBound,
+    Error, FnArg, GenericParam, Generics, Ident, ItemTrait, Lifetime, Pat, PatType, Receiver,
+    Result, ReturnType, Signature, Token, TraitBound, TraitItem, TraitItemConst, TraitItemFn,
+    TraitItemType, Type, TypeImplTrait, TypeParamBound, WhereClause,
 };
 
 struct Attrs {
@@ -119,10 +119,10 @@ fn transform_item(item: &TraitItem, bounds: &Vec<TypeParamBound>) -> TraitItem {
     //     fn stream(&self) -> impl Iterator<Item = i32> + Send;
     //     fn call(&self) -> u32;
     // }
-    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 },
@@ -134,7 +134,22 @@ 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| syn::parse2(quote! { { async move #b } }).unwrap()),
+        )
     } else {
         match &sig.output {
             ReturnType::Type(arrow, ty) => match &**ty {
@@ -143,19 +158,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()
     })
 }
@@ -184,7 +202,26 @@ fn mk_blanket_impl(attrs: &Attrs, tr: &ItemTrait) -> TokenStream {
         .items
         .iter()
         .map(|item| blanket_impl_item(item, variant, &generic_names));
-    let where_clauses = tr.generics.where_clause.as_ref().map(|wh| &wh.predicates);
+    let mut where_clauses = tr
+        .generics
+        .where_clause
+        .as_ref()
+        .map(|wh| wh.predicates.clone())
+        .unwrap_or_default();
+    let self_is_sync = tr.items.iter().any(|item| {
+        matches!(
+            item,
+            TraitItem::Fn(TraitItemFn {
+                default: Some(_),
+                ..
+            })
+        )
+    });
+
+    if self_is_sync {
+        where_clauses.push(parse_quote! { Self: Sync });
+    }
+
     quote! {
         impl<#generics #trailing_comma TraitVariantBlanketType> #orig<#generic_names>
         for TraitVariantBlanketType
@@ -249,6 +286,7 @@ fn blanket_impl_item(
             } else {
                 quote! {}
             };
+
             quote! {
                 #sig {
                     <Self as #variant<#generic_names>>::#ident(#(#args),*)#maybe_await
@@ -272,3 +310,40 @@ fn blanket_impl_item(
         _ => Error::new_spanned(item, "unsupported item type").into_compile_error(),
     }
 }
+
+fn add_receiver_bounds(sig: &mut Signature) {
+    if let Some(FnArg::Receiver(Receiver { ty, reference, .. })) = sig.inputs.first_mut() {
+        let predicate =
+            if let (Type::Reference(reference), Some((_and, lt))) = (&mut **ty, reference) {
+                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(),
+                    }),
+                );
+                reference.lifetime = Some(lifetime.clone());
+                let predicate = parse_quote! { #reference: Send };
+                *lt = Some(lifetime);
+                predicate
+            } else {
+                parse_quote! { #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);
+        }
+    }
+}