Rollup merge of #137008 - nnethercote:mv-code-into-rustc_mir_transform, r=oli-obk

Move code into `rustc_mir_transform`

I found two modules in other crates that are better placed in `rustc_mir_transform`, because that's the only crate that uses them.

r? ``@matthewjasper``

compiler/rustc_mir_transform/src/add_moves_for_packed_drops.rs

@@ -1,8 +1,8 @@
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
 use tracing::debug;
 
+use crate::patch::MirPatch;
 use crate::util;
 
 /// This pass moves values being dropped that are within a packed

compiler/rustc_mir_transform/src/add_subtyping_projections.rs

@@ -1,8 +1,9 @@
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::visit::MutVisitor;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 
+use crate::patch::MirPatch;
+
 pub(super) struct Subtyper;
 
 struct SubTypeChecker<'a, 'tcx> {

compiler/rustc_mir_transform/src/coroutine.rs

@@ -1061,9 +1061,8 @@ fn insert_switch<'tcx>(
 }
 
 fn elaborate_coroutine_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    use rustc_middle::mir::patch::MirPatch;
-    use rustc_mir_dataflow::elaborate_drops::{Unwind, elaborate_drop};
-
+    use crate::elaborate_drop::{Unwind, elaborate_drop};
+    use crate::patch::MirPatch;
     use crate::shim::DropShimElaborator;
 
     // Note that `elaborate_drops` only drops the upvars of a coroutine, and

compiler/rustc_mir_transform/src/deref_separator.rs

@@ -1,9 +1,10 @@
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::visit::NonUseContext::VarDebugInfo;
 use rustc_middle::mir::visit::{MutVisitor, PlaceContext};
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 
+use crate::patch::MirPatch;
+
 pub(super) struct Derefer;
 
 struct DerefChecker<'a, 'tcx> {

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -1,11 +1,11 @@
 use std::fmt::Debug;
 
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{Ty, TyCtxt};
 use tracing::trace;
 
 use super::simplify::simplify_cfg;
+use crate::patch::MirPatch;
 
 /// This pass optimizes something like
 /// ```ignore (syntax-highlighting-only)

compiler/rustc_mir_transform/src/elaborate_box_derefs.rs

@@ -4,12 +4,13 @@
 
 use rustc_abi::FieldIdx;
 use rustc_hir::def_id::DefId;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::visit::MutVisitor;
 use rustc_middle::mir::*;
 use rustc_middle::span_bug;
 use rustc_middle::ty::{Ty, TyCtxt};
 
+use crate::patch::MirPatch;
+
 /// Constructs the types used when accessing a Box's pointer
 fn build_ptr_tys<'tcx>(
     tcx: TyCtxt<'tcx>,

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -3,21 +3,20 @@ use std::fmt;
 use rustc_abi::{FieldIdx, VariantIdx};
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_mir_dataflow::elaborate_drops::{
-    DropElaborator, DropFlagMode, DropFlagState, DropStyle, Unwind, elaborate_drop,
-};
 use rustc_mir_dataflow::impls::{MaybeInitializedPlaces, MaybeUninitializedPlaces};
 use rustc_mir_dataflow::move_paths::{LookupResult, MoveData, MovePathIndex};
 use rustc_mir_dataflow::{
-    Analysis, MoveDataTypingEnv, ResultsCursor, on_all_children_bits, on_lookup_result_bits,
+    Analysis, DropFlagState, MoveDataTypingEnv, ResultsCursor, on_all_children_bits,
+    on_lookup_result_bits,
 };
 use rustc_span::Span;
 use tracing::{debug, instrument};
 
 use crate::deref_separator::deref_finder;
+use crate::elaborate_drop::{DropElaborator, DropFlagMode, DropStyle, Unwind, elaborate_drop};
+use crate::patch::MirPatch;
 
 /// During MIR building, Drop terminators are inserted in every place where a drop may occur.
 /// However, in this phase, the presence of these terminators does not guarantee that a destructor

compiler/rustc_mir_transform/src/lib.rs

@@ -49,10 +49,12 @@ mod check_pointers;
 mod cost_checker;
 mod cross_crate_inline;
 mod deduce_param_attrs;
+mod elaborate_drop;
 mod errors;
 mod ffi_unwind_calls;
 mod lint;
 mod lint_tail_expr_drop_order;
+mod patch;
 mod shim;
 mod ssa;
 

compiler/rustc_mir_transform/src/match_branches.rs

@@ -2,7 +2,6 @@ use std::iter;
 
 use rustc_abi::Integer;
 use rustc_index::IndexSlice;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
 use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt};
@@ -10,6 +9,7 @@ use rustc_type_ir::TyKind::*;
 use tracing::instrument;
 
 use super::simplify::simplify_cfg;
+use crate::patch::MirPatch;
 
 pub(super) struct MatchBranchSimplification;
 

compiler/rustc_mir_transform/src/patch.rs

@@ -0,0 +1,252 @@
+use rustc_index::{Idx, IndexVec};
+use rustc_middle::mir::*;
+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.
+pub(crate) struct MirPatch<'tcx> {
+    patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
+    new_blocks: Vec<BasicBlockData<'tcx>>,
+    new_statements: Vec<(Location, StatementKind<'tcx>)>,
+    new_locals: Vec<LocalDecl<'tcx>>,
+    resume_block: Option<BasicBlock>,
+    // Only for unreachable in cleanup path.
+    unreachable_cleanup_block: Option<BasicBlock>,
+    // Only for unreachable not in cleanup path.
+    unreachable_no_cleanup_block: Option<BasicBlock>,
+    // Cached block for UnwindTerminate (with reason)
+    terminate_block: Option<(BasicBlock, UnwindTerminateReason)>,
+    body_span: Span,
+    next_local: usize,
+}
+
+impl<'tcx> MirPatch<'tcx> {
+    pub(crate) fn new(body: &Body<'tcx>) -> Self {
+        let mut result = MirPatch {
+            patch_map: IndexVec::from_elem(None, &body.basic_blocks),
+            new_blocks: vec![],
+            new_statements: vec![],
+            new_locals: vec![],
+            next_local: body.local_decls.len(),
+            resume_block: None,
+            unreachable_cleanup_block: None,
+            unreachable_no_cleanup_block: None,
+            terminate_block: None,
+            body_span: body.span,
+        };
+
+        for (bb, block) in body.basic_blocks.iter_enumerated() {
+            // Check if we already have a resume block
+            if matches!(block.terminator().kind, TerminatorKind::UnwindResume)
+                && block.statements.is_empty()
+            {
+                result.resume_block = Some(bb);
+                continue;
+            }
+
+            // Check if we already have an unreachable block
+            if matches!(block.terminator().kind, TerminatorKind::Unreachable)
+                && block.statements.is_empty()
+            {
+                if block.is_cleanup {
+                    result.unreachable_cleanup_block = Some(bb);
+                } else {
+                    result.unreachable_no_cleanup_block = Some(bb);
+                }
+                continue;
+            }
+
+            // Check if we already have a terminate block
+            if let TerminatorKind::UnwindTerminate(reason) = block.terminator().kind
+                && block.statements.is_empty()
+            {
+                result.terminate_block = Some((bb, reason));
+                continue;
+            }
+        }
+
+        result
+    }
+
+    pub(crate) fn resume_block(&mut self) -> BasicBlock {
+        if let Some(bb) = self.resume_block {
+            return bb;
+        }
+
+        let bb = self.new_block(BasicBlockData {
+            statements: vec![],
+            terminator: Some(Terminator {
+                source_info: SourceInfo::outermost(self.body_span),
+                kind: TerminatorKind::UnwindResume,
+            }),
+            is_cleanup: true,
+        });
+        self.resume_block = Some(bb);
+        bb
+    }
+
+    pub(crate) fn unreachable_cleanup_block(&mut self) -> BasicBlock {
+        if let Some(bb) = self.unreachable_cleanup_block {
+            return bb;
+        }
+
+        let bb = self.new_block(BasicBlockData {
+            statements: vec![],
+            terminator: Some(Terminator {
+                source_info: SourceInfo::outermost(self.body_span),
+                kind: TerminatorKind::Unreachable,
+            }),
+            is_cleanup: true,
+        });
+        self.unreachable_cleanup_block = Some(bb);
+        bb
+    }
+
+    pub(crate) fn unreachable_no_cleanup_block(&mut self) -> BasicBlock {
+        if let Some(bb) = self.unreachable_no_cleanup_block {
+            return bb;
+        }
+
+        let bb = self.new_block(BasicBlockData {
+            statements: vec![],
+            terminator: Some(Terminator {
+                source_info: SourceInfo::outermost(self.body_span),
+                kind: TerminatorKind::Unreachable,
+            }),
+            is_cleanup: false,
+        });
+        self.unreachable_no_cleanup_block = Some(bb);
+        bb
+    }
+
+    pub(crate) fn terminate_block(&mut self, reason: UnwindTerminateReason) -> BasicBlock {
+        if let Some((cached_bb, cached_reason)) = self.terminate_block
+            && reason == cached_reason
+        {
+            return cached_bb;
+        }
+
+        let bb = self.new_block(BasicBlockData {
+            statements: vec![],
+            terminator: Some(Terminator {
+                source_info: SourceInfo::outermost(self.body_span),
+                kind: TerminatorKind::UnwindTerminate(reason),
+            }),
+            is_cleanup: true,
+        });
+        self.terminate_block = Some((bb, reason));
+        bb
+    }
+
+    pub(crate) fn is_patched(&self, bb: BasicBlock) -> bool {
+        self.patch_map[bb].is_some()
+    }
+
+    pub(crate) fn new_local_with_info(
+        &mut self,
+        ty: Ty<'tcx>,
+        span: Span,
+        local_info: LocalInfo<'tcx>,
+    ) -> Local {
+        let index = self.next_local;
+        self.next_local += 1;
+        let mut new_decl = LocalDecl::new(ty, span);
+        **new_decl.local_info.as_mut().assert_crate_local() = local_info;
+        self.new_locals.push(new_decl);
+        Local::new(index)
+    }
+
+    pub(crate) fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
+        let index = self.next_local;
+        self.next_local += 1;
+        self.new_locals.push(LocalDecl::new(ty, span));
+        Local::new(index)
+    }
+
+    pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
+        let block = BasicBlock::new(self.patch_map.len());
+        debug!("MirPatch: new_block: {:?}: {:?}", block, data);
+        self.new_blocks.push(data);
+        self.patch_map.push(None);
+        block
+    }
+
+    pub(crate) fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
+        assert!(self.patch_map[block].is_none());
+        debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
+        self.patch_map[block] = Some(new);
+    }
+
+    pub(crate) fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
+        debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
+        self.new_statements.push((loc, stmt));
+    }
+
+    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))));
+    }
+
+    pub(crate) fn apply(self, body: &mut Body<'tcx>) {
+        debug!(
+            "MirPatch: {:?} new temps, starting from index {}: {:?}",
+            self.new_locals.len(),
+            body.local_decls.len(),
+            self.new_locals
+        );
+        debug!(
+            "MirPatch: {} new blocks, starting from index {}",
+            self.new_blocks.len(),
+            body.basic_blocks.len()
+        );
+        let bbs = if self.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() {
+            if let Some(patch) = patch {
+                debug!("MirPatch: patching block {:?}", src);
+                bbs[src].terminator_mut().kind = patch;
+            }
+        }
+
+        let mut new_statements = self.new_statements;
+        new_statements.sort_by_key(|s| s.0);
+
+        let mut delta = 0;
+        let mut last_bb = START_BLOCK;
+        for (mut loc, stmt) in new_statements {
+            if loc.block != last_bb {
+                delta = 0;
+                last_bb = loc.block;
+            }
+            debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
+            loc.statement_index += delta;
+            let source_info = Self::source_info_for_index(&body[loc.block], loc);
+            body[loc.block]
+                .statements
+                .insert(loc.statement_index, Statement { source_info, kind: stmt });
+            delta += 1;
+        }
+    }
+
+    fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
+        match data.statements.get(loc.statement_index) {
+            Some(stmt) => stmt.source_info,
+            None => data.terminator().source_info,
+        }
+    }
+
+    pub(crate) fn source_info_for_location(&self, body: &Body<'tcx>, loc: Location) -> SourceInfo {
+        let data = match loc.block.index().checked_sub(body.basic_blocks.len()) {
+            Some(new) => &self.new_blocks[new],
+            None => &body[loc.block],
+        };
+        Self::source_info_for_index(data, loc)
+    }
+}

compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs

@@ -1,10 +1,11 @@
 use rustc_index::bit_set::DenseBitSet;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 use rustc_target::spec::PanicStrategy;
 use tracing::debug;
 
+use crate::patch::MirPatch;
+
 /// A pass that removes noop landing pads and replaces jumps to them with
 /// `UnwindAction::Continue`. This is important because otherwise LLVM generates
 /// terrible code for these.

compiler/rustc_mir_transform/src/shim.rs

@@ -6,7 +6,6 @@ use rustc_hir as hir;
 use rustc_hir::def_id::DefId;
 use rustc_hir::lang_items::LangItem;
 use rustc_index::{Idx, IndexVec};
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::adjustment::PointerCoercion;
@@ -14,11 +13,12 @@ use rustc_middle::ty::{
     self, CoroutineArgs, CoroutineArgsExt, EarlyBinder, GenericArgs, Ty, TyCtxt,
 };
 use rustc_middle::{bug, span_bug};
-use rustc_mir_dataflow::elaborate_drops::{self, DropElaborator, DropFlagMode, DropStyle};
 use rustc_span::source_map::Spanned;
 use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument};
 
+use crate::elaborate_drop::{DropElaborator, DropFlagMode, DropStyle, Unwind, elaborate_drop};
+use crate::patch::MirPatch;
 use crate::{
     abort_unwinding_calls, add_call_guards, add_moves_for_packed_drops, deref_separator, inline,
     instsimplify, mentioned_items, pass_manager as pm, remove_noop_landing_pads, simplify,
@@ -283,13 +283,13 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
                 DropShimElaborator { body: &body, patch: MirPatch::new(&body), tcx, typing_env };
             let dropee = tcx.mk_place_deref(dropee_ptr);
             let resume_block = elaborator.patch.resume_block();
-            elaborate_drops::elaborate_drop(
+            elaborate_drop(
                 &mut elaborator,
                 source_info,
                 dropee,
                 (),
                 return_block,
-                elaborate_drops::Unwind::To(resume_block),
+                Unwind::To(resume_block),
                 START_BLOCK,
             );
             elaborator.patch

compiler/rustc_mir_transform/src/sroa.rs

@@ -4,13 +4,14 @@ use rustc_hir::LangItem;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::{DenseBitSet, GrowableBitSet};
 use rustc_middle::bug;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_mir_dataflow::value_analysis::{excluded_locals, iter_fields};
 use tracing::{debug, instrument};
 
+use crate::patch::MirPatch;
+
 pub(super) struct ScalarReplacementOfAggregates;
 
 impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {

compiler/rustc_mir_transform/src/unreachable_enum_branching.rs

@@ -3,7 +3,6 @@
 use rustc_abi::Variants;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::bug;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::{
     BasicBlock, BasicBlockData, BasicBlocks, Body, Local, Operand, Rvalue, StatementKind,
     TerminatorKind,
@@ -12,6 +11,8 @@ use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{Ty, TyCtxt};
 use tracing::trace;
 
+use crate::patch::MirPatch;
+
 pub(super) struct UnreachableEnumBranching;
 
 fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> {

compiler/rustc_mir_transform/src/unreachable_prop.rs

@@ -6,10 +6,11 @@ use rustc_abi::Size;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::bug;
 use rustc_middle::mir::interpret::Scalar;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
 
+use crate::patch::MirPatch;
+
 pub(super) struct UnreachablePropagation;
 
 impl crate::MirPass<'_> for UnreachablePropagation {