Auto merge of #147858 - yotamofek:pr/mir/coroutine-layout-opt, r=<try>

Micro-optimization attempt in coroutine layout computation

Author: rust-bors[bot]

compiler/rustc_mir_transform/src/coroutine.rs

@@ -735,53 +735,53 @@ fn locals_live_across_suspend_points<'tcx>(
     let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
 
     for (block, data) in body.basic_blocks.iter_enumerated() {
-        if let TerminatorKind::Yield { .. } = data.terminator().kind {
-            let loc = Location { block, statement_index: data.statements.len() };
-
-            liveness.seek_to_block_end(block);
-            let mut live_locals = liveness.get().clone();
-
-            if !movable {
-                // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
-                // This is correct for movable coroutines since borrows cannot live across
-                // suspension points. However for immovable coroutines we need to account for
-                // borrows, so we conservatively assume that all borrowed locals are live until
-                // we find a StorageDead statement referencing the locals.
-                // To do this we just union our `liveness` result with `borrowed_locals`, which
-                // contains all the locals which has been borrowed before this suspension point.
-                // If a borrow is converted to a raw reference, we must also assume that it lives
-                // forever. Note that the final liveness is still bounded by the storage liveness
-                // of the local, which happens using the `intersect` operation below.
-                borrowed_locals_cursor2.seek_before_primary_effect(loc);
-                live_locals.union(borrowed_locals_cursor2.get());
-            }
+        let TerminatorKind::Yield { .. } = data.terminator().kind else { continue };
+
+        let loc = Location { block, statement_index: data.statements.len() };
+
+        liveness.seek_to_block_end(block);
+        let mut live_locals = liveness.get().clone();
+
+        if !movable {
+            // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
+            // This is correct for movable coroutines since borrows cannot live across
+            // suspension points. However for immovable coroutines we need to account for
+            // borrows, so we conservatively assume that all borrowed locals are live until
+            // we find a StorageDead statement referencing the locals.
+            // To do this we just union our `liveness` result with `borrowed_locals`, which
+            // contains all the locals which has been borrowed before this suspension point.
+            // If a borrow is converted to a raw reference, we must also assume that it lives
+            // forever. Note that the final liveness is still bounded by the storage liveness
+            // of the local, which happens using the `intersect` operation below.
+            borrowed_locals_cursor2.seek_before_primary_effect(loc);
+            live_locals.union(borrowed_locals_cursor2.get());
+        }
 
-            // Store the storage liveness for later use so we can restore the state
-            // after a suspension point
-            storage_live.seek_before_primary_effect(loc);
-            storage_liveness_map[block] = Some(storage_live.get().clone());
+        // Store the storage liveness for later use so we can restore the state
+        // after a suspension point
+        storage_live.seek_before_primary_effect(loc);
+        storage_liveness_map[block] = Some(storage_live.get().clone());
 
-            // Locals live are live at this point only if they are used across
-            // suspension points (the `liveness` variable)
-            // and their storage is required (the `storage_required` variable)
-            requires_storage_cursor.seek_before_primary_effect(loc);
-            live_locals.intersect(requires_storage_cursor.get());
+        // Locals live are live at this point only if they are used across
+        // suspension points (the `liveness` variable)
+        // and their storage is required (the `storage_required` variable)
+        requires_storage_cursor.seek_before_primary_effect(loc);
+        live_locals.intersect(requires_storage_cursor.get());
 
-            // The coroutine argument is ignored.
-            live_locals.remove(SELF_ARG);
+        // The coroutine argument is ignored.
+        live_locals.remove(SELF_ARG);
 
-            debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
+        debug!(?loc, ?live_locals);
 
-            // Add the locals live at this suspension point to the set of locals which live across
-            // any suspension points
-            live_locals_at_any_suspension_point.union(&live_locals);
+        // Add the locals live at this suspension point to the set of locals which live across
+        // any suspension points
+        live_locals_at_any_suspension_point.union(&live_locals);
 
-            live_locals_at_suspension_points.push(live_locals);
-            source_info_at_suspension_points.push(data.terminator().source_info);
-        }
+        live_locals_at_suspension_points.push(live_locals);
+        source_info_at_suspension_points.push(data.terminator().source_info);
     }
 
-    debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
+    debug!(?live_locals_at_any_suspension_point);
     let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point);
 
     // Renumber our liveness_map bitsets to include only the locals we are
@@ -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);