diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs index 26b6003e1095e..b010b37e75cdd 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 @@ use drop::{ 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; @@ -735,53 +736,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,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 {