From 0e7eca77e13b0021d3763e1470cd15a6c25bf80b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= Date: Mon, 23 May 2022 00:00:00 +0000 Subject: [PATCH] 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. --- compiler/rustc_middle/src/thir.rs | 4 +- compiler/rustc_middle/src/thir/visit.rs | 2 +- .../src/build/expr/as_place.rs | 59 +++++++++++++++---- compiler/rustc_mir_build/src/thir/cx/expr.rs | 14 ++--- .../capture-enum-field.rs | 27 +++++++++ 5 files changed, 85 insertions(+), 21 deletions(-) create mode 100644 src/test/ui/closures/2229_closure_analysis/capture-enum-field.rs diff --git a/compiler/rustc_middle/src/thir.rs b/compiler/rustc_middle/src/thir.rs index b99e7573000c9..c3b79917dda91 100644 --- a/compiler/rustc_middle/src/thir.rs +++ b/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, }, diff --git a/compiler/rustc_middle/src/thir/visit.rs b/compiler/rustc_middle/src/thir/visit.rs index f57569522d58d..8c8ebb0a6b87a 100644 --- a/compiler/rustc_middle/src/thir/visit.rs +++ b/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]); 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 045d6eb1c3021..438a68e71f054 100644 --- a/compiler/rustc_mir_build/src/build/expr/as_place.rs +++ b/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>, + prefix_projections: &[HirProjection<'tcx>], +) -> impl Iterator> { + 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 } => { diff --git a/compiler/rustc_mir_build/src/thir/cx/expr.rs b/compiler/rustc_mir_build/src/thir/cx/expr.rs index 480c5b195cc2d..bd9f599fff0a1 100644 --- a/compiler/rustc_mir_build/src/thir/cx/expr.rs +++ b/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; diff --git a/src/test/ui/closures/2229_closure_analysis/capture-enum-field.rs b/src/test/ui/closures/2229_closure_analysis/capture-enum-field.rs new file mode 100644 index 0000000000000..bbe3aa31a98df --- /dev/null +++ b/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); +}