From cbd79836a5b7ef07e33213cb979294ad1b9faccc Mon Sep 17 00:00:00 2001
From: Fabian Wolff <fabian.wolff@alumni.ethz.ch>
Date: Mon, 13 Sep 2021 19:39:30 +0200
Subject: [PATCH] Improve error message for type mismatch in generator
 arguments

---
 .../src/traits/error_reporting/mod.rs         |  5 ++-
 .../src/traits/error_reporting/suggestions.rs | 33 +++++++++++--------
 src/test/ui/generator/issue-88653.rs          | 19 +++++++++++
 src/test/ui/generator/issue-88653.stderr      | 12 +++++++
 4 files changed, 55 insertions(+), 14 deletions(-)
 create mode 100644 src/test/ui/generator/issue-88653.rs
 create mode 100644 src/test/ui/generator/issue-88653.stderr

diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
index 761b217c78f4e..6eb6164846ccc 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs
@@ -730,7 +730,10 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
                 };
 
                 let found_did = match *found_trait_ty.kind() {
-                    ty::Closure(did, _) | ty::Foreign(did) | ty::FnDef(did, _) => Some(did),
+                    ty::Closure(did, _)
+                    | ty::Foreign(did)
+                    | ty::FnDef(did, _)
+                    | ty::Generator(did, ..) => Some(did),
                     ty::Adt(def, _) => Some(def.did),
                     _ => None,
                 };
diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
index 9371ff3405eb0..362a5f47ffe30 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
@@ -1250,33 +1250,40 @@ impl<'a, 'tcx> InferCtxtExt<'tcx> for InferCtxt<'a, 'tcx> {
             trait_ref: ty::PolyTraitRef<'tcx>,
         ) -> String {
             let inputs = trait_ref.skip_binder().substs.type_at(1);
-            let sig = if let ty::Tuple(inputs) = inputs.kind() {
-                tcx.mk_fn_sig(
-                    inputs.iter().map(|k| k.expect_ty()),
-                    tcx.mk_ty_infer(ty::TyVar(ty::TyVid::from_u32(0))),
-                    false,
-                    hir::Unsafety::Normal,
-                    abi::Abi::Rust,
-                )
-            } else {
-                tcx.mk_fn_sig(
+            let sig = match inputs.kind() {
+                ty::Tuple(inputs)
+                    if tcx.fn_trait_kind_from_lang_item(trait_ref.def_id()).is_some() =>
+                {
+                    tcx.mk_fn_sig(
+                        inputs.iter().map(|k| k.expect_ty()),
+                        tcx.mk_ty_infer(ty::TyVar(ty::TyVid::from_u32(0))),
+                        false,
+                        hir::Unsafety::Normal,
+                        abi::Abi::Rust,
+                    )
+                }
+                _ => tcx.mk_fn_sig(
                     std::iter::once(inputs),
                     tcx.mk_ty_infer(ty::TyVar(ty::TyVid::from_u32(0))),
                     false,
                     hir::Unsafety::Normal,
                     abi::Abi::Rust,
-                )
+                ),
             };
             trait_ref.rebind(sig).to_string()
         }
 
-        let argument_is_closure = expected_ref.skip_binder().substs.type_at(0).is_closure();
+        let argument_kind = match expected_ref.skip_binder().substs.type_at(0) {
+            t if t.is_closure() => "closure",
+            t if t.is_generator() => "generator",
+            _ => "function",
+        };
         let mut err = struct_span_err!(
             self.tcx.sess,
             span,
             E0631,
             "type mismatch in {} arguments",
-            if argument_is_closure { "closure" } else { "function" }
+            argument_kind
         );
 
         let found_str = format!("expected signature of `{}`", build_fn_sig_string(self.tcx, found));
diff --git a/src/test/ui/generator/issue-88653.rs b/src/test/ui/generator/issue-88653.rs
new file mode 100644
index 0000000000000..ce9159b53e0f0
--- /dev/null
+++ b/src/test/ui/generator/issue-88653.rs
@@ -0,0 +1,19 @@
+// Regression test for #88653, where a confusing warning about a
+// type mismatch in generator arguments was issued.
+
+#![feature(generators, generator_trait)]
+
+use std::ops::Generator;
+
+fn foo(bar: bool) -> impl Generator<(bool,)> {
+//~^ ERROR: type mismatch in generator arguments [E0631]
+//~| NOTE: expected signature of `fn((bool,)) -> _`
+    |bar| {
+    //~^ NOTE: found signature of `fn(bool) -> _`
+        if bar {
+            yield bar;
+        }
+    }
+}
+
+fn main() {}
diff --git a/src/test/ui/generator/issue-88653.stderr b/src/test/ui/generator/issue-88653.stderr
new file mode 100644
index 0000000000000..5bd8ad129fef9
--- /dev/null
+++ b/src/test/ui/generator/issue-88653.stderr
@@ -0,0 +1,12 @@
+error[E0631]: type mismatch in generator arguments
+  --> $DIR/issue-88653.rs:8:22
+   |
+LL | fn foo(bar: bool) -> impl Generator<(bool,)> {
+   |                      ^^^^^^^^^^^^^^^^^^^^^^^ expected signature of `fn((bool,)) -> _`
+...
+LL |     |bar| {
+   |     ----- found signature of `fn(bool) -> _`
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0631`.