Rollup merge of #137203 - nnethercote:improve-MIR-modification, r=compiler-errors

Improve MIR modification

A few commits that simplify code that manipulates MIR bodies.

r? `@tmiasko`

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, SwitchTargetValue};
 use either::Either;
@@ -1365,66 +1365,6 @@ impl<'tcx> BasicBlockData<'tcx> {
         self.terminator.as_mut().expect("invalid terminator state")
     }
 
-    pub fn retain_statements<F>(&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 expand_statements<F, I>(&mut self, mut f: F)
-    where
-        F: FnMut(&mut Statement<'tcx>) -> Option<I>,
-        I: iter::TrustedLen<Item = Statement<'tcx>>,
-    {
-        // 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 }
     }

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> {

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 {
+        for s in &mut data.statements {
-            StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
+            if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = s.kind
-                !self.remap.contains(l)
+                && self.remap.contains(l)
+            {
+                s.make_nop();
+            }
         }
-            _ => true,
-        });
 
         let ret_val = match data.terminator().kind {
             TerminatorKind::Return => {

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

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 mut patch = MirPatch::new(body);
-        let local_decls = &mut body.local_decls;
 
-        for bb in blocks {
+        for (block, data) in body.basic_blocks.as_mut().iter_enumerated_mut() {
-            bb.expand_statements(|st| {
+            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) =
+                let ty = lhs.ty(&body.local_decls, tcx).ty;
-                    self.candidate(tcx, typing_env, ty, &mut alloc_cache)?;
 
-                let source_info = st.source_info;
+                let Some((adt_def, num_variants, alloc_id)) =
-                let span = source_info.span;
+                    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 size_array_local = patch.new_temp(tmp_ty, span);
-                let store_live =
+
-                    Statement { source_info, kind: StatementKind::StorageLive(size_array_local) };
+                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 =
+                let const_assign = StatementKind::Assign(Box::new((place, rval)));
-                    Statement { source_info, kind: StatementKind::Assign(Box::new((place, rval))) };
 
-                let discr_place = Place::from(
+                let discr_place =
-                    local_decls.push(LocalDecl::new(adt_def.repr().discr_type().to_ty(tcx), span)),
+                    Place::from(patch.new_temp(adt_def.repr().discr_type().to_ty(tcx), span));
-                );
+                let store_discr =
-                let store_discr = Statement {
+                    StatementKind::Assign(Box::new((discr_place, Rvalue::Discriminant(*rhs))));
-                    source_info,
-                    kind: StatementKind::Assign(Box::new((
-                        discr_place,
-                        Rvalue::Discriminant(*rhs),
-                    ))),
-                };
 
-                let discr_cast_place =
+                let discr_cast_place = Place::from(patch.new_temp(tcx.types.usize, span));
-                    Place::from(local_decls.push(LocalDecl::new(tcx.types.usize, span)));
+                let cast_discr = StatementKind::Assign(Box::new((
-                let cast_discr = Statement {
-                    source_info,
-                    kind: StatementKind::Assign(Box::new((
                     discr_cast_place,
-                        Rvalue::Cast(
+                    Rvalue::Cast(CastKind::IntToInt, Operand::Copy(discr_place), tcx.types.usize),
-                            CastKind::IntToInt,
+                )));
-                            Operand::Copy(discr_place),
-                            tcx.types.usize,
-                        ),
-                    ))),
-                };
 
-                let size_place =
+                let size_place = Place::from(patch.new_temp(tcx.types.usize, span));
-                    Place::from(local_decls.push(LocalDecl::new(tcx.types.usize, span)));
+                let store_size = StatementKind::Assign(Box::new((
-                let store_size = Statement {
-                    source_info,
-                    kind: StatementKind::Assign(Box::new((
                     size_place,
                     Rvalue::Use(Operand::Copy(Place {
                         local: size_array_local,
-                            projection: tcx
+                        projection: tcx.mk_place_elems(&[PlaceElem::Index(discr_cast_place.local)]),
-                                .mk_place_elems(&[PlaceElem::Index(discr_cast_place.local)]),
                     })),
-                    ))),
+                )));
-                };
 
-                let dst =
+                let dst = Place::from(patch.new_temp(Ty::new_mut_ptr(tcx, ty), span));
-                    Place::from(local_decls.push(LocalDecl::new(Ty::new_mut_ptr(tcx, ty), span)));
+                let dst_ptr =
-                let dst_ptr = Statement {
+                    StatementKind::Assign(Box::new((dst, Rvalue::RawPtr(RawPtrKind::Mut, *lhs))));
-                    source_info,
-                    kind: 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 =
+                let dst_cast_place = Place::from(patch.new_temp(dst_cast_ty, span));
-                    Place::from(local_decls.push(LocalDecl::new(dst_cast_ty, span)));
+                let dst_cast = StatementKind::Assign(Box::new((
-                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 src =
+                let src = Place::from(patch.new_temp(Ty::new_imm_ptr(tcx, ty), span));
-                    Place::from(local_decls.push(LocalDecl::new(Ty::new_imm_ptr(tcx, ty), span)));
+                let src_ptr =
-                let src_ptr = Statement {
+                    StatementKind::Assign(Box::new((src, Rvalue::RawPtr(RawPtrKind::Const, *rhs))));
-                    source_info,
-                    kind: 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 =
+                let src_cast_place = Place::from(patch.new_temp(src_cast_ty, span));
-                    Place::from(local_decls.push(LocalDecl::new(src_cast_ty, span)));
+                let src_cast = StatementKind::Assign(Box::new((
-                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 deinit_old =
+                let deinit_old = StatementKind::Deinit(Box::new(dst));
-                    Statement { source_info, kind: StatementKind::Deinit(Box::new(dst)) };
 
-                let copy_bytes = Statement {
+                let copy_bytes = StatementKind::Intrinsic(Box::new(
-                    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 store_dead =
+                let store_dead = StatementKind::StorageDead(size_array_local);
-                    Statement { source_info, kind: StatementKind::StorageDead(size_array_local) };
 
-                let iter = [
+                let stmts = [
                     store_live,
                     const_assign,
                     store_discr,
@@ -191,16 +151,18 @@ 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 {
         false
     }

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
+/// This struct lets you "patch" a MIR body, i.e. modify it. You can queue up
-/// new statements and basic blocks and patch over block
+/// various changes, such as the addition of new statements and basic blocks
-/// terminators.
+/// 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<BasicBlock, Option<TerminatorKind<'tcx>>>,
+    term_patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
     new_blocks: Vec<BasicBlockData<'tcx>>,
     new_statements: Vec<(Location, StatementKind<'tcx>)>,
     new_locals: Vec<LocalDecl<'tcx>>,
@@ -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 {
+    /// Has a replacement of this block's terminator been queued in this patch?
-        self.patch_map[bb].is_some()
+    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;

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 =
+            let init_statement_kind = std::mem::replace(
-                basic_blocks[init_loc.block].statements[init_loc.statement_index].replace_nop();
+                &mut basic_blocks[init_loc.block].statements[init_loc.statement_index].kind,
-            let StatementKind::Assign(place_and_rvalue) = init_statement.kind else {
+                StatementKind::Nop,
+            );
+            let StatementKind::Assign(place_and_rvalue) = init_statement_kind else {
                 bug!("No longer an assign?");
             };
             let (place, rvalue) = *place_and_rvalue;

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

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

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

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