assign upvars with saved locals for analysis

Co-authored-by: Dario Nieuwenhuis <dirbaio@dirbaio.net>

Author: dingxiangfei2009

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -27,8 +27,8 @@ use rustc_middle::ty::cast::CastTy;
 use rustc_middle::ty::fold::fold_regions;
 use rustc_middle::ty::visit::TypeVisitableExt;
 use rustc_middle::ty::{
-    self, Binder, CanonicalUserTypeAnnotation, CanonicalUserTypeAnnotations, CoroutineArgsExt,
-    Dynamic, GenericArgsRef, OpaqueHiddenType, OpaqueTypeKey, RegionVid, Ty, TyCtxt, UserArgs,
+    self, Binder, CanonicalUserTypeAnnotation, CanonicalUserTypeAnnotations, Dynamic,
+    GenericArgsRef, OpaqueHiddenType, OpaqueTypeKey, RegionVid, Ty, TyCtxt, UserArgs,
     UserTypeAnnotationIndex,
 };
 use rustc_middle::{bug, span_bug};
@@ -2174,14 +2174,15 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 }
             }
             AggregateKind::Coroutine(_, args) => {
-                // It doesn't make sense to look at a field beyond the prefix;
-                // these require a variant index, and are not initialized in
-                // aggregate rvalues.
-                match args.as_coroutine().prefix_tys().get(field_index.as_usize()) {
-                    Some(ty) => Ok(*ty),
-                    None => Err(FieldAccessError::OutOfRange {
-                        field_count: args.as_coroutine().prefix_tys().len(),
-                    }),
+                // It doesn't make sense to look at a field beyond the captured
+                // upvars.
+                // Otherwise it require a variant index, and are not initialized
+                // in aggregate rvalues.
+                let upvar_tys = &args.as_coroutine().upvar_tys();
+                if let Some(ty) = upvar_tys.get(field_index.as_usize()) {
+                    Ok(*ty)
+                } else {
+                    Err(FieldAccessError::OutOfRange { field_count: upvar_tys.len() })
                 }
             }
             AggregateKind::CoroutineClosure(_, args) => {

compiler/rustc_codegen_cranelift/src/base.rs

@@ -915,6 +915,9 @@ fn codegen_stmt<'tcx>(
                             let variant_dest = lval.downcast_variant(fx, variant_index);
                             (variant_index, variant_dest, active_field_index)
                         }
+                        mir::AggregateKind::Coroutine(_def_id, _args) => {
+                            (FIRST_VARIANT, lval.downcast_variant(fx, FIRST_VARIANT), None)
+                        }
                         _ => (FIRST_VARIANT, lval, None),
                     };
                     if active_field_index.is_some() {

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -144,6 +144,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         let variant_dest = dest.project_downcast(bx, variant_index);
                         (variant_index, variant_dest, active_field_index)
                     }
+                    mir::AggregateKind::Coroutine(_, _) => {
+                        (FIRST_VARIANT, dest.project_downcast(bx, FIRST_VARIANT), None)
+                    }
                     _ => (FIRST_VARIANT, dest, None),
                 };
                 if active_field_index.is_some() {

compiler/rustc_index/src/vec.rs

@@ -197,6 +197,11 @@ impl<I: Idx, T> IndexVec<I, T> {
     pub fn append(&mut self, other: &mut Self) {
         self.raw.append(&mut other.raw);
     }
+
+    #[inline]
+    pub fn debug_map_view(&self) -> IndexSliceMapView<'_, I, T> {
+        IndexSliceMapView(self.as_slice())
+    }
 }
 
 /// `IndexVec` is often used as a map, so it provides some map-like APIs.
@@ -220,14 +225,44 @@ impl<I: Idx, T> IndexVec<I, Option<T>> {
     pub fn contains(&self, index: I) -> bool {
         self.get(index).and_then(Option::as_ref).is_some()
     }
+
+    #[inline]
+    pub fn debug_map_view_compact(&self) -> IndexSliceMapViewCompact<'_, I, T> {
+        IndexSliceMapViewCompact(self.as_slice())
+    }
 }
 
+pub struct IndexSliceMapView<'a, I: Idx, T>(&'a IndexSlice<I, T>);
+pub struct IndexSliceMapViewCompact<'a, I: Idx, T>(&'a IndexSlice<I, Option<T>>);
+
 impl<I: Idx, T: fmt::Debug> fmt::Debug for IndexVec<I, T> {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt::Debug::fmt(&self.raw, fmt)
     }
 }
 
+impl<'a, I: Idx, T: fmt::Debug> fmt::Debug for IndexSliceMapView<'a, I, T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut entries = fmt.debug_map();
+        for (idx, val) in self.0.iter_enumerated() {
+            entries.entry(&idx, val);
+        }
+        entries.finish()
+    }
+}
+
+impl<'a, I: Idx, T: fmt::Debug> fmt::Debug for IndexSliceMapViewCompact<'a, I, T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut entries = fmt.debug_map();
+        for (idx, val) in self.0.iter_enumerated() {
+            if let Some(val) = val {
+                entries.entry(&idx, val);
+            }
+        }
+        entries.finish()
+    }
+}
+
 impl<I: Idx, T> Deref for IndexVec<I, T> {
     type Target = IndexSlice<I, T>;
 

compiler/rustc_middle/src/mir/mod.rs

@@ -368,6 +368,9 @@ pub struct Body<'tcx> {
     #[type_foldable(identity)]
     #[type_visitable(ignore)]
     pub function_coverage_info: Option<Box<coverage::FunctionCoverageInfo>>,
+
+    /// Coroutine local-upvar map
+    pub local_upvar_map: IndexVec<FieldIdx, Option<Local>>,
 }
 
 impl<'tcx> Body<'tcx> {
@@ -411,6 +414,7 @@ impl<'tcx> Body<'tcx> {
             tainted_by_errors,
             coverage_info_hi: None,
             function_coverage_info: None,
+            local_upvar_map: IndexVec::new(),
         };
         body.is_polymorphic = body.has_non_region_param();
         body
@@ -442,6 +446,7 @@ impl<'tcx> Body<'tcx> {
             tainted_by_errors: None,
             coverage_info_hi: None,
             function_coverage_info: None,
+            local_upvar_map: IndexVec::new(),
         };
         body.is_polymorphic = body.has_non_region_param();
         body

compiler/rustc_middle/src/mir/query.rs

@@ -10,6 +10,7 @@ use rustc_index::bit_set::BitMatrix;
 use rustc_index::{Idx, IndexVec};
 use rustc_macros::{HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
 use rustc_span::{Span, Symbol};
+use rustc_type_ir::data_structures::IndexMap;
 use smallvec::SmallVec;
 
 use super::{ConstValue, SourceInfo};
@@ -19,7 +20,7 @@ use crate::ty::{self, CoroutineArgsExt, OpaqueHiddenType, Ty, TyCtxt};
 rustc_index::newtype_index! {
     #[derive(HashStable)]
     #[encodable]
-    #[debug_format = "_{}"]
+    #[debug_format = "corsl_{}"]
     pub struct CoroutineSavedLocal {}
 }
 
@@ -57,6 +58,13 @@ pub struct CoroutineLayout<'tcx> {
     #[type_foldable(identity)]
     #[type_visitable(ignore)]
     pub storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
+
+    /// This map `A -> B` allows later MIR passes, error reporters
+    /// and layout calculator to relate saved locals `A` sourced from upvars
+    /// and locals `B` that upvars are moved into.
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
+    pub relocated_upvars: IndexMap<CoroutineSavedLocal, CoroutineSavedLocal>,
 }
 
 impl Debug for CoroutineLayout<'_> {

compiler/rustc_middle/src/mir/statement.rs

@@ -123,11 +123,10 @@ impl<'tcx> PlaceTy<'tcx> {
                     .get(f.index())
                     .copied()
                     .unwrap_or_else(|| bug!("field {f:?} out of range: {self:?}")),
-                // Only prefix fields (upvars and current state) are
-                // accessible without a variant index.
+                // Only upvars are accessible without a variant index.
                 ty::Coroutine(_, args) => args
                     .as_coroutine()
-                    .prefix_tys()
+                    .upvar_tys()
                     .get(f.index())
                     .copied()
                     .unwrap_or_else(|| bug!("field {f:?} out of range: {self:?}")),

compiler/rustc_middle/src/ty/layout.rs

@@ -935,9 +935,10 @@ where
                     ),
                     Variants::Multiple { tag, tag_field, .. } => {
                         if i == tag_field {
-                            return TyMaybeWithLayout::TyAndLayout(tag_layout(tag));
+                            TyMaybeWithLayout::TyAndLayout(tag_layout(tag))
+                        } else {
+                            TyMaybeWithLayout::Ty(args.as_coroutine().upvar_tys()[i])
                         }
-                        TyMaybeWithLayout::Ty(args.as_coroutine().prefix_tys()[i])
                     }
                 },
 

compiler/rustc_middle/src/ty/sty.rs

@@ -146,13 +146,6 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
             })
         })
     }
-
-    /// This is the types of the fields of a coroutine which are not stored in a
-    /// variant.
-    #[inline]
-    fn prefix_tys(self) -> &'tcx List<Ty<'tcx>> {
-        self.upvar_tys()
-    }
 }
 
 #[derive(Debug, Copy, Clone, HashStable, TypeFoldable, TypeVisitable)]

compiler/rustc_mir_build/src/builder/custom/mod.rs

@@ -61,6 +61,7 @@ pub(super) fn build_custom_mir<'tcx>(
         pass_count: 0,
         coverage_info_hi: None,
         function_coverage_info: None,
+        local_upvar_map: IndexVec::new(),
     };
 
     body.local_decls.push(LocalDecl::new(return_ty, return_ty_span));