From 819d39511ee64b059f0af651ae5d9e6bb91650a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Le=C3=B3n=20Orell=20Valerian=20Liehr?= <me@fmease.dev>
Date: Sat, 2 Mar 2024 19:29:53 +0100
Subject: [PATCH] Properly deal with GATs when looking for method chains to
 point at

---
 compiler/rustc_trait_selection/src/lib.rs     |  1 +
 .../src/traits/error_reporting/suggestions.rs | 14 ++++-----
 tests/ui/typeck/method-chain-gats.rs          | 22 ++++++++++++++
 tests/ui/typeck/method-chain-gats.stderr      | 30 +++++++++++++++++++
 4 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 tests/ui/typeck/method-chain-gats.rs
 create mode 100644 tests/ui/typeck/method-chain-gats.stderr

diff --git a/compiler/rustc_trait_selection/src/lib.rs b/compiler/rustc_trait_selection/src/lib.rs
index 00a2adccf64a5..7c1bfebef2cd0 100644
--- a/compiler/rustc_trait_selection/src/lib.rs
+++ b/compiler/rustc_trait_selection/src/lib.rs
@@ -16,6 +16,7 @@
 #![allow(internal_features)]
 #![allow(rustc::diagnostic_outside_of_impl)]
 #![allow(rustc::untranslatable_diagnostic)]
+#![feature(assert_matches)]
 #![feature(associated_type_bounds)]
 #![feature(box_patterns)]
 #![feature(control_flow_enum)]
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 16290c8dbcafb..2fdc684a81c75 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/suggestions.rs
@@ -38,6 +38,7 @@ use rustc_span::def_id::LocalDefId;
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::{BytePos, DesugaringKind, ExpnKind, MacroKind, Span, DUMMY_SP};
 use rustc_target::spec::abi;
+use std::assert_matches::debug_assert_matches;
 use std::borrow::Cow;
 use std::iter;
 
@@ -4219,18 +4220,13 @@ impl<'tcx> TypeErrCtxt<'_, 'tcx> {
             };
 
             let origin = TypeVariableOrigin { kind: TypeVariableOriginKind::TypeInference, span };
-            let trait_def_id = proj.trait_def_id(self.tcx);
             // Make `Self` be equivalent to the type of the call chain
             // expression we're looking at now, so that we can tell what
             // for example `Iterator::Item` is at this point in the chain.
-            let args = GenericArgs::for_item(self.tcx, trait_def_id, |param, _| {
-                match param.kind {
-                    ty::GenericParamDefKind::Type { .. } => {
-                        if param.index == 0 {
-                            return prev_ty.into();
-                        }
-                    }
-                    ty::GenericParamDefKind::Lifetime | ty::GenericParamDefKind::Const { .. } => {}
+            let args = GenericArgs::for_item(self.tcx, proj.def_id, |param, _| {
+                if param.index == 0 {
+                    debug_assert_matches!(param.kind, ty::GenericParamDefKind::Type { .. });
+                    return prev_ty.into();
                 }
                 self.var_for_def(span, param)
             });
diff --git a/tests/ui/typeck/method-chain-gats.rs b/tests/ui/typeck/method-chain-gats.rs
new file mode 100644
index 0000000000000..c7081c9a3b167
--- /dev/null
+++ b/tests/ui/typeck/method-chain-gats.rs
@@ -0,0 +1,22 @@
+// Regression test for issue #121898.
+
+trait Base {
+    type Base<B>;
+}
+
+trait Functor<A>: Base {
+    fn fmap<B>(self, f: impl Fn(A) -> B) -> Self::Base<B>
+    where
+        Self::Base<B>: Functor<B>;
+}
+
+fn fmap2<T, A, B, C>(input: T, f1: impl Fn(A) -> B, f2: impl Fn(B) -> C) -> T::Base<C>
+where
+    T: Functor<A>,
+    T::Base<B>: Functor<B, Base<C> = T::Base<C>>,
+{
+    input.fmap(f1).fmap(f2)
+    //~^ ERROR the trait bound `<T as Base>::Base<C>: Functor<C>` is not satisfied
+}
+
+fn main() {}
diff --git a/tests/ui/typeck/method-chain-gats.stderr b/tests/ui/typeck/method-chain-gats.stderr
new file mode 100644
index 0000000000000..b0796e07d2eda
--- /dev/null
+++ b/tests/ui/typeck/method-chain-gats.stderr
@@ -0,0 +1,30 @@
+error[E0277]: the trait bound `<T as Base>::Base<C>: Functor<C>` is not satisfied
+  --> $DIR/method-chain-gats.rs:18:20
+   |
+LL |     input.fmap(f1).fmap(f2)
+   |                    ^^^^ the trait `Functor<C>` is not implemented for `<T as Base>::Base<C>`
+   |
+note: the method call chain might not have had the expected associated types
+  --> $DIR/method-chain-gats.rs:13:29
+   |
+LL | fn fmap2<T, A, B, C>(input: T, f1: impl Fn(A) -> B, f2: impl Fn(B) -> C) -> T::Base<C>
+   |                             ^ `Base::Base` is `<T as Base>::Base<_>` here
+...
+LL |     input.fmap(f1).fmap(f2)
+   |           -------- `Base::Base` remains `_` here
+note: required by a bound in `Functor::fmap`
+  --> $DIR/method-chain-gats.rs:10:24
+   |
+LL |     fn fmap<B>(self, f: impl Fn(A) -> B) -> Self::Base<B>
+   |        ---- required by a bound in this associated function
+LL |     where
+LL |         Self::Base<B>: Functor<B>;
+   |                        ^^^^^^^^^^ required by this bound in `Functor::fmap`
+help: consider further restricting the associated type
+   |
+LL |     T::Base<B>: Functor<B, Base<C> = T::Base<C>>, <T as Base>::Base<C>: Functor<C>
+   |                                                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0277`.