collect collections instead of pushing iteratively in coroutine layout computation

Author: yotamofek

compiler/rustc_mir_transform/src/coroutine.rs

@@ -982,71 +982,70 @@ fn compute_layout<'tcx>(
     } = liveness;
 
     // Gather live local types and their indices.
-    let mut locals = IndexVec::<CoroutineSavedLocal, _>::new();
-    let mut tys = IndexVec::<CoroutineSavedLocal, _>::new();
-    for (saved_local, local) in saved_locals.iter_enumerated() {
-        debug!("coroutine saved local {:?} => {:?}", saved_local, local);
-
-        locals.push(local);
-        let decl = &body.local_decls[local];
-        debug!(?decl);
-
-        // Do not `unwrap_crate_local` here, as post-borrowck cleanup may have already cleared
-        // the information. This is alright, since `ignore_for_traits` is only relevant when
-        // this code runs on pre-cleanup MIR, and `ignore_for_traits = false` is the safer
-        // default.
-        let ignore_for_traits = match decl.local_info {
-            // Do not include raw pointers created from accessing `static` items, as those could
-            // well be re-created by another access to the same static.
-            ClearCrossCrate::Set(box LocalInfo::StaticRef { is_thread_local, .. }) => {
-                !is_thread_local
-            }
-            // Fake borrows are only read by fake reads, so do not have any reality in
-            // post-analysis MIR.
-            ClearCrossCrate::Set(box LocalInfo::FakeBorrow) => true,
-            _ => false,
-        };
-        let decl =
-            CoroutineSavedTy { ty: decl.ty, source_info: decl.source_info, ignore_for_traits };
-        debug!(?decl);
+    let (locals, tys): (IndexVec<_, _>, IndexVec<_, _>) = saved_locals
+        .iter_enumerated()
+        .map(|(saved_local, local)| {
+            debug!("coroutine saved local {saved_local:?} => {local:?}");
 
-        tys.push(decl);
-    }
+            let decl = &body.local_decls[local];
+            debug!(?decl);
+
+            // Do not `unwrap_crate_local` here, as post-borrowck cleanup may have already cleared
+            // the information. This is alright, since `ignore_for_traits` is only relevant when
+            // this code runs on pre-cleanup MIR, and `ignore_for_traits = false` is the safer
+            // default.
+            let ignore_for_traits = match decl.local_info {
+                // Do not include raw pointers created from accessing `static` items, as those could
+                // well be re-created by another access to the same static.
+                ClearCrossCrate::Set(box LocalInfo::StaticRef { is_thread_local, .. }) => {
+                    !is_thread_local
+                }
+                // Fake borrows are only read by fake reads, so do not have any reality in
+                // post-analysis MIR.
+                ClearCrossCrate::Set(box LocalInfo::FakeBorrow) => true,
+                _ => false,
+            };
+            let decl =
+                CoroutineSavedTy { ty: decl.ty, source_info: decl.source_info, ignore_for_traits };
+            debug!(?decl);
+
+            (local, decl)
+        })
+        .unzip();
 
     // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
     // 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 variant_source_info: IndexVec<VariantIdx, SourceInfo> = [
         SourceInfo::outermost(body_span.shrink_to_lo()),
         SourceInfo::outermost(body_span.shrink_to_hi()),
         SourceInfo::outermost(body_span.shrink_to_hi()),
     ]
-    .iter()
-    .copied()
+    .into_iter()
+    .chain(source_info_at_suspension_points)
     .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 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);
-            // 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]);
-    }
+    let variant_fields = iter::repeat_n(IndexVec::new(), CoroutineArgs::RESERVED_VARIANTS)
+        .chain(live_locals_at_suspension_points.into_iter().enumerate().map(
+            |(suspension_point_idx, live_locals)| {
+                let variant_index =
+                    VariantIdx::from(CoroutineArgs::RESERVED_VARIANTS + suspension_point_idx);
+                let fields = live_locals.iter().collect::<IndexVec<FieldIdx, _>>();
+                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.
+                    remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
+                }
+                fields
+            },
+        ))
+        .collect::<IndexVec<VariantIdx, _>>();
     debug!("coroutine variant_fields = {:?}", variant_fields);
     debug!("coroutine storage_conflicts = {:#?}", storage_conflicts);