Skip to content

Commit

Permalink
Auto merge of rust-lang#100968 - cjgillot:mir-upvar-vec, r=wesleywiser
Browse files Browse the repository at this point in the history 
Only compute captures once when building MIR.
  • Loading branch information
@bors
bors committed Sep 10, 2022
2 parents 5197c96 + bcbd183 commit 2e44c17
Show file tree
Hide file tree
Showing 8 changed files with 83 additions and 152 deletions.
136 changes: 34 additions & 102 deletions compiler/rustc_mir_build/src/build/expr/as_place.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
use crate::build::expr::category::Category;
use crate::build::ForGuard::{OutsideGuard, RefWithinGuard};
use crate::build::{BlockAnd, BlockAndExtension, Builder};
use crate::build::{BlockAnd, BlockAndExtension, Builder, Capture, CaptureMap};
use rustc_hir::def_id::LocalDefId;
use rustc_middle::hir::place::Projection as HirProjection;
use rustc_middle::hir::place::ProjectionKind as HirProjectionKind;
Expand Down Expand Up @@ -59,8 +59,6 @@ pub(crate) enum PlaceBase {
var_hir_id: LocalVarId,
/// DefId of the closure
closure_def_id: LocalDefId,
/// The trait closure implements, `Fn`, `FnMut`, `FnOnce`
closure_kind: ty::ClosureKind,
},
}

Expand Down Expand Up @@ -145,27 +143,6 @@ fn is_ancestor_or_same_capture(
iter::zip(proj_possible_ancestor, proj_capture).all(|(a, b)| a == b)
}

/// Computes the index of a capture within the desugared closure provided the closure's
/// `closure_min_captures` and the capture's index of the capture in the
/// `ty::MinCaptureList` of the root variable `var_hir_id`.
fn compute_capture_idx<'tcx>(
closure_min_captures: &ty::RootVariableMinCaptureList<'tcx>,
var_hir_id: LocalVarId,
root_var_idx: usize,
) -> usize {
let mut res = 0;
for (var_id, capture_list) in closure_min_captures {
if *var_id == var_hir_id.0 {
res += root_var_idx;
break;
} else {
res += capture_list.len();
}
}

res
}

/// Given a closure, returns the index of a capture within the desugared closure struct and the
/// `ty::CapturedPlace` which is the ancestor of the Place represented using the `var_hir_id`
/// and `projection`.
Expand All @@ -174,27 +151,17 @@ fn compute_capture_idx<'tcx>(
///
/// Returns None, when the ancestor is not found.
fn find_capture_matching_projections<'a, 'tcx>(
typeck_results: &'a ty::TypeckResults<'tcx>,
upvars: &'a CaptureMap<'tcx>,
var_hir_id: LocalVarId,
closure_def_id: LocalDefId,
projections: &[PlaceElem<'tcx>],
) -> Option<(usize, &'a ty::CapturedPlace<'tcx>)> {
let closure_min_captures = typeck_results.closure_min_captures.get(&closure_def_id)?;
let root_variable_min_captures = closure_min_captures.get(&var_hir_id.0)?;

) -> Option<(usize, &'a Capture<'tcx>)> {
let hir_projections = convert_to_hir_projections_and_truncate_for_capture(projections);

// If an ancestor is found, `idx` is the index within the list of captured places
// for root variable `var_hir_id` and `capture` is the `ty::CapturedPlace` itself.
let (idx, capture) = root_variable_min_captures.iter().enumerate().find(|(_, capture)| {
upvars.get_by_key_enumerated(var_hir_id.0).find(|(_, capture)| {
let possible_ancestor_proj_kinds: Vec<_> =
capture.place.projections.iter().map(|proj| proj.kind).collect();
capture.captured_place.place.projections.iter().map(|proj| proj.kind).collect();
is_ancestor_or_same_capture(&possible_ancestor_proj_kinds, &hir_projections)
})?;

// Convert index to be from the perspective of the entire closure_min_captures map
// instead of just the root variable capture list
Some((compute_capture_idx(closure_min_captures, var_hir_id, idx), capture))
})
}

/// Takes a PlaceBuilder and resolves the upvar (if any) within it, so that the
Expand All @@ -204,24 +171,15 @@ fn find_capture_matching_projections<'a, 'tcx>(
fn to_upvars_resolved_place_builder<'a, 'tcx>(
from_builder: PlaceBuilder<'tcx>,
tcx: TyCtxt<'tcx>,
typeck_results: &'a ty::TypeckResults<'tcx>,
upvars: &'a CaptureMap<'tcx>,
) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
match from_builder.base {
PlaceBase::Local(_) => Ok(from_builder),
PlaceBase::Upvar { var_hir_id, closure_def_id, closure_kind } => {
let mut upvar_resolved_place_builder = PlaceBuilder::from(ty::CAPTURE_STRUCT_LOCAL);
match closure_kind {
ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
upvar_resolved_place_builder = upvar_resolved_place_builder.deref();
}
ty::ClosureKind::FnOnce => {}
}

PlaceBase::Upvar { var_hir_id, closure_def_id } => {
let Some((capture_index, capture)) =
find_capture_matching_projections(
typeck_results,
upvars,
var_hir_id,
closure_def_id,
&from_builder.projection,
) else {
let closure_span = tcx.def_span(closure_def_id);
Expand All @@ -241,39 +199,17 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>(
return Err(from_builder);
};

// We won't be building MIR if the closure wasn't local
let closure_hir_id = tcx.hir().local_def_id_to_hir_id(closure_def_id);
let closure_ty = typeck_results.node_type(closure_hir_id);

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 `tuple_element_ty` and we can unwrap here because
// we know that the capture exists and is the `capture_index`-th capture.
let var_ty = substs.tupled_upvars_ty().tuple_fields()[capture_index];

upvar_resolved_place_builder =
upvar_resolved_place_builder.field(Field::new(capture_index), var_ty);

// If the variable is captured via ByRef(Immutable/Mutable) Borrow,
// we need to deref it
upvar_resolved_place_builder = match capture.info.capture_kind {
ty::UpvarCapture::ByRef(_) => upvar_resolved_place_builder.deref(),
ty::UpvarCapture::ByValue => upvar_resolved_place_builder,
};
let capture_info = &upvars[capture_index];

let mut upvar_resolved_place_builder = PlaceBuilder::from(capture_info.use_place);

// We used some of the projections to build the capture itself,
// now we apply the remaining to the upvar resolved place.
let remaining_projections = strip_prefix(
capture.place.base_ty,
capture.captured_place.place.base_ty,
from_builder.projection,
&capture.place.projections,
&capture.captured_place.place.projections,
);
upvar_resolved_place_builder.projection.extend(remaining_projections);

Expand Down Expand Up @@ -315,24 +251,24 @@ fn strip_prefix<'tcx>(
}

impl<'tcx> PlaceBuilder<'tcx> {
pub(crate) fn into_place<'a>(
pub(in crate::build) fn into_place<'a>(
self,
tcx: TyCtxt<'tcx>,
typeck_results: &'a ty::TypeckResults<'tcx>,
upvars: &'a CaptureMap<'tcx>,
) -> Place<'tcx> {
if let PlaceBase::Local(local) = self.base {
Place { local, projection: tcx.intern_place_elems(&self.projection) }
} else {
self.expect_upvars_resolved(tcx, typeck_results).into_place(tcx, typeck_results)
self.expect_upvars_resolved(tcx, upvars).into_place(tcx, upvars)
}
}

fn expect_upvars_resolved<'a>(
self,
tcx: TyCtxt<'tcx>,
typeck_results: &'a ty::TypeckResults<'tcx>,
upvars: &'a CaptureMap<'tcx>,
) -> PlaceBuilder<'tcx> {
to_upvars_resolved_place_builder(self, tcx, typeck_results).unwrap()
to_upvars_resolved_place_builder(self, tcx, upvars).unwrap()
}

/// Attempts to resolve the `PlaceBuilder`.
Expand All @@ -346,12 +282,12 @@ impl<'tcx> PlaceBuilder<'tcx> {
/// not captured. This can happen because the final mir that will be
/// generated doesn't require a read for this place. Failures will only
/// happen inside closures.
pub(crate) fn try_upvars_resolved<'a>(
pub(in crate::build) fn try_upvars_resolved<'a>(
self,
tcx: TyCtxt<'tcx>,
typeck_results: &'a ty::TypeckResults<'tcx>,
upvars: &'a CaptureMap<'tcx>,
) -> Result<PlaceBuilder<'tcx>, PlaceBuilder<'tcx>> {
to_upvars_resolved_place_builder(self, tcx, typeck_results)
to_upvars_resolved_place_builder(self, tcx, upvars)
}

pub(crate) fn base(&self) -> PlaceBase {
Expand Down Expand Up @@ -392,6 +328,12 @@ impl<'tcx> From<PlaceBase> for PlaceBuilder<'tcx> {
}
}

impl<'tcx> From<Place<'tcx>> for PlaceBuilder<'tcx> {
fn from(p: Place<'tcx>) -> Self {
Self { base: PlaceBase::Local(p.local), projection: p.projection.to_vec() }
}
}

impl<'a, 'tcx> Builder<'a, 'tcx> {
/// Compile `expr`, yielding a place that we can move from etc.
///
Expand All @@ -411,7 +353,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
expr: &Expr<'tcx>,
) -> BlockAnd<Place<'tcx>> {
let place_builder = unpack!(block = self.as_place_builder(block, expr));
block.and(place_builder.into_place(self.tcx, self.typeck_results))
block.and(place_builder.into_place(self.tcx, &self.upvars))
}

/// This is used when constructing a compound `Place`, so that we can avoid creating
Expand All @@ -435,7 +377,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
expr: &Expr<'tcx>,
) -> BlockAnd<Place<'tcx>> {
let place_builder = unpack!(block = self.as_read_only_place_builder(block, expr));
block.and(place_builder.into_place(self.tcx, self.typeck_results))
block.and(place_builder.into_place(self.tcx, &self.upvars))
}

/// This is used when constructing a compound `Place`, so that we can avoid creating
Expand Down Expand Up @@ -530,7 +472,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
inferred_ty: expr.ty,
});

let place = place_builder.clone().into_place(this.tcx, this.typeck_results);
let place = place_builder.clone().into_place(this.tcx, &this.upvars);
this.cfg.push(
block,
Statement {
Expand Down Expand Up @@ -629,17 +571,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
closure_def_id: LocalDefId,
var_hir_id: LocalVarId,
) -> BlockAnd<PlaceBuilder<'tcx>> {
let closure_ty =
self.typeck_results.node_type(self.tcx.hir().local_def_id_to_hir_id(closure_def_id));

let closure_kind = if let ty::Closure(_, closure_substs) = closure_ty.kind() {
self.infcx.closure_kind(closure_substs).unwrap()
} else {
// Generators are considered FnOnce.
ty::ClosureKind::FnOnce
};

block.and(PlaceBuilder::from(PlaceBase::Upvar { var_hir_id, closure_def_id, closure_kind }))
block.and(PlaceBuilder::from(PlaceBase::Upvar { var_hir_id, closure_def_id }))
}

/// Lower an index expression
Expand Down Expand Up @@ -678,7 +610,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
if is_outermost_index {
self.read_fake_borrows(block, fake_borrow_temps, source_info)
} else {
base_place = base_place.expect_upvars_resolved(self.tcx, self.typeck_results);
base_place = base_place.expect_upvars_resolved(self.tcx, &self.upvars);
self.add_fake_borrows_of_base(
&base_place,
block,
Expand Down Expand Up @@ -710,7 +642,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
block,
source_info,
len,
Rvalue::Len(slice.into_place(self.tcx, self.typeck_results)),
Rvalue::Len(slice.into_place(self.tcx, &self.upvars)),
);
// lt = idx < len
self.cfg.push_assign(
Expand Down
17 changes: 8 additions & 9 deletions compiler/rustc_mir_build/src/build/expr/as_rvalue.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,10 +328,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
unpack!(block = this.as_place_builder(block, &this.thir[*thir_place]));

if let Ok(place_builder_resolved) =
place_builder.try_upvars_resolved(this.tcx, this.typeck_results)
place_builder.try_upvars_resolved(this.tcx, &this.upvars)
{
let mir_place =
place_builder_resolved.into_place(this.tcx, this.typeck_results);
let mir_place = place_builder_resolved.into_place(this.tcx, &this.upvars);
this.cfg.push_fake_read(
block,
this.source_info(this.tcx.hir().span(*hir_id)),
Expand Down Expand Up @@ -623,7 +622,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
// is same as that of the capture in the parent closure.
PlaceBase::Upvar { .. } => {
let enclosing_upvars_resolved =
arg_place_builder.clone().into_place(this.tcx, this.typeck_results);
arg_place_builder.clone().into_place(this.tcx, &this.upvars);

match enclosing_upvars_resolved.as_ref() {
PlaceRef {
Expand All @@ -643,12 +642,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
);
// Not in a closure
debug_assert!(
this.upvar_mutbls.len() > upvar_index.index(),
"Unexpected capture place, upvar_mutbls={:#?}, upvar_index={:?}",
this.upvar_mutbls,
this.upvars.len() > upvar_index.index(),
"Unexpected capture place, upvars={:#?}, upvar_index={:?}",
this.upvars,
upvar_index
);
this.upvar_mutbls[upvar_index.index()]
this.upvars[upvar_index.index()].mutability
}
_ => bug!("Unexpected capture place"),
}
Expand All @@ -660,7 +659,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
Mutability::Mut => BorrowKind::Mut { allow_two_phase_borrow: false },
};

let arg_place = arg_place_builder.into_place(this.tcx, this.typeck_results);
let arg_place = arg_place_builder.into_place(this.tcx, &this.upvars);

this.cfg.push_assign(
block,
Expand Down
4 changes: 1 addition & 3 deletions compiler/rustc_mir_build/src/build/expr/into.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,9 +366,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
None => {
let place_builder = place_builder.clone();
this.consume_by_copy_or_move(
place_builder
.field(n, *ty)
.into_place(this.tcx, this.typeck_results),
place_builder.field(n, *ty).into_place(this.tcx, &this.upvars),
)
}
})
Expand Down
26 changes: 11 additions & 15 deletions compiler/rustc_mir_build/src/build/matches/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,9 +221,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
let source_info = self.source_info(scrutinee_span);

if let Ok(scrutinee_builder) =
scrutinee_place_builder.clone().try_upvars_resolved(self.tcx, self.typeck_results)
scrutinee_place_builder.clone().try_upvars_resolved(self.tcx, &self.upvars)
{
let scrutinee_place = scrutinee_builder.into_place(self.tcx, self.typeck_results);
let scrutinee_place = scrutinee_builder.into_place(self.tcx, &self.upvars);
self.cfg.push_fake_read(block, source_info, cause_matched_place, scrutinee_place);
}

Expand Down Expand Up @@ -348,12 +348,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
// ```
let mut opt_scrutinee_place: Option<(Option<&Place<'tcx>>, Span)> = None;
let scrutinee_place: Place<'tcx>;
if let Ok(scrutinee_builder) = scrutinee_place_builder
.clone()
.try_upvars_resolved(this.tcx, this.typeck_results)
if let Ok(scrutinee_builder) =
scrutinee_place_builder.clone().try_upvars_resolved(this.tcx, &this.upvars)
{
scrutinee_place =
scrutinee_builder.into_place(this.tcx, this.typeck_results);
scrutinee_place = scrutinee_builder.into_place(this.tcx, &this.upvars);
opt_scrutinee_place = Some((Some(&scrutinee_place), scrutinee_span));
}
let scope = this.declare_bindings(
Expand Down Expand Up @@ -620,9 +618,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
// };
// ```
if let Ok(match_pair_resolved) =
initializer.clone().try_upvars_resolved(self.tcx, self.typeck_results)
initializer.clone().try_upvars_resolved(self.tcx, &self.upvars)
{
let place = match_pair_resolved.into_place(self.tcx, self.typeck_results);
let place = match_pair_resolved.into_place(self.tcx, &self.upvars);
*match_place = Some(place);
}
}
Expand Down Expand Up @@ -1602,9 +1600,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {

// Insert a Shallow borrow of any places that is switched on.
if let Some(fb) = fake_borrows && let Ok(match_place_resolved) =
match_place.clone().try_upvars_resolved(self.tcx, self.typeck_results)
match_place.clone().try_upvars_resolved(self.tcx, &self.upvars)
{
let resolved_place = match_place_resolved.into_place(self.tcx, self.typeck_results);
let resolved_place = match_place_resolved.into_place(self.tcx, &self.upvars);
fb.insert(resolved_place);
}

Expand Down Expand Up @@ -1791,10 +1789,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
);
let mut opt_expr_place: Option<(Option<&Place<'tcx>>, Span)> = None;
let expr_place: Place<'tcx>;
if let Ok(expr_builder) =
expr_place_builder.try_upvars_resolved(self.tcx, self.typeck_results)
{
expr_place = expr_builder.into_place(self.tcx, self.typeck_results);
if let Ok(expr_builder) = expr_place_builder.try_upvars_resolved(self.tcx, &self.upvars) {
expr_place = expr_builder.into_place(self.tcx, &self.upvars);
opt_expr_place = Some((Some(&expr_place), expr_span));
}
let otherwise_post_guard_block = otherwise_candidate.pre_binding_block.unwrap();
Expand Down
Loading

0 comments on commit 2e44c17

Please sign in to comment.