Move capture lowering from THIR to MIR

This allows us to:
- Handle precise Places captured by a closure directly in MIR. Handling
captures in MIR is easier since we can rely on/ tweak PlaceBuilder to
generate `mir::Place`s that resemble how we store captures (`hir::Place`).
- Allows us to handle `let _ = x` case when feature `capture_disjoint_fields`
is enabled directly in MIR. This is required to be done in MIR since
patterns are desugared in MIR.

Author: arora-aman

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

@@ -160,6 +160,30 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 expr_span,
                 source_info,
             ),
+            ExprKind::UpvarRef { closure_def_id, var_hir_id } => {
+                let capture = this
+                    .hir
+                    .typeck_results
+                    .closure_captures
+                    .get(&closure_def_id)
+                    .and_then(|captures| captures.get_full(&var_hir_id));
+
+                if capture.is_none() {
+                    if !this.hir.tcx().features().capture_disjoint_fields {
+                        bug!(
+                            "No associated capture found for {:?} even though \
+                            capture_disjoint_fields isn't enabled",
+                            expr.kind
+                        )
+                    }
+                    // FIXME(project-rfc-2229#24): Handle this case properly
+                }
+
+                // Unwrap until the FIXME has been resolved
+                let (capture_index, _, upvar_id) = capture.unwrap();
+                this.lower_closure_capture(block, capture_index, *upvar_id)
+            }
+
             ExprKind::SelfRef => block.and(PlaceBuilder::from(Local::new(1))),
             ExprKind::VarRef { id } => {
                 let place_builder = if this.is_bound_var_in_guard(id) {
@@ -270,6 +294,61 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
     }
 
+    /// Lower a closure/generator capture by representing it as a field
+    /// access within the desugared closure/generator.
+    ///
+    /// `capture_index` is the index of the capture within the desugared
+    /// closure/generator.
+    fn lower_closure_capture(
+        &mut self,
+        block: BasicBlock,
+        capture_index: usize,
+        upvar_id: ty::UpvarId,
+    )  -> BlockAnd<PlaceBuilder<'tcx>> {
+        let closure_ty = self
+            .hir
+            .typeck_results()
+            .node_type(self.hir.tcx().hir().local_def_id_to_hir_id(upvar_id.closure_expr_id));
+
+        // Captures are represented using fields inside a structure.
+        // This represents accessing self in the closure structure
+        let mut place_builder = PlaceBuilder::from(Local::new(1));
+
+        // In case of Fn/FnMut closures we must deref to access the fields
+        // Generators are considered FnOnce, so we ignore this step for them.
+        if let ty::Closure(_, closure_substs) = closure_ty.kind() {
+            match self.hir.infcx().closure_kind(closure_substs).unwrap() {
+                ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
+                    place_builder = place_builder.deref();
+                }
+                ty::ClosureKind::FnOnce => {}
+            }
+        }
+
+        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 `upvar_tys` and we can unwrap here because
+        // we know that the capture exists and is the `capture_index`-th capture.
+        let var_ty = substs.upvar_tys().nth(capture_index).unwrap();
+        place_builder = place_builder.field(Field::new(capture_index), var_ty);
+
+        // If the variable is captured via ByRef(Immutable/Mutable) Borrow,
+        // we need to deref it
+        match self.hir.typeck_results.upvar_capture(upvar_id) {
+            ty::UpvarCapture::ByRef(_) => {
+                block.and(place_builder.deref())
+            }
+            ty::UpvarCapture::ByValue(_) => block.and(place_builder),
+        }
+    }
+
     /// Lower an index expression
     ///
     /// This has two complications;

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

@@ -251,6 +251,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             | ExprKind::Index { .. }
             | ExprKind::VarRef { .. }
             | ExprKind::SelfRef
+            | ExprKind::UpvarRef { .. }
             | ExprKind::Break { .. }
             | ExprKind::Continue { .. }
             | ExprKind::Return { .. }

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

@@ -39,6 +39,7 @@ impl Category {
             | ExprKind::Deref { .. }
             | ExprKind::Index { .. }
             | ExprKind::SelfRef
+            | ExprKind::UpvarRef { .. }
             | ExprKind::VarRef { .. }
             | ExprKind::PlaceTypeAscription { .. }
             | ExprKind::ValueTypeAscription { .. } => Some(Category::Place),

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

@@ -401,6 +401,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             // Avoid creating a temporary
             ExprKind::VarRef { .. }
             | ExprKind::SelfRef
+            | ExprKind::UpvarRef { .. }
             | ExprKind::PlaceTypeAscription { .. }
             | ExprKind::ValueTypeAscription { .. } => {
                 debug_assert!(Category::of(&expr.kind) == Some(Category::Place));

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

@@ -880,130 +880,26 @@ fn convert_path_expr<'a, 'tcx>(
             ExprKind::Deref { arg: Expr { ty, temp_lifetime, span: expr.span, kind }.to_ref() }
         }
 
-        Res::Local(var_hir_id) => convert_var(cx, expr, var_hir_id),
+        Res::Local(var_hir_id) => convert_var(cx, var_hir_id),
 
         _ => span_bug!(expr.span, "res `{:?}` not yet implemented", res),
     }
 }
 
-fn convert_var<'tcx>(
-    cx: &mut Cx<'_, 'tcx>,
-    expr: &'tcx hir::Expr<'tcx>,
-    var_hir_id: hir::HirId,
-) -> ExprKind<'tcx> {
-    let upvar_index = cx
-        .typeck_results()
-        .closure_captures
-        .get(&cx.body_owner)
-        .and_then(|upvars| upvars.get_full(&var_hir_id).map(|(i, _, _)| i));
-
-    debug!(
-        "convert_var({:?}): upvar_index={:?}, body_owner={:?}",
-        var_hir_id, upvar_index, cx.body_owner
-    );
-
-    let temp_lifetime = cx.region_scope_tree.temporary_scope(expr.hir_id.local_id);
-
-    match upvar_index {
-        None => ExprKind::VarRef { id: var_hir_id },
+fn convert_var<'tcx>(cx: &mut Cx<'_, 'tcx>, var_hir_id: hir::HirId) -> ExprKind<'tcx> {
+    // We want upvars here not captures.
+    // Captures will be handled in MIR.
+    let is_upvar = cx
+        .tcx
+        .upvars_mentioned(cx.body_owner)
+        .map_or(false, |upvars| upvars.contains_key(&var_hir_id));
 
-        Some(upvar_index) => {
-            let closure_def_id = cx.body_owner;
-            let upvar_id = ty::UpvarId {
-                var_path: ty::UpvarPath { hir_id: var_hir_id },
-                closure_expr_id: closure_def_id.expect_local(),
-            };
-            let var_ty = cx.typeck_results().node_type(var_hir_id);
+    debug!("convert_var({:?}): is_upvar={}, body_owner={:?}", var_hir_id, is_upvar, cx.body_owner);
 
-            // FIXME free regions in closures are not right
-            let closure_ty = cx
-                .typeck_results()
-                .node_type(cx.tcx.hir().local_def_id_to_hir_id(upvar_id.closure_expr_id));
-
-            // FIXME we're just hard-coding the idea that the
-            // signature will be &self or &mut self and hence will
-            // have a bound region with number 0
-            let region = ty::ReFree(ty::FreeRegion {
-                scope: closure_def_id,
-                bound_region: ty::BoundRegion::BrAnon(0),
-            });
-            let region = cx.tcx.mk_region(region);
-
-            let self_expr = if let ty::Closure(_, closure_substs) = closure_ty.kind() {
-                match cx.infcx.closure_kind(closure_substs).unwrap() {
-                    ty::ClosureKind::Fn => {
-                        let ref_closure_ty = cx.tcx.mk_ref(
-                            region,
-                            ty::TypeAndMut { ty: closure_ty, mutbl: hir::Mutability::Not },
-                        );
-                        Expr {
-                            ty: closure_ty,
-                            temp_lifetime,
-                            span: expr.span,
-                            kind: ExprKind::Deref {
-                                arg: Expr {
-                                    ty: ref_closure_ty,
-                                    temp_lifetime,
-                                    span: expr.span,
-                                    kind: ExprKind::SelfRef,
-                                }
-                                .to_ref(),
-                            },
-                        }
-                    }
-                    ty::ClosureKind::FnMut => {
-                        let ref_closure_ty = cx.tcx.mk_ref(
-                            region,
-                            ty::TypeAndMut { ty: closure_ty, mutbl: hir::Mutability::Mut },
-                        );
-                        Expr {
-                            ty: closure_ty,
-                            temp_lifetime,
-                            span: expr.span,
-                            kind: ExprKind::Deref {
-                                arg: Expr {
-                                    ty: ref_closure_ty,
-                                    temp_lifetime,
-                                    span: expr.span,
-                                    kind: ExprKind::SelfRef,
-                                }
-                                .to_ref(),
-                            },
-                        }
-                    }
-                    ty::ClosureKind::FnOnce => Expr {
-                        ty: closure_ty,
-                        temp_lifetime,
-                        span: expr.span,
-                        kind: ExprKind::SelfRef,
-                    },
-                }
-            } else {
-                Expr { ty: closure_ty, temp_lifetime, span: expr.span, kind: ExprKind::SelfRef }
-            };
-
-            // at this point we have `self.n`, which loads up the upvar
-            let field_kind =
-                ExprKind::Field { lhs: self_expr.to_ref(), name: Field::new(upvar_index) };
-
-            // ...but the upvar might be an `&T` or `&mut T` capture, at which
-            // point we need an implicit deref
-            match cx.typeck_results().upvar_capture(upvar_id) {
-                ty::UpvarCapture::ByValue(_) => field_kind,
-                ty::UpvarCapture::ByRef(borrow) => ExprKind::Deref {
-                    arg: Expr {
-                        temp_lifetime,
-                        ty: cx.tcx.mk_ref(
-                            borrow.region,
-                            ty::TypeAndMut { ty: var_ty, mutbl: borrow.kind.to_mutbl_lossy() },
-                        ),
-                        span: expr.span,
-                        kind: field_kind,
-                    }
-                    .to_ref(),
-                },
-            }
-        }
+    if is_upvar {
+        ExprKind::UpvarRef { closure_def_id: cx.body_owner, var_hir_id }
+    } else {
+        ExprKind::VarRef { id: var_hir_id }
     }
 }
 
@@ -1102,7 +998,7 @@ fn capture_upvar<'tcx>(
         temp_lifetime,
         ty: var_ty,
         span: closure_expr.span,
-        kind: convert_var(cx, closure_expr, var_hir_id),
+        kind: convert_var(cx, var_hir_id),
     };
     match upvar_capture {
         ty::UpvarCapture::ByValue(_) => captured_var.to_ref(),

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

@@ -186,6 +186,10 @@ impl<'a, 'tcx> Cx<'a, 'tcx> {
         ty.needs_drop(self.tcx, self.param_env)
     }
 
+    crate fn infcx(&self) -> &'a InferCtxt<'a, 'tcx> {
+        self.infcx
+    }
+
     crate fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }

compiler/rustc_mir_build/src/thir/mod.rs

@@ -211,6 +211,14 @@ crate enum ExprKind<'tcx> {
     VarRef {
         id: hir::HirId,
     },
+    /// Used to represent upvars mentioned in a closure/generator
+    UpvarRef {
+        /// DefId of the closure/generator
+        closure_def_id: DefId,
+
+        /// HirId of the root variable
+        var_hir_id: hir::HirId,
+    },
     /// first argument, used for self in a closure
     SelfRef,
     Borrow {