Rollup merge of #143258 - compiler-errors:disambiguator-state, r=oli-obk

Don't recompute `DisambiguatorState` for every RPITIT in trait definition

The `associated_type_for_impl_trait_in_trait` currently needs to rerun the `RPITVisitor` for every RPITIT to compute its disambiguator.

Instead of synthesizing all of the RPITITs def ids one at a time in different queries, just synthesize them inside of the `associated_types_for_impl_traits_in_associated_fn` query. There we can just share the same `DisambiguatorState` for all the RPITITs in one function signature.

r? ``````@Zoxc`````` or ``````@oli-obk`````` cc rust-lang/rust#140453

3 files changed, 47 insertions, 46 deletions

diff --git a/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs b/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
index 6b2854d96af..d2566cda7bf 100644
--- a/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
+++ b/compiler/rustc_hir_analysis/src/hir_ty_lowering/mod.rs
@@ -2459,13 +2459,15 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
                 // type a projection.
                 let in_trait = match opaque_ty.origin {
                     hir::OpaqueTyOrigin::FnReturn {
+                        parent,
                         in_trait_or_impl: Some(hir::RpitContext::Trait),
                         ..
                     }
                     | hir::OpaqueTyOrigin::AsyncFn {
+                        parent,
                         in_trait_or_impl: Some(hir::RpitContext::Trait),
                         ..
-                    } => true,
+                    } => Some(parent),
                     hir::OpaqueTyOrigin::FnReturn {
                         in_trait_or_impl: None | Some(hir::RpitContext::TraitImpl),
                         ..
@@ -2474,7 +2476,7 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
                         in_trait_or_impl: None | Some(hir::RpitContext::TraitImpl),
                         ..
                     }
-                    | hir::OpaqueTyOrigin::TyAlias { .. } => false,
+                    | hir::OpaqueTyOrigin::TyAlias { .. } => None,
                 };
 
                 self.lower_opaque_ty(opaque_ty.def_id, in_trait)
@@ -2594,17 +2596,25 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
 
     /// Lower an opaque type (i.e., an existential impl-Trait type) from the HIR.
     #[instrument(level = "debug", skip(self), ret)]
-    fn lower_opaque_ty(&self, def_id: LocalDefId, in_trait: bool) -> Ty<'tcx> {
+    fn lower_opaque_ty(&self, def_id: LocalDefId, in_trait: Option<LocalDefId>) -> Ty<'tcx> {
         let tcx = self.tcx();
 
         let lifetimes = tcx.opaque_captured_lifetimes(def_id);
         debug!(?lifetimes);
 
-        // If this is an RPITIT and we are using the new RPITIT lowering scheme, we
-        // generate the def_id of an associated type for the trait and return as
-        // type a projection.
-        let def_id = if in_trait {
-            tcx.associated_type_for_impl_trait_in_trait(def_id).to_def_id()
+        // If this is an RPITIT and we are using the new RPITIT lowering scheme,
+        // do a linear search to map this to the synthetic associated type that
+        // it will be lowered to.
+        let def_id = if let Some(parent_def_id) = in_trait {
+            *tcx.associated_types_for_impl_traits_in_associated_fn(parent_def_id)
+                .iter()
+                .find(|rpitit| match tcx.opt_rpitit_info(**rpitit) {
+                    Some(ty::ImplTraitInTraitData::Trait { opaque_def_id, .. }) => {
+                        opaque_def_id.expect_local() == def_id
+                    }
+                    _ => unreachable!(),
+                })
+                .unwrap()
         } else {
             def_id.to_def_id()
         };
@@ -2627,7 +2637,7 @@ impl<'tcx> dyn HirTyLowerer<'tcx> + '_ {
         });
         debug!(?args);
 
-        if in_trait {
+        if in_trait.is_some() {
             Ty::new_projection_from_args(tcx, def_id, args)
         } else {
             Ty::new_opaque(tcx, def_id, args)
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 3e68f0f784e..05154633a38 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1093,13 +1093,6 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    /// Given an impl trait in trait `opaque_ty_def_id`, create and return the corresponding
-    /// associated item.
-    query associated_type_for_impl_trait_in_trait(opaque_ty_def_id: LocalDefId) -> LocalDefId {
-        desc { |tcx| "creating the associated item corresponding to the opaque type `{}`", tcx.def_path_str(opaque_ty_def_id.to_def_id()) }
-        cache_on_disk_if { true }
-    }
-
     /// Given an `impl_id`, return the trait it implements along with some header information.
     /// Return `None` if this is an inherent impl.
     query impl_trait_header(impl_id: DefId) -> Option<ty::ImplTraitHeader<'tcx>> {
diff --git a/compiler/rustc_ty_utils/src/assoc.rs b/compiler/rustc_ty_utils/src/assoc.rs
index f14a45aa1e3..a65f9b347dc 100644
--- a/compiler/rustc_ty_utils/src/assoc.rs
+++ b/compiler/rustc_ty_utils/src/assoc.rs
@@ -1,9 +1,8 @@
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::definitions::{DefPathData, DisambiguatorState};
 use rustc_hir::intravisit::{self, Visitor};
-use rustc_hir::{self as hir, AmbigArg};
 use rustc_middle::query::Providers;
 use rustc_middle::ty::{self, ImplTraitInTraitData, TyCtxt};
 use rustc_middle::{bug, span_bug};
@@ -14,7 +13,6 @@ pub(crate) fn provide(providers: &mut Providers) {
         associated_item_def_ids,
         associated_items,
         associated_types_for_impl_traits_in_associated_fn,
-        associated_type_for_impl_trait_in_trait,
         impl_item_implementor_ids,
         ..*providers
     };
@@ -160,20 +158,22 @@ fn associated_item_from_impl_item_ref(impl_item_ref: &hir::ImplItemRef) -> ty::A
         container: ty::AssocItemContainer::Impl,
     }
 }
-struct RPITVisitor {
-    rpits: FxIndexSet<LocalDefId>,
+struct RPITVisitor<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    synthetics: Vec<LocalDefId>,
+    data: DefPathData,
+    disambiguator: DisambiguatorState,
 }
 
-impl<'tcx> Visitor<'tcx> for RPITVisitor {
-    fn visit_ty(&mut self, ty: &'tcx hir::Ty<'tcx, AmbigArg>) {
-        if let hir::TyKind::OpaqueDef(opaq) = ty.kind
-            && self.rpits.insert(opaq.def_id)
-        {
-            for bound in opaq.bounds {
-                intravisit::walk_param_bound(self, bound);
-            }
-        }
-        intravisit::walk_ty(self, ty)
+impl<'tcx> Visitor<'tcx> for RPITVisitor<'tcx> {
+    fn visit_opaque_ty(&mut self, opaque: &'tcx hir::OpaqueTy<'tcx>) -> Self::Result {
+        self.synthetics.push(associated_type_for_impl_trait_in_trait(
+            self.tcx,
+            opaque.def_id,
+            self.data,
+            &mut self.disambiguator,
+        ));
+        intravisit::walk_opaque_ty(self, opaque)
     }
 }
 
@@ -194,14 +194,18 @@ fn associated_types_for_impl_traits_in_associated_fn(
 
     match tcx.def_kind(parent_def_id) {
         DefKind::Trait => {
-            let mut visitor = RPITVisitor { rpits: FxIndexSet::default() };
-
             if let Some(output) = tcx.hir_get_fn_output(fn_def_id) {
+                let data = DefPathData::AnonAssocTy(tcx.item_name(fn_def_id.to_def_id()));
+                let mut visitor = RPITVisitor {
+                    tcx,
+                    synthetics: vec![],
+                    data,
+                    disambiguator: DisambiguatorState::with(parent_def_id, data, 0),
+                };
                 visitor.visit_fn_ret_ty(output);
-
-                tcx.arena.alloc_from_iter(visitor.rpits.iter().map(|opaque_ty_def_id| {
-                    tcx.associated_type_for_impl_trait_in_trait(opaque_ty_def_id).to_def_id()
-                }))
+                tcx.arena.alloc_from_iter(
+                    visitor.synthetics.into_iter().map(|def_id| def_id.to_def_id()),
+                )
             } else {
                 &[]
             }
@@ -211,7 +215,6 @@ fn associated_types_for_impl_traits_in_associated_fn(
             let Some(trait_fn_def_id) = tcx.associated_item(fn_def_id).trait_item_def_id else {
                 return &[];
             };
-
             tcx.arena.alloc_from_iter(
                 tcx.associated_types_for_impl_traits_in_associated_fn(trait_fn_def_id).iter().map(
                     move |&trait_assoc_def_id| {
@@ -236,6 +239,8 @@ fn associated_types_for_impl_traits_in_associated_fn(
 fn associated_type_for_impl_trait_in_trait(
     tcx: TyCtxt<'_>,
     opaque_ty_def_id: LocalDefId,
+    data: DefPathData,
+    disambiguator: &mut DisambiguatorState,
 ) -> LocalDefId {
     let (hir::OpaqueTyOrigin::FnReturn { parent: fn_def_id, .. }
     | hir::OpaqueTyOrigin::AsyncFn { parent: fn_def_id, .. }) =
@@ -246,22 +251,15 @@ fn associated_type_for_impl_trait_in_trait(
     let trait_def_id = tcx.local_parent(fn_def_id);
     assert_eq!(tcx.def_kind(trait_def_id), DefKind::Trait);
 
-    // Collect all opaque types in return position for the method and use
-    // the index as the disambiguator to make an unique def path.
-    let mut visitor = RPITVisitor { rpits: FxIndexSet::default() };
-    visitor.visit_fn_ret_ty(tcx.hir_get_fn_output(fn_def_id).unwrap());
-    let disambiguator = visitor.rpits.get_index_of(&opaque_ty_def_id).unwrap().try_into().unwrap();
-
     let span = tcx.def_span(opaque_ty_def_id);
     // Also use the method name to create an unique def path.
-    let data = DefPathData::AnonAssocTy(tcx.item_name(fn_def_id.to_def_id()));
     let trait_assoc_ty = tcx.at(span).create_def(
         trait_def_id,
         // No name because this is an anonymous associated type.
         None,
         DefKind::AssocTy,
         Some(data),
-        &mut DisambiguatorState::with(trait_def_id, data, disambiguator),
+        disambiguator,
     );
 
     let local_def_id = trait_assoc_ty.def_id();