Double check that hidden types match the expected hidden type

compiler/rustc_hir_analysis/src/check/check.rs

@@ -19,11 +19,13 @@ use rustc_lint_defs::builtin::REPR_TRANSPARENT_EXTERNAL_PRIVATE_FIELDS;
 use rustc_middle::hir::nested_filter;
 use rustc_middle::middle::stability::EvalResult;
 use rustc_middle::traits::DefiningAnchor;
+use rustc_middle::ty::fold::BottomUpFolder;
 use rustc_middle::ty::layout::{LayoutError, MAX_SIMD_LANES};
 use rustc_middle::ty::util::{Discr, IntTypeExt};
 use rustc_middle::ty::GenericArgKind;
 use rustc_middle::ty::{
-    self, AdtDef, ParamEnv, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitableExt,
+    self, AdtDef, ParamEnv, RegionKind, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable,
+    TypeVisitableExt,
 };
 use rustc_session::lint::builtin::{UNINHABITED_STATIC, UNSUPPORTED_CALLING_CONVENTIONS};
 use rustc_span::symbol::sym;
@@ -34,6 +36,7 @@ use rustc_trait_selection::traits::error_reporting::on_unimplemented::OnUnimplem
 use rustc_trait_selection::traits::error_reporting::TypeErrCtxtExt as _;
 use rustc_trait_selection::traits::outlives_bounds::InferCtxtExt as _;
 use rustc_trait_selection::traits::{self, ObligationCtxt, TraitEngine, TraitEngineExt as _};
+use rustc_type_ir::fold::TypeFoldable;
 
 use std::ops::ControlFlow;
 
@@ -437,7 +440,7 @@ fn check_opaque_meets_bounds<'tcx>(
     // hidden type is well formed even without those bounds.
     let predicate =
         ty::Binder::dummy(ty::PredicateKind::Clause(ty::ClauseKind::WellFormed(hidden_ty.into())));
-    ocx.register_obligation(Obligation::new(tcx, misc_cause, param_env, predicate));
+    ocx.register_obligation(Obligation::new(tcx, misc_cause.clone(), param_env, predicate));
 
     // Check that all obligations are satisfied by the implementation's
     // version.
@@ -464,11 +467,143 @@ fn check_opaque_meets_bounds<'tcx>(
             ocx.resolve_regions_and_report_errors(defining_use_anchor, &outlives_env)?;
         }
     }
-    // Clean up after ourselves
-    let _ = infcx.take_opaque_types();
+    // Check that any hidden types found during wf checking match the hidden types that `type_of` sees.
+    for (key, mut ty) in infcx.take_opaque_types() {
+        ty.hidden_type.ty = infcx.resolve_vars_if_possible(ty.hidden_type.ty);
+        sanity_check_found_hidden_type(tcx, key, ty.hidden_type, defining_use_anchor, origin)?;
+    }
     Ok(())
 }
 
+fn sanity_check_found_hidden_type<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    key: ty::OpaqueTypeKey<'tcx>,
+    mut ty: ty::OpaqueHiddenType<'tcx>,
+    defining_use_anchor: LocalDefId,
+    origin: &hir::OpaqueTyOrigin,
+) -> Result<(), ErrorGuaranteed> {
+    if ty.ty.is_ty_var() {
+        // Nothing was actually constrained.
+        return Ok(());
+    }
+    if let ty::Alias(ty::Opaque, alias) = ty.ty.kind() {
+        if alias.def_id == key.def_id.to_def_id() && alias.args == key.args {
+            // Nothing was actually constrained, this is an opaque usage that was
+            // only discovered to be opaque after inference vars resolved.
+            return Ok(());
+        }
+    }
+    // Closures frequently end up containing erased lifetimes in their final representation.
+    // These correspond to lifetime variables that never got resolved, so we patch this up here.
+    ty.ty = ty.ty.fold_with(&mut BottomUpFolder {
+        tcx,
+        ty_op: |t| t,
+        ct_op: |c| c,
+        lt_op: |l| match l.kind() {
+            RegionKind::ReVar(_) => tcx.lifetimes.re_erased,
+            _ => l,
+        },
+    });
+    // Get the hidden type, and in case it is in a nested opaque type, find that opaque type's
+    // usage in the function signature and use the generic arguments from the usage site.
+    let mut hidden_ty = tcx.type_of(key.def_id).instantiate(tcx, key.args);
+    if let hir::OpaqueTyOrigin::FnReturn(..) | hir::OpaqueTyOrigin::AsyncFn(..) = origin {
+        if hidden_ty != ty.ty {
+            hidden_ty = find_and_apply_rpit_substs(
+                tcx,
+                hidden_ty,
+                defining_use_anchor.to_def_id(),
+                key.def_id.to_def_id(),
+            )?;
+        }
+    }
+    // If the hidden types differ, emit a type mismatch diagnostic.
+    if hidden_ty == ty.ty {
+        Ok(())
+    } else {
+        let span = tcx.def_span(key.def_id);
+        let other = ty::OpaqueHiddenType { ty: hidden_ty, span };
+        Err(ty.report_mismatch(&other, key.def_id, tcx).emit())
+    }
+}
+
+fn find_and_apply_rpit_substs<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mut hidden_ty: Ty<'tcx>,
+    function: DefId,
+    opaque: DefId,
+) -> Result<Ty<'tcx>, ErrorGuaranteed> {
+    // Find use of the RPIT in the function signature and thus find the right substs to
+    // convert it into the parameter space of the function signature. This is needed,
+    // because that's what `type_of` returns, against which we compare later.
+    let ret = tcx.fn_sig(function).instantiate_identity().output();
+    struct Visitor<'tcx> {
+        tcx: TyCtxt<'tcx>,
+        opaque: DefId,
+        function: DefId,
+        seen: FxHashSet<DefId>,
+    }
+    impl<'tcx> ty::TypeVisitor<TyCtxt<'tcx>> for Visitor<'tcx> {
+        type BreakTy = GenericArgsRef<'tcx>;
+
+        #[instrument(level = "trace", skip(self), ret)]
+        fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
+            trace!("{:#?}", t.kind());
+            match t.kind() {
+                ty::Alias(ty::Opaque, alias) => {
+                    trace!(?alias.def_id);
+                    if alias.def_id == self.opaque {
+                        return ControlFlow::Break(alias.args);
+                    } else if self.seen.insert(alias.def_id) {
+                        for clause in self
+                            .tcx
+                            .explicit_item_bounds(alias.def_id)
+                            .iter_instantiated_copied(self.tcx, alias.args)
+                        {
+                            trace!(?clause);
+                            clause.visit_with(self)?;
+                        }
+                    }
+                }
+                ty::Alias(ty::Projection, alias) => {
+                    if self.tcx.is_impl_trait_in_trait(alias.def_id)
+                        && self.tcx.impl_trait_in_trait_parent_fn(alias.def_id) == self.function
+                    {
+                        // If we're lowering to associated item, install the opaque type which is just
+                        // the `type_of` of the trait's associated item. If we're using the old lowering
+                        // strategy, then just reinterpret the associated type like an opaque :^)
+                        self.tcx
+                            .type_of(alias.def_id)
+                            .instantiate(self.tcx, alias.args)
+                            .visit_with(self)?;
+                    }
+                }
+                ty::Alias(ty::Weak, alias) => {
+                    self.tcx
+                        .type_of(alias.def_id)
+                        .instantiate(self.tcx, alias.args)
+                        .visit_with(self)?;
+                }
+                _ => (),
+            }
+
+            t.super_visit_with(self)
+        }
+    }
+    if let ControlFlow::Break(args) =
+        ret.visit_with(&mut Visitor { tcx, function, opaque, seen: Default::default() })
+    {
+        trace!(?args);
+        trace!("expected: {hidden_ty:#?}");
+        hidden_ty = ty::EarlyBinder::bind(hidden_ty).instantiate(tcx, args);
+        trace!("expected: {hidden_ty:#?}");
+    } else {
+        tcx.sess
+            .delay_span_bug(tcx.def_span(function), format!("{ret:?} does not contain {opaque:?}"));
+    }
+    Ok(hidden_ty)
+}
+
 fn is_enum_of_nonnullable_ptr<'tcx>(
     tcx: TyCtxt<'tcx>,
     adt_def: AdtDef<'tcx>,