Introduce PlaceBuilder::resolve_upvar by ref

1 files changed, 72 insertions, 63 deletions

diff --git a/compiler/rustc_mir_build/src/build/expr/as_place.rs b/compiler/rustc_mir_build/src/build/expr/as_place.rs
index 396782d45d2..fd6ade42b4f 100644
--- a/compiler/rustc_mir_build/src/build/expr/as_place.rs
+++ b/compiler/rustc_mir_build/src/build/expr/as_place.rs
@@ -167,59 +167,54 @@ fn find_capture_matching_projections<'a, 'tcx>(
     })
 }
 
-/// Takes a PlaceBuilder and resolves the upvar (if any) within it, so that the
-/// `PlaceBuilder` now starts from `PlaceBase::Local`.
-///
-/// Returns a Result with the error being the PlaceBuilder (`from_builder`) that was not found.
+/// Takes an upvar place and tries to resolve it into a `PlaceBuilder`
+/// with `PlaceBase::Local`
 #[instrument(level = "trace", skip(cx), ret)]
 fn to_upvars_resolved_place_builder<'tcx>(
-    from_builder: PlaceBuilder<'tcx>,
     cx: &Builder<'_, 'tcx>,
-) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
-    match from_builder.base {
-        PlaceBase::Local(_) => Ok(from_builder),
-        PlaceBase::Upvar { var_hir_id, closure_def_id } => {
-            let Some((capture_index, capture)) =
-                find_capture_matching_projections(
-                    &cx.upvars,
-                    var_hir_id,
-                    &from_builder.projection,
-                ) else {
-                let closure_span = cx.tcx.def_span(closure_def_id);
-                if !enable_precise_capture(cx.tcx, closure_span) {
-                    bug!(
-                        "No associated capture found for {:?}[{:#?}] even though \
-                            capture_disjoint_fields isn't enabled",
-                        var_hir_id,
-                        from_builder.projection
-                    )
-                } else {
-                    debug!(
-                        "No associated capture found for {:?}[{:#?}]",
-                        var_hir_id, from_builder.projection,
-                    );
-                }
-                return Err(from_builder);
-            };
+    var_hir_id: LocalVarId,
+    closure_def_id: LocalDefId,
+    projection: &[PlaceElem<'tcx>],
+) -> Option<PlaceBuilder<'tcx>> {
+    let Some((capture_index, capture)) =
+        find_capture_matching_projections(
+            &cx.upvars,
+            var_hir_id,
+            &projection,
+        ) else {
+        let closure_span = cx.tcx.def_span(closure_def_id);
+        if !enable_precise_capture(cx.tcx, closure_span) {
+            bug!(
+                "No associated capture found for {:?}[{:#?}] even though \
+                    capture_disjoint_fields isn't enabled",
+                var_hir_id,
+                projection
+            )
+        } else {
+            debug!(
+                "No associated capture found for {:?}[{:#?}]",
+                var_hir_id, projection,
+            );
+        }
+        return None;
+    };
 
-            // Access the capture by accessing the field within the Closure struct.
-            let capture_info = &cx.upvars[capture_index];
+    // Access the capture by accessing the field within the Closure struct.
+    let capture_info = &cx.upvars[capture_index];
 
-            let mut upvar_resolved_place_builder = PlaceBuilder::from(capture_info.use_place);
+    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.
-            trace!(?capture.captured_place, ?from_builder.projection);
-            let remaining_projections = strip_prefix(
-                capture.captured_place.place.base_ty,
-                from_builder.projection,
-                &capture.captured_place.place.projections,
-            );
-            upvar_resolved_place_builder.projection.extend(remaining_projections);
+    // We used some of the projections to build the capture itself,
+    // now we apply the remaining to the upvar resolved place.
+    trace!(?capture.captured_place, ?projection);
+    let remaining_projections = strip_prefix(
+        capture.captured_place.place.base_ty,
+        projection,
+        &capture.captured_place.place.projections,
+    );
+    upvar_resolved_place_builder.projection.extend(remaining_projections);
 
-            Ok(upvar_resolved_place_builder)
-        }
-    }
+    Some(upvar_resolved_place_builder)
 }
 
 /// Returns projections remaining after stripping an initial prefix of HIR
@@ -228,13 +223,14 @@ fn to_upvars_resolved_place_builder<'tcx>(
 /// Supports only HIR projection kinds that represent a path that might be
 /// captured by a closure or a generator, i.e., an `Index` or a `Subslice`
 /// projection kinds are unsupported.
-fn strip_prefix<'tcx>(
+fn strip_prefix<'a, 'tcx>(
     mut base_ty: Ty<'tcx>,
-    projections: Vec<PlaceElem<'tcx>>,
+    projections: &'a [PlaceElem<'tcx>],
     prefix_projections: &[HirProjection<'tcx>],
-) -> impl Iterator<Item = PlaceElem<'tcx>> {
+) -> impl Iterator<Item = PlaceElem<'tcx>> + 'a {
     let mut iter = projections
-        .into_iter()
+        .iter()
+        .copied()
         // Filter out opaque casts, they are unnecessary in the prefix.
         .filter(|elem| !matches!(elem, ProjectionElem::OpaqueCast(..)));
     for projection in prefix_projections {
@@ -258,21 +254,31 @@ fn strip_prefix<'tcx>(
 }
 
 impl<'tcx> PlaceBuilder<'tcx> {
-    pub(in crate::build) fn into_place(self, cx: &Builder<'_, 'tcx>) -> Place<'tcx> {
-        if let PlaceBase::Local(local) = self.base {
-            Place { local, projection: cx.tcx.intern_place_elems(&self.projection) }
-        } else {
-            self.expect_upvars_resolved(cx).into_place(cx)
-        }
+    pub(in crate::build) fn into_place(mut self, cx: &Builder<'_, 'tcx>) -> Place<'tcx> {
+        self = self.resolve_upvar(cx).unwrap_or(self);
+        let PlaceBase::Local(local) = self.base else { panic!("expected local") };
+        Place { local, projection: cx.tcx.intern_place_elems(&self.projection) }
     }
 
     fn expect_upvars_resolved(self, cx: &Builder<'_, 'tcx>) -> PlaceBuilder<'tcx> {
-        to_upvars_resolved_place_builder(self, cx).unwrap()
+        match self.base {
+            PlaceBase::Local(_) => self,
+            PlaceBase::Upvar {..} => self.resolve_upvar(cx).unwrap(),
+        }
+    }
+
+    pub(in crate::build) fn try_upvars_resolved(
+        self,
+        cx: &Builder<'_, 'tcx>,
+    ) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
+        match self.base {
+            PlaceBase::Local(_) => Ok(self),
+            PlaceBase::Upvar { .. } => self.resolve_upvar(cx).ok_or(self),
+        }
     }
 
     /// Attempts to resolve the `PlaceBuilder`.
-    /// On success, it will return the resolved `PlaceBuilder`.
-    /// On failure, it will return itself.
+    /// Returns `None` if this is not an upvar.
     ///
     /// Upvars resolve may fail for a `PlaceBuilder` when attempting to
     /// resolve a disjoint field whose root variable is not captured
@@ -281,11 +287,14 @@ 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(in crate::build) fn try_upvars_resolved(
-        self,
+    pub(in crate::build) fn resolve_upvar(
+        &self,
         cx: &Builder<'_, 'tcx>,
-    ) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
-        to_upvars_resolved_place_builder(self, cx)
+    ) -> Option<PlaceBuilder<'tcx>> {
+        let PlaceBase::Upvar { var_hir_id, closure_def_id } = self.base else {
+            return None;
+        };
+        to_upvars_resolved_place_builder(cx, var_hir_id, closure_def_id, &self.projection)
     }
 
     pub(crate) fn base(&self) -> PlaceBase {