Remember where a type was kept in MIR.

Author: cjgillot

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_middle/src/mir/query.rs

@@ -135,11 +135,20 @@ rustc_index::newtype_index! {
     pub struct GeneratorSavedLocal {}
 }
 
+#[derive(Clone, Debug, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub struct GeneratorSavedTy<'tcx> {
+    pub ty: Ty<'tcx>,
+    /// Source info corresponding to the local in the original MIR body.
+    pub source_info: SourceInfo,
+    /// Whether the local was introduced as a raw pointer to a static.
+    pub is_static_ptr: bool,
+}
+
 /// The layout of generator state.
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
 pub struct GeneratorLayout<'tcx> {
     /// The type of every local stored inside the generator.
-    pub field_tys: IndexVec<GeneratorSavedLocal, Ty<'tcx>>,
+    pub field_tys: IndexVec<GeneratorSavedLocal, GeneratorSavedTy<'tcx>>,
 
     /// Which of the above fields are in each variant. Note that one field may
     /// be stored in multiple variants.

compiler/rustc_middle/src/ty/sty.rs

@@ -571,9 +571,9 @@ impl<'tcx> GeneratorSubsts<'tcx> {
     ) -> impl Iterator<Item = impl Iterator<Item = Ty<'tcx>> + Captures<'tcx>> {
         let layout = tcx.generator_layout(def_id).unwrap();
         layout.variant_fields.iter().map(move |variant| {
-            variant
-                .iter()
-                .map(move |field| ty::EarlyBinder(layout.field_tys[*field]).subst(tcx, self.substs))
+            variant.iter().map(move |field| {
+                ty::EarlyBinder(layout.field_tys[*field].ty).subst(tcx, self.substs)
+            })
         })
     }
 

compiler/rustc_mir_transform/src/generator.rs

@@ -916,7 +916,15 @@ fn compute_layout<'tcx>(
     let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
     for (saved_local, local) in saved_locals.iter_enumerated() {
         locals.push(local);
-        tys.push(body.local_decls[local].ty);
+        let decl = &body.local_decls[local];
+        let decl = GeneratorSavedTy {
+            ty: decl.ty,
+            source_info: decl.source_info,
+            is_static_ptr: decl.internal
+                && decl.ty.is_unsafe_ptr()
+                && matches!(decl.local_info.as_deref(), Some(LocalInfo::StaticRef { .. })),
+        };
+        tys.push(decl);
         debug!("generator saved local {:?} => {:?}", saved_local, local);
     }
 
@@ -947,7 +955,7 @@ fn compute_layout<'tcx>(
             // just use the first one here. That's fine; fields do not move
             // around inside generators, so it doesn't matter which variant
             // index we access them by.
-            remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx));
+            remap.entry(locals[saved_local]).or_insert((tys[saved_local].ty, variant_index, idx));
         }
         variant_fields.push(fields);
         variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]);
@@ -957,6 +965,7 @@ fn compute_layout<'tcx>(
 
     let layout =
         GeneratorLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
+    debug!(?layout);
 
     (remap, layout, storage_liveness)
 }

compiler/rustc_mir_transform/src/inline.rs

@@ -947,12 +947,12 @@ impl<'tcx> Visitor<'tcx> for CostChecker<'_, 'tcx> {
                             return;
                         };
 
-                        let Some(&f_ty) = layout.field_tys.get(local) else {
+                        let Some(f_ty) = layout.field_tys.get(local) else {
                             self.validation = Err("malformed MIR");
                             return;
                         };
 
-                        f_ty
+                        f_ty.ty
                     } else {
                         let Some(f_ty) = substs.as_generator().prefix_tys().nth(f.index()) else {
                             self.validation = Err("malformed MIR");

compiler/rustc_ty_utils/src/layout.rs

@@ -643,7 +643,7 @@ fn generator_layout<'tcx>(
 
     let promoted_layouts = ineligible_locals
         .iter()
-        .map(|local| subst_field(info.field_tys[local]))
+        .map(|local| subst_field(info.field_tys[local].ty))
         .map(|ty| tcx.mk_maybe_uninit(ty))
         .map(|ty| cx.layout_of(ty));
     let prefix_layouts = substs
@@ -713,7 +713,7 @@ fn generator_layout<'tcx>(
                     Assigned(_) => bug!("assignment does not match variant"),
                     Ineligible(_) => false,
                 })
-                .map(|local| subst_field(info.field_tys[*local]));
+                .map(|local| subst_field(info.field_tys[*local].ty));
 
             let mut variant = univariant_uninterned(
                 cx,

tests/mir-opt/generator_drop_cleanup.main-{closure#0}.generator_drop.0.mir

@@ -1,7 +1,14 @@
 // MIR for `main::{closure#0}` 0 generator_drop
 /* generator_layout = GeneratorLayout {
     field_tys: {
-        _0: std::string::String,
+        _0: GeneratorSavedTy {
+            ty: std::string::String,
+            source_info: SourceInfo {
+                span: $DIR/generator_drop_cleanup.rs:11:13: 11:15 (#0),
+                scope: scope[0],
+            },
+            is_static_ptr: false,
+        },
     },
     variant_fields: {
         Unresumed(0): [],

tests/mir-opt/generator_tiny.main-{closure#0}.generator_resume.0.mir

@@ -1,7 +1,14 @@
 // MIR for `main::{closure#0}` 0 generator_resume
 /* generator_layout = GeneratorLayout {
     field_tys: {
-        _0: HasDrop,
+        _0: GeneratorSavedTy {
+            ty: HasDrop,
+            source_info: SourceInfo {
+                span: $DIR/generator_tiny.rs:20:13: 20:15 (#0),
+                scope: scope[0],
+            },
+            is_static_ptr: false,
+        },
     },
     variant_fields: {
         Unresumed(0): [],