Refactor generator interior types calculation

Author: Aaron1011

src/librustc/query/mod.rs

@@ -108,6 +108,11 @@ rustc_queries! {
         /// unreachable code.
         query mir_built(_: DefId) -> &'tcx Steal<mir::BodyCache<'tcx>> {}
 
+        /// Compute the generator interior types for a given `DefId`
+        /// (if it corresponds to a generator), for use in determining
+        /// generator auto trait implementation
+        query mir_generator_interior(_: DefId) -> &'tcx Steal<mir::BodyCache<'tcx>> {}
+
         /// Fetch the MIR for a given `DefId` up till the point where it is
         /// ready for const evaluation.
         ///

src/librustc/traits/select.rs

@@ -2807,10 +2807,12 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
 
                 let mut used_types = Vec::new();
 
-                // Note - while we use `optimized_mir`, the .generator_interior_tys
-                // field is only set during construction of the original MIR,
-                // and is unaffected by optimizations
-                let interior_tys = self.tcx().optimized_mir(did).generator_interior_tys
+                // We use `mir_validated` since earlier queries (e.g. `mir_const`)
+                // may be been stolen by the time this code runs. However, `generator_interior_tys`
+                // is computed early on and never modified, so it's fine to use
+                // a later query.
+                let mir = self.tcx().mir_validated(did).0.borrow();
+                let interior_tys = mir.generator_interior_tys
                     .as_ref().expect("Missing generator interior types!");
 
                 for ty in interior_tys {

src/librustc_mir/transform/mod.rs

@@ -46,6 +46,7 @@ pub(crate) fn provide(providers: &mut Providers<'_>) {
         mir_const,
         mir_const_qualif,
         mir_validated,
+        mir_generator_interior,
         optimized_mir,
         is_mir_available,
         promoted_mir,
@@ -98,7 +99,12 @@ fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> &DefIdSet {
 }
 
 fn mir_built(tcx: TyCtxt<'_>, def_id: DefId) -> &Steal<BodyCache<'_>> {
-    let mut mir = build::mir_build(tcx, def_id);
+    let mir = build::mir_build(tcx, def_id);
+    tcx.alloc_steal_mir(mir)
+}
+
+fn mir_generator_interior(tcx: TyCtxt<'_>, def_id: DefId) -> &Steal<BodyCache<'_>> {
+    let mut mir = tcx.mir_const(def_id).steal();
     if let ty::Generator(..) = tcx.type_of(def_id).kind {
         let interior_types = generator::generator_interior_tys(tcx, def_id, &mir);
         mir.generator_interior_tys = Some(interior_types);
@@ -248,7 +254,7 @@ fn mir_validated(
     // this point, before we steal the mir-const result.
     let _ = tcx.mir_const_qualif(def_id);
 
-    let mut body = tcx.mir_const(def_id).steal();
+    let mut body = tcx.mir_generator_interior(def_id).steal();
     let promote_pass = promote_consts::PromoteTemps::default();
     run_passes(tcx, &mut body, InstanceDef::Item(def_id), None, MirPhase::Validated, &[
         // What we need to run borrowck etc.