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);