From c71a0113a463693cf7817e603f4e89072357e14b Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Sat, 2 Mar 2024 19:56:16 +0000
Subject: [PATCH] Introduce trait_obj_ty query

This query computes the trait object, complete with associated type
projections for its supertraits, from a trait ref.

This is intended for use by CFI labeling, to allow it to determine the
abstract type that will be used by a receiver, based on the instantiated
trait that defines the vtable.
---
 compiler/rustc_middle/src/query/mod.rs        |  9 ++++++
 .../src/traits/vtable.rs                      | 29 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 3984b3b61c294..b305dfd87b535 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -2257,6 +2257,15 @@ rustc_queries! {
     query find_field((def_id, ident): (DefId, rustc_span::symbol::Ident)) -> Option<rustc_target::abi::FieldIdx> {
         desc { |tcx| "find the index of maybe nested field `{ident}` in `{}`", tcx.def_path_str(def_id) }
     }
+
+    /// Construct a type for a trait object corresponding to `trait_ref`. This type will have all
+    /// associated types for it and its supertraits expanded and resolved as additional predicates.
+    ///
+    /// The provided `trait_ref` must be sufficiently instantiated that all associated types can be
+    /// successfully resolved.
+    query trait_object_ty(trait_ref: ty::PolyTraitRef<'tcx>) -> Ty<'tcx> {
+        desc { "Compute the trait object type for calling a method on a trait" }
+    }
 }
 
 rustc_query_append! { define_callbacks! }
diff --git a/compiler/rustc_trait_selection/src/traits/vtable.rs b/compiler/rustc_trait_selection/src/traits/vtable.rs
index 46a68508753c0..7d2a20cba4547 100644
--- a/compiler/rustc_trait_selection/src/traits/vtable.rs
+++ b/compiler/rustc_trait_selection/src/traits/vtable.rs
@@ -1,4 +1,5 @@
 use crate::errors::DumpVTableEntries;
+use crate::traits;
 use crate::traits::{impossible_predicates, is_vtable_safe_method};
 use rustc_hir::def_id::DefId;
 use rustc_hir::lang_items::LangItem;
@@ -13,6 +14,7 @@ use rustc_span::{sym, Span};
 use smallvec::SmallVec;
 
 use std::fmt::Debug;
+use std::iter;
 use std::ops::ControlFlow;
 
 #[derive(Clone, Debug)]
@@ -232,6 +234,32 @@ fn own_existential_vtable_entries_iter(
     own_entries
 }
 
+fn trait_object_ty<'tcx>(tcx: TyCtxt<'tcx>, poly_trait_ref: ty::PolyTraitRef<'tcx>) -> Ty<'tcx> {
+    let principal_pred = poly_trait_ref.map_bound(|trait_ref| {
+        ty::ExistentialPredicate::Trait(ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref))
+    });
+    let assoc_preds = traits::supertraits(tcx, poly_trait_ref).flat_map(|super_poly_trait_ref| {
+        tcx.associated_items(super_poly_trait_ref.def_id())
+            .in_definition_order()
+            .filter(|item| item.kind == ty::AssocKind::Type)
+            .map(move |assoc_ty| {
+                super_poly_trait_ref.map_bound(|super_trait_ref| {
+                    let alias_ty = ty::AliasTy::new(tcx, assoc_ty.def_id, super_trait_ref.args);
+                    let resolved = tcx
+                        .normalize_erasing_regions(ty::ParamEnv::reveal_all(), alias_ty.to_ty(tcx));
+                    ty::ExistentialPredicate::Projection(ty::ExistentialProjection {
+                        def_id: assoc_ty.def_id,
+                        args: ty::ExistentialTraitRef::erase_self_ty(tcx, super_trait_ref).args,
+                        term: resolved.into(),
+                    })
+                })
+            })
+    });
+    let preds =
+        tcx.mk_poly_existential_predicates_from_iter(iter::once(principal_pred).chain(assoc_preds));
+    Ty::new_dynamic(tcx, preds, tcx.lifetimes.re_erased, ty::Dyn)
+}
+
 /// Given a trait `trait_ref`, iterates the vtable entries
 /// that come from `trait_ref`, including its supertraits.
 fn vtable_entries<'tcx>(
@@ -404,6 +432,7 @@ pub(super) fn provide(providers: &mut Providers) {
         own_existential_vtable_entries,
         vtable_entries,
         vtable_trait_upcasting_coercion_new_vptr_slot,
+        trait_object_ty,
         ..*providers
     };
 }