Allocate vec capacity upfront in coroutine layout computation

This commit is contained in:
Yotam Ofek
2025-10-18 22:49:38 +03:00
parent c6efb9019b
commit f0c98c2caa
+23 -21
View File
@@ -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 @@
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 {