From a1daa34ad005d0b34d30c878cb4e2e995346d300 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 18 Feb 2025 11:24:57 +1100 Subject: [PATCH 1/5] Use `MirPatch` in `EnumSizeOpt`. Instead of `expand_statements`. This makes the code shorter and consistent with other MIR transform passes. The tests require updating because there is a slight change in MIR output: - the old code replaced the original statement with twelve new statements. - the new code inserts converts the original statement to a `nop` and then insert twelve new statements in front of it. I.e. we now end up with an extra `nop`, which doesn't matter at all. --- .../rustc_mir_transform/src/large_enums.rs | 186 +++++++----------- .../enum_opt.cand.EnumSizeOpt.32bit.diff | 2 + .../enum_opt.cand.EnumSizeOpt.64bit.diff | 2 + .../enum_opt.unin.EnumSizeOpt.32bit.diff | 2 + .../enum_opt.unin.EnumSizeOpt.64bit.diff | 2 + 5 files changed, 82 insertions(+), 112 deletions(-) diff --git a/compiler/rustc_mir_transform/src/large_enums.rs b/compiler/rustc_mir_transform/src/large_enums.rs index 1e546bfbeb303..47cb478fe33ee 100644 --- a/compiler/rustc_mir_transform/src/large_enums.rs +++ b/compiler/rustc_mir_transform/src/large_enums.rs @@ -6,6 +6,8 @@ use rustc_middle::ty::util::IntTypeExt; use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt}; use rustc_session::Session; +use crate::patch::MirPatch; + /// A pass that seeks to optimize unnecessary moves of large enum types, if there is a large /// enough discrepancy between them. /// @@ -41,31 +43,34 @@ impl<'tcx> crate::MirPass<'tcx> for EnumSizeOpt { let mut alloc_cache = FxHashMap::default(); let typing_env = body.typing_env(tcx); - let blocks = body.basic_blocks.as_mut(); - let local_decls = &mut body.local_decls; + let mut patch = MirPatch::new(body); - for bb in blocks { - bb.expand_statements(|st| { + for (block, data) in body.basic_blocks.as_mut().iter_enumerated_mut() { + for (statement_index, st) in data.statements.iter_mut().enumerate() { let StatementKind::Assign(box ( lhs, Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)), )) = &st.kind else { - return None; + continue; }; - let ty = lhs.ty(local_decls, tcx).ty; + let location = Location { block, statement_index }; - let (adt_def, num_variants, alloc_id) = - self.candidate(tcx, typing_env, ty, &mut alloc_cache)?; + let ty = lhs.ty(&body.local_decls, tcx).ty; - let source_info = st.source_info; - let span = source_info.span; + let Some((adt_def, num_variants, alloc_id)) = + self.candidate(tcx, typing_env, ty, &mut alloc_cache) + else { + continue; + }; + + let span = st.source_info.span; let tmp_ty = Ty::new_array(tcx, tcx.types.usize, num_variants as u64); - let size_array_local = local_decls.push(LocalDecl::new(tmp_ty, span)); - let store_live = - Statement { source_info, kind: StatementKind::StorageLive(size_array_local) }; + let size_array_local = patch.new_temp(tmp_ty, span); + + let store_live = StatementKind::StorageLive(size_array_local); let place = Place::from(size_array_local); let constant_vals = ConstOperand { @@ -77,108 +82,63 @@ impl<'tcx> crate::MirPass<'tcx> for EnumSizeOpt { ), }; let rval = Rvalue::Use(Operand::Constant(Box::new(constant_vals))); - let const_assign = - Statement { source_info, kind: StatementKind::Assign(Box::new((place, rval))) }; - - let discr_place = Place::from( - local_decls.push(LocalDecl::new(adt_def.repr().discr_type().to_ty(tcx), span)), - ); - let store_discr = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - discr_place, - Rvalue::Discriminant(*rhs), - ))), - }; - - let discr_cast_place = - Place::from(local_decls.push(LocalDecl::new(tcx.types.usize, span))); - let cast_discr = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - discr_cast_place, - Rvalue::Cast( - CastKind::IntToInt, - Operand::Copy(discr_place), - tcx.types.usize, - ), - ))), - }; - - let size_place = - Place::from(local_decls.push(LocalDecl::new(tcx.types.usize, span))); - let store_size = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - size_place, - Rvalue::Use(Operand::Copy(Place { - local: size_array_local, - projection: tcx - .mk_place_elems(&[PlaceElem::Index(discr_cast_place.local)]), - })), - ))), - }; - - let dst = - Place::from(local_decls.push(LocalDecl::new(Ty::new_mut_ptr(tcx, ty), span))); - let dst_ptr = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - dst, - Rvalue::RawPtr(RawPtrKind::Mut, *lhs), - ))), - }; + let const_assign = StatementKind::Assign(Box::new((place, rval))); + + let discr_place = + Place::from(patch.new_temp(adt_def.repr().discr_type().to_ty(tcx), span)); + let store_discr = + StatementKind::Assign(Box::new((discr_place, Rvalue::Discriminant(*rhs)))); + + let discr_cast_place = Place::from(patch.new_temp(tcx.types.usize, span)); + let cast_discr = StatementKind::Assign(Box::new(( + discr_cast_place, + Rvalue::Cast(CastKind::IntToInt, Operand::Copy(discr_place), tcx.types.usize), + ))); + + let size_place = Place::from(patch.new_temp(tcx.types.usize, span)); + let store_size = StatementKind::Assign(Box::new(( + size_place, + Rvalue::Use(Operand::Copy(Place { + local: size_array_local, + projection: tcx.mk_place_elems(&[PlaceElem::Index(discr_cast_place.local)]), + })), + ))); + + let dst = Place::from(patch.new_temp(Ty::new_mut_ptr(tcx, ty), span)); + let dst_ptr = + StatementKind::Assign(Box::new((dst, Rvalue::RawPtr(RawPtrKind::Mut, *lhs)))); let dst_cast_ty = Ty::new_mut_ptr(tcx, tcx.types.u8); - let dst_cast_place = - Place::from(local_decls.push(LocalDecl::new(dst_cast_ty, span))); - let dst_cast = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - dst_cast_place, - Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(dst), dst_cast_ty), - ))), - }; + let dst_cast_place = Place::from(patch.new_temp(dst_cast_ty, span)); + let dst_cast = StatementKind::Assign(Box::new(( + dst_cast_place, + Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(dst), dst_cast_ty), + ))); - let src = - Place::from(local_decls.push(LocalDecl::new(Ty::new_imm_ptr(tcx, ty), span))); - let src_ptr = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - src, - Rvalue::RawPtr(RawPtrKind::Const, *rhs), - ))), - }; + let src = Place::from(patch.new_temp(Ty::new_imm_ptr(tcx, ty), span)); + let src_ptr = + StatementKind::Assign(Box::new((src, Rvalue::RawPtr(RawPtrKind::Const, *rhs)))); let src_cast_ty = Ty::new_imm_ptr(tcx, tcx.types.u8); - let src_cast_place = - Place::from(local_decls.push(LocalDecl::new(src_cast_ty, span))); - let src_cast = Statement { - source_info, - kind: StatementKind::Assign(Box::new(( - src_cast_place, - Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(src), src_cast_ty), - ))), - }; + let src_cast_place = Place::from(patch.new_temp(src_cast_ty, span)); + let src_cast = StatementKind::Assign(Box::new(( + src_cast_place, + Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(src), src_cast_ty), + ))); - let deinit_old = - Statement { source_info, kind: StatementKind::Deinit(Box::new(dst)) }; - - let copy_bytes = Statement { - source_info, - kind: StatementKind::Intrinsic(Box::new( - NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping { - src: Operand::Copy(src_cast_place), - dst: Operand::Copy(dst_cast_place), - count: Operand::Copy(size_place), - }), - )), - }; + let deinit_old = StatementKind::Deinit(Box::new(dst)); + + let copy_bytes = StatementKind::Intrinsic(Box::new( + NonDivergingIntrinsic::CopyNonOverlapping(CopyNonOverlapping { + src: Operand::Copy(src_cast_place), + dst: Operand::Copy(dst_cast_place), + count: Operand::Copy(size_place), + }), + )); - let store_dead = - Statement { source_info, kind: StatementKind::StorageDead(size_array_local) }; + let store_dead = StatementKind::StorageDead(size_array_local); - let iter = [ + let stmts = [ store_live, const_assign, store_discr, @@ -191,14 +151,16 @@ impl<'tcx> crate::MirPass<'tcx> for EnumSizeOpt { deinit_old, copy_bytes, store_dead, - ] - .into_iter(); + ]; + for stmt in stmts { + patch.add_statement(location, stmt); + } st.make_nop(); - - Some(iter) - }); + } } + + patch.apply(body); } fn is_required(&self) -> bool { diff --git a/tests/mir-opt/enum_opt.cand.EnumSizeOpt.32bit.diff b/tests/mir-opt/enum_opt.cand.EnumSizeOpt.32bit.diff index 727efe4b0d95d..267a4c1cf6beb 100644 --- a/tests/mir-opt/enum_opt.cand.EnumSizeOpt.32bit.diff +++ b/tests/mir-opt/enum_opt.cand.EnumSizeOpt.32bit.diff @@ -47,6 +47,7 @@ + Deinit(_8); + copy_nonoverlapping(dst = copy _9, src = copy _11, count = copy _7); + StorageDead(_4); ++ nop; StorageDead(_2); - _0 = move _1; + StorageLive(_12); @@ -61,6 +62,7 @@ + Deinit(_16); + copy_nonoverlapping(dst = copy _17, src = copy _19, count = copy _15); + StorageDead(_12); ++ nop; StorageDead(_1); return; } diff --git a/tests/mir-opt/enum_opt.cand.EnumSizeOpt.64bit.diff b/tests/mir-opt/enum_opt.cand.EnumSizeOpt.64bit.diff index 8d0cd97f78661..8e5c403cd7e6b 100644 --- a/tests/mir-opt/enum_opt.cand.EnumSizeOpt.64bit.diff +++ b/tests/mir-opt/enum_opt.cand.EnumSizeOpt.64bit.diff @@ -47,6 +47,7 @@ + Deinit(_8); + copy_nonoverlapping(dst = copy _9, src = copy _11, count = copy _7); + StorageDead(_4); ++ nop; StorageDead(_2); - _0 = move _1; + StorageLive(_12); @@ -61,6 +62,7 @@ + Deinit(_16); + copy_nonoverlapping(dst = copy _17, src = copy _19, count = copy _15); + StorageDead(_12); ++ nop; StorageDead(_1); return; } diff --git a/tests/mir-opt/enum_opt.unin.EnumSizeOpt.32bit.diff b/tests/mir-opt/enum_opt.unin.EnumSizeOpt.32bit.diff index 6d1e2a72fdb73..96c5aadd85fd4 100644 --- a/tests/mir-opt/enum_opt.unin.EnumSizeOpt.32bit.diff +++ b/tests/mir-opt/enum_opt.unin.EnumSizeOpt.32bit.diff @@ -47,6 +47,7 @@ + Deinit(_8); + copy_nonoverlapping(dst = copy _9, src = copy _11, count = copy _7); + StorageDead(_4); ++ nop; StorageDead(_2); - _0 = move _1; + StorageLive(_12); @@ -61,6 +62,7 @@ + Deinit(_16); + copy_nonoverlapping(dst = copy _17, src = copy _19, count = copy _15); + StorageDead(_12); ++ nop; StorageDead(_1); return; } diff --git a/tests/mir-opt/enum_opt.unin.EnumSizeOpt.64bit.diff b/tests/mir-opt/enum_opt.unin.EnumSizeOpt.64bit.diff index 4b1406d0d623b..d20e2e08eaafd 100644 --- a/tests/mir-opt/enum_opt.unin.EnumSizeOpt.64bit.diff +++ b/tests/mir-opt/enum_opt.unin.EnumSizeOpt.64bit.diff @@ -47,6 +47,7 @@ + Deinit(_8); + copy_nonoverlapping(dst = copy _9, src = copy _11, count = copy _7); + StorageDead(_4); ++ nop; StorageDead(_2); - _0 = move _1; + StorageLive(_12); @@ -61,6 +62,7 @@ + Deinit(_16); + copy_nonoverlapping(dst = copy _17, src = copy _19, count = copy _15); + StorageDead(_12); ++ nop; StorageDead(_1); return; } From e3316ae453a86eed28840a85b12df2ea1917aac7 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 18 Feb 2025 12:48:27 +1100 Subject: [PATCH 2/5] Improve `MirPatch` documentation and naming. It's currently lacking comments. This commit adds some, which is useful because there are some methods with non-obvious behaviour. The commit also renames two things: - `patch_map` becomes `term_patch_map`, because it's only about terminators. - `is_patched` becomes `is_term_patched`, for the same reason. (I would guess that originally `MirPatch` only handled terminators, and then over time it expanded to allow other modifications, but these names weren't updated.) --- .../src/elaborate_drops.rs | 6 +-- compiler/rustc_mir_transform/src/patch.rs | 51 ++++++++++++++----- 2 files changed, 41 insertions(+), 16 deletions(-) diff --git a/compiler/rustc_mir_transform/src/elaborate_drops.rs b/compiler/rustc_mir_transform/src/elaborate_drops.rs index ab6aafab446bb..530c72ca549a6 100644 --- a/compiler/rustc_mir_transform/src/elaborate_drops.rs +++ b/compiler/rustc_mir_transform/src/elaborate_drops.rs @@ -417,7 +417,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> { .. } = data.terminator().kind { - assert!(!self.patch.is_patched(bb)); + assert!(!self.patch.is_term_patched(bb)); let loc = Location { block: tgt, statement_index: 0 }; let path = self.move_data().rev_lookup.find(destination.as_ref()); @@ -462,7 +462,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> { // a Goto; see `MirPatch::new`). } _ => { - assert!(!self.patch.is_patched(bb)); + assert!(!self.patch.is_term_patched(bb)); } } } @@ -486,7 +486,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> { .. } = data.terminator().kind { - assert!(!self.patch.is_patched(bb)); + assert!(!self.patch.is_term_patched(bb)); let loc = Location { block: bb, statement_index: data.statements.len() }; let path = self.move_data().rev_lookup.find(destination.as_ref()); diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs index b4f6fa514a487..d3d181f6cb2b1 100644 --- a/compiler/rustc_mir_transform/src/patch.rs +++ b/compiler/rustc_mir_transform/src/patch.rs @@ -4,11 +4,12 @@ use rustc_middle::ty::Ty; use rustc_span::Span; use tracing::debug; -/// This struct represents a patch to MIR, which can add -/// new statements and basic blocks and patch over block -/// terminators. +/// This struct lets you "patch" a MIR body, i.e. modify it. You can queue up +/// various changes, such as the addition of new statements and basic blocks +/// and replacement of terminators, and then apply the queued changes all at +/// once with `apply`. This is useful for MIR transformation passes. pub(crate) struct MirPatch<'tcx> { - patch_map: IndexVec>>, + term_patch_map: IndexVec>>, new_blocks: Vec>, new_statements: Vec<(Location, StatementKind<'tcx>)>, new_locals: Vec>, @@ -24,9 +25,10 @@ pub(crate) struct MirPatch<'tcx> { } impl<'tcx> MirPatch<'tcx> { + /// Creates a new, empty patch. pub(crate) fn new(body: &Body<'tcx>) -> Self { let mut result = MirPatch { - patch_map: IndexVec::from_elem(None, &body.basic_blocks), + term_patch_map: IndexVec::from_elem(None, &body.basic_blocks), new_blocks: vec![], new_statements: vec![], new_locals: vec![], @@ -141,10 +143,12 @@ impl<'tcx> MirPatch<'tcx> { bb } - pub(crate) fn is_patched(&self, bb: BasicBlock) -> bool { - self.patch_map[bb].is_some() + /// Has a replacement of this block's terminator been queued in this patch? + pub(crate) fn is_term_patched(&self, bb: BasicBlock) -> bool { + self.term_patch_map[bb].is_some() } + /// Queues the addition of a new temporary with additional local info. pub(crate) fn new_local_with_info( &mut self, ty: Ty<'tcx>, @@ -159,6 +163,7 @@ impl<'tcx> MirPatch<'tcx> { Local::new(index) } + /// Queues the addition of a new temporary. pub(crate) fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local { let index = self.next_local; self.next_local += 1; @@ -174,29 +179,46 @@ impl<'tcx> MirPatch<'tcx> { self.new_locals[new_local_idx].ty } + /// Queues the addition of a new basic block. pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock { - let block = BasicBlock::new(self.patch_map.len()); + let block = BasicBlock::new(self.term_patch_map.len()); debug!("MirPatch: new_block: {:?}: {:?}", block, data); self.new_blocks.push(data); - self.patch_map.push(None); + self.term_patch_map.push(None); block } + /// Queues the replacement of a block's terminator. pub(crate) fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) { - assert!(self.patch_map[block].is_none()); + assert!(self.term_patch_map[block].is_none()); debug!("MirPatch: patch_terminator({:?}, {:?})", block, new); - self.patch_map[block] = Some(new); + self.term_patch_map[block] = Some(new); } + /// Queues the insertion of a statement at a given location. The statement + /// currently at that location, and all statements that follow, are shifted + /// down. If multiple statements are queued for addition at the same + /// location, the final statement order after calling `apply` will match + /// the queue insertion order. + /// + /// E.g. if we have `s0` at location `loc` and do these calls: + /// + /// p.add_statement(loc, s1); + /// p.add_statement(loc, s2); + /// p.apply(body); + /// + /// then the final order will be `s1, s2, s0`, with `s1` at `loc`. pub(crate) fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) { debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt); self.new_statements.push((loc, stmt)); } + /// Like `add_statement`, but specialized for assignments. pub(crate) fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) { self.add_statement(loc, StatementKind::Assign(Box::new((place, rv)))); } + /// Applies the queued changes. pub(crate) fn apply(self, body: &mut Body<'tcx>) { debug!( "MirPatch: {:?} new temps, starting from index {}: {:?}", @@ -209,14 +231,14 @@ impl<'tcx> MirPatch<'tcx> { self.new_blocks.len(), body.basic_blocks.len() ); - let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() { + let bbs = if self.term_patch_map.is_empty() && self.new_blocks.is_empty() { body.basic_blocks.as_mut_preserves_cfg() } else { body.basic_blocks.as_mut() }; bbs.extend(self.new_blocks); body.local_decls.extend(self.new_locals); - for (src, patch) in self.patch_map.into_iter_enumerated() { + for (src, patch) in self.term_patch_map.into_iter_enumerated() { if let Some(patch) = patch { debug!("MirPatch: patching block {:?}", src); bbs[src].terminator_mut().kind = patch; @@ -224,6 +246,9 @@ impl<'tcx> MirPatch<'tcx> { } let mut new_statements = self.new_statements; + + // This must be a stable sort to provide the ordering described in the + // comment for `add_statement`. new_statements.sort_by_key(|s| s.0); let mut delta = 0; From 627e08c909168aa451c613570a7f624e05f72766 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 18 Feb 2025 13:05:41 +1100 Subject: [PATCH 3/5] Remove `BasicBlockData::expand_statements`. The previous commit removed its single use. `MirPatch` is a more flexible alternative. --- compiler/rustc_middle/src/mir/mod.rs | 49 ---------------------------- 1 file changed, 49 deletions(-) diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs index 795cfcef2d36d..c78cde82e15d1 100644 --- a/compiler/rustc_middle/src/mir/mod.rs +++ b/compiler/rustc_middle/src/mir/mod.rs @@ -1376,55 +1376,6 @@ impl<'tcx> BasicBlockData<'tcx> { } } - pub fn expand_statements(&mut self, mut f: F) - where - F: FnMut(&mut Statement<'tcx>) -> Option, - I: iter::TrustedLen>, - { - // Gather all the iterators we'll need to splice in, and their positions. - let mut splices: Vec<(usize, I)> = vec![]; - let mut extra_stmts = 0; - for (i, s) in self.statements.iter_mut().enumerate() { - if let Some(mut new_stmts) = f(s) { - if let Some(first) = new_stmts.next() { - // We can already store the first new statement. - *s = first; - - // Save the other statements for optimized splicing. - let remaining = new_stmts.size_hint().0; - if remaining > 0 { - splices.push((i + 1 + extra_stmts, new_stmts)); - extra_stmts += remaining; - } - } else { - s.make_nop(); - } - } - } - - // Splice in the new statements, from the end of the block. - // FIXME(eddyb) This could be more efficient with a "gap buffer" - // where a range of elements ("gap") is left uninitialized, with - // splicing adding new elements to the end of that gap and moving - // existing elements from before the gap to the end of the gap. - // For now, this is safe code, emulating a gap but initializing it. - let mut gap = self.statements.len()..self.statements.len() + extra_stmts; - self.statements.resize( - gap.end, - Statement { source_info: SourceInfo::outermost(DUMMY_SP), kind: StatementKind::Nop }, - ); - for (splice_start, new_stmts) in splices.into_iter().rev() { - let splice_end = splice_start + new_stmts.size_hint().0; - while gap.end > splice_end { - gap.start -= 1; - gap.end -= 1; - self.statements.swap(gap.start, gap.end); - } - self.statements.splice(splice_start..splice_end, new_stmts); - gap.end = splice_start; - } - } - pub fn visitable(&self, index: usize) -> &dyn MirVisitable<'tcx> { if index < self.statements.len() { &self.statements[index] } else { &self.terminator } } From 69f5e342bf179c991e325587fb9b16a75851b372 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 18 Feb 2025 13:31:08 +1100 Subject: [PATCH 4/5] Inline and remove `BasicBlockData::retain_statements`. It has a single call site, and the code is clearer this way. --- compiler/rustc_middle/src/mir/mod.rs | 11 ----------- compiler/rustc_mir_transform/src/coroutine.rs | 11 ++++++----- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs index c78cde82e15d1..582941e7e0493 100644 --- a/compiler/rustc_middle/src/mir/mod.rs +++ b/compiler/rustc_middle/src/mir/mod.rs @@ -1365,17 +1365,6 @@ impl<'tcx> BasicBlockData<'tcx> { self.terminator.as_mut().expect("invalid terminator state") } - pub fn retain_statements(&mut self, mut f: F) - where - F: FnMut(&mut Statement<'_>) -> bool, - { - for s in &mut self.statements { - if !f(s) { - s.make_nop(); - } - } - } - pub fn visitable(&self, index: usize) -> &dyn MirVisitable<'tcx> { if index < self.statements.len() { &self.statements[index] } else { &self.terminator } } diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs index afc49c5cc54af..f3f3a65cd805d 100644 --- a/compiler/rustc_mir_transform/src/coroutine.rs +++ b/compiler/rustc_mir_transform/src/coroutine.rs @@ -393,12 +393,13 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) { // Remove StorageLive and StorageDead statements for remapped locals - data.retain_statements(|s| match s.kind { - StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => { - !self.remap.contains(l) + for s in &mut data.statements { + if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = s.kind + && self.remap.contains(l) + { + s.make_nop(); } - _ => true, - }); + } let ret_val = match data.terminator().kind { TerminatorKind::Return => { From 04eeda47abd07a5ba6f3f93e586ecf75d5574545 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 18 Feb 2025 13:40:58 +1100 Subject: [PATCH 5/5] Inline and replace `Statement::replace_nop`. It has a single call site, and doesn't seem worth having as an API function. --- compiler/rustc_middle/src/mir/mod.rs | 2 +- compiler/rustc_middle/src/mir/statement.rs | 9 --------- compiler/rustc_mir_transform/src/single_use_consts.rs | 8 +++++--- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs index 582941e7e0493..74ecff7208244 100644 --- a/compiler/rustc_middle/src/mir/mod.rs +++ b/compiler/rustc_middle/src/mir/mod.rs @@ -4,8 +4,8 @@ use std::borrow::Cow; use std::fmt::{self, Debug, Formatter}; +use std::iter; use std::ops::{Index, IndexMut}; -use std::{iter, mem}; pub use basic_blocks::BasicBlocks; use either::Either; diff --git a/compiler/rustc_middle/src/mir/statement.rs b/compiler/rustc_middle/src/mir/statement.rs index d345c99f902f2..690a907f9a335 100644 --- a/compiler/rustc_middle/src/mir/statement.rs +++ b/compiler/rustc_middle/src/mir/statement.rs @@ -19,15 +19,6 @@ impl Statement<'_> { pub fn make_nop(&mut self) { self.kind = StatementKind::Nop } - - /// Changes a statement to a nop and returns the original statement. - #[must_use = "If you don't need the statement, use `make_nop` instead"] - pub fn replace_nop(&mut self) -> Self { - Statement { - source_info: self.source_info, - kind: mem::replace(&mut self.kind, StatementKind::Nop), - } - } } impl<'tcx> StatementKind<'tcx> { diff --git a/compiler/rustc_mir_transform/src/single_use_consts.rs b/compiler/rustc_mir_transform/src/single_use_consts.rs index c5e951eb8b2c2..02caa92ad3fc8 100644 --- a/compiler/rustc_mir_transform/src/single_use_consts.rs +++ b/compiler/rustc_mir_transform/src/single_use_consts.rs @@ -48,9 +48,11 @@ impl<'tcx> crate::MirPass<'tcx> for SingleUseConsts { // We're only changing an operand, not the terminator kinds or successors let basic_blocks = body.basic_blocks.as_mut_preserves_cfg(); - let init_statement = - basic_blocks[init_loc.block].statements[init_loc.statement_index].replace_nop(); - let StatementKind::Assign(place_and_rvalue) = init_statement.kind else { + let init_statement_kind = std::mem::replace( + &mut basic_blocks[init_loc.block].statements[init_loc.statement_index].kind, + StatementKind::Nop, + ); + let StatementKind::Assign(place_and_rvalue) = init_statement_kind else { bug!("No longer an assign?"); }; let (place, rvalue) = *place_and_rvalue;