Suggest returning closure as impl Fn

6 files changed, 86 insertions, 24 deletions

diff --git a/compiler/rustc_middle/src/ty/print/pretty.rs b/compiler/rustc_middle/src/ty/print/pretty.rs
index cc55b7e8611..4e1bba6f1fa 100644
--- a/compiler/rustc_middle/src/ty/print/pretty.rs
+++ b/compiler/rustc_middle/src/ty/print/pretty.rs
@@ -1536,6 +1536,34 @@ pub trait PrettyPrinter<'tcx>:
         }
         Ok(self)
     }
+
+    fn pretty_closure_as_impl(
+        mut self,
+        closure: ty::ClosureSubsts<'tcx>,
+    ) -> Result<Self::Const, Self::Error> {
+        let sig = closure.sig();
+        let kind = closure.kind_ty().to_opt_closure_kind().unwrap_or(ty::ClosureKind::Fn);
+
+        write!(self, "impl ")?;
+        self.wrap_binder(&sig, |sig, mut cx| {
+            define_scoped_cx!(cx);
+
+            p!(print(kind), "(");
+            for (i, arg) in sig.inputs()[0].tuple_fields().iter().enumerate() {
+                if i > 0 {
+                    p!(", ");
+                }
+                p!(print(arg));
+            }
+            p!(")");
+
+            if !sig.output().is_unit() {
+                p!(" -> ", print(sig.output()));
+            }
+
+            Ok(cx)
+        })
+    }
 }
 
 // HACK(eddyb) boxed to avoid moving around a large struct by-value.
@@ -2450,6 +2478,11 @@ impl<'tcx> ty::PolyTraitPredicate<'tcx> {
     }
 }
 
+#[derive(Debug, Copy, Clone, TypeFoldable, TypeVisitable, Lift)]
+pub struct PrintClosureAsImpl<'tcx> {
+    pub closure: ty::ClosureSubsts<'tcx>,
+}
+
 forward_display_to_print! {
     ty::Region<'tcx>,
     Ty<'tcx>,
@@ -2542,6 +2575,10 @@ define_print_and_forward_display! {
         p!(print(self.0.trait_ref.print_only_trait_path()));
     }
 
+    PrintClosureAsImpl<'tcx> {
+        p!(pretty_closure_as_impl(self.closure))
+    }
+
     ty::ParamTy {
         p!(write("{}", self.name))
     }
diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
index 0070575f213..80354a3f8a2 100644
--- a/compiler/rustc_middle/src/ty/sty.rs
+++ b/compiler/rustc_middle/src/ty/sty.rs
@@ -325,6 +325,10 @@ impl<'tcx> ClosureSubsts<'tcx> {
             _ => bug!("closure_sig_as_fn_ptr_ty is not a fn-ptr: {:?}", ty.kind()),
         }
     }
+
+    pub fn print_as_impl_trait(self) -> ty::print::PrintClosureAsImpl<'tcx> {
+        ty::print::PrintClosureAsImpl { closure: self }
+    }
 }
 
 /// Similar to `ClosureSubsts`; see the above documentation for more.
diff --git a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
index 57771e0969b..64d261285c5 100644
--- a/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
+++ b/compiler/rustc_typeck/src/check/fn_ctxt/suggestions.rs
@@ -506,30 +506,30 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             self.resolve_numeric_literals_with_default(self.resolve_vars_if_possible(found));
         // Only suggest changing the return type for methods that
         // haven't set a return type at all (and aren't `fn main()` or an impl).
-        match (
-            &fn_decl.output,
-            found.is_suggestable(self.tcx, false),
-            can_suggest,
-            expected.is_unit(),
-        ) {
-            (&hir::FnRetTy::DefaultReturn(span), true, true, true) => {
-                err.subdiagnostic(AddReturnTypeSuggestion::Add { span, found });
-                true
-            }
-            (&hir::FnRetTy::DefaultReturn(span), false, true, true) => {
-                // FIXME: if `found` could be `impl Iterator` or `impl Fn*`, we should suggest
-                // that.
-                err.subdiagnostic(AddReturnTypeSuggestion::MissingHere { span });
-                true
-            }
-            (&hir::FnRetTy::DefaultReturn(span), _, false, true) => {
+        match &fn_decl.output {
+            &hir::FnRetTy::DefaultReturn(span) if expected.is_unit() && !can_suggest => {
                 // `fn main()` must return `()`, do not suggest changing return type
                 err.subdiagnostic(ExpectedReturnTypeLabel::Unit { span });
-                true
+                return true;
             }
-            // expectation was caused by something else, not the default return
-            (&hir::FnRetTy::DefaultReturn(_), _, _, false) => false,
-            (&hir::FnRetTy::Return(ref ty), _, _, _) => {
+            &hir::FnRetTy::DefaultReturn(span) if expected.is_unit() => {
+                if found.is_suggestable(self.tcx, false) {
+                    err.subdiagnostic(AddReturnTypeSuggestion::Add { span, found: found.to_string() });
+                    return true;
+                } else if let ty::Closure(_, substs) = found.kind()
+                    // FIXME(compiler-errors): Get better at printing binders...
+                    && let closure = substs.as_closure()
+                    && closure.sig().is_suggestable(self.tcx, false)
+                {
+                    err.subdiagnostic(AddReturnTypeSuggestion::Add { span, found: closure.print_as_impl_trait().to_string() });
+                    return true;
+                } else {
+                    // FIXME: if `found` could be `impl Iterator` we should suggest that.
+                    err.subdiagnostic(AddReturnTypeSuggestion::MissingHere { span });
+                    return true
+                }
+            }
+            &hir::FnRetTy::Return(ref ty) => {
                 // Only point to return type if the expected type is the return type, as if they
                 // are not, the expectation must have been caused by something else.
                 debug!("suggest_missing_return_type: return type {:?} node {:?}", ty, ty.kind);
@@ -546,9 +546,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     self.try_suggest_return_impl_trait(err, expected, ty, fn_id);
                     return true;
                 }
-                false
             }
+            _ => {}
         }
+        false
     }
 
     /// check whether the return type is a generic type with a trait bound
diff --git a/compiler/rustc_typeck/src/errors.rs b/compiler/rustc_typeck/src/errors.rs
index 2214fc2ced8..14c0558cdde 100644
--- a/compiler/rustc_typeck/src/errors.rs
+++ b/compiler/rustc_typeck/src/errors.rs
@@ -195,7 +195,7 @@ pub struct AddressOfTemporaryTaken {
 }
 
 #[derive(SessionSubdiagnostic)]
-pub enum AddReturnTypeSuggestion<'tcx> {
+pub enum AddReturnTypeSuggestion {
     #[suggestion(
         typeck::add_return_type_add,
         code = "-> {found} ",
@@ -204,7 +204,7 @@ pub enum AddReturnTypeSuggestion<'tcx> {
     Add {
         #[primary_span]
         span: Span,
-        found: Ty<'tcx>,
+        found: String,
     },
     #[suggestion(
         typeck::add_return_type_missing_here,
diff --git a/src/test/ui/suggestions/return-closures.rs b/src/test/ui/suggestions/return-closures.rs
new file mode 100644
index 00000000000..eb593a03965
--- /dev/null
+++ b/src/test/ui/suggestions/return-closures.rs
@@ -0,0 +1,6 @@
+fn foo() {
+    |x: &i32| 1i32
+    //~^ ERROR mismatched types
+}
+
+fn main() {}
diff --git a/src/test/ui/suggestions/return-closures.stderr b/src/test/ui/suggestions/return-closures.stderr
new file mode 100644
index 00000000000..6689f697088
--- /dev/null
+++ b/src/test/ui/suggestions/return-closures.stderr
@@ -0,0 +1,14 @@
+error[E0308]: mismatched types
+  --> $DIR/return-closures.rs:2:5
+   |
+LL | fn foo() {
+   |          - help: try adding a return type: `-> impl for<'r> Fn(&'r i32) -> i32`
+LL |     |x: &i32| 1i32
+   |     ^^^^^^^^^^^^^^ expected `()`, found closure
+   |
+   = note: expected unit type `()`
+                found closure `[closure@$DIR/return-closures.rs:2:5: 2:14]`
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0308`.