Rollup merge of #97325 - tmiasko:capture-enum-field, r=arora-aman

Fix precise field capture of univariant enums

When constructing a MIR from a THIR field expression, introduce an
additional downcast projection before accessing a field of an enum.

When rebasing a place builder on top of a captured place, account for
the fact that a single HIR enum field projection corresponds to two MIR
projection elements: a downcast element and a field element.

Fixes #95271.
Fixes #96299.
Fixes #96512.
Fixes #97378.

r? ``@nikomatsakis`` ``@arora-aman``

Author: Dylan-DPC

compiler/rustc_middle/src/thir.rs

@@ -317,9 +317,11 @@ pub enum ExprKind<'tcx> {
         lhs: ExprId,
         rhs: ExprId,
     },
-    /// Access to a struct or tuple field.
+    /// Access to a field of a struct, a tuple, an union, or an enum.
     Field {
         lhs: ExprId,
+        /// Variant containing the field.
+        variant_index: VariantIdx,
         /// This can be a named (`.foo`) or unnamed (`.0`) field.
         name: Field,
     },

compiler/rustc_middle/src/thir/visit.rs

@@ -80,7 +80,7 @@ pub fn walk_expr<'a, 'tcx: 'a, V: Visitor<'a, 'tcx>>(visitor: &mut V, expr: &Exp
             visitor.visit_expr(&visitor.thir()[lhs]);
             visitor.visit_expr(&visitor.thir()[rhs]);
         }
-        Field { lhs, name: _ } => visitor.visit_expr(&visitor.thir()[lhs]),
+        Field { lhs, variant_index: _, name: _ } => visitor.visit_expr(&visitor.thir()[lhs]),
         Index { lhs, index } => {
             visitor.visit_expr(&visitor.thir()[lhs]);
             visitor.visit_expr(&visitor.thir()[index]);

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

@@ -5,6 +5,7 @@ use crate::build::ForGuard::{OutsideGuard, RefWithinGuard};
 use crate::build::{BlockAnd, BlockAndExtension, Builder};
 use rustc_hir::def_id::DefId;
 use rustc_hir::HirId;
+use rustc_middle::hir::place::Projection as HirProjection;
 use rustc_middle::hir::place::ProjectionKind as HirProjectionKind;
 use rustc_middle::middle::region;
 use rustc_middle::mir::AssertKind::BoundsCheck;
@@ -268,20 +269,52 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>(
                 ty::UpvarCapture::ByValue => upvar_resolved_place_builder,
             };
 
-            let next_projection = capture.place.projections.len();
-            let mut curr_projections = from_builder.projection;
-
             // We used some of the projections to build the capture itself,
             // now we apply the remaining to the upvar resolved place.
-            upvar_resolved_place_builder
-                .projection
-                .extend(curr_projections.drain(next_projection..));
+            let remaining_projections = strip_prefix(
+                capture.place.base_ty,
+                from_builder.projection,
+                &capture.place.projections,
+            );
+            upvar_resolved_place_builder.projection.extend(remaining_projections);
 
             Ok(upvar_resolved_place_builder)
         }
     }
 }
 
+/// Returns projections remaining after stripping an initial prefix of HIR
+/// projections.
+///
+/// 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>(
+    mut base_ty: Ty<'tcx>,
+    projections: Vec<PlaceElem<'tcx>>,
+    prefix_projections: &[HirProjection<'tcx>],
+) -> impl Iterator<Item = PlaceElem<'tcx>> {
+    let mut iter = projections.into_iter();
+    for projection in prefix_projections {
+        match projection.kind {
+            HirProjectionKind::Deref => {
+                assert!(matches!(iter.next(), Some(ProjectionElem::Deref)));
+            }
+            HirProjectionKind::Field(..) => {
+                if base_ty.is_enum() {
+                    assert!(matches!(iter.next(), Some(ProjectionElem::Downcast(..))));
+                }
+                assert!(matches!(iter.next(), Some(ProjectionElem::Field(..))));
+            }
+            HirProjectionKind::Index | HirProjectionKind::Subslice => {
+                bug!("unexpected projection kind: {:?}", projection);
+            }
+        }
+        base_ty = projection.ty;
+    }
+    iter
+}
+
 impl<'tcx> PlaceBuilder<'tcx> {
     pub(crate) fn into_place<'a>(
         self,
@@ -438,11 +471,15 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     this.expr_as_place(block, &this.thir[value], mutability, fake_borrow_temps)
                 })
             }
-            ExprKind::Field { lhs, name } => {
-                let place_builder = unpack!(
-                    block =
-                        this.expr_as_place(block, &this.thir[lhs], mutability, fake_borrow_temps,)
-                );
+            ExprKind::Field { lhs, variant_index, name } => {
+                let lhs = &this.thir[lhs];
+                let mut place_builder =
+                    unpack!(block = this.expr_as_place(block, lhs, mutability, fake_borrow_temps,));
+                if let ty::Adt(adt_def, _) = lhs.ty.kind() {
+                    if adt_def.is_enum() {
+                        place_builder = place_builder.downcast(*adt_def, variant_index);
+                    }
+                }
                 block.and(place_builder.field(name, expr.ty))
             }
             ExprKind::Deref { arg } => {

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

@@ -591,6 +591,7 @@ impl<'tcx> Cx<'tcx> {
             }
             hir::ExprKind::Field(ref source, ..) => ExprKind::Field {
                 lhs: self.mirror_expr(source),
+                variant_index: VariantIdx::new(0),
                 name: Field::new(tcx.field_index(expr.hir_id, self.typeck_results)),
             },
             hir::ExprKind::Cast(ref source, ref cast_ty) => {
@@ -994,14 +995,11 @@ impl<'tcx> Cx<'tcx> {
                 HirProjectionKind::Deref => {
                     ExprKind::Deref { arg: self.thir.exprs.push(captured_place_expr) }
                 }
-                HirProjectionKind::Field(field, ..) => {
-                    // Variant index will always be 0, because for multi-variant
-                    // enums, we capture the enum entirely.
-                    ExprKind::Field {
-                        lhs: self.thir.exprs.push(captured_place_expr),
-                        name: Field::new(field as usize),
-                    }
-                }
+                HirProjectionKind::Field(field, variant_index) => ExprKind::Field {
+                    lhs: self.thir.exprs.push(captured_place_expr),
+                    variant_index,
+                    name: Field::new(field as usize),
+                },
                 HirProjectionKind::Index | HirProjectionKind::Subslice => {
                     // We don't capture these projections, so we can ignore them here
                     continue;

src/test/ui/closures/2229_closure_analysis/capture-enum-field.rs

@@ -0,0 +1,27 @@
+// edition:2021
+// run-pass
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum Color {
+    RGB(u8, u8, u8),
+}
+
+fn main() {
+    let mut color = Color::RGB(0, 0, 0);
+    let mut red = |v| {
+        let Color::RGB(ref mut r, _, _) = color;
+        *r = v;
+    };
+    let mut green = |v| {
+        let Color::RGB(_, ref mut g, _) = color;
+        *g = v;
+    };
+    let mut blue = |v| {
+        let Color::RGB(_, _, ref mut b) = color;
+        *b = v;
+    };
+    red(1);
+    green(2);
+    blue(3);
+    assert_eq!(Color::RGB(1, 2, 3), color);
+}