Auto merge of #101692 - cjgillot:generator-lazy-witness, r=oli-obk

Compute generator saved locals on MIR

Generators are currently type-checked by introducing a `witness` type variable, which is unified with a `GeneratorWitness(captured types)` whose purpose is to ensure that the auto traits correctly migrate from the captured types to the `witness` type.  This requires computing the captured types on HIR during type-checking, only to re-do it on MIR later.

This PR proposes to drop the HIR-based computation, and only keep the MIR one.  This is done in 3 steps.
1. During type-checking, the `witness` type variable is never unified.  This allows to stall all the obligations that depend on it until the end of type-checking.  Then, the stalled obligations are marked as successful, and saved into the typeck results for later verification.
2. At type-checking writeback, `witness` is replaced by `GeneratorWitnessMIR(def_id, substs)`.  From this point on, all trait selection involving `GeneratorWitnessMIR` will fetch the MIR-computed locals, similar to what opaque types do.  There is no lifetime to be preserved here: we consider all the lifetimes appearing in this witness type to be higher-ranked.
3. After borrowck, the stashed obligations are verified against the actually computed types, in the `check_generator_obligations` query.  If any obligation was wrongly marked as fulfilled in step 1, it should be reported here.

There are still many issues:
- ~I am not too happy having to filter out some locals from the checked bounds, I think this is MIR building that introduces raw pointers polluting the analysis;~ solved by a check specific to static variables.
- the diagnostics for captured types don't show where they are used/dropped;
- I do not attempt to support chalk.

cc `@eholk` `@jyn514` for the drop-tracking work
r? `@oli-obk` as you warned me of potential unsoundness

Author: bors

compiler/rustc_codegen_ssa/src/debuginfo/type_names.rs

@@ -414,6 +414,7 @@ fn push_debuginfo_type_name<'tcx>(
         | ty::Placeholder(..)
         | ty::Alias(..)
         | ty::Bound(..)
+        | ty::GeneratorWitnessMIR(..)
         | ty::GeneratorWitness(..) => {
             bug!(
                 "debuginfo: Trying to create type name for \

compiler/rustc_const_eval/src/const_eval/valtrees.rs

@@ -151,7 +151,7 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
         // FIXME(oli-obk): we can probably encode closures just like structs
         | ty::Closure(..)
         | ty::Generator(..)
-        | ty::GeneratorWitness(..) => Err(ValTreeCreationError::NonSupportedType),
+        | ty::GeneratorWitness(..) |ty::GeneratorWitnessMIR(..)=> Err(ValTreeCreationError::NonSupportedType),
     }
 }
 
@@ -314,6 +314,7 @@ pub fn valtree_to_const_value<'tcx>(
         | ty::Closure(..)
         | ty::Generator(..)
         | ty::GeneratorWitness(..)
+        | ty::GeneratorWitnessMIR(..)
         | ty::FnPtr(_)
         | ty::RawPtr(_)
         | ty::Str

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -101,6 +101,7 @@ pub(crate) fn eval_nullary_intrinsic<'tcx>(
             | ty::Closure(_, _)
             | ty::Generator(_, _, _)
             | ty::GeneratorWitness(_)
+            | ty::GeneratorWitnessMIR(_, _)
             | ty::Never
             | ty::Tuple(_)
             | ty::Error(_) => ConstValue::from_machine_usize(0u64, &tcx),

compiler/rustc_const_eval/src/interpret/validity.rs

@@ -602,6 +602,7 @@ impl<'rt, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> ValidityVisitor<'rt, 'mir, '
             | ty::Bound(..)
             | ty::Param(..)
             | ty::Alias(..)
+            | ty::GeneratorWitnessMIR(..)
             | ty::GeneratorWitness(..) => bug!("Encountered invalid type {:?}", ty),
         }
     }

compiler/rustc_const_eval/src/transform/validate.rs

@@ -372,12 +372,12 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                                 return;
                             };
 
-                            let Some(&f_ty) = layout.field_tys.get(local) else {
+                            let Some(f_ty) = layout.field_tys.get(local) else {
                                 self.fail(location, format!("Out of bounds local {:?} for {:?}", local, parent_ty));
                                 return;
                             };
 
-                            f_ty
+                            f_ty.ty
                         } else {
                             let Some(f_ty) = substs.as_generator().prefix_tys().nth(f.index()) else {
                                 fail_out_of_bounds(self, location);

compiler/rustc_const_eval/src/util/type_name.rs

@@ -64,6 +64,7 @@ impl<'tcx> Printer<'tcx> for AbsolutePathPrinter<'tcx> {
             ty::Foreign(def_id) => self.print_def_path(def_id, &[]),
 
             ty::GeneratorWitness(_) => bug!("type_name: unexpected `GeneratorWitness`"),
+            ty::GeneratorWitnessMIR(..) => bug!("type_name: unexpected `GeneratorWitnessMIR`"),
         }
     }
 

compiler/rustc_hir/src/hir.rs

@@ -2106,8 +2106,8 @@ pub enum LocalSource {
 }
 
 /// Hints at the original code for a `match _ { .. }`.
-#[derive(Copy, Clone, PartialEq, Eq, Encodable, Hash, Debug)]
-#[derive(HashStable_Generic)]
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
+#[derive(HashStable_Generic, Encodable, Decodable)]
 pub enum MatchSource {
     /// A `match _ { .. }`.
     Normal,

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_analysis/src/coherence/inherent_impls.rs

@@ -240,6 +240,7 @@ impl<'tcx> InherentCollect<'tcx> {
             | ty::Closure(..)
             | ty::Generator(..)
             | ty::GeneratorWitness(..)
+            | ty::GeneratorWitnessMIR(..)
             | ty::Bound(..)
             | ty::Placeholder(_)
             | ty::Infer(_) => {

compiler/rustc_hir_analysis/src/variance/constraints.rs

@@ -295,12 +295,12 @@ impl<'a, 'tcx> ConstraintContext<'a, 'tcx> {
                 // types, where we use Error as the Self type
             }
 
-            ty::Placeholder(..) | ty::GeneratorWitness(..) | ty::Bound(..) | ty::Infer(..) => {
-                bug!(
-                    "unexpected type encountered in \
-                      variance inference: {}",
-                    ty
-                );
+            ty::Placeholder(..)
+            | ty::GeneratorWitness(..)
+            | ty::GeneratorWitnessMIR(..)
+            | ty::Bound(..)
+            | ty::Infer(..) => {
+                bug!("unexpected type encountered in variance inference: {}", ty);
             }
         }
     }

compiler/rustc_hir_typeck/src/cast.rs

@@ -130,6 +130,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             | ty::Float(_)
             | ty::Array(..)
             | ty::GeneratorWitness(..)
+            | ty::GeneratorWitnessMIR(..)
             | ty::RawPtr(_)
             | ty::Ref(..)
             | ty::FnDef(..)

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,16 +517,73 @@ 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);
         }
     }
 
     #[instrument(skip(self), level = "debug")]
-    pub(in super::super) fn select_all_obligations_or_error(&self) {
-        let mut errors = self.fulfillment_cx.borrow_mut().select_all_or_error(&self);
+    pub(in super::super) fn report_ambiguity_errors(&self) {
+        let mut errors = self.fulfillment_cx.borrow_mut().collect_remaining_errors();
 
         if !errors.is_empty() {
             self.adjust_fulfillment_errors_for_expr_obligation(&mut errors);

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,14 +294,24 @@ 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);
             fcx.require_type_is_sized(ty, span, code);
         }
 
-        fcx.select_all_obligations_or_error();
+        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();
+        }
 
         if let None = fcx.infcx.tainted_by_errors() {
             fcx.check_transmutes();

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/infer/canonical/canonicalizer.rs

@@ -435,6 +435,7 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for Canonicalizer<'cx, 'tcx> {
             ty::Closure(..)
             | ty::Generator(..)
             | ty::GeneratorWitness(..)
+            | ty::GeneratorWitnessMIR(..)
             | ty::Bool
             | ty::Char
             | ty::Int(..)

compiler/rustc_infer/src/infer/canonical/query_response.rs

@@ -17,7 +17,7 @@ use crate::infer::region_constraints::{Constraint, RegionConstraintData};
 use crate::infer::{InferCtxt, InferOk, InferResult, NllRegionVariableOrigin};
 use crate::traits::query::{Fallible, NoSolution};
 use crate::traits::{Obligation, ObligationCause, PredicateObligation};
-use crate::traits::{PredicateObligations, TraitEngine};
+use crate::traits::{PredicateObligations, TraitEngine, TraitEngineExt};
 use rustc_data_structures::captures::Captures;
 use rustc_index::vec::Idx;
 use rustc_index::vec::IndexVec;

compiler/rustc_infer/src/infer/freshen.rs

@@ -209,6 +209,7 @@ impl<'a, 'tcx> TypeFolder<'tcx> for TypeFreshener<'a, 'tcx> {
             | ty::Foreign(..)
             | ty::Param(..)
             | ty::Closure(..)
+            | ty::GeneratorWitnessMIR(..)
             | ty::GeneratorWitness(..) => t.super_fold_with(self),
 
             ty::Placeholder(..) | ty::Bound(..) => bug!("unexpected type {:?}", t),

compiler/rustc_infer/src/infer/outlives/components.rs

@@ -112,7 +112,7 @@ fn compute_components<'tcx>(
             }
 
             // All regions are bound inside a witness
-            ty::GeneratorWitness(..) => (),
+            ty::GeneratorWitness(..) | ty::GeneratorWitnessMIR(..) => (),
 
             // OutlivesTypeParameterEnv -- the actual checking that `X:'a`
             // is implied by the environment is done in regionck.

compiler/rustc_infer/src/traits/engine.rs

@@ -36,11 +36,19 @@ pub trait TraitEngine<'tcx>: 'tcx {
         obligation: PredicateObligation<'tcx>,
     );
 
-    fn select_all_or_error(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>>;
-
     fn select_where_possible(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'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> {
@@ -49,6 +57,8 @@ pub trait TraitEngineExt<'tcx> {
         infcx: &InferCtxt<'tcx>,
         obligations: impl IntoIterator<Item = PredicateObligation<'tcx>>,
     );
+
+    fn select_all_or_error(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>>;
 }
 
 impl<'tcx, T: ?Sized + TraitEngine<'tcx>> TraitEngineExt<'tcx> for T {
@@ -61,4 +71,13 @@ impl<'tcx, T: ?Sized + TraitEngine<'tcx>> TraitEngineExt<'tcx> for T {
             self.register_predicate_obligation(infcx, obligation);
         }
     }
+
+    fn select_all_or_error(&mut self, infcx: &InferCtxt<'tcx>) -> Vec<FulfillmentError<'tcx>> {
+        let errors = self.select_where_possible(infcx);
+        if !errors.is_empty() {
+            return errors;
+        }
+
+        self.collect_remaining_errors()
+    }
 }