diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs index 26b6003e1095..46b4e8b86e51 100644 --- a/compiler/rustc_mir_transform/src/coroutine.rs +++ b/compiler/rustc_mir_transform/src/coroutine.rs @@ -52,7 +52,7 @@ mod by_move_body; mod drop; -use std::{iter, ops}; +use std::ops; pub(super) use by_move_body::coroutine_by_move_body_def_id; use drop::{ @@ -60,6 +60,7 @@ create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, expand_async_drops, has_expandable_async_drops, insert_clean_drop, }; +use itertools::izip; use rustc_abi::{FieldIdx, VariantIdx}; use rustc_data_structures::fx::FxHashSet; use rustc_errors::pluralize; @@ -982,8 +983,8 @@ fn compute_layout<'tcx>( } = liveness; // Gather live local types and their indices. - let mut locals = IndexVec::::new(); - let mut tys = IndexVec::::new(); + let mut locals = IndexVec::::with_capacity(saved_locals.domain_size()); + let mut tys = IndexVec::::with_capacity(saved_locals.domain_size()); for (saved_local, local) in saved_locals.iter_enumerated() { debug!("coroutine saved local {:?} => {:?}", saved_local, local); @@ -1017,38 +1018,39 @@ fn compute_layout<'tcx>( // In debuginfo, these will correspond to the beginning (UNRESUMED) or end // (RETURNED, POISONED) of the function. let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span; - let mut variant_source_info: IndexVec = [ + let mut variant_source_info: IndexVec = IndexVec::with_capacity( + CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(), + ); + variant_source_info.extend([ SourceInfo::outermost(body_span.shrink_to_lo()), SourceInfo::outermost(body_span.shrink_to_hi()), SourceInfo::outermost(body_span.shrink_to_hi()), - ] - .iter() - .copied() - .collect(); + ]); // Build the coroutine variant field list. // Create a map from local indices to coroutine struct indices. - let mut variant_fields: IndexVec> = - iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect(); + let mut variant_fields: IndexVec = IndexVec::from_elem_n( + IndexVec::new(), + CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(), + ); let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size()); - for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() { - let variant_index = - VariantIdx::from(CoroutineArgs::RESERVED_VARIANTS + suspension_point_idx); - let mut fields = IndexVec::new(); - for (idx, saved_local) in live_locals.iter().enumerate() { - fields.push(saved_local); + for (live_locals, &source_info_at_suspension_point, (variant_index, fields)) in izip!( + &live_locals_at_suspension_points, + &source_info_at_suspension_points, + variant_fields.iter_enumerated_mut().skip(CoroutineArgs::RESERVED_VARIANTS) + ) { + *fields = live_locals.iter().collect(); + for (idx, &saved_local) in fields.iter_enumerated() { // Note that if a field is included in multiple variants, we will // just use the first one here. That's fine; fields do not move // around inside coroutines, so it doesn't matter which variant // index we access them by. - let idx = FieldIdx::from_usize(idx); remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx)); } - variant_fields.push(fields); - variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]); + variant_source_info.push(source_info_at_suspension_point); } - debug!("coroutine variant_fields = {:?}", variant_fields); - debug!("coroutine storage_conflicts = {:#?}", storage_conflicts); + debug!(?variant_fields); + debug!(?storage_conflicts); let mut field_names = IndexVec::from_elem(None, &tys); for var in &body.var_debug_info {