Support impl Trait in return type

dtolnay · Jan 3, 2025 · bf8f641 · bf8f641
1 parent aff365f
commit bf8f641
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/src/expand.rs b/src/expand.rs
@@ -12,7 +12,7 @@ use syn::visit_mut::{self, VisitMut};
 use syn::{
     parse_quote, parse_quote_spanned, Attribute, Block, FnArg, GenericArgument, GenericParam,
     Generics, Ident, ImplItem, Lifetime, LifetimeParam, Pat, PatIdent, PathArguments, Receiver,
-    ReturnType, Signature, Token, TraitItem, Type, TypePath, WhereClause,
+    ReturnType, Signature, Token, TraitItem, Type, TypeInfer, TypePath, WhereClause,
 };
 
 impl ToTokens for Item {
@@ -410,6 +410,8 @@ fn transform_block(context: Context, sig: &mut Signature, block: &mut Block) {
                     quote!(#(#decls)* { #(#stmts)* })
                 }
             } else {
+                let mut ret = ret.clone();
+                replace_impl_trait_with_infer(&mut ret);
                 quote! {
                     if let ::core::option::Option::Some(__ret) = ::core::option::Option::None::<#ret> {
                         #[allow(unreachable_code)]
@@ -475,3 +477,20 @@ fn where_clause_or_default(clause: &mut Option<WhereClause>) -> &mut WhereClause
         predicates: Punctuated::new(),
     })
 }
+
+fn replace_impl_trait_with_infer(ty: &mut Type) {
+    struct ReplaceImplTraitWithInfer;
+
+    impl VisitMut for ReplaceImplTraitWithInfer {
+        fn visit_type_mut(&mut self, ty: &mut Type) {
+            if let Type::ImplTrait(impl_trait) = ty {
+                *ty = Type::Infer(TypeInfer {
+                    underscore_token: Token![_](impl_trait.impl_token.span),
+                });
+            }
+            visit_mut::visit_type_mut(self, ty);
+        }
+    }
+
+    ReplaceImplTraitWithInfer.visit_type_mut(ty);
+}