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> {

compiler/rustc_middle/src/ty/typeck_results.rs

@@ -1,6 +1,7 @@
 use crate::{
     hir::place::Place as HirPlace,
     infer::canonical::Canonical,
+    traits::ObligationCause,
     ty::{
         self, tls, BindingMode, BoundVar, CanonicalPolyFnSig, ClosureSizeProfileData,
         GenericArgKind, InternalSubsts, SubstsRef, Ty, UserSubsts,
@@ -193,6 +194,11 @@ pub struct TypeckResults<'tcx> {
     /// that are live across the yield of this generator (if a generator).
     pub generator_interior_types: ty::Binder<'tcx, Vec<GeneratorInteriorTypeCause<'tcx>>>,
 
+    /// Stores the predicates that apply on generator witness types.
+    /// formatting modified file tests/ui/generator/retain-resume-ref.rs
+    pub generator_interior_predicates:
+        FxHashMap<LocalDefId, Vec<(ty::Predicate<'tcx>, ObligationCause<'tcx>)>>,
+
     /// We sometimes treat byte string literals (which are of type `&[u8; N]`)
     /// as `&[u8]`, depending on the pattern in which they are used.
     /// This hashset records all instances where we behave
@@ -271,6 +277,7 @@ impl<'tcx> TypeckResults<'tcx> {
             closure_fake_reads: Default::default(),
             rvalue_scopes: Default::default(),
             generator_interior_types: ty::Binder::dummy(Default::default()),
+            generator_interior_predicates: Default::default(),
             treat_byte_string_as_slice: Default::default(),
             closure_size_eval: Default::default(),
         }

compiler/rustc_middle/src/ty/util.rs

@@ -615,6 +615,36 @@ impl<'tcx> TyCtxt<'tcx> {
         }
     }
 
+    /// Return the set of types that should be taken into accound when checking
+    /// trait bounds on a generator's internal state.
+    pub fn generator_hidden_types(
+        self,
+        def_id: DefId,
+    ) -> impl Iterator<Item = ty::EarlyBinder<Ty<'tcx>>> {
+        let generator_layout = &self.mir_generator_witnesses(def_id);
+        generator_layout
+            .field_tys
+            .iter()
+            .filter(|decl| !decl.is_static_ptr)
+            .map(|decl| ty::EarlyBinder(decl.ty))
+    }
+
+    /// Normalizes all opaque types in the given value, replacing them
+    /// with their underlying types.
+    pub fn expand_opaque_types(self, val: Ty<'tcx>) -> Ty<'tcx> {
+        let mut visitor = OpaqueTypeExpander {
+            seen_opaque_tys: FxHashSet::default(),
+            expanded_cache: FxHashMap::default(),
+            primary_def_id: None,
+            found_recursion: false,
+            found_any_recursion: false,
+            check_recursion: false,
+            expand_generators: false,
+            tcx: self,
+        };
+        val.fold_with(&mut visitor)
+    }
+
     /// Expands the given impl trait type, stopping if the type is recursive.
     #[instrument(skip(self), level = "debug", ret)]
     pub fn try_expand_impl_trait_type(
@@ -629,6 +659,7 @@ impl<'tcx> TyCtxt<'tcx> {
             found_recursion: false,
             found_any_recursion: false,
             check_recursion: true,
+            expand_generators: true,
             tcx: self,
         };
 
@@ -741,6 +772,7 @@ struct OpaqueTypeExpander<'tcx> {
     primary_def_id: Option<DefId>,
     found_recursion: bool,
     found_any_recursion: bool,
+    expand_generators: bool,
     /// Whether or not to check for recursive opaque types.
     /// This is `true` when we're explicitly checking for opaque type
     /// recursion, and 'false' otherwise to avoid unnecessary work.
@@ -777,6 +809,37 @@ impl<'tcx> OpaqueTypeExpander<'tcx> {
             None
         }
     }
+
+    fn expand_generator(&mut self, def_id: DefId, substs: SubstsRef<'tcx>) -> Option<Ty<'tcx>> {
+        if self.found_any_recursion {
+            return None;
+        }
+        let substs = substs.fold_with(self);
+        if !self.check_recursion || self.seen_opaque_tys.insert(def_id) {
+            let expanded_ty = match self.expanded_cache.get(&(def_id, substs)) {
+                Some(expanded_ty) => *expanded_ty,
+                None => {
+                    for bty in self.tcx.generator_hidden_types(def_id) {
+                        let hidden_ty = bty.subst(self.tcx, substs);
+                        self.fold_ty(hidden_ty);
+                    }
+                    let expanded_ty = self.tcx.mk_generator_witness_mir(def_id, substs);
+                    self.expanded_cache.insert((def_id, substs), expanded_ty);
+                    expanded_ty
+                }
+            };
+            if self.check_recursion {
+                self.seen_opaque_tys.remove(&def_id);
+            }
+            Some(expanded_ty)
+        } else {
+            // If another opaque type that we contain is recursive, then it
+            // will report the error, so we don't have to.
+            self.found_any_recursion = true;
+            self.found_recursion = def_id == *self.primary_def_id.as_ref().unwrap();
+            None
+        }
+    }
 }
 
 impl<'tcx> TypeFolder<'tcx> for OpaqueTypeExpander<'tcx> {
@@ -785,13 +848,19 @@ impl<'tcx> TypeFolder<'tcx> for OpaqueTypeExpander<'tcx> {
     }
 
     fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
-        if let ty::Alias(ty::Opaque, ty::AliasTy { def_id, substs, .. }) = *t.kind() {
+        let mut t = if let ty::Alias(ty::Opaque, ty::AliasTy { def_id, substs, .. }) = *t.kind() {
             self.expand_opaque_ty(def_id, substs).unwrap_or(t)
-        } else if t.has_opaque_types() {
+        } else if t.has_opaque_types() || t.has_generators() {
             t.super_fold_with(self)
         } else {
             t
+        };
+        if self.expand_generators {
+            if let ty::GeneratorWitnessMIR(def_id, substs) = *t.kind() {
+                t = self.expand_generator(def_id, substs).unwrap_or(t);
+            }
         }
+        t
     }
 }
 
@@ -1299,6 +1368,7 @@ pub fn reveal_opaque_types_in_bounds<'tcx>(
         found_recursion: false,
         found_any_recursion: false,
         check_recursion: false,
+        expand_generators: false,
         tcx,
     };
     val.fold_with(&mut visitor)

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -139,4 +139,11 @@ impl<'tcx> TraitEngine<'tcx> for FulfillmentCtxt<'tcx> {
     fn pending_obligations(&self) -> Vec<PredicateObligation<'tcx>> {
         self.obligations.clone()
     }
+
+    fn drain_unstalled_obligations(
+        &mut self,
+        _: &InferCtxt<'tcx>,
+    ) -> Vec<PredicateObligation<'tcx>> {
+        unimplemented!()
+    }
 }