Skip to content

Micro-optimization attempt in coroutine layout computation #147858

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 30, 2025
Merged

Micro-optimization attempt in coroutine layout computation #147858

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0c98c2
Allocate vec capacity upfront in coroutine layout computation
yotamofek Oct 18, 2025
a67c615
Reduce rightward drift in `locals_live_across_suspend_points`
yotamofek Oct 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 62 additions & 60 deletions compiler/rustc_mir_transform/src/coroutine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,15 @@

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::{
cleanup_async_drops, create_coroutine_drop_shim, create_coroutine_drop_shim_async,
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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
Copy link
Contributor

@cjgillot cjgillot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we do simpler? The original code but using IndexVec::with_capacity?

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

Expand Down Expand Up @@ -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();
Copy link
Contributor

@cjgillot cjgillot Oct 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of pushing, could we use IndexVec::from_elem_n(IndexVec::new(), CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len())? And then assign to each element?

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 {
Expand Down
Loading