Fix suggestion for `_` on return type for fn in impl for Trait

7 files changed, 129 insertions, 73 deletions

diff --git a/compiler/rustc_typeck/src/astconv/mod.rs b/compiler/rustc_typeck/src/astconv/mod.rs
index 52de1cb88f4..1cd0ace8adb 100644
--- a/compiler/rustc_typeck/src/astconv/mod.rs
+++ b/compiler/rustc_typeck/src/astconv/mod.rs
@@ -2574,7 +2574,8 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
             .map(|(i, a)| {
                 if let hir::TyKind::Infer = a.kind && !self.allow_ty_infer() {
                     if let Some(suggested_ty) =
-                      self.suggest_trait_fn_ty_for_impl_fn_infer(hir_id, i) {
+                        self.suggest_trait_fn_ty_for_impl_fn_infer(hir_id, Some(i))
+                    {
                         infer_replacements.push((a.span, suggested_ty.to_string()));
                         return suggested_ty;
                     }
@@ -2588,8 +2589,17 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
 
         let output_ty = match decl.output {
             hir::FnRetTy::Return(output) => {
-                visitor.visit_ty(output);
-                self.ast_ty_to_ty(output)
+                if let hir::TyKind::Infer = output.kind
+                    && !self.allow_ty_infer()
+                    && let Some(suggested_ty) =
+                        self.suggest_trait_fn_ty_for_impl_fn_infer(hir_id, None)
+                {
+                    infer_replacements.push((output.span, suggested_ty.to_string()));
+                    suggested_ty
+                } else {
+                    visitor.visit_ty(output);
+                    self.ast_ty_to_ty(output)
+                }
             }
             hir::FnRetTy::DefaultReturn(..) => tcx.mk_unit(),
         };
@@ -2619,7 +2629,7 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
             if !infer_replacements.is_empty() {
                 diag.multipart_suggestion(&format!(
                     "try replacing `_` with the type{} in the corresponding trait method signature",
-                    if infer_replacements.len() > 1 { "s" } else { "" }
+                    rustc_errors::pluralize!(infer_replacements.len()),
                 ), infer_replacements, Applicability::MachineApplicable);
             }
 
@@ -2653,10 +2663,12 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
 
     /// Given a fn_hir_id for a impl function, suggest the type that is found on the
     /// corresponding function in the trait that the impl implements, if it exists.
+    /// If arg_idx is Some, then it corresponds to an input type index, otherwise it
+    /// corresponds to the return type.
     fn suggest_trait_fn_ty_for_impl_fn_infer(
         &self,
         fn_hir_id: hir::HirId,
-        arg_idx: usize,
+        arg_idx: Option<usize>,
     ) -> Option<Ty<'tcx>> {
         let tcx = self.tcx();
         let hir = tcx.hir();
@@ -2664,7 +2676,7 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
         let hir::Node::ImplItem(hir::ImplItem { kind: hir::ImplItemKind::Fn(..), ident, .. }) =
             hir.get(fn_hir_id) else { return None };
         let hir::Node::Item(hir::Item { kind: hir::ItemKind::Impl(i), .. }) =
-                hir.get(hir.get_parent_node(fn_hir_id)) else { return None };
+                hir.get(hir.get_parent_node(fn_hir_id)) else { bug!("ImplItem should have Impl parent") };
 
         let trait_ref =
             self.instantiate_mono_trait_ref(i.of_trait.as_ref()?, self.ast_ty_to_ty(i.self_ty));
@@ -2681,7 +2693,9 @@ impl<'o, 'tcx> dyn AstConv<'tcx> + 'o {
             trait_ref.substs.extend_to(tcx, x.def_id, |param, _| tcx.mk_param_from_def(param)),
         );
 
-        Some(tcx.erase_late_bound_regions(fn_sig.input(arg_idx)))
+        let ty = if let Some(arg_idx) = arg_idx { fn_sig.input(arg_idx) } else { fn_sig.output() };
+
+        Some(tcx.erase_late_bound_regions(ty))
     }
 
     fn validate_late_bound_regions(
diff --git a/compiler/rustc_typeck/src/collect.rs b/compiler/rustc_typeck/src/collect.rs
index 9df20aa60f9..be77bdb0bf5 100644
--- a/compiler/rustc_typeck/src/collect.rs
+++ b/compiler/rustc_typeck/src/collect.rs
@@ -1897,50 +1897,17 @@ fn fn_sig(tcx: TyCtxt<'_>, def_id: DefId) -> ty::PolyFnSig<'_> {
             generics,
             ..
         })
-        | ImplItem(hir::ImplItem { kind: ImplItemKind::Fn(sig, _), ident, generics, .. })
         | Item(hir::Item { kind: ItemKind::Fn(sig, generics, _), ident, .. }) => {
-            match get_infer_ret_ty(&sig.decl.output) {
-                Some(ty) => {
-                    let fn_sig = tcx.typeck(def_id).liberated_fn_sigs()[hir_id];
-                    // Typeck doesn't expect erased regions to be returned from `type_of`.
-                    let fn_sig = tcx.fold_regions(fn_sig, &mut false, |r, _| match *r {
-                        ty::ReErased => tcx.lifetimes.re_static,
-                        _ => r,
-                    });
-                    let fn_sig = ty::Binder::dummy(fn_sig);
-
-                    let mut visitor = HirPlaceholderCollector::default();
-                    visitor.visit_ty(ty);
-                    let mut diag = bad_placeholder(tcx, visitor.0, "return type");
-                    let ret_ty = fn_sig.skip_binder().output();
-                    if !ret_ty.references_error() {
-                        if !ret_ty.is_closure() {
-                            let ret_ty_str = match ret_ty.kind() {
-                                // Suggest a function pointer return type instead of a unique function definition
-                                // (e.g. `fn() -> i32` instead of `fn() -> i32 { f }`, the latter of which is invalid
-                                // syntax)
-                                ty::FnDef(..) => ret_ty.fn_sig(tcx).to_string(),
-                                _ => ret_ty.to_string(),
-                            };
-                            diag.span_suggestion(
-                                ty.span,
-                                "replace with the correct return type",
-                                ret_ty_str,
-                                Applicability::MaybeIncorrect,
-                            );
-                        } else {
-                            // We're dealing with a closure, so we should suggest using `impl Fn` or trait bounds
-                            // to prevent the user from getting a papercut while trying to use the unique closure
-                            // syntax (e.g. `[closure@src/lib.rs:2:5: 2:9]`).
-                            diag.help("consider using an `Fn`, `FnMut`, or `FnOnce` trait bound");
-                            diag.note("for more information on `Fn` traits and closure types, see https://doc.rust-lang.org/book/ch13-01-closures.html");
-                        }
-                    }
-                    diag.emit();
+            infer_return_ty_for_fn_sig(tcx, sig, *ident, generics, def_id, &icx)
+        }
 
-                    fn_sig
-                }
-                None => <dyn AstConv<'_>>::ty_of_fn(
+        ImplItem(hir::ImplItem { kind: ImplItemKind::Fn(sig, _), ident, generics, .. }) => {
+            // Do not try to inference the return type for a impl method coming from a trait
+            if let Item(hir::Item { kind: ItemKind::Impl(i), .. }) =
+                tcx.hir().get(tcx.hir().get_parent_node(hir_id))
+                && i.of_trait.is_some()
+            {
+                <dyn AstConv<'_>>::ty_of_fn(
                     &icx,
                     hir_id,
                     sig.header.unsafety,
@@ -1949,7 +1916,9 @@ fn fn_sig(tcx: TyCtxt<'_>, def_id: DefId) -> ty::PolyFnSig<'_> {
                     generics,
                     Some(ident.span),
                     None,
-                ),
+                )
+            } else {
+                infer_return_ty_for_fn_sig(tcx, sig, *ident, generics, def_id, &icx)
             }
         }
 
@@ -2011,6 +1980,70 @@ fn fn_sig(tcx: TyCtxt<'_>, def_id: DefId) -> ty::PolyFnSig<'_> {
     }
 }
 
+fn infer_return_ty_for_fn_sig<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    sig: &hir::FnSig<'_>,
+    ident: Ident,
+    generics: &hir::Generics<'_>,
+    def_id: LocalDefId,
+    icx: &ItemCtxt<'tcx>,
+) -> ty::PolyFnSig<'tcx> {
+    let hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
+
+    match get_infer_ret_ty(&sig.decl.output) {
+        Some(ty) => {
+            let fn_sig = tcx.typeck(def_id).liberated_fn_sigs()[hir_id];
+            // Typeck doesn't expect erased regions to be returned from `type_of`.
+            let fn_sig = tcx.fold_regions(fn_sig, &mut false, |r, _| match *r {
+                ty::ReErased => tcx.lifetimes.re_static,
+                _ => r,
+            });
+            let fn_sig = ty::Binder::dummy(fn_sig);
+
+            let mut visitor = HirPlaceholderCollector::default();
+            visitor.visit_ty(ty);
+            let mut diag = bad_placeholder(tcx, visitor.0, "return type");
+            let ret_ty = fn_sig.skip_binder().output();
+            if !ret_ty.references_error() {
+                if !ret_ty.is_closure() {
+                    let ret_ty_str = match ret_ty.kind() {
+                        // Suggest a function pointer return type instead of a unique function definition
+                        // (e.g. `fn() -> i32` instead of `fn() -> i32 { f }`, the latter of which is invalid
+                        // syntax)
+                        ty::FnDef(..) => ret_ty.fn_sig(tcx).to_string(),
+                        _ => ret_ty.to_string(),
+                    };
+                    diag.span_suggestion(
+                        ty.span,
+                        "replace with the correct return type",
+                        ret_ty_str,
+                        Applicability::MaybeIncorrect,
+                    );
+                } else {
+                    // We're dealing with a closure, so we should suggest using `impl Fn` or trait bounds
+                    // to prevent the user from getting a papercut while trying to use the unique closure
+                    // syntax (e.g. `[closure@src/lib.rs:2:5: 2:9]`).
+                    diag.help("consider using an `Fn`, `FnMut`, or `FnOnce` trait bound");
+                    diag.note("for more information on `Fn` traits and closure types, see https://doc.rust-lang.org/book/ch13-01-closures.html");
+                }
+            }
+            diag.emit();
+
+            fn_sig
+        }
+        None => <dyn AstConv<'_>>::ty_of_fn(
+            icx,
+            hir_id,
+            sig.header.unsafety,
+            sig.header.abi,
+            sig.decl,
+            generics,
+            Some(ident.span),
+            None,
+        ),
+    }
+}
+
 fn impl_trait_ref(tcx: TyCtxt<'_>, def_id: DefId) -> Option<ty::TraitRef<'_>> {
     let icx = ItemCtxt::new(tcx, def_id);
     match tcx.hir().expect_item(def_id.expect_local()).kind {
diff --git a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.fixed b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.fixed
index 20e762a98c3..4963790c35d 100644
--- a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.fixed
+++ b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.fixed
@@ -2,12 +2,14 @@
 #![allow(unused)]
 
 trait Foo<T>: Sized {
-    fn bar(i: i32, t: T, s: &Self) {}
+    fn bar(i: i32, t: T, s: &Self) -> (T, i32);
 }
 
 impl Foo<usize> for () {
-    fn bar(i: i32, t: usize, s: &()) {}
-    //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
+    fn bar(i: i32, t: usize, s: &()) -> (usize, i32) {
+        //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
+        (1, 2)
+    }
 }
 
 fn main() {}
diff --git a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.rs b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.rs
index ae71fd4e390..ddf39c9c861 100644
--- a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.rs
+++ b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.rs
@@ -2,12 +2,14 @@
 #![allow(unused)]
 
 trait Foo<T>: Sized {
-    fn bar(i: i32, t: T, s: &Self) {}
+    fn bar(i: i32, t: T, s: &Self) -> (T, i32);
 }
 
 impl Foo<usize> for () {
-    fn bar(i: _, t: _, s: _) {}
-    //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
+    fn bar(i: _, t: _, s: _) -> _ {
+        //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
+        (1, 2)
+    }
 }
 
 fn main() {}
diff --git a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.stderr b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.stderr
index 56044e4670b..730836a40c2 100644
--- a/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.stderr
+++ b/src/test/ui/did_you_mean/replace-impl-infer-ty-from-trait.stderr
@@ -1,16 +1,17 @@
 error[E0121]: the placeholder `_` is not allowed within types on item signatures for functions
   --> $DIR/replace-impl-infer-ty-from-trait.rs:9:15
    |
-LL |     fn bar(i: _, t: _, s: _) {}
-   |               ^     ^     ^ not allowed in type signatures
-   |               |     |
+LL |     fn bar(i: _, t: _, s: _) -> _ {
+   |               ^     ^     ^     ^ not allowed in type signatures
+   |               |     |     |
+   |               |     |     not allowed in type signatures
    |               |     not allowed in type signatures
    |               not allowed in type signatures
    |
 help: try replacing `_` with the types in the corresponding trait method signature
    |
-LL |     fn bar(i: i32, t: usize, s: &()) {}
-   |               ~~~     ~~~~~     ~~~
+LL |     fn bar(i: i32, t: usize, s: &()) -> (usize, i32) {
+   |               ~~~     ~~~~~     ~~~     ~~~~~~~~~~~~
 
 error: aborting due to previous error
 
diff --git a/src/test/ui/typeck/typeck_type_placeholder_item.rs b/src/test/ui/typeck/typeck_type_placeholder_item.rs
index ca0876be58d..22fedb22d66 100644
--- a/src/test/ui/typeck/typeck_type_placeholder_item.rs
+++ b/src/test/ui/typeck/typeck_type_placeholder_item.rs
@@ -57,7 +57,7 @@ unsafe fn test12(x: *const usize) -> *const *const _ {
 
 impl Clone for Test9 {
     fn clone(&self) -> _ { Test9 }
-    //~^ ERROR the placeholder `_` is not allowed within types on item signatures for return types
+    //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
 
     fn clone_from(&mut self, other: _) { *self = Test9; }
     //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
@@ -113,7 +113,7 @@ pub fn main() {
 
     impl Clone for FnTest9 {
         fn clone(&self) -> _ { FnTest9 }
-        //~^ ERROR the placeholder `_` is not allowed within types on item signatures for return types
+        //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
 
         fn clone_from(&mut self, other: _) { *self = FnTest9; }
         //~^ ERROR the placeholder `_` is not allowed within types on item signatures for functions
diff --git a/src/test/ui/typeck/typeck_type_placeholder_item.stderr b/src/test/ui/typeck/typeck_type_placeholder_item.stderr
index 64c7a306e5d..3ea317dfb1a 100644
--- a/src/test/ui/typeck/typeck_type_placeholder_item.stderr
+++ b/src/test/ui/typeck/typeck_type_placeholder_item.stderr
@@ -545,14 +545,16 @@ help: use type parameters instead
 LL |     fn test10<T>(&self, _x : T) { }
    |              +++             ~
 
-error[E0121]: the placeholder `_` is not allowed within types on item signatures for return types
+error[E0121]: the placeholder `_` is not allowed within types on item signatures for functions
   --> $DIR/typeck_type_placeholder_item.rs:59:24
    |
 LL |     fn clone(&self) -> _ { Test9 }
-   |                        ^
-   |                        |
-   |                        not allowed in type signatures
-   |                        help: replace with the correct return type: `Test9`
+   |                        ^ not allowed in type signatures
+   |
+help: try replacing `_` with the type in the corresponding trait method signature
+   |
+LL |     fn clone(&self) -> Test9 { Test9 }
+   |                        ~~~~~
 
 error[E0121]: the placeholder `_` is not allowed within types on item signatures for functions
   --> $DIR/typeck_type_placeholder_item.rs:62:37
@@ -585,14 +587,16 @@ help: use type parameters instead
 LL |         fn fn_test10<T>(&self, _x : T) { }
    |                     +++             ~
 
-error[E0121]: the placeholder `_` is not allowed within types on item signatures for return types
+error[E0121]: the placeholder `_` is not allowed within types on item signatures for functions
   --> $DIR/typeck_type_placeholder_item.rs:115:28
    |
 LL |         fn clone(&self) -> _ { FnTest9 }
-   |                            ^
-   |                            |
-   |                            not allowed in type signatures
-   |                            help: replace with the correct return type: `FnTest9`
+   |                            ^ not allowed in type signatures
+   |
+help: try replacing `_` with the type in the corresponding trait method signature
+   |
+LL |         fn clone(&self) -> FnTest9 { FnTest9 }
+   |                            ~~~~~~~
 
 error[E0121]: the placeholder `_` is not allowed within types on item signatures for functions
   --> $DIR/typeck_type_placeholder_item.rs:118:41