Skip to content

Commit

Permalink
Fix precise field capture of univariant enums
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
tmiasko committed May 23, 2022
1 parent 6534637 commit 0e7eca7
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 21 deletions.
4 changes: 3 additions & 1 deletion compiler/rustc_middle/src/thir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_middle/src/thir/visit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
Expand Down
59 changes: 48 additions & 11 deletions compiler/rustc_mir_build/src/build/expr/as_place.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 } => {
Expand Down
14 changes: 6 additions & 8 deletions compiler/rustc_mir_build/src/thir/cx/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down Expand Up @@ -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;
Expand Down
27 changes: 27 additions & 0 deletions src/test/ui/closures/2229_closure_analysis/capture-enum-field.rs
Original file line number Diff line number Diff line change
@@ -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);
}

0 comments on commit 0e7eca7

Please sign in to comment.