Skip to content

Commit

Permalink
Prefer Struct-of-array
Browse files Browse the repository at this point in the history
  • Loading branch information
tesuji committed Jun 16, 2024
1 parent c036594 commit 45aa179
Showing 1 changed file with 23 additions and 29 deletions.
52 changes: 23 additions & 29 deletions compiler/rustc_mir_transform/src/promote_consts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,10 @@ impl<'tcx> MirPass<'tcx> for PromoteTemps<'tcx> {
}

let ccx = ConstCx::new(tcx, body);
let (mut temps, all_candidates) = collect_temps_and_candidates(&ccx);
let (mut temps, mut all_candidates) = collect_temps_and_candidates(&ccx);

let promotable_candidates = validate_candidates(&ccx, &mut temps, &all_candidates);
validate_candidates(&ccx, &mut temps, &mut all_candidates);
let promotable_candidates = &all_candidates.0[..];

let promoted = promote_candidates(body, tcx, temps, promotable_candidates);
self.promoted_fragments.set(promoted);
Expand All @@ -83,18 +84,16 @@ enum TempState {
PromotedOut,
}

/// A "root candidate" for promotion, which will become the
/// returned value in a promoted MIR, unless it's a subset
/// of a larger candidate.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct Candidate {
location: Location,
}
/// Each `[Location]` is a "root candidate" for promotion, which will become
/// the returned value in a promoted MIR, unless it's a subset of a larger
/// candidate.
#[derive(Clone, Debug)]
struct Candidates(Vec<Location>);

struct Collector<'a, 'tcx> {
ccx: &'a ConstCx<'a, 'tcx>,
temps: IndexVec<Local, TempState>,
candidates: Vec<Candidate>,
candidates: Candidates,
}

impl<'tcx> Visitor<'tcx> for Collector<'_, 'tcx> {
Expand Down Expand Up @@ -141,24 +140,23 @@ impl<'tcx> Visitor<'tcx> for Collector<'_, 'tcx> {
}
TempState::Unpromotable | TempState::PromotedOut => TempState::Unpromotable,
};
debug!(?temp);
}

fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
self.super_rvalue(rvalue, location);

if let Rvalue::Ref(..) = *rvalue {
self.candidates.push(Candidate { location });
self.candidates.0.push(location);
}
}
}

fn collect_temps_and_candidates<'tcx>(
ccx: &ConstCx<'_, 'tcx>,
) -> (IndexVec<Local, TempState>, Vec<Candidate>) {
) -> (IndexVec<Local, TempState>, Candidates) {
let mut collector = Collector {
temps: IndexVec::from_elem(TempState::Undefined, &ccx.body.local_decls),
candidates: vec![],
candidates: Candidates(vec![]),
ccx,
};
for (bb, data) in traversal::reverse_postorder(ccx.body) {
Expand Down Expand Up @@ -192,8 +190,8 @@ impl<'a, 'tcx> std::ops::Deref for Validator<'a, 'tcx> {
struct Unpromotable;

impl<'tcx> Validator<'_, 'tcx> {
fn validate_candidate(&mut self, candidate: Candidate) -> Result<(), Unpromotable> {
let Left(statement) = self.body.stmt_at(candidate.location) else { bug!() };
fn validate_candidate(&mut self, candidate: Location) -> Result<(), Unpromotable> {
let Left(statement) = self.body.stmt_at(candidate) else { bug!() };
let Some((_, Rvalue::Ref(_, kind, place))) = statement.kind.as_assign() else { bug!() };

// We can only promote interior borrows of promotable temps (non-temps
Expand Down Expand Up @@ -691,15 +689,11 @@ impl<'tcx> Validator<'_, 'tcx> {
fn validate_candidates(
ccx: &ConstCx<'_, '_>,
temps: &mut IndexSlice<Local, TempState>,
candidates: &[Candidate],
) -> Vec<Candidate> {
candidates: &mut Candidates,
) {
let mut validator = Validator { ccx, temps, promotion_safe_blocks: None };

candidates
.iter()
.copied()
.filter(|&candidate| validator.validate_candidate(candidate).is_ok())
.collect()
candidates.0.retain(|&candidate| validator.validate_candidate(candidate).is_ok());
}

struct Promoter<'a, 'tcx> {
Expand Down Expand Up @@ -854,7 +848,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
new_temp
}

fn promote_candidate(mut self, candidate: Candidate, next_promoted_id: usize) -> Body<'tcx> {
fn promote_candidate(mut self, candidate: Location, next_promoted_id: usize) -> Body<'tcx> {
let def = self.source.source.def_id();
let (mut rvalue, promoted_op) = {
let promoted = &mut self.promoted;
Expand All @@ -871,7 +865,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {

let blocks = self.source.basic_blocks.as_mut();
let local_decls = &mut self.source.local_decls;
let loc = candidate.location;
let loc = candidate;
let statement = &mut blocks[loc.block].statements[loc.statement_index];
let StatementKind::Assign(box (_, Rvalue::Ref(region, borrow_kind, place))) =
&mut statement.kind
Expand Down Expand Up @@ -965,7 +959,7 @@ fn promote_candidates<'tcx>(
body: &mut Body<'tcx>,
tcx: TyCtxt<'tcx>,
mut temps: IndexVec<Local, TempState>,
candidates: Vec<Candidate>,
candidates: &[Location],
) -> IndexVec<Promoted, Body<'tcx>> {
// Visit candidates in reverse, in case they're nested.
debug!(promote_candidates = ?candidates);
Expand All @@ -978,8 +972,8 @@ fn promote_candidates<'tcx>(
let mut promotions = IndexVec::new();

let mut extra_statements = vec![];
for candidate in candidates.into_iter().rev() {
let Location { block, statement_index } = candidate.location;
for candidate in candidates.iter().copied().rev() {
let Location { block, statement_index } = candidate;
if let StatementKind::Assign(box (place, _)) = &body[block].statements[statement_index].kind
{
if let Some(local) = place.as_local() {
Expand All @@ -993,7 +987,7 @@ fn promote_candidates<'tcx>(
// Declare return place local so that `mir::Body::new` doesn't complain.
let initial_locals = iter::once(LocalDecl::new(tcx.types.never, body.span)).collect();

let mut scope = body.source_scopes[body.source_info(candidate.location).scope].clone();
let mut scope = body.source_scopes[body.source_info(candidate).scope].clone();
scope.parent_scope = None;

let mut promoted = Body::new(
Expand Down

0 comments on commit 45aa179

Please sign in to comment.