From a10406318e0888b1356e88ec8866fe260a55e6aa Mon Sep 17 00:00:00 2001
From: Santiago Pastorino <spastorino@gmail.com>
Date: Thu, 29 Jun 2023 23:08:32 -0300
Subject: [PATCH] Properly implement variances_of for RPITIT GAT

---
 .../rustc_hir_analysis/src/variance/mod.rs    | 24 ++++++++++++-------
 .../impl-trait/in-trait/variances-of-gat.rs   | 19 +++++++++++++++
 2 files changed, 34 insertions(+), 9 deletions(-)
 create mode 100644 tests/ui/impl-trait/in-trait/variances-of-gat.rs

diff --git a/compiler/rustc_hir_analysis/src/variance/mod.rs b/compiler/rustc_hir_analysis/src/variance/mod.rs
index 23d8da88a454b..49aee6b59a28c 100644
--- a/compiler/rustc_hir_analysis/src/variance/mod.rs
+++ b/compiler/rustc_hir_analysis/src/variance/mod.rs
@@ -7,7 +7,7 @@ use rustc_arena::DroplessArena;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_middle::query::Providers;
-use rustc_middle::ty::{self, CrateVariancesMap, SubstsRef, Ty, TyCtxt};
+use rustc_middle::ty::{self, CrateVariancesMap, ImplTraitInTraitData, SubstsRef, Ty, TyCtxt};
 use rustc_middle::ty::{TypeSuperVisitable, TypeVisitable};
 use std::ops::ControlFlow;
 
@@ -51,20 +51,26 @@ fn variances_of(tcx: TyCtxt<'_>, item_def_id: LocalDefId) -> &[ty::Variance] {
         | DefKind::Struct
         | DefKind::Union
         | DefKind::Variant
-        | DefKind::Ctor(..) => {}
+        | DefKind::Ctor(..) => {
+            // These are inferred.
+            let crate_map = tcx.crate_variances(());
+            return crate_map.variances.get(&item_def_id.to_def_id()).copied().unwrap_or(&[]);
+        }
         DefKind::OpaqueTy | DefKind::ImplTraitPlaceholder => {
             return variance_of_opaque(tcx, item_def_id);
         }
-        _ => {
-            // Variance not relevant.
-            span_bug!(tcx.def_span(item_def_id), "asked to compute variance for wrong kind of item")
+        DefKind::AssocTy => {
+            if let Some(ImplTraitInTraitData::Trait { .. }) =
+                tcx.opt_rpitit_info(item_def_id.to_def_id())
+            {
+                return variance_of_opaque(tcx, item_def_id);
+            }
         }
+        _ => {}
     }
 
-    // Everything else must be inferred.
-
-    let crate_map = tcx.crate_variances(());
-    crate_map.variances.get(&item_def_id.to_def_id()).copied().unwrap_or(&[])
+    // Variance not relevant.
+    span_bug!(tcx.def_span(item_def_id), "asked to compute variance for wrong kind of item");
 }
 
 #[instrument(level = "trace", skip(tcx), ret)]
diff --git a/tests/ui/impl-trait/in-trait/variances-of-gat.rs b/tests/ui/impl-trait/in-trait/variances-of-gat.rs
new file mode 100644
index 0000000000000..4008ece94da91
--- /dev/null
+++ b/tests/ui/impl-trait/in-trait/variances-of-gat.rs
@@ -0,0 +1,19 @@
+// check-pass
+// [next] compile-flags: -Zlower-impl-trait-in-trait-to-assoc-ty
+// revisions: current next
+
+#![feature(return_position_impl_trait_in_trait)]
+
+trait Foo {}
+
+impl Foo for () {}
+
+trait ThreeCellFragment {
+    fn ext_cells<'a>(&'a self) -> impl Foo + 'a {
+        self.ext_adjacent_cells()
+    }
+
+    fn ext_adjacent_cells<'a>(&'a self) -> impl Foo + 'a;
+}
+
+fn main() {}