Rollup merge of #108141 - spastorino:add_rpitit_queries, r=compiler-errors

Add rpitit queries

This is part of the changes we are making to lower RPITITs as an associated type. The rest of the stuff will follow under a `-Z` flag.

I still need to add comments to the code, explain stuff and also I'd need to avoid encoding in metadata when rpitit queries return `&[]`

r? `@compiler-errors`

compiler/rustc_ty_utils/src/assoc.rs

@@ -1,13 +1,18 @@
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
-use rustc_hir::def_id::DefId;
-use rustc_middle::ty::{self, TyCtxt};
+use rustc_hir::def::DefKind;
+use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_hir::definitions::DefPathData;
+use rustc_hir::intravisit::{self, Visitor};
+use rustc_middle::ty::{self, DefIdTree, TyCtxt};
 
 pub fn provide(providers: &mut ty::query::Providers) {
     *providers = ty::query::Providers {
         associated_item,
         associated_item_def_ids,
         associated_items,
+        associated_items_for_impl_trait_in_trait,
+        associated_item_for_impl_trait_in_trait,
         impl_item_implementor_ids,
         ..*providers
     };
@@ -112,3 +117,97 @@ fn associated_item_from_impl_item_ref(impl_item_ref: &hir::ImplItemRef) -> ty::A
         fn_has_self_parameter: has_self,
     }
 }
+
+/// Given an `fn_def_id` of a trait or of an impl that implements a given trait:
+/// if `fn_def_id` is the def id of a function defined inside a trait, then it creates and returns
+/// the associated items that correspond to each impl trait in return position for that trait.
+/// if `fn_def_id` is the def id of a function defined inside an impl that implements a trait, then it
+/// creates and returns the associated items that correspond to each impl trait in return position
+/// of the implemented trait.
+fn associated_items_for_impl_trait_in_trait(tcx: TyCtxt<'_>, fn_def_id: DefId) -> &'_ [DefId] {
+    let parent_def_id = tcx.parent(fn_def_id);
+
+    match tcx.def_kind(parent_def_id) {
+        DefKind::Trait => {
+            struct RPITVisitor {
+                rpits: Vec<LocalDefId>,
+            }
+
+            impl<'v> Visitor<'v> for RPITVisitor {
+                fn visit_ty(&mut self, ty: &'v hir::Ty<'v>) {
+                    if let hir::TyKind::OpaqueDef(item_id, _, _) = ty.kind {
+                        self.rpits.push(item_id.owner_id.def_id)
+                    }
+                    intravisit::walk_ty(self, ty)
+                }
+            }
+
+            let mut visitor = RPITVisitor { rpits: Vec::new() };
+
+            if let Some(output) = tcx.hir().get_fn_output(fn_def_id.expect_local()) {
+                visitor.visit_fn_ret_ty(output);
+
+                tcx.arena.alloc_from_iter(visitor.rpits.iter().map(|opaque_ty_def_id| {
+                    tcx.associated_item_for_impl_trait_in_trait(opaque_ty_def_id).to_def_id()
+                }))
+            } else {
+                &[]
+            }
+        }
+
+        DefKind::Impl { .. } => {
+            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_items_for_impl_trait_in_trait(trait_fn_def_id).iter().map(
+                    move |trait_assoc_def_id| {
+                        impl_associated_item_for_impl_trait_in_trait(
+                            tcx,
+                            trait_assoc_def_id.expect_local(),
+                            fn_def_id.expect_local(),
+                        )
+                        .to_def_id()
+                    },
+                ),
+            )
+        }
+
+        def_kind => bug!(
+            "associated_items_for_impl_trait_in_trait: {:?} should be Trait or Impl but is {:?}",
+            parent_def_id,
+            def_kind
+        ),
+    }
+}
+
+/// Given an `opaque_ty_def_id` corresponding to an impl trait in trait, create and return the
+/// corresponding associated item.
+fn associated_item_for_impl_trait_in_trait(
+    tcx: TyCtxt<'_>,
+    opaque_ty_def_id: LocalDefId,
+) -> LocalDefId {
+    let fn_def_id = tcx.impl_trait_in_trait_parent(opaque_ty_def_id.to_def_id());
+    let trait_def_id = tcx.parent(fn_def_id);
+    assert_eq!(tcx.def_kind(trait_def_id), DefKind::Trait);
+
+    let span = tcx.def_span(opaque_ty_def_id);
+    let trait_assoc_ty =
+        tcx.at(span).create_def(trait_def_id.expect_local(), DefPathData::ImplTraitAssocTy);
+    trait_assoc_ty.def_id()
+}
+
+/// Given an `trait_assoc_def_id` that corresponds to a previously synthethized impl trait in trait
+/// into an associated type and an `impl_def_id` corresponding to an impl block, create and return
+/// the corresponding associated item inside the impl block.
+fn impl_associated_item_for_impl_trait_in_trait(
+    tcx: TyCtxt<'_>,
+    trait_assoc_def_id: LocalDefId,
+    impl_fn_def_id: LocalDefId,
+) -> LocalDefId {
+    let impl_def_id = tcx.local_parent(impl_fn_def_id);
+
+    let span = tcx.def_span(trait_assoc_def_id);
+    let impl_assoc_ty = tcx.at(span).create_def(impl_def_id, DefPathData::ImplTraitAssocTy);
+
+    impl_assoc_ty.def_id()
+}