From 6a1d0699a484ee875c87394f70cb37f09acadd88 Mon Sep 17 00:00:00 2001 From: Aman Arora Date: Thu, 26 Nov 2020 00:05:18 -0500 Subject: [PATCH] Use Places for captures in MIR - Use closure_min_capture maps to capture precise paths - PlaceBuilder now searches for ancestors in min_capture list - Add API to `Ty` to allow access to the n-th element in a tuple in O(1) time. Co-authored-by: Roxane Fruytier --- compiler/rustc_middle/src/ty/sty.rs | 9 ++ .../src/build/expr/as_place.rs | 150 +++++++++++++++--- 2 files changed, 137 insertions(+), 22 deletions(-) diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs index 78994c6e1c7..a90807de99d 100644 --- a/compiler/rustc_middle/src/ty/sty.rs +++ b/compiler/rustc_middle/src/ty/sty.rs @@ -2174,6 +2174,15 @@ pub fn tuple_fields(&self) -> impl DoubleEndedIterator> { } } + /// Get the `i`-th element of a tuple. + /// Panics when called on anything but a tuple. + pub fn tuple_element_ty(&self, i: usize) -> Option> { + match self.kind() { + Tuple(substs) => substs.iter().nth(i).map(|field| field.expect_ty()), + _ => bug!("tuple_fields called on non-tuple"), + } + } + /// If the type contains variants, returns the valid range of variant indices. // // FIXME: This requires the optimized MIR in the case of generators. 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 2ab733fabb8..b3957fa7222 100644 --- a/compiler/rustc_mir_build/src/build/expr/as_place.rs +++ b/compiler/rustc_mir_build/src/build/expr/as_place.rs @@ -7,10 +7,12 @@ use rustc_hir::def_id::DefId; use rustc_hir::HirId; use rustc_middle::middle::region; +use rustc_middle::hir::place::ProjectionKind as HirProjectionKind; use rustc_middle::mir::AssertKind::BoundsCheck; use rustc_middle::mir::*; use rustc_middle::ty::{self, CanonicalUserTypeAnnotation, Ty, TyCtxt, Variance}; use rustc_span::Span; +use rustc_target::abi::VariantIdx; use rustc_index::vec::Idx; @@ -70,28 +72,129 @@ struct PlaceBuilder<'tcx> { projection: Vec>, } -fn capture_matching_projections<'a, 'tcx>( +/// Given a list of MIR projections, convert them to list of HIR ProjectionKind. +/// The projections are truncated to represent a path that might be captured by a +/// closure/generator. This implies the vector returned from this function doesn't contain +/// ProjectionElems `Downcast`, `ConstantIndex`, `Index`, or `Subslice` because those will never be +/// part of a path that is captued by a closure. We stop applying projections once we see the first +/// projection that isn't captured by a closure. +fn convert_to_hir_projections_and_truncate_for_capture<'tcx>( + mir_projections: &Vec>, +) -> Vec { + + let mut hir_projections = Vec::new(); + + for mir_projection in mir_projections { + let hir_projection = match mir_projection { + ProjectionElem::Deref => HirProjectionKind::Deref, + ProjectionElem::Field(field, _) => { + // We will never encouter this for multivariant enums, + // read the comment for `Downcast`. + HirProjectionKind::Field(field.index() as u32, VariantIdx::new(0)) + }, + ProjectionElem::Downcast(..) => { + // This projections exist only for enums that have + // multiple variants. Since such enums that are captured + // completely, we can stop here. + break + }, + ProjectionElem::Index(..) + | ProjectionElem::ConstantIndex { .. } + | ProjectionElem::Subslice { .. } => { + // We don't capture array-access projections. + // We can stop here as arrays are captured completely. + break + }, + }; + + hir_projections.push(hir_projection); + } + + hir_projections +} + +/// Return true if the `proj_possible_ancestor` represents an ancestor path +/// to `proj_capture` or `proj_possible_ancestor` is same as `proj_capture`, +/// assuming they both start off of the same root variable. +/// +/// **Note:** It's the caller's responsibility to ensure that both lists of projections +/// start off of the same root variable. +/// +/// Eg: 1. `foo.x` which is represented using `projections=[Field(x)]` is an ancestor of +/// `foo.x.y` which is represented using `projections=[Field(x), Field(y)]`. +/// Note both `foo.x` and `foo.x.y` start off of the same root variable `foo`. +/// 2. Since we only look at the projections here function will return `bar.x` as an a valid +/// ancestor of `foo.x.y`. It's the caller's responsibility to ensure that both projections +/// list are being applied to the same root variable. +fn is_ancestor_or_same_capture( + proj_possible_ancestor: &Vec, + proj_capture: &Vec, +) -> bool { + // We want to make sure `is_ancestor_or_same_capture("x.0.0", "x.0")` to return false. + // Therefore we can't just check if all projections are same in the zipped iterator below. + if proj_possible_ancestor.len() > proj_capture.len() { + return false; + } + + proj_possible_ancestor.iter().zip(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: HirId, + root_var_idx: usize, +) -> usize { + let mut res = 0; + for (var_id, capture_list) in closure_min_captures { + if *var_id == var_hir_id { + 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`. +/// +/// Note there will be at most one ancestor for any given Place. +/// +/// Returns None, when the ancestor is not found. +fn find_capture_matching_projections<'a, 'tcx>( typeck_results: &'a ty::TypeckResults<'tcx>, var_hir_id: HirId, closure_def_id: DefId, - _projections: &Vec>, -) -> Option<(usize, ty::UpvarCapture<'tcx>)> { - typeck_results - .closure_captures - .get(&closure_def_id) - .and_then(|captures| captures.get_full(&var_hir_id)) - .and_then(|(capture_index, _, _)|{ - let upvar_id = ty::UpvarId::new(var_hir_id, closure_def_id.expect_local()); - let capture_kind = typeck_results.upvar_capture(upvar_id); - Some((capture_index, capture_kind)) - }) + projections: &Vec>, +) -> 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)?; + + 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)| { + let possible_ancestor_proj_kinds = + capture.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 presepective 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 PlaceBuilder now starts from PlaceBase::Local. +/// Takes a PlaceBuilder and resolves the upvar (if any) within it, so that the +/// `PlaceBuilder` now starts from `PlaceBase::Local`. /// -/// Returns a Result with the error being the HirId of the -/// Upvar that was not found. +/// Returns a Result with the error being the HirId of the Upvar that was not found. fn to_upvars_resolved_place_builder<'a, 'tcx>( from_builder: PlaceBuilder<'tcx>, tcx: TyCtxt<'tcx>, @@ -110,8 +213,8 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>( ty::ClosureKind::FnOnce => {} } - let (capture_index, capture_kind) = - if let Some(capture_details) = capture_matching_projections( + let (capture_index, capture) = + if let Some(capture_details) = find_capture_matching_projections( typeck_results, var_hir_id, closure_def_id, @@ -149,21 +252,24 @@ fn to_upvars_resolved_place_builder<'a, 'tcx>( // 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 + // 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.upvar_tys().nth(capture_index).unwrap(); + let var_ty = substs.tupled_upvars_ty().tuple_element_ty(capture_index).unwrap(); 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_kind { + 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 next_projection = 0; + 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..));