Allocate vec capacity upfront in coroutine layout computation

Author: yotamofek

compiler/rustc_mir_transform/src/coroutine.rs

@@ -52,14 +52,15 @@
 
 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::{
     cleanup_async_drops, create_coroutine_drop_shim, create_coroutine_drop_shim_async,
     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::<CoroutineSavedLocal, _>::new();
-    let mut tys = IndexVec::<CoroutineSavedLocal, _>::new();
+    let mut locals = IndexVec::<CoroutineSavedLocal, _>::with_capacity(saved_locals.domain_size());
+    let mut tys = IndexVec::<CoroutineSavedLocal, _>::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<VariantIdx, SourceInfo> = [
+    let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = 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<VariantIdx, IndexVec<FieldIdx, CoroutineSavedLocal>> =
-        iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect();
+    let mut variant_fields: IndexVec<VariantIdx, _> = 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 {