Auto merge of #100968 - cjgillot:mir-upvar-vec, r=wesleywiser

Only compute captures once when building MIR.

Author: bors

compiler/rustc_mir_build/src/build/expr/as_place.rs

@@ -2,7 +2,7 @@
 
 use crate::build::expr::category::Category;
 use crate::build::ForGuard::{OutsideGuard, RefWithinGuard};
-use crate::build::{BlockAnd, BlockAndExtension, Builder};
+use crate::build::{BlockAnd, BlockAndExtension, Builder, Capture, CaptureMap};
 use rustc_hir::def_id::LocalDefId;
 use rustc_middle::hir::place::Projection as HirProjection;
 use rustc_middle::hir::place::ProjectionKind as HirProjectionKind;
@@ -59,8 +59,6 @@ pub(crate) enum PlaceBase {
         var_hir_id: LocalVarId,
         /// DefId of the closure
         closure_def_id: LocalDefId,
-        /// The trait closure implements, `Fn`, `FnMut`, `FnOnce`
-        closure_kind: ty::ClosureKind,
     },
 }
 
@@ -145,27 +143,6 @@ fn is_ancestor_or_same_capture(
     iter::zip(proj_possible_ancestor, proj_capture).all(|(a, b)| a == b)
 }
 
-/// Computes the index of a capture within the desugared closure provided the closure's
-/// `closure_min_captures` and the capture's index of the capture in the
-/// `ty::MinCaptureList` of the root variable `var_hir_id`.
-fn compute_capture_idx<'tcx>(
-    closure_min_captures: &ty::RootVariableMinCaptureList<'tcx>,
-    var_hir_id: LocalVarId,
-    root_var_idx: usize,
-) -> usize {
-    let mut res = 0;
-    for (var_id, capture_list) in closure_min_captures {
-        if *var_id == var_hir_id.0 {
-            res += root_var_idx;
-            break;
-        } else {
-            res += capture_list.len();
-        }
-    }
-
-    res
-}
-
 /// Given a closure, returns the index of a capture within the desugared closure struct and the
 /// `ty::CapturedPlace` which is the ancestor of the Place represented using the `var_hir_id`
 /// and `projection`.
@@ -174,27 +151,17 @@ fn compute_capture_idx<'tcx>(
 ///
 /// Returns None, when the ancestor is not found.
 fn find_capture_matching_projections<'a, 'tcx>(
-    typeck_results: &'a ty::TypeckResults<'tcx>,
+    upvars: &'a CaptureMap<'tcx>,
     var_hir_id: LocalVarId,
-    closure_def_id: LocalDefId,
     projections: &[PlaceElem<'tcx>],
-) -> Option<(usize, &'a ty::CapturedPlace<'tcx>)> {
-    let closure_min_captures = typeck_results.closure_min_captures.get(&closure_def_id)?;
-    let root_variable_min_captures = closure_min_captures.get(&var_hir_id.0)?;
-
+) -> Option<(usize, &'a Capture<'tcx>)> {
     let hir_projections = convert_to_hir_projections_and_truncate_for_capture(projections);
 
-    // If an ancestor is found, `idx` is the index within the list of captured places
-    // for root variable `var_hir_id` and `capture` is the `ty::CapturedPlace` itself.
-    let (idx, capture) = root_variable_min_captures.iter().enumerate().find(|(_, capture)| {
+    upvars.get_by_key_enumerated(var_hir_id.0).find(|(_, capture)| {
         let possible_ancestor_proj_kinds: Vec<_> =
-            capture.place.projections.iter().map(|proj| proj.kind).collect();
+            capture.captured_place.place.projections.iter().map(|proj| proj.kind).collect();
         is_ancestor_or_same_capture(&possible_ancestor_proj_kinds, &hir_projections)
-    })?;
-
-    // Convert index to be from the perspective of the entire closure_min_captures map
-    // instead of just the root variable capture list
-    Some((compute_capture_idx(closure_min_captures, var_hir_id, idx), capture))
+    })
 }
 
 /// Takes a PlaceBuilder and resolves the upvar (if any) within it, so that the
@@ -204,24 +171,15 @@ fn find_capture_matching_projections<'a, 'tcx>(
 fn to_upvars_resolved_place_builder<'a, 'tcx>(
     from_builder: PlaceBuilder<'tcx>,
     tcx: TyCtxt<'tcx>,
-    typeck_results: &'a ty::TypeckResults<'tcx>,
+    upvars: &'a CaptureMap<'tcx>,
 ) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
     match from_builder.base {
         PlaceBase::Local(_) => Ok(from_builder),
-        PlaceBase::Upvar { var_hir_id, closure_def_id, closure_kind } => {
-            let mut upvar_resolved_place_builder = PlaceBuilder::from(ty::CAPTURE_STRUCT_LOCAL);
-            match closure_kind {
-                ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
-                    upvar_resolved_place_builder = upvar_resolved_place_builder.deref();
-                }
-                ty::ClosureKind::FnOnce => {}
-            }
-
+        PlaceBase::Upvar { var_hir_id, closure_def_id } => {
             let Some((capture_index, capture)) =
                 find_capture_matching_projections(
-                    typeck_results,
+                    upvars,
                     var_hir_id,
-                    closure_def_id,
                     &from_builder.projection,
                 ) else {
                 let closure_span = tcx.def_span(closure_def_id);
@@ -241,39 +199,17 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>(
                 return Err(from_builder);
             };
 
-            // We won't be building MIR if the closure wasn't local
-            let closure_hir_id = tcx.hir().local_def_id_to_hir_id(closure_def_id);
-            let closure_ty = typeck_results.node_type(closure_hir_id);
-
-            let substs = match closure_ty.kind() {
-                ty::Closure(_, substs) => ty::UpvarSubsts::Closure(substs),
-                ty::Generator(_, substs, _) => ty::UpvarSubsts::Generator(substs),
-                _ => bug!("Lowering capture for non-closure type {:?}", closure_ty),
-            };
-
             // Access the capture by accessing the field within the Closure struct.
-            //
-            // We must have inferred the capture types since we are building MIR, therefore
-            // it's safe to call `tuple_element_ty` and we can unwrap here because
-            // we know that the capture exists and is the `capture_index`-th capture.
-            let var_ty = substs.tupled_upvars_ty().tuple_fields()[capture_index];
-
-            upvar_resolved_place_builder =
-                upvar_resolved_place_builder.field(Field::new(capture_index), var_ty);
-
-            // If the variable is captured via ByRef(Immutable/Mutable) Borrow,
-            // we need to deref it
-            upvar_resolved_place_builder = match capture.info.capture_kind {
-                ty::UpvarCapture::ByRef(_) => upvar_resolved_place_builder.deref(),
-                ty::UpvarCapture::ByValue => upvar_resolved_place_builder,
-            };
+            let capture_info = &upvars[capture_index];
+
+            let mut upvar_resolved_place_builder = PlaceBuilder::from(capture_info.use_place);
 
             // We used some of the projections to build the capture itself,
             // now we apply the remaining to the upvar resolved place.
             let remaining_projections = strip_prefix(
-                capture.place.base_ty,
+                capture.captured_place.place.base_ty,
                 from_builder.projection,
-                &capture.place.projections,
+                &capture.captured_place.place.projections,
             );
             upvar_resolved_place_builder.projection.extend(remaining_projections);
 
@@ -315,24 +251,24 @@ fn strip_prefix<'tcx>(
 }
 
 impl<'tcx> PlaceBuilder<'tcx> {
-    pub(crate) fn into_place<'a>(
+    pub(in crate::build) fn into_place<'a>(
         self,
         tcx: TyCtxt<'tcx>,
-        typeck_results: &'a ty::TypeckResults<'tcx>,
+        upvars: &'a CaptureMap<'tcx>,
     ) -> Place<'tcx> {
         if let PlaceBase::Local(local) = self.base {
             Place { local, projection: tcx.intern_place_elems(&self.projection) }
         } else {
-            self.expect_upvars_resolved(tcx, typeck_results).into_place(tcx, typeck_results)
+            self.expect_upvars_resolved(tcx, upvars).into_place(tcx, upvars)
         }
     }
 
     fn expect_upvars_resolved<'a>(
         self,
         tcx: TyCtxt<'tcx>,
-        typeck_results: &'a ty::TypeckResults<'tcx>,
+        upvars: &'a CaptureMap<'tcx>,
     ) -> PlaceBuilder<'tcx> {
-        to_upvars_resolved_place_builder(self, tcx, typeck_results).unwrap()
+        to_upvars_resolved_place_builder(self, tcx, upvars).unwrap()
     }
 
     /// Attempts to resolve the `PlaceBuilder`.
@@ -346,12 +282,12 @@ impl<'tcx> PlaceBuilder<'tcx> {
     /// not captured. This can happen because the final mir that will be
     /// generated doesn't require a read for this place. Failures will only
     /// happen inside closures.
-    pub(crate) fn try_upvars_resolved<'a>(
+    pub(in crate::build) fn try_upvars_resolved<'a>(
         self,
         tcx: TyCtxt<'tcx>,
-        typeck_results: &'a ty::TypeckResults<'tcx>,
+        upvars: &'a CaptureMap<'tcx>,
     ) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
-        to_upvars_resolved_place_builder(self, tcx, typeck_results)
+        to_upvars_resolved_place_builder(self, tcx, upvars)
     }
 
     pub(crate) fn base(&self) -> PlaceBase {
@@ -392,6 +328,12 @@ impl<'tcx> From<PlaceBase> for PlaceBuilder<'tcx> {
     }
 }
 
+impl<'tcx> From<Place<'tcx>> for PlaceBuilder<'tcx> {
+    fn from(p: Place<'tcx>) -> Self {
+        Self { base: PlaceBase::Local(p.local), projection: p.projection.to_vec() }
+    }
+}
+
 impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// Compile `expr`, yielding a place that we can move from etc.
     ///
@@ -411,7 +353,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         expr: &Expr<'tcx>,
     ) -> BlockAnd<Place<'tcx>> {
         let place_builder = unpack!(block = self.as_place_builder(block, expr));
-        block.and(place_builder.into_place(self.tcx, self.typeck_results))
+        block.and(place_builder.into_place(self.tcx, &self.upvars))
     }
 
     /// This is used when constructing a compound `Place`, so that we can avoid creating
@@ -435,7 +377,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         expr: &Expr<'tcx>,
     ) -> BlockAnd<Place<'tcx>> {
         let place_builder = unpack!(block = self.as_read_only_place_builder(block, expr));
-        block.and(place_builder.into_place(self.tcx, self.typeck_results))
+        block.and(place_builder.into_place(self.tcx, &self.upvars))
     }
 
     /// This is used when constructing a compound `Place`, so that we can avoid creating
@@ -530,7 +472,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                             inferred_ty: expr.ty,
                         });
 
-                    let place = place_builder.clone().into_place(this.tcx, this.typeck_results);
+                    let place = place_builder.clone().into_place(this.tcx, &this.upvars);
                     this.cfg.push(
                         block,
                         Statement {
@@ -629,17 +571,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         closure_def_id: LocalDefId,
         var_hir_id: LocalVarId,
     ) -> BlockAnd<PlaceBuilder<'tcx>> {
-        let closure_ty =
-            self.typeck_results.node_type(self.tcx.hir().local_def_id_to_hir_id(closure_def_id));
-
-        let closure_kind = if let ty::Closure(_, closure_substs) = closure_ty.kind() {
-            self.infcx.closure_kind(closure_substs).unwrap()
-        } else {
-            // Generators are considered FnOnce.
-            ty::ClosureKind::FnOnce
-        };
-
-        block.and(PlaceBuilder::from(PlaceBase::Upvar { var_hir_id, closure_def_id, closure_kind }))
+        block.and(PlaceBuilder::from(PlaceBase::Upvar { var_hir_id, closure_def_id }))
     }
 
     /// Lower an index expression
@@ -678,7 +610,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         if is_outermost_index {
             self.read_fake_borrows(block, fake_borrow_temps, source_info)
         } else {
-            base_place = base_place.expect_upvars_resolved(self.tcx, self.typeck_results);
+            base_place = base_place.expect_upvars_resolved(self.tcx, &self.upvars);
             self.add_fake_borrows_of_base(
                 &base_place,
                 block,
@@ -710,7 +642,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             block,
             source_info,
             len,
-            Rvalue::Len(slice.into_place(self.tcx, self.typeck_results)),
+            Rvalue::Len(slice.into_place(self.tcx, &self.upvars)),
         );
         // lt = idx < len
         self.cfg.push_assign(

compiler/rustc_mir_build/src/build/expr/as_rvalue.rs

@@ -328,10 +328,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         unpack!(block = this.as_place_builder(block, &this.thir[*thir_place]));
 
                     if let Ok(place_builder_resolved) =
-                        place_builder.try_upvars_resolved(this.tcx, this.typeck_results)
+                        place_builder.try_upvars_resolved(this.tcx, &this.upvars)
                     {
-                        let mir_place =
-                            place_builder_resolved.into_place(this.tcx, this.typeck_results);
+                        let mir_place = place_builder_resolved.into_place(this.tcx, &this.upvars);
                         this.cfg.push_fake_read(
                             block,
                             this.source_info(this.tcx.hir().span(*hir_id)),
@@ -623,7 +622,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             // is same as that of the capture in the parent closure.
             PlaceBase::Upvar { .. } => {
                 let enclosing_upvars_resolved =
-                    arg_place_builder.clone().into_place(this.tcx, this.typeck_results);
+                    arg_place_builder.clone().into_place(this.tcx, &this.upvars);
 
                 match enclosing_upvars_resolved.as_ref() {
                     PlaceRef {
@@ -643,12 +642,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         );
                         // Not in a closure
                         debug_assert!(
-                            this.upvar_mutbls.len() > upvar_index.index(),
-                            "Unexpected capture place, upvar_mutbls={:#?}, upvar_index={:?}",
-                            this.upvar_mutbls,
+                            this.upvars.len() > upvar_index.index(),
+                            "Unexpected capture place, upvars={:#?}, upvar_index={:?}",
+                            this.upvars,
                             upvar_index
                         );
-                        this.upvar_mutbls[upvar_index.index()]
+                        this.upvars[upvar_index.index()].mutability
                     }
                     _ => bug!("Unexpected capture place"),
                 }
@@ -660,7 +659,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             Mutability::Mut => BorrowKind::Mut { allow_two_phase_borrow: false },
         };
 
-        let arg_place = arg_place_builder.into_place(this.tcx, this.typeck_results);
+        let arg_place = arg_place_builder.into_place(this.tcx, &this.upvars);
 
         this.cfg.push_assign(
             block,

compiler/rustc_mir_build/src/build/expr/into.rs

@@ -366,9 +366,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                             None => {
                                 let place_builder = place_builder.clone();
                                 this.consume_by_copy_or_move(
-                                    place_builder
-                                        .field(n, *ty)
-                                        .into_place(this.tcx, this.typeck_results),
+                                    place_builder.field(n, *ty).into_place(this.tcx, &this.upvars),
                                 )
                             }
                         })

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -221,9 +221,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         let source_info = self.source_info(scrutinee_span);
 
         if let Ok(scrutinee_builder) =
-            scrutinee_place_builder.clone().try_upvars_resolved(self.tcx, self.typeck_results)
+            scrutinee_place_builder.clone().try_upvars_resolved(self.tcx, &self.upvars)
         {
-            let scrutinee_place = scrutinee_builder.into_place(self.tcx, self.typeck_results);
+            let scrutinee_place = scrutinee_builder.into_place(self.tcx, &self.upvars);
             self.cfg.push_fake_read(block, source_info, cause_matched_place, scrutinee_place);
         }
 
@@ -348,12 +348,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     // ```
                     let mut opt_scrutinee_place: Option<(Option<&Place<'tcx>>, Span)> = None;
                     let scrutinee_place: Place<'tcx>;
-                    if let Ok(scrutinee_builder) = scrutinee_place_builder
-                        .clone()
-                        .try_upvars_resolved(this.tcx, this.typeck_results)
+                    if let Ok(scrutinee_builder) =
+                        scrutinee_place_builder.clone().try_upvars_resolved(this.tcx, &this.upvars)
                     {
-                        scrutinee_place =
-                            scrutinee_builder.into_place(this.tcx, this.typeck_results);
+                        scrutinee_place = scrutinee_builder.into_place(this.tcx, &this.upvars);
                         opt_scrutinee_place = Some((Some(&scrutinee_place), scrutinee_span));
                     }
                     let scope = this.declare_bindings(
@@ -620,9 +618,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     // };
                     // ```
                     if let Ok(match_pair_resolved) =
-                        initializer.clone().try_upvars_resolved(self.tcx, self.typeck_results)
+                        initializer.clone().try_upvars_resolved(self.tcx, &self.upvars)
                     {
-                        let place = match_pair_resolved.into_place(self.tcx, self.typeck_results);
+                        let place = match_pair_resolved.into_place(self.tcx, &self.upvars);
                         *match_place = Some(place);
                     }
                 }
@@ -1602,9 +1600,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         // Insert a Shallow borrow of any places that is switched on.
         if let Some(fb) = fake_borrows && let Ok(match_place_resolved) =
-            match_place.clone().try_upvars_resolved(self.tcx, self.typeck_results)
+            match_place.clone().try_upvars_resolved(self.tcx, &self.upvars)
         {
-            let resolved_place = match_place_resolved.into_place(self.tcx, self.typeck_results);
+            let resolved_place = match_place_resolved.into_place(self.tcx, &self.upvars);
             fb.insert(resolved_place);
         }
 
@@ -1791,10 +1789,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         );
         let mut opt_expr_place: Option<(Option<&Place<'tcx>>, Span)> = None;
         let expr_place: Place<'tcx>;
-        if let Ok(expr_builder) =
-            expr_place_builder.try_upvars_resolved(self.tcx, self.typeck_results)
-        {
-            expr_place = expr_builder.into_place(self.tcx, self.typeck_results);
+        if let Ok(expr_builder) = expr_place_builder.try_upvars_resolved(self.tcx, &self.upvars) {
+            expr_place = expr_builder.into_place(self.tcx, &self.upvars);
             opt_expr_place = Some((Some(&expr_place), expr_span));
         }
         let otherwise_post_guard_block = otherwise_candidate.pre_binding_block.unwrap();