relocate upvars to Unresumed state and make coroutine prefix trivial

Co-authored-by: Dario Nieuwenhuis <[email protected]>

Author: dingxiangfei2009

Diff for: compiler/rustc_borrowck/src/lib.rs

@@ -26,6 +26,7 @@ use std::rc::Rc;
 use consumers::{BodyWithBorrowckFacts, ConsumerOptions};
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_data_structures::graph::dominators::Dominators;
+use rustc_data_structures::unord::UnordMap;
 use rustc_errors::Diag;
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
@@ -292,6 +293,7 @@ fn do_mir_borrowck<'tcx>(
             regioncx: regioncx.clone(),
             used_mut: Default::default(),
             used_mut_upvars: SmallVec::new(),
+            local_from_upvars: UnordMap::default(),
             borrow_set: Rc::clone(&borrow_set),
             upvars: &[],
             local_names: IndexVec::from_elem(None, &promoted_body.local_decls),
@@ -318,6 +320,12 @@ fn do_mir_borrowck<'tcx>(
         }
     }
 
+    let mut local_from_upvars = UnordMap::default();
+    for (field, &local) in body.local_upvar_map.iter_enumerated() {
+        let Some(local) = local else { continue };
+        local_from_upvars.insert(local, field);
+    }
+    debug!(?local_from_upvars, "dxf");
     let mut mbcx = MirBorrowckCtxt {
         infcx: &infcx,
         param_env,
@@ -333,6 +341,7 @@ fn do_mir_borrowck<'tcx>(
         regioncx: Rc::clone(&regioncx),
         used_mut: Default::default(),
         used_mut_upvars: SmallVec::new(),
+        local_from_upvars,
         borrow_set: Rc::clone(&borrow_set),
         upvars: tcx.closure_captures(def),
         local_names,
@@ -568,6 +577,9 @@ struct MirBorrowckCtxt<'a, 'mir, 'infcx, 'tcx> {
     /// If the function we're checking is a closure, then we'll need to report back the list of
     /// mutable upvars that have been used. This field keeps track of them.
     used_mut_upvars: SmallVec<[FieldIdx; 8]>,
+    /// Since upvars are moved to real locals, we need to map mutations to the locals back to
+    /// the upvars, so that used_mut_upvars is up-to-date.
+    local_from_upvars: UnordMap<Local, FieldIdx>,
     /// Region inference context. This contains the results from region inference and lets us e.g.
     /// find out which CFG points are contained in each borrow region.
     regioncx: Rc<RegionInferenceContext<'tcx>>,
@@ -2230,16 +2242,19 @@ impl<'mir, 'tcx> MirBorrowckCtxt<'_, 'mir, '_, 'tcx> {
     }
 
     /// Adds the place into the used mutable variables set
+    #[instrument(level = "debug", skip(self, flow_state))]
     fn add_used_mut(&mut self, root_place: RootPlace<'tcx>, flow_state: &Flows<'_, 'mir, 'tcx>) {
         match root_place {
             RootPlace { place_local: local, place_projection: [], is_local_mutation_allowed } => {
                 // If the local may have been initialized, and it is now currently being
                 // mutated, then it is justified to be annotated with the `mut`
                 // keyword, since the mutation may be a possible reassignment.
-                if is_local_mutation_allowed != LocalMutationIsAllowed::Yes
-                    && self.is_local_ever_initialized(local, flow_state).is_some()
-                {
-                    self.used_mut.insert(local);
+                if !matches!(is_local_mutation_allowed, LocalMutationIsAllowed::Yes) {
+                    if self.is_local_ever_initialized(local, flow_state).is_some() {
+                        self.used_mut.insert(local);
+                    } else if let Some(&field) = self.local_from_upvars.get(&local) {
+                        self.used_mut_upvars.push(field);
+                    }
                 }
             }
             RootPlace {
@@ -2257,6 +2272,8 @@ impl<'mir, 'tcx> MirBorrowckCtxt<'_, 'mir, '_, 'tcx> {
                     projection: place_projection,
                 }) {
                     self.used_mut_upvars.push(field);
+                } else if let Some(&field) = self.local_from_upvars.get(&place_local) {
+                    self.used_mut_upvars.push(field);
                 }
             }
         }

Diff for: compiler/rustc_borrowck/src/type_check/mod.rs

@@ -811,15 +811,15 @@ impl<'a, 'b, 'tcx> TypeVerifier<'a, 'b, 'tcx> {
                         }),
                     };
                 }
-                ty::Coroutine(_, args) => {
+                ty::Coroutine(_def_id, args) => {
                     // Only prefix fields (upvars and current state) are
                     // accessible without a variant index.
-                    return match args.as_coroutine().prefix_tys().get(field.index()) {
-                        Some(ty) => Ok(*ty),
-                        None => Err(FieldAccessError::OutOfRange {
-                            field_count: args.as_coroutine().prefix_tys().len(),
-                        }),
-                    };
+                    let upvar_tys = args.as_coroutine().upvar_tys();
+                    if let Some(ty) = upvar_tys.get(field.index()) {
+                        return Ok(*ty);
+                    } else {
+                        return Err(FieldAccessError::OutOfRange { field_count: upvar_tys.len() });
+                    }
                 }
                 ty::Tuple(tys) => {
                     return match tys.get(field.index()) {
@@ -1837,11 +1837,11 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 // 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(),
-                    }),
+                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) => {

Diff for: compiler/rustc_codegen_cranelift/src/base.rs

@@ -880,6 +880,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() {

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -13,8 +13,7 @@ use rustc_hir::def_id::{DefId, LOCAL_CRATE};
 use rustc_middle::bug;
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::{
-    self, AdtKind, CoroutineArgsExt, Instance, ParamEnv, PolyExistentialTraitRef, Ty, TyCtxt,
-    Visibility,
+    self, AdtKind, Instance, ParamEnv, PolyExistentialTraitRef, Ty, TyCtxt, Visibility,
 };
 use rustc_session::config::{self, DebugInfo, Lto};
 use rustc_span::symbol::Symbol;
@@ -1110,7 +1109,7 @@ fn build_upvar_field_di_nodes<'ll, 'tcx>(
     closure_or_coroutine_di_node: &'ll DIType,
 ) -> SmallVec<&'ll DIType> {
     let (&def_id, up_var_tys) = match closure_or_coroutine_ty.kind() {
-        ty::Coroutine(def_id, args) => (def_id, args.as_coroutine().prefix_tys()),
+        ty::Coroutine(def_id, args) => (def_id, args.as_coroutine().upvar_tys()),
         ty::Closure(def_id, args) => (def_id, args.as_closure().upvar_tys()),
         ty::CoroutineClosure(def_id, args) => (def_id, args.as_coroutine_closure().upvar_tys()),
         _ => {

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs

@@ -672,7 +672,6 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
     let coroutine_layout =
         cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.kind_ty()).unwrap();
 
-    let common_upvar_names = cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
     let variant_range = coroutine_args.variant_range(coroutine_def_id, cx.tcx);
     let variant_count = (variant_range.start.as_u32()..variant_range.end.as_u32()).len();
 
@@ -707,7 +706,6 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
                 coroutine_type_and_layout,
                 coroutine_type_di_node,
                 coroutine_layout,
-                common_upvar_names,
             );
 
             let span = coroutine_layout.variant_source_info[variant_index].span;

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs

@@ -3,12 +3,10 @@ use std::borrow::Cow;
 use rustc_codegen_ssa::debuginfo::type_names::{compute_debuginfo_type_name, cpp_like_debuginfo};
 use rustc_codegen_ssa::debuginfo::{tag_base_type, wants_c_like_enum_debuginfo};
 use rustc_hir::def::CtorKind;
-use rustc_index::IndexSlice;
 use rustc_middle::bug;
 use rustc_middle::mir::CoroutineLayout;
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, AdtDef, CoroutineArgs, CoroutineArgsExt, Ty, VariantDef};
-use rustc_span::Symbol;
 use rustc_target::abi::{FieldIdx, TagEncoding, VariantIdx, Variants};
 
 use super::type_map::{DINodeCreationResult, UniqueTypeId};
@@ -263,7 +261,6 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
     coroutine_type_and_layout: TyAndLayout<'tcx>,
     coroutine_type_di_node: &'ll DIType,
     coroutine_layout: &CoroutineLayout<'tcx>,
-    common_upvar_names: &IndexSlice<FieldIdx, Symbol>,
 ) -> &'ll DIType {
     let variant_name = CoroutineArgs::variant_name(variant_index);
     let unique_type_id = UniqueTypeId::for_enum_variant_struct_type(
@@ -274,11 +271,6 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
 
     let variant_layout = coroutine_type_and_layout.for_variant(cx, variant_index);
 
-    let coroutine_args = match coroutine_type_and_layout.ty.kind() {
-        ty::Coroutine(_, args) => args.as_coroutine(),
-        _ => unreachable!(),
-    };
-
     type_map::build_type_with_children(
         cx,
         type_map::stub(
@@ -292,7 +284,7 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
         ),
         |cx, variant_struct_type_di_node| {
             // Fields that just belong to this variant/state
-            let state_specific_fields: SmallVec<_> = (0..variant_layout.fields.count())
+            (0..variant_layout.fields.count())
                 .map(|field_index| {
                     let coroutine_saved_local = coroutine_layout.variant_fields[variant_index]
                         [FieldIdx::from_usize(field_index)];
@@ -314,28 +306,7 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
                         type_di_node(cx, field_type),
                     )
                 })
-                .collect();
-
-            // Fields that are common to all states
-            let common_fields: SmallVec<_> = coroutine_args
-                .prefix_tys()
-                .iter()
-                .zip(common_upvar_names)
-                .enumerate()
-                .map(|(index, (upvar_ty, upvar_name))| {
-                    build_field_di_node(
-                        cx,
-                        variant_struct_type_di_node,
-                        upvar_name.as_str(),
-                        cx.size_and_align_of(upvar_ty),
-                        coroutine_type_and_layout.fields.offset(index),
-                        DIFlags::FlagZero,
-                        type_di_node(cx, upvar_ty),
-                    )
-                })
-                .collect();
-
-            state_specific_fields.into_iter().chain(common_fields).collect()
+                .collect()
         },
         |cx| build_generic_type_param_di_nodes(cx, coroutine_type_and_layout.ty),
     )

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs

@@ -160,9 +160,6 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
                 )
             };
 
-            let common_upvar_names =
-                cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
-
             // Build variant struct types
             let variant_struct_type_di_nodes: SmallVec<_> = variants
                 .indices()
@@ -190,7 +187,6 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
                                 coroutine_type_and_layout,
                                 coroutine_type_di_node,
                                 coroutine_layout,
-                                common_upvar_names,
                             ),
                         source_info,
                     }

Diff for: compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -9,7 +9,7 @@ use rustc_middle::mir::{self, ConstValue};
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::Ty;
 use rustc_target::abi::{self, Abi, Align, Size};
-use tracing::debug;
+use tracing::{debug, instrument};
 
 use super::place::{PlaceRef, PlaceValue};
 use super::{FunctionCx, LocalRef};
@@ -550,13 +550,12 @@ impl<'a, 'tcx, V: CodegenObject> OperandValue<V> {
 }
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
+    #[instrument(level = "debug", skip(self, bx), ret)]
     fn maybe_codegen_consume_direct(
         &mut self,
         bx: &mut Bx,
         place_ref: mir::PlaceRef<'tcx>,
     ) -> Option<OperandRef<'tcx, Bx::Value>> {
-        debug!("maybe_codegen_consume_direct(place_ref={:?})", place_ref);
-
         match self.locals[place_ref.local] {
             LocalRef::Operand(mut o) => {
                 // Moves out of scalar and scalar pair fields are trivial.
@@ -599,13 +598,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
     }
 
+    #[instrument(level = "debug", skip(self, bx), ret)]
     pub fn codegen_consume(
         &mut self,
         bx: &mut Bx,
         place_ref: mir::PlaceRef<'tcx>,
     ) -> OperandRef<'tcx, Bx::Value> {
-        debug!("codegen_consume(place_ref={:?})", place_ref);
-
         let ty = self.monomorphized_place_ty(place_ref);
         let layout = bx.cx().layout_of(ty);
 
@@ -624,13 +622,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx.load_operand(place)
     }
 
+    #[instrument(level = "debug", skip(self, bx), ret)]
     pub fn codegen_operand(
         &mut self,
         bx: &mut Bx,
         operand: &mir::Operand<'tcx>,
     ) -> OperandRef<'tcx, Bx::Value> {
-        debug!("codegen_operand(operand={:?})", operand);
-
         match *operand {
             mir::Operand::Copy(ref place) | mir::Operand::Move(ref place) => {
                 self.codegen_consume(bx, place.as_ref())

Diff for: compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -129,6 +129,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() {

Diff for: compiler/rustc_const_eval/src/interpret/step.rs

@@ -294,6 +294,9 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 let variant_dest = self.project_downcast(dest, variant_index)?;
                 (variant_index, variant_dest, active_field_index)
             }
+            mir::AggregateKind::Coroutine(_def_id, _args) => {
+                (FIRST_VARIANT, self.project_downcast(dest, FIRST_VARIANT)?, None)
+            }
             mir::AggregateKind::RawPtr(..) => {
                 // Pointers don't have "fields" in the normal sense, so the
                 // projection-based code below would either fail in projection

Diff for: compiler/rustc_middle/src/mir/mod.rs

@@ -443,6 +443,9 @@ pub struct Body<'tcx> {
     /// If `-Cinstrument-coverage` is not active, or if an individual function
     /// is not eligible for coverage, then this should always be `None`.
     pub function_coverage_info: Option<Box<coverage::FunctionCoverageInfo>>,
+
+    /// Coroutine local-upvar map
+    pub local_upvar_map: IndexVec<FieldIdx, Option<Local>>,
 }
 
 impl<'tcx> Body<'tcx> {
@@ -486,6 +489,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
@@ -517,6 +521,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

Diff for: compiler/rustc_middle/src/mir/patch.rs

@@ -155,10 +155,14 @@ impl<'tcx> MirPatch<'tcx> {
         ty: Ty<'tcx>,
         span: Span,
         local_info: LocalInfo<'tcx>,
+        immutable: bool,
     ) -> Local {
         let index = self.next_local;
         self.next_local += 1;
         let mut new_decl = LocalDecl::new(ty, span);
+        if immutable {
+            new_decl = new_decl.immutable();
+        }
         **new_decl.local_info.as_mut().assert_crate_local() = local_info;
         self.new_locals.push(new_decl);
         Local::new(index)

Diff for: compiler/rustc_middle/src/ty/layout.rs

@@ -907,9 +907,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])
                     }
                 },
 

Diff for: compiler/rustc_middle/src/ty/sty.rs

@@ -147,13 +147,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)]