Use precise places when lowering Closures in THIR

- Closures now use closure_min_captures to figure out captured paths
- Build upvar_mutbls using closure_min_captures
- Change logic in limit_capture_mutability to differentiate b/w
  capturing parent's local variable or capturing a variable that is
  captured by the parent (in case of nested closure) using PlaceBase.

Co-authored-by: Roxane Fruytier <[email protected]>

Author: arora-aman

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

@@ -17,7 +17,7 @@ use rustc_target::abi::VariantIdx;
 use rustc_index::vec::Idx;
 
 #[derive(Copy, Clone)]
-pub enum PlaceBase {
+crate enum PlaceBase {
     Local(Local),
     Upvar { var_hir_id: HirId, closure_def_id: DefId, closure_kind: ty::ClosureKind },
 }
@@ -29,7 +29,7 @@ pub enum PlaceBase {
 /// This is used internally when building a place for an expression like `a.b.c`. The fields `b`
 /// and `c` can be progressively pushed onto the place builder that is created when converting `a`.
 #[derive(Clone)]
-struct PlaceBuilder<'tcx> {
+crate struct PlaceBuilder<'tcx> {
     base: PlaceBase,
     projection: Vec<PlaceElem<'tcx>>,
 }
@@ -241,7 +241,7 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>(
 }
 
 impl<'tcx> PlaceBuilder<'tcx> {
-    fn into_place<'a>(
+    crate fn into_place<'a>(
         self,
         tcx: TyCtxt<'tcx>,
         typeck_results: &'a ty::TypeckResults<'tcx>,
@@ -261,6 +261,10 @@ impl<'tcx> PlaceBuilder<'tcx> {
         to_upvars_resolved_place_builder(self, tcx, typeck_results).unwrap()
     }
 
+    crate fn base(&self) -> PlaceBase {
+        self.base
+    }
+
     fn field(self, f: Field, ty: Ty<'tcx>) -> Self {
         self.project(PlaceElem::Field(f, ty))
     }
@@ -308,7 +312,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
     /// This is used when constructing a compound `Place`, so that we can avoid creating
     /// intermediate `Place` values until we know the full set of projections.
-    fn as_place_builder<M>(&mut self, block: BasicBlock, expr: M) -> BlockAnd<PlaceBuilder<'tcx>>
+    crate fn as_place_builder<M>(&mut self, block: BasicBlock, expr: M) -> BlockAnd<PlaceBuilder<'tcx>>
     where
         M: Mirror<'tcx, Output = Expr<'tcx>>,
     {

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

@@ -4,6 +4,7 @@ use rustc_index::vec::Idx;
 
 use crate::build::expr::category::{Category, RvalueFunc};
 use crate::build::{BlockAnd, BlockAndExtension, Builder};
+use crate::build::expr::as_place::PlaceBase;
 use crate::thir::*;
 use rustc_middle::middle::region;
 use rustc_middle::mir::AssertKind;
@@ -381,51 +382,51 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         this.cfg.push(block, Statement { source_info, kind: StatementKind::StorageLive(temp) });
 
-        let arg_place = unpack!(block = this.as_place(block, arg));
-
-        let mutability = match arg_place.as_ref() {
-            PlaceRef { local, projection: &[] } => this.local_decls[local].mutability,
-            PlaceRef { local, projection: &[ProjectionElem::Deref] } => {
-                debug_assert!(
-                    this.local_decls[local].is_ref_for_guard(),
-                    "Unexpected capture place",
-                );
-                this.local_decls[local].mutability
-            }
-            PlaceRef {
-                local,
-                projection: &[ref proj_base @ .., ProjectionElem::Field(upvar_index, _)],
-            }
-            | PlaceRef {
-                local,
-                projection:
-                    &[ref proj_base @ .., ProjectionElem::Field(upvar_index, _), ProjectionElem::Deref],
-            } => {
-                let place = PlaceRef { local, projection: proj_base };
-
-                // Not projected from the implicit `self` in a closure.
-                debug_assert!(
-                    match place.local_or_deref_local() {
-                        Some(local) => local == Local::new(1),
-                        None => false,
-                    },
-                    "Unexpected capture place"
-                );
-                // Not in a closure
-                debug_assert!(
-                    this.upvar_mutbls.len() > upvar_index.index(),
-                    "Unexpected capture place"
-                );
-                this.upvar_mutbls[upvar_index.index()]
+        let arg_place_builder = unpack!(block = this.as_place_builder(block, arg));
+
+        let mutability = match arg_place_builder.base() {
+            // Capture parent's local variable
+            PlaceBase::Local(local) =>  this.local_decls[local].mutability,
+            // Parent is a closure and we capturing a variable that is captured
+            // by the parent itself.
+            PlaceBase::Upvar { .. } => {
+                let enclosing_upvars_resolved = arg_place_builder.clone().into_place(
+                    this.hir.tcx(),
+                    this.hir.typeck_results());
+
+                match enclosing_upvars_resolved.as_ref() {
+                    PlaceRef { local, projection: &[ProjectionElem::Field(upvar_index, _), ..] }
+                    | PlaceRef {
+                        local,
+                        projection: &[ProjectionElem::Deref, ProjectionElem::Field(upvar_index, _), ..] } => {
+                            // Not in a closure
+                            debug_assert!(
+                                local == Local::new(1),
+                                "Expected local to be Local(1), found {:?}",
+                                local
+                            );
+                            // Not in a closure
+                            debug_assert!(
+                                this.upvar_mutbls.len() > upvar_index.index(),
+                                "Unexpected capture place, upvar_mutbls={:#?}, upvar_index={:?}",
+                                this.upvar_mutbls, upvar_index
+                            );
+                            this.upvar_mutbls[upvar_index.index()]
+                        }
+                    _ => bug!("Unexpected capture place"),
+                }
             }
-            _ => bug!("Unexpected capture place"),
         };
 
         let borrow_kind = match mutability {
             Mutability::Not => BorrowKind::Unique,
             Mutability::Mut => BorrowKind::Mut { allow_two_phase_borrow: false },
         };
 
+        let arg_place = arg_place_builder.into_place(
+                    this.hir.tcx(),
+                    this.hir.typeck_results());
+
         this.cfg.push_assign(
             block,
             source_info,

compiler/rustc_mir_build/src/build/mod.rs

@@ -10,6 +10,7 @@ use rustc_hir::lang_items::LangItem;
 use rustc_hir::{GeneratorKind, HirIdMap, Node};
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_infer::infer::TyCtxtInferExt;
+use rustc_middle::hir::place::PlaceBase as HirPlaceBase;
 use rustc_middle::middle::region;
 use rustc_middle::mir::*;
 use rustc_middle::ty::subst::Subst;
@@ -811,7 +812,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         // with the closure's DefId. Here, we run through that vec of UpvarIds for
         // the given closure and use the necessary information to create upvar
         // debuginfo and to fill `self.upvar_mutbls`.
-        if let Some(upvars) = hir_typeck_results.closure_captures.get(&fn_def_id) {
+        if hir_typeck_results.closure_min_captures.get(&fn_def_id).is_some() {
             let closure_env_arg = Local::new(1);
             let mut closure_env_projs = vec![];
             let mut closure_ty = self.local_decls[closure_env_arg].ty;
@@ -824,12 +825,19 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 ty::Generator(_, substs, _) => ty::UpvarSubsts::Generator(substs),
                 _ => span_bug!(self.fn_span, "upvars with non-closure env ty {:?}", closure_ty),
             };
-            let upvar_tys = upvar_substs.upvar_tys();
-            let upvars_with_tys = upvars.iter().zip(upvar_tys);
-            self.upvar_mutbls = upvars_with_tys
+            let capture_tys = upvar_substs.upvar_tys();
+            let captures_with_tys = hir_typeck_results
+                .closure_min_captures_flattened(fn_def_id)
+                .zip(capture_tys);
+
+            self.upvar_mutbls = captures_with_tys
                 .enumerate()
-                .map(|(i, ((&var_id, &upvar_id), ty))| {
-                    let capture = hir_typeck_results.upvar_capture(upvar_id);
+                .map(|(i, (captured_place, ty))| {
+                    let capture = captured_place.info.capture_kind;
+                    let var_id = match captured_place.place.base {
+                        HirPlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
+                        _ => bug!("Expected an upvar")
+                    };
 
                     let mut mutability = Mutability::Not;
 

compiler/rustc_mir_build/src/thir/cx/expr.rs

@@ -6,6 +6,8 @@ use crate::thir::*;
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind, Res};
 use rustc_index::vec::Idx;
+use rustc_middle::hir::place::PlaceBase as HirPlaceBase;
+use rustc_middle::hir::place::ProjectionKind as HirProjectionKind;
 use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::BorrowKind;
 use rustc_middle::ty::adjustment::{
@@ -386,14 +388,12 @@ fn make_mirror_unadjusted<'a, 'tcx>(
                     span_bug!(expr.span, "closure expr w/o closure type: {:?}", closure_ty);
                 }
             };
+
             let upvars = cx
                 .typeck_results()
-                .closure_captures
-                .get(&def_id)
-                .iter()
-                .flat_map(|upvars| upvars.iter())
+                .closure_min_captures_flattened(def_id)
                 .zip(substs.upvar_tys())
-                .map(|((&var_hir_id, _), ty)| capture_upvar(cx, expr, var_hir_id, ty))
+                .map(|(captured_place, ty)| capture_upvar(cx, expr, captured_place, ty))
                 .collect();
             ExprKind::Closure { closure_id: def_id, substs, upvars, movability }
         }
@@ -981,27 +981,50 @@ fn overloaded_place<'a, 'tcx>(
     ExprKind::Deref { arg: ref_expr.to_ref() }
 }
 
-fn capture_upvar<'tcx>(
+fn capture_upvar<'a, 'tcx>(
     cx: &mut Cx<'_, 'tcx>,
     closure_expr: &'tcx hir::Expr<'tcx>,
-    var_hir_id: hir::HirId,
+    captured_place: &'a ty::CapturedPlace<'tcx>,
     upvar_ty: Ty<'tcx>,
 ) -> ExprRef<'tcx> {
-    let upvar_id = ty::UpvarId {
-        var_path: ty::UpvarPath { hir_id: var_hir_id },
-        closure_expr_id: cx.tcx.hir().local_def_id(closure_expr.hir_id),
-    };
-    let upvar_capture = cx.typeck_results().upvar_capture(upvar_id);
+    let upvar_capture = captured_place.info.capture_kind;
     let temp_lifetime = cx.region_scope_tree.temporary_scope(closure_expr.hir_id.local_id);
-    let var_ty = cx.typeck_results().node_type(var_hir_id);
-    let captured_var = Expr {
+    let var_ty = captured_place.place.base_ty;
+
+    let var_hir_id = match captured_place.place.base {
+        HirPlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
+        base => bug!("Expected an upvar, found {:?}", base),
+    };
+
+    let mut captured_place_expr = Expr {
         temp_lifetime,
         ty: var_ty,
         span: closure_expr.span,
         kind: convert_var(cx, var_hir_id),
     };
+
+    for proj in captured_place.place.projections.iter() {
+        let kind = match proj.kind {
+            HirProjectionKind::Deref => ExprKind::Deref { arg: captured_place_expr.to_ref() },
+            HirProjectionKind::Field(field, ..) => {
+                // Variant index will always be 0, because for multi-variant
+                // enums, we capture the enum entirely.
+                ExprKind::Field {
+                    lhs: captured_place_expr.to_ref(),
+                    name: Field::new(field as usize),
+                }
+            }
+            HirProjectionKind::Index | HirProjectionKind::Subslice => {
+                // We don't capture these projections, so we can ignore them here
+                continue;
+            }
+        };
+
+        captured_place_expr = Expr { temp_lifetime, ty: proj.ty, span: closure_expr.span, kind };
+    }
+
     match upvar_capture {
-        ty::UpvarCapture::ByValue(_) => captured_var.to_ref(),
+        ty::UpvarCapture::ByValue(_) => captured_place_expr.to_ref(),
         ty::UpvarCapture::ByRef(upvar_borrow) => {
             let borrow_kind = match upvar_borrow.kind {
                 ty::BorrowKind::ImmBorrow => BorrowKind::Shared,
@@ -1012,7 +1035,7 @@ fn capture_upvar<'tcx>(
                 temp_lifetime,
                 ty: upvar_ty,
                 span: closure_expr.span,
-                kind: ExprKind::Borrow { borrow_kind, arg: captured_var.to_ref() },
+                kind: ExprKind::Borrow { borrow_kind, arg: captured_place_expr.to_ref() },
             }
             .to_ref()
         }