Compute generator saved locals on MIR.

Author: cjgillot

compiler/rustc_hir_analysis/src/check/check.rs

@@ -14,7 +14,7 @@ use rustc_hir::{ItemKind, Node, PathSegment};
 use rustc_infer::infer::opaque_types::ConstrainOpaqueTypeRegionVisitor;
 use rustc_infer::infer::outlives::env::OutlivesEnvironment;
 use rustc_infer::infer::{DefiningAnchor, RegionVariableOrigin, TyCtxtInferExt};
-use rustc_infer::traits::Obligation;
+use rustc_infer::traits::{Obligation, TraitEngineExt as _};
 use rustc_lint::builtin::REPR_TRANSPARENT_EXTERNAL_PRIVATE_FIELDS;
 use rustc_middle::hir::nested_filter;
 use rustc_middle::middle::stability::EvalResult;
@@ -28,7 +28,7 @@ use rustc_span::{self, Span};
 use rustc_target::spec::abi::Abi;
 use rustc_trait_selection::traits::error_reporting::on_unimplemented::OnUnimplementedDirective;
 use rustc_trait_selection::traits::error_reporting::TypeErrCtxtExt as _;
-use rustc_trait_selection::traits::{self, ObligationCtxt};
+use rustc_trait_selection::traits::{self, ObligationCtxt, TraitEngine, TraitEngineExt as _};
 
 use std::ops::ControlFlow;
 
@@ -1460,7 +1460,8 @@ fn opaque_type_cycle_error(
                 for def_id in visitor.opaques {
                     let ty_span = tcx.def_span(def_id);
                     if !seen.contains(&ty_span) {
-                        err.span_label(ty_span, &format!("returning this opaque type `{ty}`"));
+                        let descr = if ty.is_impl_trait() { "opaque " } else { "" };
+                        err.span_label(ty_span, &format!("returning this {descr}type `{ty}`"));
                         seen.insert(ty_span);
                     }
                     err.span_label(sp, &format!("returning here with type `{ty}`"));
@@ -1507,3 +1508,34 @@ fn opaque_type_cycle_error(
     }
     err.emit()
 }
+
+pub(super) fn check_generator_obligations(tcx: TyCtxt<'_>, def_id: LocalDefId) {
+    debug_assert!(tcx.sess.opts.unstable_opts.drop_tracking_mir);
+    debug_assert!(matches!(tcx.def_kind(def_id), DefKind::Generator));
+
+    let typeck = tcx.typeck(def_id);
+    let param_env = tcx.param_env(def_id);
+
+    let generator_interior_predicates = &typeck.generator_interior_predicates[&def_id];
+    debug!(?generator_interior_predicates);
+
+    let infcx = tcx
+        .infer_ctxt()
+        // typeck writeback gives us predicates with their regions erased.
+        // As borrowck already has checked lifetimes, we do not need to do it again.
+        .ignoring_regions()
+        // Bind opaque types to `def_id` as they should have been checked by borrowck.
+        .with_opaque_type_inference(DefiningAnchor::Bind(def_id))
+        .build();
+
+    let mut fulfillment_cx = <dyn TraitEngine<'_>>::new(infcx.tcx);
+    for (predicate, cause) in generator_interior_predicates {
+        let obligation = Obligation::new(tcx, cause.clone(), param_env, *predicate);
+        fulfillment_cx.register_predicate_obligation(&infcx, obligation);
+    }
+    let errors = fulfillment_cx.select_all_or_error(&infcx);
+    debug!(?errors);
+    if !errors.is_empty() {
+        infcx.err_ctxt().report_fulfillment_errors(&errors, None);
+    }
+}

compiler/rustc_hir_analysis/src/check/mod.rs

@@ -105,6 +105,7 @@ pub fn provide(providers: &mut Providers) {
         region_scope_tree,
         collect_return_position_impl_trait_in_trait_tys,
         compare_impl_const: compare_impl_item::compare_impl_const_raw,
+        check_generator_obligations: check::check_generator_obligations,
         ..*providers
     };
 }

compiler/rustc_hir_typeck/src/check.rs

@@ -130,7 +130,7 @@ pub(super) fn check_fn<'a, 'tcx>(
     let gen_ty = if let (Some(_), Some(gen_kind)) = (can_be_generator, body.generator_kind) {
         let interior = fcx
             .next_ty_var(TypeVariableOrigin { kind: TypeVariableOriginKind::MiscVariable, span });
-        fcx.deferred_generator_interiors.borrow_mut().push((body.id(), interior, gen_kind));
+        fcx.deferred_generator_interiors.borrow_mut().push((fn_id, body.id(), interior, gen_kind));
 
         let (resume_ty, yield_ty) = fcx.resume_yield_tys.unwrap();
         Some(GeneratorTypes {

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -517,10 +517,67 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     }
 
     pub(in super::super) fn resolve_generator_interiors(&self, def_id: DefId) {
+        if self.tcx.sess.opts.unstable_opts.drop_tracking_mir {
+            self.save_generator_interior_predicates(def_id);
+            return;
+        }
+
+        self.select_obligations_where_possible(|_| {});
+
         let mut generators = self.deferred_generator_interiors.borrow_mut();
-        for (body_id, interior, kind) in generators.drain(..) {
-            self.select_obligations_where_possible(|_| {});
+        for (_, body_id, interior, kind) in generators.drain(..) {
             crate::generator_interior::resolve_interior(self, def_id, body_id, interior, kind);
+            self.select_obligations_where_possible(|_| {});
+        }
+    }
+
+    /// Unify the inference variables corresponding to generator witnesses, and save all the
+    /// predicates that were stalled on those inference variables.
+    ///
+    /// This process allows to conservatively save all predicates that do depend on the generator
+    /// interior types, for later processing by `check_generator_obligations`.
+    ///
+    /// We must not attempt to select obligations after this method has run, or risk query cycle
+    /// ICE.
+    #[instrument(level = "debug", skip(self))]
+    fn save_generator_interior_predicates(&self, def_id: DefId) {
+        // Try selecting all obligations that are not blocked on inference variables.
+        // Once we start unifying generator witnesses, trying to select obligations on them will
+        // trigger query cycle ICEs, as doing so requires MIR.
+        self.select_obligations_where_possible(|_| {});
+
+        let generators = std::mem::take(&mut *self.deferred_generator_interiors.borrow_mut());
+        debug!(?generators);
+
+        for &(expr_hir_id, body_id, interior, _) in generators.iter() {
+            let expr_def_id = self.tcx.hir().local_def_id(expr_hir_id);
+            debug!(?expr_def_id);
+
+            // Create the `GeneratorWitness` type that we will unify with `interior`.
+            let substs = ty::InternalSubsts::identity_for_item(
+                self.tcx,
+                self.tcx.typeck_root_def_id(expr_def_id.to_def_id()),
+            );
+            let witness = self.tcx.mk_generator_witness_mir(expr_def_id.to_def_id(), substs);
+
+            // Unify `interior` with `witness` and collect all the resulting obligations.
+            let span = self.tcx.hir().body(body_id).value.span;
+            let ok = self
+                .at(&self.misc(span), self.param_env)
+                .eq(interior, witness)
+                .expect("Failed to unify generator interior type");
+            let mut obligations = ok.obligations;
+
+            // Also collect the obligations that were unstalled by this unification.
+            obligations
+                .extend(self.fulfillment_cx.borrow_mut().drain_unstalled_obligations(&self.infcx));
+
+            let obligations = obligations.into_iter().map(|o| (o.predicate, o.cause)).collect();
+            debug!(?obligations);
+            self.typeck_results
+                .borrow_mut()
+                .generator_interior_predicates
+                .insert(expr_def_id, obligations);
         }
     }
 

compiler/rustc_hir_typeck/src/inherited.rs

@@ -56,7 +56,7 @@ pub struct Inherited<'tcx> {
     pub(super) deferred_asm_checks: RefCell<Vec<(&'tcx hir::InlineAsm<'tcx>, hir::HirId)>>,
 
     pub(super) deferred_generator_interiors:
-        RefCell<Vec<(hir::BodyId, Ty<'tcx>, hir::GeneratorKind)>>,
+        RefCell<Vec<(hir::HirId, hir::BodyId, Ty<'tcx>, hir::GeneratorKind)>>,
 
     pub(super) body_id: Option<hir::BodyId>,
 

compiler/rustc_hir_typeck/src/lib.rs

@@ -294,7 +294,6 @@ fn typeck_with_fallback<'tcx>(
         // Before the generator analysis, temporary scopes shall be marked to provide more
         // precise information on types to be captured.
         fcx.resolve_rvalue_scopes(def_id.to_def_id());
-        fcx.resolve_generator_interiors(def_id.to_def_id());
 
         for (ty, span, code) in fcx.deferred_sized_obligations.borrow_mut().drain(..) {
             let ty = fcx.normalize(span, ty);
@@ -303,6 +302,13 @@ fn typeck_with_fallback<'tcx>(
 
         fcx.select_obligations_where_possible(|_| {});
 
+        debug!(pending_obligations = ?fcx.fulfillment_cx.borrow().pending_obligations());
+
+        // This must be the last thing before `report_ambiguity_errors`.
+        fcx.resolve_generator_interiors(def_id.to_def_id());
+
+        debug!(pending_obligations = ?fcx.fulfillment_cx.borrow().pending_obligations());
+
         if let None = fcx.infcx.tainted_by_errors() {
             fcx.report_ambiguity_errors();
         }

compiler/rustc_hir_typeck/src/writeback.rs

@@ -545,6 +545,10 @@ impl<'cx, 'tcx> WritebackCx<'cx, 'tcx> {
         assert_eq!(fcx_typeck_results.hir_owner, self.typeck_results.hir_owner);
         self.typeck_results.generator_interior_types =
             fcx_typeck_results.generator_interior_types.clone();
+        for (&expr_def_id, predicates) in fcx_typeck_results.generator_interior_predicates.iter() {
+            let predicates = self.resolve(predicates.clone(), &self.fcx.tcx.def_span(expr_def_id));
+            self.typeck_results.generator_interior_predicates.insert(expr_def_id, predicates);
+        }
     }
 
     #[instrument(skip(self), level = "debug")]

compiler/rustc_infer/src/traits/engine.rs

@@ -41,6 +41,14 @@ pub trait TraitEngine<'tcx>: 'tcx {
     fn collect_remaining_errors(&mut self) -> Vec<FulfillmentError<'tcx>>;
 
     fn pending_obligations(&self) -> Vec<PredicateObligation<'tcx>>;
+
+    /// Among all pending obligations, collect those are stalled on a inference variable which has
+    /// changed since the last call to `select_where_possible`.  Those obligations are marked as
+    /// successful and returned.
+    fn drain_unstalled_obligations(
+        &mut self,
+        infcx: &InferCtxt<'tcx>,
+    ) -> Vec<PredicateObligation<'tcx>>;
 }
 
 pub trait TraitEngineExt<'tcx> {

compiler/rustc_interface/src/passes.rs

@@ -893,6 +893,15 @@ fn analysis(tcx: TyCtxt<'_>, (): ()) -> Result<()> {
         }
     });
 
+    if tcx.sess.opts.unstable_opts.drop_tracking_mir {
+        tcx.hir().par_body_owners(|def_id| {
+            if let rustc_hir::def::DefKind::Generator = tcx.def_kind(def_id) {
+                tcx.ensure().mir_generator_witnesses(def_id);
+                tcx.ensure().check_generator_obligations(def_id);
+            }
+        });
+    }
+
     sess.time("layout_testing", || layout_test::test_layout(tcx));
 
     // Avoid overwhelming user with errors if borrow checking failed.

compiler/rustc_middle/src/query/mod.rs

@@ -478,6 +478,10 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    query check_generator_obligations(key: LocalDefId) {
+        desc { |tcx| "verify auto trait bounds for generator interior type `{}`", tcx.def_path_str(key.to_def_id()) }
+    }
+
     /// MIR after our optimization passes have run. This is MIR that is ready
     /// for codegen. This is also the only query that can fetch non-local MIR, at present.
     query optimized_mir(key: DefId) -> &'tcx mir::Body<'tcx> {