diff options
Diffstat (limited to 'compiler/rustc_mir_transform/src')
72 files changed, 23207 insertions, 0 deletions
diff --git a/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs b/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs new file mode 100644 index 00000000000..dfc7a9891f9 --- /dev/null +++ b/compiler/rustc_mir_transform/src/abort_unwinding_calls.rs @@ -0,0 +1,117 @@ +use rustc_ast::InlineAsmOptions; +use rustc_middle::mir::*; +use rustc_middle::ty::layout; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_target::spec::abi::Abi; +use rustc_target::spec::PanicStrategy; + +/// A pass that runs which is targeted at ensuring that codegen guarantees about +/// unwinding are upheld for compilations of panic=abort programs. +/// +/// When compiling with panic=abort codegen backends generally want to assume +/// that all Rust-defined functions do not unwind, and it's UB if they actually +/// do unwind. Foreign functions, however, can be declared as "may unwind" via +/// their ABI (e.g. `extern "C-unwind"`). To uphold the guarantees that +/// Rust-defined functions never unwind a well-behaved Rust program needs to +/// catch unwinding from foreign functions and force them to abort. +/// +/// This pass walks over all functions calls which may possibly unwind, +/// and if any are found sets their cleanup to a block that aborts the process. +/// This forces all unwinds, in panic=abort mode happening in foreign code, to +/// trigger a process abort. +#[derive(PartialEq)] +pub struct AbortUnwindingCalls; + +impl<'tcx> MirPass<'tcx> for AbortUnwindingCalls { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + let kind = tcx.def_kind(def_id); + + // We don't simplify the MIR of constants at this time because that + // namely results in a cyclic query when we call `tcx.type_of` below. + if !kind.is_fn_like() { + return; + } + + // Here we test for this function itself whether its ABI allows + // unwinding or not. + let body_ty = tcx.type_of(def_id).skip_binder(); + let body_abi = match body_ty.kind() { + ty::FnDef(..) => body_ty.fn_sig(tcx).abi(), + ty::Closure(..) => Abi::RustCall, + ty::Coroutine(..) => Abi::Rust, + _ => span_bug!(body.span, "unexpected body ty: {:?}", body_ty), + }; + let body_can_unwind = layout::fn_can_unwind(tcx, Some(def_id), body_abi); + + // Look in this function body for any basic blocks which are terminated + // with a function call, and whose function we're calling may unwind. + // This will filter to functions with `extern "C-unwind"` ABIs, for + // example. + let mut calls_to_terminate = Vec::new(); + let mut cleanups_to_remove = Vec::new(); + for (id, block) in body.basic_blocks.iter_enumerated() { + if block.is_cleanup { + continue; + } + let Some(terminator) = &block.terminator else { continue }; + let span = terminator.source_info.span; + + let call_can_unwind = match &terminator.kind { + TerminatorKind::Call { func, .. } => { + let ty = func.ty(body, tcx); + let sig = ty.fn_sig(tcx); + let fn_def_id = match ty.kind() { + ty::FnPtr(_) => None, + &ty::FnDef(def_id, _) => Some(def_id), + _ => span_bug!(span, "invalid callee of type {:?}", ty), + }; + layout::fn_can_unwind(tcx, fn_def_id, sig.abi()) + } + TerminatorKind::Drop { .. } => { + tcx.sess.opts.unstable_opts.panic_in_drop == PanicStrategy::Unwind + && layout::fn_can_unwind(tcx, None, Abi::Rust) + } + TerminatorKind::Assert { .. } | TerminatorKind::FalseUnwind { .. } => { + layout::fn_can_unwind(tcx, None, Abi::Rust) + } + TerminatorKind::InlineAsm { options, .. } => { + options.contains(InlineAsmOptions::MAY_UNWIND) + } + _ if terminator.unwind().is_some() => { + span_bug!(span, "unexpected terminator that may unwind {:?}", terminator) + } + _ => continue, + }; + + // If this function call can't unwind, then there's no need for it + // to have a landing pad. This means that we can remove any cleanup + // registered for it. + if !call_can_unwind { + cleanups_to_remove.push(id); + continue; + } + + // Otherwise if this function can unwind, then if the outer function + // can also unwind there's nothing to do. If the outer function + // can't unwind, however, we need to change the landing pad for this + // function call to one that aborts. + if !body_can_unwind { + calls_to_terminate.push(id); + } + } + + for id in calls_to_terminate { + let cleanup = body.basic_blocks_mut()[id].terminator_mut().unwind_mut().unwrap(); + *cleanup = UnwindAction::Terminate(UnwindTerminateReason::Abi); + } + + for id in cleanups_to_remove { + let cleanup = body.basic_blocks_mut()[id].terminator_mut().unwind_mut().unwrap(); + *cleanup = UnwindAction::Unreachable; + } + + // We may have invalidated some `cleanup` blocks so clean those up now. + super::simplify::remove_dead_blocks(body); + } +} diff --git a/compiler/rustc_mir_transform/src/add_call_guards.rs b/compiler/rustc_mir_transform/src/add_call_guards.rs new file mode 100644 index 00000000000..a47c8d94bba --- /dev/null +++ b/compiler/rustc_mir_transform/src/add_call_guards.rs @@ -0,0 +1,83 @@ +use rustc_index::{Idx, IndexVec}; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +#[derive(PartialEq)] +pub enum AddCallGuards { + AllCallEdges, + CriticalCallEdges, +} +pub use self::AddCallGuards::*; + +/** + * Breaks outgoing critical edges for call terminators in the MIR. + * + * Critical edges are edges that are neither the only edge leaving a + * block, nor the only edge entering one. + * + * When you want something to happen "along" an edge, you can either + * do at the end of the predecessor block, or at the start of the + * successor block. Critical edges have to be broken in order to prevent + * "edge actions" from affecting other edges. We need this for calls that are + * codegened to LLVM invoke instructions, because invoke is a block terminator + * in LLVM so we can't insert any code to handle the call's result into the + * block that performs the call. + * + * This function will break those edges by inserting new blocks along them. + * + * NOTE: Simplify CFG will happily undo most of the work this pass does. + * + */ + +impl<'tcx> MirPass<'tcx> for AddCallGuards { + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + self.add_call_guards(body); + } +} + +impl AddCallGuards { + pub fn add_call_guards(&self, body: &mut Body<'_>) { + let mut pred_count: IndexVec<_, _> = + body.basic_blocks.predecessors().iter().map(|ps| ps.len()).collect(); + pred_count[START_BLOCK] += 1; + + // We need a place to store the new blocks generated + let mut new_blocks = Vec::new(); + + let cur_len = body.basic_blocks.len(); + + for block in body.basic_blocks_mut() { + match block.terminator { + Some(Terminator { + kind: TerminatorKind::Call { target: Some(ref mut destination), unwind, .. }, + source_info, + }) if pred_count[*destination] > 1 + && (matches!( + unwind, + UnwindAction::Cleanup(_) | UnwindAction::Terminate(_) + ) || self == &AllCallEdges) => + { + // It's a critical edge, break it + let call_guard = BasicBlockData { + statements: vec![], + is_cleanup: block.is_cleanup, + terminator: Some(Terminator { + source_info, + kind: TerminatorKind::Goto { target: *destination }, + }), + }; + + // Get the index it will be when inserted into the MIR + let idx = cur_len + new_blocks.len(); + new_blocks.push(call_guard); + *destination = BasicBlock::new(idx); + } + _ => {} + } + } + + debug!("Broke {} N edges", new_blocks.len()); + + body.basic_blocks_mut().extend(new_blocks); + } +} diff --git a/compiler/rustc_mir_transform/src/add_moves_for_packed_drops.rs b/compiler/rustc_mir_transform/src/add_moves_for_packed_drops.rs new file mode 100644 index 00000000000..de6d20ae3e8 --- /dev/null +++ b/compiler/rustc_mir_transform/src/add_moves_for_packed_drops.rs @@ -0,0 +1,107 @@ +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +use crate::util; +use rustc_middle::mir::patch::MirPatch; + +/// This pass moves values being dropped that are within a packed +/// struct to a separate local before dropping them, to ensure that +/// they are dropped from an aligned address. +/// +/// For example, if we have something like +/// ```ignore (illustrative) +/// #[repr(packed)] +/// struct Foo { +/// dealign: u8, +/// data: Vec<u8> +/// } +/// +/// let foo = ...; +/// ``` +/// +/// We want to call `drop_in_place::<Vec<u8>>` on `data` from an aligned +/// address. This means we can't simply drop `foo.data` directly, because +/// its address is not aligned. +/// +/// Instead, we move `foo.data` to a local and drop that: +/// ```ignore (illustrative) +/// storage.live(drop_temp) +/// drop_temp = foo.data; +/// drop(drop_temp) -> next +/// next: +/// storage.dead(drop_temp) +/// ``` +/// +/// The storage instructions are required to avoid stack space +/// blowup. +pub struct AddMovesForPackedDrops; + +impl<'tcx> MirPass<'tcx> for AddMovesForPackedDrops { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!("add_moves_for_packed_drops({:?} @ {:?})", body.source, body.span); + add_moves_for_packed_drops(tcx, body); + } +} + +pub fn add_moves_for_packed_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let patch = add_moves_for_packed_drops_patch(tcx, body); + patch.apply(body); +} + +fn add_moves_for_packed_drops_patch<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> MirPatch<'tcx> { + let def_id = body.source.def_id(); + let mut patch = MirPatch::new(body); + let param_env = tcx.param_env(def_id); + + for (bb, data) in body.basic_blocks.iter_enumerated() { + let loc = Location { block: bb, statement_index: data.statements.len() }; + let terminator = data.terminator(); + + match terminator.kind { + TerminatorKind::Drop { place, .. } + if util::is_disaligned(tcx, body, param_env, place) => + { + add_move_for_packed_drop(tcx, body, &mut patch, terminator, loc, data.is_cleanup); + } + _ => {} + } + } + + patch +} + +fn add_move_for_packed_drop<'tcx>( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + patch: &mut MirPatch<'tcx>, + terminator: &Terminator<'tcx>, + loc: Location, + is_cleanup: bool, +) { + debug!("add_move_for_packed_drop({:?} @ {:?})", terminator, loc); + let TerminatorKind::Drop { ref place, target, unwind, replace } = terminator.kind else { + unreachable!(); + }; + + let source_info = terminator.source_info; + let ty = place.ty(body, tcx).ty; + let temp = patch.new_temp(ty, terminator.source_info.span); + + let storage_dead_block = patch.new_block(BasicBlockData { + statements: vec![Statement { source_info, kind: StatementKind::StorageDead(temp) }], + terminator: Some(Terminator { source_info, kind: TerminatorKind::Goto { target } }), + is_cleanup, + }); + + patch.add_statement(loc, StatementKind::StorageLive(temp)); + patch.add_assign(loc, Place::from(temp), Rvalue::Use(Operand::Move(*place))); + patch.patch_terminator( + loc.block, + TerminatorKind::Drop { + place: Place::from(temp), + target: storage_dead_block, + unwind, + replace, + }, + ); +} diff --git a/compiler/rustc_mir_transform/src/add_retag.rs b/compiler/rustc_mir_transform/src/add_retag.rs new file mode 100644 index 00000000000..94077c63057 --- /dev/null +++ b/compiler/rustc_mir_transform/src/add_retag.rs @@ -0,0 +1,156 @@ +//! This pass adds validation calls (AcquireValid, ReleaseValid) where appropriate. +//! It has to be run really early, before transformations like inlining, because +//! introducing these calls *adds* UB -- so, conceptually, this pass is actually part +//! of MIR building, and only after this pass we think of the program has having the +//! normal MIR semantics. + +use rustc_middle::mir::*; +use rustc_middle::ty::{self, Ty, TyCtxt}; + +pub struct AddRetag; + +/// Determine whether this type may contain a reference (or box), and thus needs retagging. +/// We will only recurse `depth` times into Tuples/ADTs to bound the cost of this. +fn may_contain_reference<'tcx>(ty: Ty<'tcx>, depth: u32, tcx: TyCtxt<'tcx>) -> bool { + match ty.kind() { + // Primitive types that are not references + ty::Bool + | ty::Char + | ty::Float(_) + | ty::Int(_) + | ty::Uint(_) + | ty::RawPtr(..) + | ty::FnPtr(..) + | ty::Str + | ty::FnDef(..) + | ty::Never => false, + // References + ty::Ref(..) => true, + ty::Adt(..) if ty.is_box() => true, + ty::Adt(adt, _) if Some(adt.did()) == tcx.lang_items().ptr_unique() => true, + // Compound types: recurse + ty::Array(ty, _) | ty::Slice(ty) => { + // This does not branch so we keep the depth the same. + may_contain_reference(*ty, depth, tcx) + } + ty::Tuple(tys) => { + depth == 0 || tys.iter().any(|ty| may_contain_reference(ty, depth - 1, tcx)) + } + ty::Adt(adt, subst) => { + depth == 0 + || adt.variants().iter().any(|v| { + v.fields.iter().any(|f| may_contain_reference(f.ty(tcx, subst), depth - 1, tcx)) + }) + } + // Conservative fallback + _ => true, + } +} + +impl<'tcx> MirPass<'tcx> for AddRetag { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.opts.unstable_opts.mir_emit_retag + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // We need an `AllCallEdges` pass before we can do any work. + super::add_call_guards::AllCallEdges.run_pass(tcx, body); + + let basic_blocks = body.basic_blocks.as_mut(); + let local_decls = &body.local_decls; + let needs_retag = |place: &Place<'tcx>| { + !place.is_indirect_first_projection() // we're not really interested in stores to "outside" locations, they are hard to keep track of anyway + && may_contain_reference(place.ty(&*local_decls, tcx).ty, /*depth*/ 3, tcx) + && !local_decls[place.local].is_deref_temp() + }; + + // PART 1 + // Retag arguments at the beginning of the start block. + { + // Gather all arguments, skip return value. + let places = local_decls.iter_enumerated().skip(1).take(body.arg_count).filter_map( + |(local, decl)| { + let place = Place::from(local); + needs_retag(&place).then_some((place, decl.source_info)) + }, + ); + + // Emit their retags. + basic_blocks[START_BLOCK].statements.splice( + 0..0, + places.map(|(place, source_info)| Statement { + source_info, + kind: StatementKind::Retag(RetagKind::FnEntry, Box::new(place)), + }), + ); + } + + // PART 2 + // Retag return values of functions. + // We collect the return destinations because we cannot mutate while iterating. + let returns = basic_blocks + .iter_mut() + .filter_map(|block_data| { + match block_data.terminator().kind { + TerminatorKind::Call { target: Some(target), destination, .. } + if needs_retag(&destination) => + { + // Remember the return destination for later + Some((block_data.terminator().source_info, destination, target)) + } + + // `Drop` is also a call, but it doesn't return anything so we are good. + TerminatorKind::Drop { .. } => None, + // Not a block ending in a Call -> ignore. + _ => None, + } + }) + .collect::<Vec<_>>(); + // Now we go over the returns we collected to retag the return values. + for (source_info, dest_place, dest_block) in returns { + basic_blocks[dest_block].statements.insert( + 0, + Statement { + source_info, + kind: StatementKind::Retag(RetagKind::Default, Box::new(dest_place)), + }, + ); + } + + // PART 3 + // Add retag after assignments where data "enters" this function: the RHS is behind a deref and the LHS is not. + for block_data in basic_blocks { + // We want to insert statements as we iterate. To this end, we + // iterate backwards using indices. + for i in (0..block_data.statements.len()).rev() { + let (retag_kind, place) = match block_data.statements[i].kind { + // Retag after assignments of reference type. + StatementKind::Assign(box (ref place, ref rvalue)) if needs_retag(place) => { + let add_retag = match rvalue { + // Ptr-creating operations already do their own internal retagging, no + // need to also add a retag statement. + Rvalue::Ref(..) | Rvalue::AddressOf(..) => false, + _ => true, + }; + if add_retag { + (RetagKind::Default, *place) + } else { + continue; + } + } + // Do nothing for the rest + _ => continue, + }; + // Insert a retag after the statement. + let source_info = block_data.statements[i].source_info; + block_data.statements.insert( + i + 1, + Statement { + source_info, + kind: StatementKind::Retag(retag_kind, Box::new(place)), + }, + ); + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/add_subtyping_projections.rs b/compiler/rustc_mir_transform/src/add_subtyping_projections.rs new file mode 100644 index 00000000000..04204c68f7b --- /dev/null +++ b/compiler/rustc_mir_transform/src/add_subtyping_projections.rs @@ -0,0 +1,69 @@ +use rustc_index::IndexVec; +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::visit::MutVisitor; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +pub struct Subtyper; + +pub struct SubTypeChecker<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + patcher: MirPatch<'tcx>, + local_decls: &'a IndexVec<Local, LocalDecl<'tcx>>, +} + +impl<'a, 'tcx> MutVisitor<'tcx> for SubTypeChecker<'a, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_assign( + &mut self, + place: &mut Place<'tcx>, + rvalue: &mut Rvalue<'tcx>, + location: Location, + ) { + // We don't need to do anything for deref temps as they are + // not part of the source code, but used for desugaring purposes. + if self.local_decls[place.local].is_deref_temp() { + return; + } + let mut place_ty = place.ty(self.local_decls, self.tcx).ty; + let mut rval_ty = rvalue.ty(self.local_decls, self.tcx); + // Not erasing this causes `Free Regions` errors in validator, + // when rval is `ReStatic`. + rval_ty = self.tcx.erase_regions_ty(rval_ty); + place_ty = self.tcx.erase_regions(place_ty); + if place_ty != rval_ty { + let temp = self + .patcher + .new_temp(rval_ty, self.local_decls[place.as_ref().local].source_info.span); + let new_place = Place::from(temp); + self.patcher.add_assign(location, new_place, rvalue.clone()); + let subtyped = new_place.project_deeper(&[ProjectionElem::Subtype(place_ty)], self.tcx); + *rvalue = Rvalue::Use(Operand::Move(subtyped)); + } + } +} + +// Aim here is to do this kind of transformation: +// +// let place: place_ty = rval; +// // gets transformed to +// let temp: rval_ty = rval; +// let place: place_ty = temp as place_ty; +pub fn subtype_finder<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let patch = MirPatch::new(body); + let mut checker = SubTypeChecker { tcx, patcher: patch, local_decls: &body.local_decls }; + + for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() { + checker.visit_basic_block_data(bb, data); + } + checker.patcher.apply(body); +} + +impl<'tcx> MirPass<'tcx> for Subtyper { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + subtype_finder(tcx, body); + } +} diff --git a/compiler/rustc_mir_transform/src/check_alignment.rs b/compiler/rustc_mir_transform/src/check_alignment.rs new file mode 100644 index 00000000000..9eec724ef21 --- /dev/null +++ b/compiler/rustc_mir_transform/src/check_alignment.rs @@ -0,0 +1,245 @@ +use rustc_hir::lang_items::LangItem; +use rustc_index::IndexVec; +use rustc_middle::mir::*; +use rustc_middle::mir::{ + interpret::Scalar, + visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor}, +}; +use rustc_middle::ty::{self, ParamEnv, Ty, TyCtxt, TypeAndMut}; +use rustc_session::Session; + +pub struct CheckAlignment; + +impl<'tcx> MirPass<'tcx> for CheckAlignment { + fn is_enabled(&self, sess: &Session) -> bool { + // FIXME(#112480) MSVC and rustc disagree on minimum stack alignment on x86 Windows + if sess.target.llvm_target == "i686-pc-windows-msvc" { + return false; + } + sess.opts.debug_assertions + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // This pass emits new panics. If for whatever reason we do not have a panic + // implementation, running this pass may cause otherwise-valid code to not compile. + if tcx.lang_items().get(LangItem::PanicImpl).is_none() { + return; + } + + let basic_blocks = body.basic_blocks.as_mut(); + let local_decls = &mut body.local_decls; + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + + // This pass inserts new blocks. Each insertion changes the Location for all + // statements/blocks after. Iterating or visiting the MIR in order would require updating + // our current location after every insertion. By iterating backwards, we dodge this issue: + // The only Locations that an insertion changes have already been handled. + for block in (0..basic_blocks.len()).rev() { + let block = block.into(); + for statement_index in (0..basic_blocks[block].statements.len()).rev() { + let location = Location { block, statement_index }; + let statement = &basic_blocks[block].statements[statement_index]; + let source_info = statement.source_info; + + let mut finder = + PointerFinder { tcx, local_decls, param_env, pointers: Vec::new() }; + finder.visit_statement(statement, location); + + for (local, ty) in finder.pointers { + debug!("Inserting alignment check for {:?}", ty); + let new_block = split_block(basic_blocks, location); + insert_alignment_check( + tcx, + local_decls, + &mut basic_blocks[block], + local, + ty, + source_info, + new_block, + ); + } + } + } + } +} + +struct PointerFinder<'tcx, 'a> { + tcx: TyCtxt<'tcx>, + local_decls: &'a mut LocalDecls<'tcx>, + param_env: ParamEnv<'tcx>, + pointers: Vec<(Place<'tcx>, Ty<'tcx>)>, +} + +impl<'tcx, 'a> Visitor<'tcx> for PointerFinder<'tcx, 'a> { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) { + // We want to only check reads and writes to Places, so we specifically exclude + // Borrows and AddressOf. + match context { + PlaceContext::MutatingUse( + MutatingUseContext::Store + | MutatingUseContext::AsmOutput + | MutatingUseContext::Call + | MutatingUseContext::Yield + | MutatingUseContext::Drop, + ) => {} + PlaceContext::NonMutatingUse( + NonMutatingUseContext::Copy | NonMutatingUseContext::Move, + ) => {} + _ => { + return; + } + } + + if !place.is_indirect() { + return; + } + + // Since Deref projections must come first and only once, the pointer for an indirect place + // is the Local that the Place is based on. + let pointer = Place::from(place.local); + let pointer_ty = self.local_decls[place.local].ty; + + // We only want to check places based on unsafe pointers + if !pointer_ty.is_unsafe_ptr() { + trace!("Indirect, but not based on an unsafe ptr, not checking {:?}", place); + return; + } + + let pointee_ty = + pointer_ty.builtin_deref(true).expect("no builtin_deref for an unsafe pointer").ty; + // Ideally we'd support this in the future, but for now we are limited to sized types. + if !pointee_ty.is_sized(self.tcx, self.param_env) { + debug!("Unsafe pointer, but pointee is not known to be sized: {:?}", pointer_ty); + return; + } + + // Try to detect types we are sure have an alignment of 1 and skip the check + // We don't need to look for str and slices, we already rejected unsized types above + let element_ty = match pointee_ty.kind() { + ty::Array(ty, _) => *ty, + _ => pointee_ty, + }; + if [self.tcx.types.bool, self.tcx.types.i8, self.tcx.types.u8].contains(&element_ty) { + debug!("Trivially aligned place type: {:?}", pointee_ty); + return; + } + + // Ensure that this place is based on an aligned pointer. + self.pointers.push((pointer, pointee_ty)); + + self.super_place(place, context, location); + } +} + +fn split_block( + basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'_>>, + location: Location, +) -> BasicBlock { + let block_data = &mut basic_blocks[location.block]; + + // Drain every statement after this one and move the current terminator to a new basic block + let new_block = BasicBlockData { + statements: block_data.statements.split_off(location.statement_index), + terminator: block_data.terminator.take(), + is_cleanup: block_data.is_cleanup, + }; + + basic_blocks.push(new_block) +} + +fn insert_alignment_check<'tcx>( + tcx: TyCtxt<'tcx>, + local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, + block_data: &mut BasicBlockData<'tcx>, + pointer: Place<'tcx>, + pointee_ty: Ty<'tcx>, + source_info: SourceInfo, + new_block: BasicBlock, +) { + // Cast the pointer to a *const () + let const_raw_ptr = Ty::new_ptr(tcx, TypeAndMut { ty: tcx.types.unit, mutbl: Mutability::Not }); + let rvalue = Rvalue::Cast(CastKind::PtrToPtr, Operand::Copy(pointer), const_raw_ptr); + let thin_ptr = local_decls.push(LocalDecl::with_source_info(const_raw_ptr, source_info)).into(); + block_data + .statements + .push(Statement { source_info, kind: StatementKind::Assign(Box::new((thin_ptr, rvalue))) }); + + // Transmute the pointer to a usize (equivalent to `ptr.addr()`) + let rvalue = Rvalue::Cast(CastKind::Transmute, Operand::Copy(thin_ptr), tcx.types.usize); + let addr = local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); + block_data + .statements + .push(Statement { source_info, kind: StatementKind::Assign(Box::new((addr, rvalue))) }); + + // Get the alignment of the pointee + let alignment = + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); + let rvalue = Rvalue::NullaryOp(NullOp::AlignOf, pointee_ty); + block_data.statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new((alignment, rvalue))), + }); + + // Subtract 1 from the alignment to get the alignment mask + let alignment_mask = + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); + let one = Operand::Constant(Box::new(ConstOperand { + span: source_info.span, + user_ty: None, + const_: Const::Val(ConstValue::Scalar(Scalar::from_target_usize(1, &tcx)), tcx.types.usize), + })); + block_data.statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + alignment_mask, + Rvalue::BinaryOp(BinOp::Sub, Box::new((Operand::Copy(alignment), one))), + ))), + }); + + // BitAnd the alignment mask with the pointer + let alignment_bits = + local_decls.push(LocalDecl::with_source_info(tcx.types.usize, source_info)).into(); + block_data.statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + alignment_bits, + Rvalue::BinaryOp( + BinOp::BitAnd, + Box::new((Operand::Copy(addr), Operand::Copy(alignment_mask))), + ), + ))), + }); + + // Check if the alignment bits are all zero + let is_ok = local_decls.push(LocalDecl::with_source_info(tcx.types.bool, source_info)).into(); + let zero = Operand::Constant(Box::new(ConstOperand { + span: source_info.span, + user_ty: None, + const_: Const::Val(ConstValue::Scalar(Scalar::from_target_usize(0, &tcx)), tcx.types.usize), + })); + block_data.statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + is_ok, + Rvalue::BinaryOp(BinOp::Eq, Box::new((Operand::Copy(alignment_bits), zero.clone()))), + ))), + }); + + // Set this block's terminator to our assert, continuing to new_block if we pass + block_data.terminator = Some(Terminator { + source_info, + kind: TerminatorKind::Assert { + cond: Operand::Copy(is_ok), + expected: true, + target: new_block, + msg: Box::new(AssertKind::MisalignedPointerDereference { + required: Operand::Copy(alignment), + found: Operand::Copy(addr), + }), + // This calls panic_misaligned_pointer_dereference, which is #[rustc_nounwind]. + // We never want to insert an unwind into unsafe code, because unwinding could + // make a failing UB check turn into much worse UB when we start unwinding. + unwind: UnwindAction::Unreachable, + }, + }); +} diff --git a/compiler/rustc_mir_transform/src/check_const_item_mutation.rs b/compiler/rustc_mir_transform/src/check_const_item_mutation.rs new file mode 100644 index 00000000000..3195cd3622d --- /dev/null +++ b/compiler/rustc_mir_transform/src/check_const_item_mutation.rs @@ -0,0 +1,159 @@ +use rustc_hir::HirId; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_session::lint::builtin::CONST_ITEM_MUTATION; +use rustc_span::def_id::DefId; +use rustc_span::Span; + +use crate::{errors, MirLint}; + +pub struct CheckConstItemMutation; + +impl<'tcx> MirLint<'tcx> for CheckConstItemMutation { + fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) { + let mut checker = ConstMutationChecker { body, tcx, target_local: None }; + checker.visit_body(body); + } +} + +struct ConstMutationChecker<'a, 'tcx> { + body: &'a Body<'tcx>, + tcx: TyCtxt<'tcx>, + target_local: Option<Local>, +} + +impl<'tcx> ConstMutationChecker<'_, 'tcx> { + fn is_const_item(&self, local: Local) -> Option<DefId> { + if let LocalInfo::ConstRef { def_id } = *self.body.local_decls[local].local_info() { + Some(def_id) + } else { + None + } + } + + fn is_const_item_without_destructor(&self, local: Local) -> Option<DefId> { + let def_id = self.is_const_item(local)?; + + // We avoid linting mutation of a const item if the const's type has a + // Drop impl. The Drop logic observes the mutation which was performed. + // + // pub struct Log { msg: &'static str } + // pub const LOG: Log = Log { msg: "" }; + // impl Drop for Log { + // fn drop(&mut self) { println!("{}", self.msg); } + // } + // + // LOG.msg = "wow"; // prints "wow" + // + // FIXME(https://github.com/rust-lang/rust/issues/77425): + // Drop this exception once there is a stable attribute to suppress the + // const item mutation lint for a single specific const only. Something + // equivalent to: + // + // #[const_mutation_allowed] + // pub const LOG: Log = Log { msg: "" }; + match self.tcx.calculate_dtor(def_id, |_, _| Ok(())) { + Some(_) => None, + None => Some(def_id), + } + } + + /// If we should lint on this usage, return the [`HirId`], source [`Span`] + /// and [`Span`] of the const item to use in the lint. + fn should_lint_const_item_usage( + &self, + place: &Place<'tcx>, + const_item: DefId, + location: Location, + ) -> Option<(HirId, Span, Span)> { + // Don't lint on borrowing/assigning when a dereference is involved. + // If we 'leave' the temporary via a dereference, we must + // be modifying something else + // + // `unsafe { *FOO = 0; *BAR.field = 1; }` + // `unsafe { &mut *FOO }` + // `unsafe { (*ARRAY)[0] = val; }` + if !place.projection.iter().any(|p| matches!(p, PlaceElem::Deref)) { + let source_info = self.body.source_info(location); + let lint_root = self.body.source_scopes[source_info.scope] + .local_data + .as_ref() + .assert_crate_local() + .lint_root; + + Some((lint_root, source_info.span, self.tcx.def_span(const_item))) + } else { + None + } + } +} + +impl<'tcx> Visitor<'tcx> for ConstMutationChecker<'_, 'tcx> { + fn visit_statement(&mut self, stmt: &Statement<'tcx>, loc: Location) { + if let StatementKind::Assign(box (lhs, _)) = &stmt.kind { + // Check for assignment to fields of a constant + // Assigning directly to a constant (e.g. `FOO = true;`) is a hard error, + // so emitting a lint would be redundant. + if !lhs.projection.is_empty() { + if let Some(def_id) = self.is_const_item_without_destructor(lhs.local) + && let Some((lint_root, span, item)) = + self.should_lint_const_item_usage(lhs, def_id, loc) + { + self.tcx.emit_spanned_lint( + CONST_ITEM_MUTATION, + lint_root, + span, + errors::ConstMutate::Modify { konst: item }, + ); + } + } + // We are looking for MIR of the form: + // + // ``` + // _1 = const FOO; + // _2 = &mut _1; + // method_call(_2, ..) + // ``` + // + // Record our current LHS, so that we can detect this + // pattern in `visit_rvalue` + self.target_local = lhs.as_local(); + } + self.super_statement(stmt, loc); + self.target_local = None; + } + fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, loc: Location) { + if let Rvalue::Ref(_, BorrowKind::Mut { .. }, place) = rvalue { + let local = place.local; + if let Some(def_id) = self.is_const_item(local) { + // If this Rvalue is being used as the right-hand side of a + // `StatementKind::Assign`, see if it ends up getting used as + // the `self` parameter of a method call (as the terminator of our current + // BasicBlock). If so, we emit a more specific lint. + let method_did = self.target_local.and_then(|target_local| { + rustc_middle::util::find_self_call(self.tcx, self.body, target_local, loc.block) + }); + let lint_loc = + if method_did.is_some() { self.body.terminator_loc(loc.block) } else { loc }; + + let method_call = if let Some((method_did, _)) = method_did { + Some(self.tcx.def_span(method_did)) + } else { + None + }; + if let Some((lint_root, span, item)) = + self.should_lint_const_item_usage(place, def_id, lint_loc) + { + self.tcx.emit_spanned_lint( + CONST_ITEM_MUTATION, + lint_root, + span, + errors::ConstMutate::MutBorrow { method_call, konst: item }, + ); + } + } + } + self.super_rvalue(rvalue, loc); + } +} diff --git a/compiler/rustc_mir_transform/src/check_packed_ref.rs b/compiler/rustc_mir_transform/src/check_packed_ref.rs new file mode 100644 index 00000000000..77bcba50a3c --- /dev/null +++ b/compiler/rustc_mir_transform/src/check_packed_ref.rs @@ -0,0 +1,61 @@ +use rustc_middle::mir::visit::{PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; + +use crate::MirLint; +use crate::{errors, util}; + +pub struct CheckPackedRef; + +impl<'tcx> MirLint<'tcx> for CheckPackedRef { + fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) { + let param_env = tcx.param_env(body.source.def_id()); + let source_info = SourceInfo::outermost(body.span); + let mut checker = PackedRefChecker { body, tcx, param_env, source_info }; + checker.visit_body(body); + } +} + +struct PackedRefChecker<'a, 'tcx> { + body: &'a Body<'tcx>, + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + source_info: SourceInfo, +} + +impl<'tcx> Visitor<'tcx> for PackedRefChecker<'_, 'tcx> { + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + // Make sure we know where in the MIR we are. + self.source_info = terminator.source_info; + self.super_terminator(terminator, location); + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + // Make sure we know where in the MIR we are. + self.source_info = statement.source_info; + self.super_statement(statement, location); + } + + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + if context.is_borrow() { + if util::is_disaligned(self.tcx, self.body, self.param_env, *place) { + let def_id = self.body.source.instance.def_id(); + if let Some(impl_def_id) = self.tcx.impl_of_method(def_id) + && self.tcx.is_builtin_derived(impl_def_id) + { + // If we ever reach here it means that the generated derive + // code is somehow doing an unaligned reference, which it + // shouldn't do. + span_bug!( + self.source_info.span, + "builtin derive created an unaligned reference" + ); + } else { + self.tcx + .sess + .emit_err(errors::UnalignedPackedRef { span: self.source_info.span }); + } + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/check_unsafety.rs b/compiler/rustc_mir_transform/src/check_unsafety.rs new file mode 100644 index 00000000000..f246de55ca8 --- /dev/null +++ b/compiler/rustc_mir_transform/src/check_unsafety.rs @@ -0,0 +1,611 @@ +use rustc_data_structures::unord::{ExtendUnord, UnordItems, UnordSet}; +use rustc_hir as hir; +use rustc_hir::def::DefKind; +use rustc_hir::def_id::{DefId, LocalDefId}; +use rustc_hir::hir_id::HirId; +use rustc_hir::intravisit; +use rustc_hir::{BlockCheckMode, ExprKind, Node}; +use rustc_middle::mir::visit::{MutatingUseContext, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::query::Providers; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_session::lint::builtin::{UNSAFE_OP_IN_UNSAFE_FN, UNUSED_UNSAFE}; +use rustc_session::lint::Level; + +use std::ops::Bound; + +use crate::errors; + +pub struct UnsafetyChecker<'a, 'tcx> { + body: &'a Body<'tcx>, + body_did: LocalDefId, + violations: Vec<UnsafetyViolation>, + source_info: SourceInfo, + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + + /// Used `unsafe` blocks in this function. This is used for the "unused_unsafe" lint. + used_unsafe_blocks: UnordSet<HirId>, +} + +impl<'a, 'tcx> UnsafetyChecker<'a, 'tcx> { + fn new( + body: &'a Body<'tcx>, + body_did: LocalDefId, + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + ) -> Self { + Self { + body, + body_did, + violations: vec![], + source_info: SourceInfo::outermost(body.span), + tcx, + param_env, + used_unsafe_blocks: Default::default(), + } + } +} + +impl<'tcx> Visitor<'tcx> for UnsafetyChecker<'_, 'tcx> { + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + self.source_info = terminator.source_info; + match terminator.kind { + TerminatorKind::Goto { .. } + | TerminatorKind::SwitchInt { .. } + | TerminatorKind::Drop { .. } + | TerminatorKind::Yield { .. } + | TerminatorKind::Assert { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } => { + // safe (at least as emitted during MIR construction) + } + + TerminatorKind::Call { ref func, .. } => { + let func_ty = func.ty(self.body, self.tcx); + let func_id = + if let ty::FnDef(func_id, _) = func_ty.kind() { Some(func_id) } else { None }; + let sig = func_ty.fn_sig(self.tcx); + if let hir::Unsafety::Unsafe = sig.unsafety() { + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::CallToUnsafeFunction, + ) + } + + if let Some(func_id) = func_id { + self.check_target_features(*func_id); + } + } + + TerminatorKind::InlineAsm { .. } => self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::UseOfInlineAssembly, + ), + } + self.super_terminator(terminator, location); + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + self.source_info = statement.source_info; + match statement.kind { + StatementKind::Assign(..) + | StatementKind::FakeRead(..) + | StatementKind::SetDiscriminant { .. } + | StatementKind::Deinit(..) + | StatementKind::StorageLive(..) + | StatementKind::StorageDead(..) + | StatementKind::Retag { .. } + | StatementKind::PlaceMention(..) + | StatementKind::Coverage(..) + | StatementKind::Intrinsic(..) + | StatementKind::ConstEvalCounter + | StatementKind::Nop => { + // safe (at least as emitted during MIR construction) + } + // `AscribeUserType` just exists to help MIR borrowck. + // It has no semantics, and everything is already reported by `PlaceMention`. + StatementKind::AscribeUserType(..) => return, + } + self.super_statement(statement, location); + } + + fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) { + match rvalue { + Rvalue::Aggregate(box ref aggregate, _) => match aggregate { + &AggregateKind::Array(..) | &AggregateKind::Tuple => {} + &AggregateKind::Adt(adt_did, ..) => { + match self.tcx.layout_scalar_valid_range(adt_did) { + (Bound::Unbounded, Bound::Unbounded) => {} + _ => self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::InitializingTypeWith, + ), + } + } + &AggregateKind::Closure(def_id, _) | &AggregateKind::Coroutine(def_id, _, _) => { + let def_id = def_id.expect_local(); + let UnsafetyCheckResult { violations, used_unsafe_blocks, .. } = + self.tcx.unsafety_check_result(def_id); + self.register_violations(violations, used_unsafe_blocks.items().copied()); + } + }, + _ => {} + } + self.super_rvalue(rvalue, location); + } + + fn visit_operand(&mut self, op: &Operand<'tcx>, location: Location) { + if let Operand::Constant(constant) = op { + let maybe_uneval = match constant.const_ { + Const::Val(..) | Const::Ty(_) => None, + Const::Unevaluated(uv, _) => Some(uv), + }; + + if let Some(uv) = maybe_uneval { + if uv.promoted.is_none() { + let def_id = uv.def; + if self.tcx.def_kind(def_id) == DefKind::InlineConst { + let local_def_id = def_id.expect_local(); + let UnsafetyCheckResult { violations, used_unsafe_blocks, .. } = + self.tcx.unsafety_check_result(local_def_id); + self.register_violations(violations, used_unsafe_blocks.items().copied()); + } + } + } + } + self.super_operand(op, location); + } + + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + // On types with `scalar_valid_range`, prevent + // * `&mut x.field` + // * `x.field = y;` + // * `&x.field` if `field`'s type has interior mutability + // because either of these would allow modifying the layout constrained field and + // insert values that violate the layout constraints. + if context.is_mutating_use() || context.is_borrow() { + self.check_mut_borrowing_layout_constrained_field(*place, context.is_mutating_use()); + } + + // Some checks below need the extra meta info of the local declaration. + let decl = &self.body.local_decls[place.local]; + + // Check the base local: it might be an unsafe-to-access static. We only check derefs of the + // temporary holding the static pointer to avoid duplicate errors + // <https://github.com/rust-lang/rust/pull/78068#issuecomment-731753506>. + if place.projection.first() == Some(&ProjectionElem::Deref) { + // If the projection root is an artificial local that we introduced when + // desugaring `static`, give a more specific error message + // (avoid the general "raw pointer" clause below, that would only be confusing). + if let LocalInfo::StaticRef { def_id, .. } = *decl.local_info() { + if self.tcx.is_mutable_static(def_id) { + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::UseOfMutableStatic, + ); + return; + } else if self.tcx.is_foreign_item(def_id) { + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::UseOfExternStatic, + ); + return; + } + } + } + + // Check for raw pointer `Deref`. + for (base, proj) in place.iter_projections() { + if proj == ProjectionElem::Deref { + let base_ty = base.ty(self.body, self.tcx).ty; + if base_ty.is_unsafe_ptr() { + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::DerefOfRawPointer, + ) + } + } + } + + // Check for union fields. For this we traverse right-to-left, as the last `Deref` changes + // whether we *read* the union field or potentially *write* to it (if this place is being assigned to). + let mut saw_deref = false; + for (base, proj) in place.iter_projections().rev() { + if proj == ProjectionElem::Deref { + saw_deref = true; + continue; + } + + let base_ty = base.ty(self.body, self.tcx).ty; + if base_ty.is_union() { + // If we did not hit a `Deref` yet and the overall place use is an assignment, the + // rules are different. + let assign_to_field = !saw_deref + && matches!( + context, + PlaceContext::MutatingUse( + MutatingUseContext::Store + | MutatingUseContext::Drop + | MutatingUseContext::AsmOutput + ) + ); + // If this is just an assignment, determine if the assigned type needs dropping. + if assign_to_field { + // We have to check the actual type of the assignment, as that determines if the + // old value is being dropped. + let assigned_ty = place.ty(&self.body.local_decls, self.tcx).ty; + if assigned_ty.needs_drop(self.tcx, self.param_env) { + // This would be unsafe, but should be outright impossible since we reject such unions. + self.tcx.sess.span_delayed_bug( + self.source_info.span, + format!("union fields that need dropping should be impossible: {assigned_ty}") + ); + } + } else { + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::AccessToUnionField, + ) + } + } + } + } +} + +impl<'tcx> UnsafetyChecker<'_, 'tcx> { + fn require_unsafe(&mut self, kind: UnsafetyViolationKind, details: UnsafetyViolationDetails) { + // Violations can turn out to be `UnsafeFn` during analysis, but they should not start out as such. + assert_ne!(kind, UnsafetyViolationKind::UnsafeFn); + + let source_info = self.source_info; + let lint_root = self.body.source_scopes[self.source_info.scope] + .local_data + .as_ref() + .assert_crate_local() + .lint_root; + self.register_violations( + [&UnsafetyViolation { source_info, lint_root, kind, details }], + UnordItems::empty(), + ); + } + + fn register_violations<'a>( + &mut self, + violations: impl IntoIterator<Item = &'a UnsafetyViolation>, + new_used_unsafe_blocks: UnordItems<HirId, impl Iterator<Item = HirId>>, + ) { + let safety = self.body.source_scopes[self.source_info.scope] + .local_data + .as_ref() + .assert_crate_local() + .safety; + match safety { + // `unsafe` blocks are required in safe code + Safety::Safe => violations.into_iter().for_each(|violation| { + match violation.kind { + UnsafetyViolationKind::General => {} + UnsafetyViolationKind::UnsafeFn => { + bug!("`UnsafetyViolationKind::UnsafeFn` in an `Safe` context") + } + } + if !self.violations.contains(violation) { + self.violations.push(violation.clone()) + } + }), + // With the RFC 2585, no longer allow `unsafe` operations in `unsafe fn`s + Safety::FnUnsafe => violations.into_iter().for_each(|violation| { + let mut violation = violation.clone(); + violation.kind = UnsafetyViolationKind::UnsafeFn; + if !self.violations.contains(&violation) { + self.violations.push(violation) + } + }), + Safety::BuiltinUnsafe => {} + Safety::ExplicitUnsafe(hir_id) => violations.into_iter().for_each(|_violation| { + self.used_unsafe_blocks.insert(hir_id); + }), + }; + + self.used_unsafe_blocks.extend_unord(new_used_unsafe_blocks); + } + fn check_mut_borrowing_layout_constrained_field( + &mut self, + place: Place<'tcx>, + is_mut_use: bool, + ) { + for (place_base, elem) in place.iter_projections().rev() { + match elem { + // Modifications behind a dereference don't affect the value of + // the pointer. + ProjectionElem::Deref => return, + ProjectionElem::Field(..) => { + let ty = place_base.ty(&self.body.local_decls, self.tcx).ty; + if let ty::Adt(def, _) = ty.kind() { + if self.tcx.layout_scalar_valid_range(def.did()) + != (Bound::Unbounded, Bound::Unbounded) + { + let details = if is_mut_use { + UnsafetyViolationDetails::MutationOfLayoutConstrainedField + + // Check `is_freeze` as late as possible to avoid cycle errors + // with opaque types. + } else if !place + .ty(self.body, self.tcx) + .ty + .is_freeze(self.tcx, self.param_env) + { + UnsafetyViolationDetails::BorrowOfLayoutConstrainedField + } else { + continue; + }; + self.require_unsafe(UnsafetyViolationKind::General, details); + } + } + } + _ => {} + } + } + } + + /// Checks whether calling `func_did` needs an `unsafe` context or not, i.e. whether + /// the called function has target features the calling function hasn't. + fn check_target_features(&mut self, func_did: DefId) { + // Unsafety isn't required on wasm targets. For more information see + // the corresponding check in typeck/src/collect.rs + if self.tcx.sess.target.options.is_like_wasm { + return; + } + + let callee_features = &self.tcx.codegen_fn_attrs(func_did).target_features; + // The body might be a constant, so it doesn't have codegen attributes. + let self_features = &self.tcx.body_codegen_attrs(self.body_did.to_def_id()).target_features; + + // Is `callee_features` a subset of `calling_features`? + if !callee_features.iter().all(|feature| self_features.contains(feature)) { + let missing: Vec<_> = callee_features + .iter() + .copied() + .filter(|feature| !self_features.contains(feature)) + .collect(); + let build_enabled = self + .tcx + .sess + .target_features + .iter() + .copied() + .filter(|feature| missing.contains(feature)) + .collect(); + self.require_unsafe( + UnsafetyViolationKind::General, + UnsafetyViolationDetails::CallToFunctionWith { missing, build_enabled }, + ) + } + } +} + +pub(crate) fn provide(providers: &mut Providers) { + *providers = Providers { unsafety_check_result, ..*providers }; +} + +/// Context information for [`UnusedUnsafeVisitor`] traversal, +/// saves (innermost) relevant context +#[derive(Copy, Clone, Debug)] +enum Context { + Safe, + /// in an `unsafe fn` + UnsafeFn, + /// in a *used* `unsafe` block + /// (i.e. a block without unused-unsafe warning) + UnsafeBlock(HirId), +} + +struct UnusedUnsafeVisitor<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + used_unsafe_blocks: &'a UnordSet<HirId>, + context: Context, + unused_unsafes: &'a mut Vec<(HirId, UnusedUnsafe)>, +} + +impl<'tcx> intravisit::Visitor<'tcx> for UnusedUnsafeVisitor<'_, 'tcx> { + fn visit_block(&mut self, block: &'tcx hir::Block<'tcx>) { + if let hir::BlockCheckMode::UnsafeBlock(hir::UnsafeSource::UserProvided) = block.rules { + let used = match self.tcx.lint_level_at_node(UNUSED_UNSAFE, block.hir_id) { + (Level::Allow, _) => true, + _ => self.used_unsafe_blocks.contains(&block.hir_id), + }; + let unused_unsafe = match (self.context, used) { + (_, false) => UnusedUnsafe::Unused, + (Context::Safe, true) | (Context::UnsafeFn, true) => { + let previous_context = self.context; + self.context = Context::UnsafeBlock(block.hir_id); + intravisit::walk_block(self, block); + self.context = previous_context; + return; + } + (Context::UnsafeBlock(hir_id), true) => UnusedUnsafe::InUnsafeBlock(hir_id), + }; + self.unused_unsafes.push((block.hir_id, unused_unsafe)); + } + intravisit::walk_block(self, block); + } + + fn visit_inline_const(&mut self, c: &'tcx hir::ConstBlock) { + self.visit_body(self.tcx.hir().body(c.body)) + } + + fn visit_fn( + &mut self, + fk: intravisit::FnKind<'tcx>, + _fd: &'tcx hir::FnDecl<'tcx>, + b: hir::BodyId, + _s: rustc_span::Span, + _id: LocalDefId, + ) { + if matches!(fk, intravisit::FnKind::Closure) { + self.visit_body(self.tcx.hir().body(b)) + } + } +} + +fn check_unused_unsafe( + tcx: TyCtxt<'_>, + def_id: LocalDefId, + used_unsafe_blocks: &UnordSet<HirId>, +) -> Vec<(HirId, UnusedUnsafe)> { + let body_id = tcx.hir().maybe_body_owned_by(def_id); + + let Some(body_id) = body_id else { + debug!("check_unused_unsafe({:?}) - no body found", def_id); + return vec![]; + }; + + let body = tcx.hir().body(body_id); + let hir_id = tcx.local_def_id_to_hir_id(def_id); + let context = match tcx.hir().fn_sig_by_hir_id(hir_id) { + Some(sig) if sig.header.unsafety == hir::Unsafety::Unsafe => Context::UnsafeFn, + _ => Context::Safe, + }; + + debug!( + "check_unused_unsafe({:?}, context={:?}, body={:?}, used_unsafe_blocks={:?})", + def_id, body, context, used_unsafe_blocks + ); + + let mut unused_unsafes = vec![]; + + let mut visitor = UnusedUnsafeVisitor { + tcx, + used_unsafe_blocks, + context, + unused_unsafes: &mut unused_unsafes, + }; + intravisit::Visitor::visit_body(&mut visitor, body); + + unused_unsafes +} + +fn unsafety_check_result(tcx: TyCtxt<'_>, def: LocalDefId) -> &UnsafetyCheckResult { + debug!("unsafety_violations({:?})", def); + + // N.B., this borrow is valid because all the consumers of + // `mir_built` force this. + let body = &tcx.mir_built(def).borrow(); + + if body.is_custom_mir() || body.tainted_by_errors.is_some() { + return tcx.arena.alloc(UnsafetyCheckResult { + violations: Vec::new(), + used_unsafe_blocks: Default::default(), + unused_unsafes: Some(Vec::new()), + }); + } + + let param_env = tcx.param_env(def); + + let mut checker = UnsafetyChecker::new(body, def, tcx, param_env); + checker.visit_body(body); + + let unused_unsafes = (!tcx.is_typeck_child(def.to_def_id())) + .then(|| check_unused_unsafe(tcx, def, &checker.used_unsafe_blocks)); + + tcx.arena.alloc(UnsafetyCheckResult { + violations: checker.violations, + used_unsafe_blocks: checker.used_unsafe_blocks, + unused_unsafes, + }) +} + +fn report_unused_unsafe(tcx: TyCtxt<'_>, kind: UnusedUnsafe, id: HirId) { + let span = tcx.sess.source_map().guess_head_span(tcx.hir().span(id)); + let nested_parent = if let UnusedUnsafe::InUnsafeBlock(id) = kind { + Some(tcx.sess.source_map().guess_head_span(tcx.hir().span(id))) + } else { + None + }; + tcx.emit_spanned_lint(UNUSED_UNSAFE, id, span, errors::UnusedUnsafe { span, nested_parent }); +} + +pub fn check_unsafety(tcx: TyCtxt<'_>, def_id: LocalDefId) { + debug!("check_unsafety({:?})", def_id); + + // closures and inline consts are handled by their parent fn. + if tcx.is_typeck_child(def_id.to_def_id()) { + return; + } + + let UnsafetyCheckResult { violations, unused_unsafes, .. } = tcx.unsafety_check_result(def_id); + // Only suggest wrapping the entire function body in an unsafe block once + let mut suggest_unsafe_block = true; + + for &UnsafetyViolation { source_info, lint_root, kind, ref details } in violations.iter() { + let details = + errors::RequiresUnsafeDetail { violation: details.clone(), span: source_info.span }; + + match kind { + UnsafetyViolationKind::General => { + let op_in_unsafe_fn_allowed = unsafe_op_in_unsafe_fn_allowed(tcx, lint_root); + let note_non_inherited = tcx.hir().parent_iter(lint_root).find(|(id, node)| { + if let Node::Expr(block) = node + && let ExprKind::Block(block, _) = block.kind + && let BlockCheckMode::UnsafeBlock(_) = block.rules + { + true + } else if let Some(sig) = tcx.hir().fn_sig_by_hir_id(*id) + && sig.header.is_unsafe() + { + true + } else { + false + } + }); + let enclosing = if let Some((id, _)) = note_non_inherited { + Some(tcx.sess.source_map().guess_head_span(tcx.hir().span(id))) + } else { + None + }; + tcx.sess.emit_err(errors::RequiresUnsafe { + span: source_info.span, + enclosing, + details, + op_in_unsafe_fn_allowed, + }); + } + UnsafetyViolationKind::UnsafeFn => { + tcx.emit_spanned_lint( + UNSAFE_OP_IN_UNSAFE_FN, + lint_root, + source_info.span, + errors::UnsafeOpInUnsafeFn { + details, + suggest_unsafe_block: suggest_unsafe_block.then(|| { + let hir_id = tcx.local_def_id_to_hir_id(def_id); + let fn_sig = tcx + .hir() + .fn_sig_by_hir_id(hir_id) + .expect("this violation only occurs in fn"); + let body = tcx.hir().body_owned_by(def_id); + let body_span = tcx.hir().body(body).value.span; + let start = tcx.sess.source_map().start_point(body_span).shrink_to_hi(); + let end = tcx.sess.source_map().end_point(body_span).shrink_to_lo(); + (start, end, fn_sig.span) + }), + }, + ); + suggest_unsafe_block = false; + } + } + } + + for &(block_id, kind) in unused_unsafes.as_ref().unwrap() { + report_unused_unsafe(tcx, kind, block_id); + } +} + +fn unsafe_op_in_unsafe_fn_allowed(tcx: TyCtxt<'_>, id: HirId) -> bool { + tcx.lint_level_at_node(UNSAFE_OP_IN_UNSAFE_FN, id).0 == Level::Allow +} diff --git a/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs b/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs new file mode 100644 index 00000000000..5b4bc4fa134 --- /dev/null +++ b/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs @@ -0,0 +1,48 @@ +//! This module provides a pass that removes parts of MIR that are no longer relevant after +//! analysis phase and borrowck. In particular, it removes false edges, user type annotations and +//! replaces following statements with [`Nop`]s: +//! +//! - [`AscribeUserType`] +//! - [`FakeRead`] +//! - [`Assign`] statements with a [`Fake`] borrow +//! +//! [`AscribeUserType`]: rustc_middle::mir::StatementKind::AscribeUserType +//! [`Assign`]: rustc_middle::mir::StatementKind::Assign +//! [`FakeRead`]: rustc_middle::mir::StatementKind::FakeRead +//! [`Nop`]: rustc_middle::mir::StatementKind::Nop +//! [`Fake`]: rustc_middle::mir::BorrowKind::Fake + +use crate::MirPass; +use rustc_middle::mir::{Body, BorrowKind, Rvalue, StatementKind, TerminatorKind}; +use rustc_middle::ty::TyCtxt; + +pub struct CleanupPostBorrowck; + +impl<'tcx> MirPass<'tcx> for CleanupPostBorrowck { + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + for basic_block in body.basic_blocks.as_mut() { + for statement in basic_block.statements.iter_mut() { + match statement.kind { + StatementKind::AscribeUserType(..) + | StatementKind::Assign(box (_, Rvalue::Ref(_, BorrowKind::Fake, _))) + | StatementKind::FakeRead(..) => statement.make_nop(), + _ => (), + } + } + let terminator = basic_block.terminator_mut(); + match terminator.kind { + TerminatorKind::FalseEdge { real_target, .. } + | TerminatorKind::FalseUnwind { real_target, .. } => { + terminator.kind = TerminatorKind::Goto { target: real_target }; + } + _ => {} + } + } + + body.user_type_annotations.raw.clear(); + + for decl in &mut body.local_decls { + decl.user_ty = None; + } + } +} diff --git a/compiler/rustc_mir_transform/src/const_debuginfo.rs b/compiler/rustc_mir_transform/src/const_debuginfo.rs new file mode 100644 index 00000000000..e4e4270c499 --- /dev/null +++ b/compiler/rustc_mir_transform/src/const_debuginfo.rs @@ -0,0 +1,102 @@ +//! Finds locals which are assigned once to a const and unused except for debuginfo and converts +//! their debuginfo to use the const directly, allowing the local to be removed. + +use rustc_middle::{ + mir::{ + visit::{PlaceContext, Visitor}, + Body, ConstOperand, Local, Location, Operand, Rvalue, StatementKind, VarDebugInfoContents, + }, + ty::TyCtxt, +}; + +use crate::MirPass; +use rustc_index::{bit_set::BitSet, IndexVec}; + +pub struct ConstDebugInfo; + +impl<'tcx> MirPass<'tcx> for ConstDebugInfo { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("running ConstDebugInfo on {:?}", body.source); + + for (local, constant) in find_optimization_opportunities(body) { + for debuginfo in &mut body.var_debug_info { + if let VarDebugInfoContents::Place(p) = debuginfo.value { + if p.local == local && p.projection.is_empty() { + trace!( + "changing debug info for {:?} from place {:?} to constant {:?}", + debuginfo.name, + p, + constant + ); + debuginfo.value = VarDebugInfoContents::Const(constant); + } + } + } + } + } +} + +struct LocalUseVisitor { + local_mutating_uses: IndexVec<Local, u8>, + local_assignment_locations: IndexVec<Local, Option<Location>>, +} + +fn find_optimization_opportunities<'tcx>(body: &Body<'tcx>) -> Vec<(Local, ConstOperand<'tcx>)> { + let mut visitor = LocalUseVisitor { + local_mutating_uses: IndexVec::from_elem(0, &body.local_decls), + local_assignment_locations: IndexVec::from_elem(None, &body.local_decls), + }; + + visitor.visit_body(body); + + let mut locals_to_debuginfo = BitSet::new_empty(body.local_decls.len()); + for debuginfo in &body.var_debug_info { + if let VarDebugInfoContents::Place(p) = debuginfo.value + && let Some(l) = p.as_local() + { + locals_to_debuginfo.insert(l); + } + } + + let mut eligible_locals = Vec::new(); + for (local, mutating_uses) in visitor.local_mutating_uses.drain_enumerated(..) { + if mutating_uses != 1 || !locals_to_debuginfo.contains(local) { + continue; + } + + if let Some(location) = visitor.local_assignment_locations[local] { + let bb = &body[location.block]; + + // The value is assigned as the result of a call, not a constant + if bb.statements.len() == location.statement_index { + continue; + } + + if let StatementKind::Assign(box (p, Rvalue::Use(Operand::Constant(box c)))) = + &bb.statements[location.statement_index].kind + { + if let Some(local) = p.as_local() { + eligible_locals.push((local, *c)); + } + } + } + } + + eligible_locals +} + +impl Visitor<'_> for LocalUseVisitor { + fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) { + if context.is_mutating_use() { + self.local_mutating_uses[local] = self.local_mutating_uses[local].saturating_add(1); + + if context.is_place_assignment() { + self.local_assignment_locations[local] = Some(location); + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/const_goto.rs b/compiler/rustc_mir_transform/src/const_goto.rs new file mode 100644 index 00000000000..3884346076e --- /dev/null +++ b/compiler/rustc_mir_transform/src/const_goto.rs @@ -0,0 +1,128 @@ +//! This pass optimizes the following sequence +//! ```rust,ignore (example) +//! bb2: { +//! _2 = const true; +//! goto -> bb3; +//! } +//! +//! bb3: { +//! switchInt(_2) -> [false: bb4, otherwise: bb5]; +//! } +//! ``` +//! into +//! ```rust,ignore (example) +//! bb2: { +//! _2 = const true; +//! goto -> bb5; +//! } +//! ``` + +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_middle::{mir::visit::Visitor, ty::ParamEnv}; + +use super::simplify::{simplify_cfg, simplify_locals}; + +pub struct ConstGoto; + +impl<'tcx> MirPass<'tcx> for ConstGoto { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // This pass participates in some as-of-yet untested unsoundness found + // in https://github.com/rust-lang/rust/issues/112460 + sess.mir_opt_level() >= 2 && sess.opts.unstable_opts.unsound_mir_opts + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("Running ConstGoto on {:?}", body.source); + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + let mut opt_finder = + ConstGotoOptimizationFinder { tcx, body, optimizations: vec![], param_env }; + opt_finder.visit_body(body); + let should_simplify = !opt_finder.optimizations.is_empty(); + for opt in opt_finder.optimizations { + let block = &mut body.basic_blocks_mut()[opt.bb_with_goto]; + block.statements.extend(opt.stmts_move_up); + let terminator = block.terminator_mut(); + let new_goto = TerminatorKind::Goto { target: opt.target_to_use_in_goto }; + debug!("SUCCESS: replacing `{:?}` with `{:?}`", terminator.kind, new_goto); + terminator.kind = new_goto; + } + + // if we applied optimizations, we potentially have some cfg to cleanup to + // make it easier for further passes + if should_simplify { + simplify_cfg(tcx, body); + simplify_locals(body, tcx); + } + } +} + +impl<'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'_, 'tcx> { + fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) { + if data.is_cleanup { + // Because of the restrictions around control flow in cleanup blocks, we don't perform + // this optimization at all in such blocks. + return; + } + self.super_basic_block_data(block, data); + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + let _: Option<_> = try { + let target = terminator.kind.as_goto()?; + // We only apply this optimization if the last statement is a const assignment + let last_statement = self.body.basic_blocks[location.block].statements.last()?; + + if let (place, Rvalue::Use(Operand::Constant(_const))) = + last_statement.kind.as_assign()? + { + // We found a constant being assigned to `place`. + // Now check that the target of this Goto switches on this place. + let target_bb = &self.body.basic_blocks[target]; + + // The `StorageDead(..)` statement does not affect the functionality of mir. + // We can move this part of the statement up to the predecessor. + let mut stmts_move_up = Vec::new(); + for stmt in &target_bb.statements { + if let StatementKind::StorageDead(..) = stmt.kind { + stmts_move_up.push(stmt.clone()) + } else { + None?; + } + } + + let target_bb_terminator = target_bb.terminator(); + let (discr, targets) = target_bb_terminator.kind.as_switch()?; + if discr.place() == Some(*place) { + let switch_ty = place.ty(self.body.local_decls(), self.tcx).ty; + debug_assert_eq!(switch_ty, _const.ty()); + // We now know that the Switch matches on the const place, and it is statementless + // Now find which value in the Switch matches the const value. + let const_value = _const.const_.try_eval_bits(self.tcx, self.param_env)?; + let target_to_use_in_goto = targets.target_for_value(const_value); + self.optimizations.push(OptimizationToApply { + bb_with_goto: location.block, + target_to_use_in_goto, + stmts_move_up, + }); + } + } + Some(()) + }; + + self.super_terminator(terminator, location); + } +} + +struct OptimizationToApply<'tcx> { + bb_with_goto: BasicBlock, + target_to_use_in_goto: BasicBlock, + stmts_move_up: Vec<Statement<'tcx>>, +} + +pub struct ConstGotoOptimizationFinder<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + body: &'a Body<'tcx>, + param_env: ParamEnv<'tcx>, + optimizations: Vec<OptimizationToApply<'tcx>>, +} diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs new file mode 100644 index 00000000000..e66d5e0a9f9 --- /dev/null +++ b/compiler/rustc_mir_transform/src/const_prop.rs @@ -0,0 +1,830 @@ +//! Propagates constants for early reporting of statically known +//! assertion failures + +use either::Right; +use rustc_const_eval::ReportErrorExt; +use rustc_data_structures::fx::FxHashSet; +use rustc_hir::def::DefKind; +use rustc_index::bit_set::BitSet; +use rustc_index::{IndexSlice, IndexVec}; +use rustc_middle::mir::visit::{ + MutVisitor, MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor, +}; +use rustc_middle::mir::*; +use rustc_middle::query::TyCtxtAt; +use rustc_middle::ty::layout::{LayoutError, LayoutOf, LayoutOfHelpers, TyAndLayout}; +use rustc_middle::ty::{self, GenericArgs, Instance, ParamEnv, Ty, TyCtxt, TypeVisitableExt}; +use rustc_span::{def_id::DefId, Span}; +use rustc_target::abi::{self, HasDataLayout, Size, TargetDataLayout}; +use rustc_target::spec::abi::Abi as CallAbi; + +use crate::dataflow_const_prop::Patch; +use rustc_const_eval::interpret::{ + self, compile_time_machine, AllocId, ConstAllocation, FnArg, Frame, ImmTy, Immediate, InterpCx, + InterpResult, MemoryKind, OpTy, PlaceTy, Pointer, Scalar, StackPopCleanup, +}; + +/// The maximum number of bytes that we'll allocate space for a local or the return value. +/// Needed for #66397, because otherwise we eval into large places and that can cause OOM or just +/// Severely regress performance. +const MAX_ALLOC_LIMIT: u64 = 1024; + +/// Macro for machine-specific `InterpError` without allocation. +/// (These will never be shown to the user, but they help diagnose ICEs.) +pub(crate) macro throw_machine_stop_str($($tt:tt)*) {{ + // We make a new local type for it. The type itself does not carry any information, + // but its vtable (for the `MachineStopType` trait) does. + #[derive(Debug)] + struct Zst; + // Printing this type shows the desired string. + impl std::fmt::Display for Zst { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, $($tt)*) + } + } + + impl rustc_middle::mir::interpret::MachineStopType for Zst { + fn diagnostic_message(&self) -> rustc_errors::DiagnosticMessage { + self.to_string().into() + } + + fn add_args( + self: Box<Self>, + _: &mut dyn FnMut(std::borrow::Cow<'static, str>, rustc_errors::DiagnosticArgValue<'static>), + ) {} + } + throw_machine_stop!(Zst) +}} + +pub struct ConstProp; + +impl<'tcx> MirPass<'tcx> for ConstProp { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 2 + } + + #[instrument(skip(self, tcx), level = "debug")] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // will be evaluated by miri and produce its errors there + if body.source.promoted.is_some() { + return; + } + + let def_id = body.source.def_id().expect_local(); + let def_kind = tcx.def_kind(def_id); + let is_fn_like = def_kind.is_fn_like(); + let is_assoc_const = def_kind == DefKind::AssocConst; + + // Only run const prop on functions, methods, closures and associated constants + if !is_fn_like && !is_assoc_const { + // skip anon_const/statics/consts because they'll be evaluated by miri anyway + trace!("ConstProp skipped for {:?}", def_id); + return; + } + + // FIXME(welseywiser) const prop doesn't work on coroutines because of query cycles + // computing their layout. + if tcx.is_coroutine(def_id.to_def_id()) { + trace!("ConstProp skipped for coroutine {:?}", def_id); + return; + } + + trace!("ConstProp starting for {:?}", def_id); + + // FIXME(oli-obk, eddyb) Optimize locals (or even local paths) to hold + // constants, instead of just checking for const-folding succeeding. + // That would require a uniform one-def no-mutation analysis + // and RPO (or recursing when needing the value of a local). + let mut optimization_finder = ConstPropagator::new(body, tcx); + + // Traverse the body in reverse post-order, to ensure that `FullConstProp` locals are + // assigned before being read. + for &bb in body.basic_blocks.reverse_postorder() { + let data = &body.basic_blocks[bb]; + optimization_finder.visit_basic_block_data(bb, data); + } + + let mut patch = optimization_finder.patch; + patch.visit_body_preserves_cfg(body); + + trace!("ConstProp done for {:?}", def_id); + } +} + +pub struct ConstPropMachine<'mir, 'tcx> { + /// The virtual call stack. + stack: Vec<Frame<'mir, 'tcx>>, + pub written_only_inside_own_block_locals: FxHashSet<Local>, + pub can_const_prop: IndexVec<Local, ConstPropMode>, +} + +impl ConstPropMachine<'_, '_> { + pub fn new(can_const_prop: IndexVec<Local, ConstPropMode>) -> Self { + Self { + stack: Vec::new(), + written_only_inside_own_block_locals: Default::default(), + can_const_prop, + } + } +} + +impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx> { + compile_time_machine!(<'mir, 'tcx>); + + const PANIC_ON_ALLOC_FAIL: bool = true; // all allocations are small (see `MAX_ALLOC_LIMIT`) + + const POST_MONO_CHECKS: bool = false; // this MIR is still generic! + + type MemoryKind = !; + + #[inline(always)] + fn enforce_alignment(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool { + false // no reason to enforce alignment + } + + #[inline(always)] + fn enforce_validity(_ecx: &InterpCx<'mir, 'tcx, Self>, _layout: TyAndLayout<'tcx>) -> bool { + false // for now, we don't enforce validity + } + + fn load_mir( + _ecx: &InterpCx<'mir, 'tcx, Self>, + _instance: ty::InstanceDef<'tcx>, + ) -> InterpResult<'tcx, &'tcx Body<'tcx>> { + throw_machine_stop_str!("calling functions isn't supported in ConstProp") + } + + fn panic_nounwind(_ecx: &mut InterpCx<'mir, 'tcx, Self>, _msg: &str) -> InterpResult<'tcx> { + throw_machine_stop_str!("panicking isn't supported in ConstProp") + } + + fn find_mir_or_eval_fn( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _instance: ty::Instance<'tcx>, + _abi: CallAbi, + _args: &[FnArg<'tcx>], + _destination: &PlaceTy<'tcx>, + _target: Option<BasicBlock>, + _unwind: UnwindAction, + ) -> InterpResult<'tcx, Option<(&'mir Body<'tcx>, ty::Instance<'tcx>)>> { + Ok(None) + } + + fn call_intrinsic( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _instance: ty::Instance<'tcx>, + _args: &[OpTy<'tcx>], + _destination: &PlaceTy<'tcx>, + _target: Option<BasicBlock>, + _unwind: UnwindAction, + ) -> InterpResult<'tcx> { + throw_machine_stop_str!("calling intrinsics isn't supported in ConstProp") + } + + fn assert_panic( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _msg: &rustc_middle::mir::AssertMessage<'tcx>, + _unwind: rustc_middle::mir::UnwindAction, + ) -> InterpResult<'tcx> { + bug!("panics terminators are not evaluated in ConstProp") + } + + fn binary_ptr_op( + _ecx: &InterpCx<'mir, 'tcx, Self>, + _bin_op: BinOp, + _left: &ImmTy<'tcx>, + _right: &ImmTy<'tcx>, + ) -> InterpResult<'tcx, (ImmTy<'tcx>, bool)> { + // We can't do this because aliasing of memory can differ between const eval and llvm + throw_machine_stop_str!("pointer arithmetic or comparisons aren't supported in ConstProp") + } + + fn before_access_local_mut<'a>( + ecx: &'a mut InterpCx<'mir, 'tcx, Self>, + frame: usize, + local: Local, + ) -> InterpResult<'tcx> { + assert_eq!(frame, 0); + match ecx.machine.can_const_prop[local] { + ConstPropMode::NoPropagation => { + throw_machine_stop_str!( + "tried to write to a local that is marked as not propagatable" + ) + } + ConstPropMode::OnlyInsideOwnBlock => { + ecx.machine.written_only_inside_own_block_locals.insert(local); + } + ConstPropMode::FullConstProp => {} + } + Ok(()) + } + + fn before_access_global( + _tcx: TyCtxtAt<'tcx>, + _machine: &Self, + _alloc_id: AllocId, + alloc: ConstAllocation<'tcx>, + _static_def_id: Option<DefId>, + is_write: bool, + ) -> InterpResult<'tcx> { + if is_write { + throw_machine_stop_str!("can't write to global"); + } + // If the static allocation is mutable, then we can't const prop it as its content + // might be different at runtime. + if alloc.inner().mutability.is_mut() { + throw_machine_stop_str!("can't access mutable globals in ConstProp"); + } + + Ok(()) + } + + #[inline(always)] + fn expose_ptr(_ecx: &mut InterpCx<'mir, 'tcx, Self>, _ptr: Pointer) -> InterpResult<'tcx> { + throw_machine_stop_str!("exposing pointers isn't supported in ConstProp") + } + + #[inline(always)] + fn init_frame_extra( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + frame: Frame<'mir, 'tcx>, + ) -> InterpResult<'tcx, Frame<'mir, 'tcx>> { + Ok(frame) + } + + #[inline(always)] + fn stack<'a>( + ecx: &'a InterpCx<'mir, 'tcx, Self>, + ) -> &'a [Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>] { + &ecx.machine.stack + } + + #[inline(always)] + fn stack_mut<'a>( + ecx: &'a mut InterpCx<'mir, 'tcx, Self>, + ) -> &'a mut Vec<Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>> { + &mut ecx.machine.stack + } +} + +/// Finds optimization opportunities on the MIR. +struct ConstPropagator<'mir, 'tcx> { + ecx: InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + local_decls: &'mir IndexSlice<Local, LocalDecl<'tcx>>, + patch: Patch<'tcx>, +} + +impl<'tcx> LayoutOfHelpers<'tcx> for ConstPropagator<'_, 'tcx> { + type LayoutOfResult = Result<TyAndLayout<'tcx>, LayoutError<'tcx>>; + + #[inline] + fn handle_layout_err(&self, err: LayoutError<'tcx>, _: Span, _: Ty<'tcx>) -> LayoutError<'tcx> { + err + } +} + +impl HasDataLayout for ConstPropagator<'_, '_> { + #[inline] + fn data_layout(&self) -> &TargetDataLayout { + &self.tcx.data_layout + } +} + +impl<'tcx> ty::layout::HasTyCtxt<'tcx> for ConstPropagator<'_, 'tcx> { + #[inline] + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } +} + +impl<'tcx> ty::layout::HasParamEnv<'tcx> for ConstPropagator<'_, 'tcx> { + #[inline] + fn param_env(&self) -> ty::ParamEnv<'tcx> { + self.param_env + } +} + +impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> { + fn new(body: &'mir Body<'tcx>, tcx: TyCtxt<'tcx>) -> ConstPropagator<'mir, 'tcx> { + let def_id = body.source.def_id(); + let args = &GenericArgs::identity_for_item(tcx, def_id); + let param_env = tcx.param_env_reveal_all_normalized(def_id); + + let can_const_prop = CanConstProp::check(tcx, param_env, body); + let mut ecx = InterpCx::new( + tcx, + tcx.def_span(def_id), + param_env, + ConstPropMachine::new(can_const_prop), + ); + + let ret_layout = ecx + .layout_of(body.bound_return_ty().instantiate(tcx, args)) + .ok() + // Don't bother allocating memory for large values. + // I don't know how return types can seem to be unsized but this happens in the + // `type/type-unsatisfiable.rs` test. + .filter(|ret_layout| { + ret_layout.is_sized() && ret_layout.size < Size::from_bytes(MAX_ALLOC_LIMIT) + }) + .unwrap_or_else(|| ecx.layout_of(tcx.types.unit).unwrap()); + + let ret = ecx + .allocate(ret_layout, MemoryKind::Stack) + .expect("couldn't perform small allocation") + .into(); + + ecx.push_stack_frame( + Instance::new(def_id, args), + body, + &ret, + StackPopCleanup::Root { cleanup: false }, + ) + .expect("failed to push initial stack frame"); + + for local in body.local_decls.indices() { + // Mark everything initially live. + // This is somewhat dicey since some of them might be unsized and it is incoherent to + // mark those as live... We rely on `local_to_place`/`local_to_op` in the interpreter + // stopping us before those unsized immediates can cause issues deeper in the + // interpreter. + ecx.frame_mut().locals[local].make_live_uninit(); + } + + let patch = Patch::new(tcx); + ConstPropagator { ecx, tcx, param_env, local_decls: &body.local_decls, patch } + } + + fn get_const(&self, place: Place<'tcx>) -> Option<OpTy<'tcx>> { + let op = match self.ecx.eval_place_to_op(place, None) { + Ok(op) => { + if op + .as_mplace_or_imm() + .right() + .is_some_and(|imm| matches!(*imm, Immediate::Uninit)) + { + // Make sure nobody accidentally uses this value. + return None; + } + op + } + Err(e) => { + trace!("get_const failed: {:?}", e.into_kind().debug()); + return None; + } + }; + + // Try to read the local as an immediate so that if it is representable as a scalar, we can + // handle it as such, but otherwise, just return the value as is. + Some(match self.ecx.read_immediate_raw(&op) { + Ok(Right(imm)) => imm.into(), + _ => op, + }) + } + + /// Remove `local` from the pool of `Locals`. Allows writing to them, + /// but not reading from them anymore. + fn remove_const(ecx: &mut InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, local: Local) { + ecx.frame_mut().locals[local].make_live_uninit(); + ecx.machine.written_only_inside_own_block_locals.remove(&local); + } + + fn check_rvalue(&mut self, rvalue: &Rvalue<'tcx>) -> Option<()> { + // Perform any special handling for specific Rvalue types. + // Generally, checks here fall into one of two categories: + // 1. Additional checking to provide useful lints to the user + // - In this case, we will do some validation and then fall through to the + // end of the function which evals the assignment. + // 2. Working around bugs in other parts of the compiler + // - In this case, we'll return `None` from this function to stop evaluation. + match rvalue { + // Do not try creating references (#67862) + Rvalue::AddressOf(_, place) | Rvalue::Ref(_, _, place) => { + trace!("skipping AddressOf | Ref for {:?}", place); + + // This may be creating mutable references or immutable references to cells. + // If that happens, the pointed to value could be mutated via that reference. + // Since we aren't tracking references, the const propagator loses track of what + // value the local has right now. + // Thus, all locals that have their reference taken + // must not take part in propagation. + Self::remove_const(&mut self.ecx, place.local); + + return None; + } + Rvalue::ThreadLocalRef(def_id) => { + trace!("skipping ThreadLocalRef({:?})", def_id); + + return None; + } + // There's no other checking to do at this time. + Rvalue::Aggregate(..) + | Rvalue::Use(..) + | Rvalue::CopyForDeref(..) + | Rvalue::Repeat(..) + | Rvalue::Len(..) + | Rvalue::Cast(..) + | Rvalue::ShallowInitBox(..) + | Rvalue::Discriminant(..) + | Rvalue::NullaryOp(..) + | Rvalue::UnaryOp(..) + | Rvalue::BinaryOp(..) + | Rvalue::CheckedBinaryOp(..) => {} + } + + // FIXME we need to revisit this for #67176 + if rvalue.has_param() { + trace!("skipping, has param"); + return None; + } + if !rvalue + .ty(&self.ecx.frame().body.local_decls, *self.ecx.tcx) + .is_sized(*self.ecx.tcx, self.param_env) + { + // the interpreter doesn't support unsized locals (only unsized arguments), + // but rustc does (in a kinda broken way), so we have to skip them here + return None; + } + + Some(()) + } + + // Attempt to use algebraic identities to eliminate constant expressions + fn eval_rvalue_with_identities( + &mut self, + rvalue: &Rvalue<'tcx>, + place: Place<'tcx>, + ) -> Option<()> { + match rvalue { + Rvalue::BinaryOp(op, box (left, right)) + | Rvalue::CheckedBinaryOp(op, box (left, right)) => { + let l = self.ecx.eval_operand(left, None).and_then(|x| self.ecx.read_immediate(&x)); + let r = + self.ecx.eval_operand(right, None).and_then(|x| self.ecx.read_immediate(&x)); + + let const_arg = match (l, r) { + (Ok(x), Err(_)) | (Err(_), Ok(x)) => x, // exactly one side is known + (Err(_), Err(_)) => return None, // neither side is known + (Ok(_), Ok(_)) => return self.ecx.eval_rvalue_into_place(rvalue, place).ok(), // both sides are known + }; + + if !matches!(const_arg.layout.abi, abi::Abi::Scalar(..)) { + // We cannot handle Scalar Pair stuff. + // No point in calling `eval_rvalue_into_place`, since only one side is known + return None; + } + + let arg_value = const_arg.to_scalar().to_bits(const_arg.layout.size).ok()?; + let dest = self.ecx.eval_place(place).ok()?; + + match op { + BinOp::BitAnd if arg_value == 0 => { + self.ecx.write_immediate(*const_arg, &dest).ok() + } + BinOp::BitOr + if arg_value == const_arg.layout.size.truncate(u128::MAX) + || (const_arg.layout.ty.is_bool() && arg_value == 1) => + { + self.ecx.write_immediate(*const_arg, &dest).ok() + } + BinOp::Mul if const_arg.layout.ty.is_integral() && arg_value == 0 => { + if let Rvalue::CheckedBinaryOp(_, _) = rvalue { + let val = Immediate::ScalarPair( + const_arg.to_scalar(), + Scalar::from_bool(false), + ); + self.ecx.write_immediate(val, &dest).ok() + } else { + self.ecx.write_immediate(*const_arg, &dest).ok() + } + } + _ => None, + } + } + _ => self.ecx.eval_rvalue_into_place(rvalue, place).ok(), + } + } + + fn replace_with_const(&mut self, place: Place<'tcx>) -> Option<Const<'tcx>> { + // This will return None if the above `const_prop` invocation only "wrote" a + // type whose creation requires no write. E.g. a coroutine whose initial state + // consists solely of uninitialized memory (so it doesn't capture any locals). + let value = self.get_const(place)?; + if !self.tcx.consider_optimizing(|| format!("ConstantPropagation - {value:?}")) { + return None; + } + trace!("replacing {:?} with {:?}", place, value); + + // FIXME: figure out what to do when read_immediate_raw fails + let imm = self.ecx.read_immediate_raw(&value).ok()?; + + let Right(imm) = imm else { return None }; + match *imm { + Immediate::Scalar(scalar) if scalar.try_to_int().is_ok() => { + Some(Const::from_scalar(self.tcx, scalar, value.layout.ty)) + } + Immediate::ScalarPair(l, r) if l.try_to_int().is_ok() && r.try_to_int().is_ok() => { + let alloc_id = self + .ecx + .intern_with_temp_alloc(value.layout, |ecx, dest| { + ecx.write_immediate(*imm, dest) + }) + .ok()?; + + Some(Const::Val( + ConstValue::Indirect { alloc_id, offset: Size::ZERO }, + value.layout.ty, + )) + } + // Scalars or scalar pairs that contain undef values are assumed to not have + // successfully evaluated and are thus not propagated. + _ => None, + } + } + + fn ensure_not_propagated(&self, local: Local) { + if cfg!(debug_assertions) { + assert!( + self.get_const(local.into()).is_none() + || self + .layout_of(self.local_decls[local].ty) + .map_or(true, |layout| layout.is_zst()), + "failed to remove values for `{local:?}`, value={:?}", + self.get_const(local.into()), + ) + } + } +} + +/// The mode that `ConstProp` is allowed to run in for a given `Local`. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum ConstPropMode { + /// The `Local` can be propagated into and reads of this `Local` can also be propagated. + FullConstProp, + /// The `Local` can only be propagated into and from its own block. + OnlyInsideOwnBlock, + /// The `Local` cannot be part of propagation at all. Any statement + /// referencing it either for reading or writing will not get propagated. + NoPropagation, +} + +pub struct CanConstProp { + can_const_prop: IndexVec<Local, ConstPropMode>, + // False at the beginning. Once set, no more assignments are allowed to that local. + found_assignment: BitSet<Local>, +} + +impl CanConstProp { + /// Returns true if `local` can be propagated + pub fn check<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + body: &Body<'tcx>, + ) -> IndexVec<Local, ConstPropMode> { + let mut cpv = CanConstProp { + can_const_prop: IndexVec::from_elem(ConstPropMode::FullConstProp, &body.local_decls), + found_assignment: BitSet::new_empty(body.local_decls.len()), + }; + for (local, val) in cpv.can_const_prop.iter_enumerated_mut() { + let ty = body.local_decls[local].ty; + match tcx.layout_of(param_env.and(ty)) { + Ok(layout) if layout.size < Size::from_bytes(MAX_ALLOC_LIMIT) => {} + // Either the layout fails to compute, then we can't use this local anyway + // or the local is too large, then we don't want to. + _ => { + *val = ConstPropMode::NoPropagation; + continue; + } + } + } + // Consider that arguments are assigned on entry. + for arg in body.args_iter() { + cpv.found_assignment.insert(arg); + } + cpv.visit_body(body); + cpv.can_const_prop + } +} + +impl<'tcx> Visitor<'tcx> for CanConstProp { + fn visit_place(&mut self, place: &Place<'tcx>, mut context: PlaceContext, loc: Location) { + use rustc_middle::mir::visit::PlaceContext::*; + + // Dereferencing just read the addess of `place.local`. + if place.projection.first() == Some(&PlaceElem::Deref) { + context = NonMutatingUse(NonMutatingUseContext::Copy); + } + + self.visit_local(place.local, context, loc); + self.visit_projection(place.as_ref(), context, loc); + } + + fn visit_local(&mut self, local: Local, context: PlaceContext, _: Location) { + use rustc_middle::mir::visit::PlaceContext::*; + match context { + // These are just stores, where the storing is not propagatable, but there may be later + // mutations of the same local via `Store` + | MutatingUse(MutatingUseContext::Call) + | MutatingUse(MutatingUseContext::AsmOutput) + | MutatingUse(MutatingUseContext::Deinit) + // Actual store that can possibly even propagate a value + | MutatingUse(MutatingUseContext::Store) + | MutatingUse(MutatingUseContext::SetDiscriminant) => { + if !self.found_assignment.insert(local) { + match &mut self.can_const_prop[local] { + // If the local can only get propagated in its own block, then we don't have + // to worry about multiple assignments, as we'll nuke the const state at the + // end of the block anyway, and inside the block we overwrite previous + // states as applicable. + ConstPropMode::OnlyInsideOwnBlock => {} + ConstPropMode::NoPropagation => {} + other @ ConstPropMode::FullConstProp => { + trace!( + "local {:?} can't be propagated because of multiple assignments. Previous state: {:?}", + local, other, + ); + *other = ConstPropMode::OnlyInsideOwnBlock; + } + } + } + } + // Reading constants is allowed an arbitrary number of times + NonMutatingUse(NonMutatingUseContext::Copy) + | NonMutatingUse(NonMutatingUseContext::Move) + | NonMutatingUse(NonMutatingUseContext::Inspect) + | NonMutatingUse(NonMutatingUseContext::PlaceMention) + | NonUse(_) => {} + + // These could be propagated with a smarter analysis or just some careful thinking about + // whether they'd be fine right now. + MutatingUse(MutatingUseContext::Yield) + | MutatingUse(MutatingUseContext::Drop) + | MutatingUse(MutatingUseContext::Retag) + // These can't ever be propagated under any scheme, as we can't reason about indirect + // mutation. + | NonMutatingUse(NonMutatingUseContext::SharedBorrow) + | NonMutatingUse(NonMutatingUseContext::FakeBorrow) + | NonMutatingUse(NonMutatingUseContext::AddressOf) + | MutatingUse(MutatingUseContext::Borrow) + | MutatingUse(MutatingUseContext::AddressOf) => { + trace!("local {:?} can't be propagated because it's used: {:?}", local, context); + self.can_const_prop[local] = ConstPropMode::NoPropagation; + } + MutatingUse(MutatingUseContext::Projection) + | NonMutatingUse(NonMutatingUseContext::Projection) => bug!("visit_place should not pass {context:?} for {local:?}"), + } + } +} + +impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> { + fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) { + self.super_operand(operand, location); + if let Some(place) = operand.place() + && let Some(value) = self.replace_with_const(place) + { + self.patch.before_effect.insert((location, place), value); + } + } + + fn visit_projection_elem( + &mut self, + _: PlaceRef<'tcx>, + elem: PlaceElem<'tcx>, + _: PlaceContext, + location: Location, + ) { + if let PlaceElem::Index(local) = elem + && let Some(value) = self.replace_with_const(local.into()) + { + self.patch.before_effect.insert((location, local.into()), value); + } + } + + fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, location: Location) { + self.super_assign(place, rvalue, location); + + let Some(()) = self.check_rvalue(rvalue) else { + trace!("rvalue check failed, removing const"); + Self::remove_const(&mut self.ecx, place.local); + return; + }; + + match self.ecx.machine.can_const_prop[place.local] { + // Do nothing if the place is indirect. + _ if place.is_indirect() => {} + ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local), + ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => { + if let Some(()) = self.eval_rvalue_with_identities(rvalue, *place) { + // If this was already an evaluated constant, keep it. + if let Rvalue::Use(Operand::Constant(c)) = rvalue + && let Const::Val(..) = c.const_ + { + trace!( + "skipping replace of Rvalue::Use({:?} because it is already a const", + c + ); + } else if let Some(operand) = self.replace_with_const(*place) { + self.patch.assignments.insert(location, operand); + } + } else { + // Const prop failed, so erase the destination, ensuring that whatever happens + // from here on, does not know about the previous value. + // This is important in case we have + // ```rust + // let mut x = 42; + // x = SOME_MUTABLE_STATIC; + // // x must now be uninit + // ``` + // FIXME: we overzealously erase the entire local, because that's easier to + // implement. + trace!( + "propagation into {:?} failed. + Nuking the entire site from orbit, it's the only way to be sure", + place, + ); + Self::remove_const(&mut self.ecx, place.local); + } + } + } + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + trace!("visit_statement: {:?}", statement); + + // We want to evaluate operands before any change to the assigned-to value, + // so we recurse first. + self.super_statement(statement, location); + + match statement.kind { + StatementKind::SetDiscriminant { ref place, .. } => { + match self.ecx.machine.can_const_prop[place.local] { + // Do nothing if the place is indirect. + _ if place.is_indirect() => {} + ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local), + ConstPropMode::FullConstProp | ConstPropMode::OnlyInsideOwnBlock => { + if self.ecx.statement(statement).is_ok() { + trace!("propped discriminant into {:?}", place); + } else { + Self::remove_const(&mut self.ecx, place.local); + } + } + } + } + StatementKind::StorageLive(local) => { + Self::remove_const(&mut self.ecx, local); + } + // We do not need to mark dead locals as such. For `FullConstProp` locals, + // this allows to propagate the single assigned value in this case: + // ``` + // let x = SOME_CONST; + // if a { + // f(copy x); + // StorageDead(x); + // } else { + // g(copy x); + // StorageDead(x); + // } + // ``` + // + // This may propagate a constant where the local would be uninit or dead. + // In both cases, this does not matter, as those reads would be UB anyway. + _ => {} + } + } + + fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) { + self.super_basic_block_data(block, data); + + // We remove all Locals which are restricted in propagation to their containing blocks and + // which were modified in the current block. + // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`. + let mut written_only_inside_own_block_locals = + std::mem::take(&mut self.ecx.machine.written_only_inside_own_block_locals); + + // This loop can get very hot for some bodies: it check each local in each bb. + // To avoid this quadratic behaviour, we only clear the locals that were modified inside + // the current block. + for local in written_only_inside_own_block_locals.drain() { + debug_assert_eq!( + self.ecx.machine.can_const_prop[local], + ConstPropMode::OnlyInsideOwnBlock + ); + Self::remove_const(&mut self.ecx, local); + } + self.ecx.machine.written_only_inside_own_block_locals = + written_only_inside_own_block_locals; + + if cfg!(debug_assertions) { + for (local, &mode) in self.ecx.machine.can_const_prop.iter_enumerated() { + match mode { + ConstPropMode::FullConstProp => {} + ConstPropMode::NoPropagation | ConstPropMode::OnlyInsideOwnBlock => { + self.ensure_not_propagated(local); + } + } + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/const_prop_lint.rs b/compiler/rustc_mir_transform/src/const_prop_lint.rs new file mode 100644 index 00000000000..99eecb567f2 --- /dev/null +++ b/compiler/rustc_mir_transform/src/const_prop_lint.rs @@ -0,0 +1,694 @@ +//! Propagates constants for early reporting of statically known +//! assertion failures + +use std::fmt::Debug; + +use either::Left; + +use rustc_const_eval::interpret::Immediate; +use rustc_const_eval::interpret::{ + InterpCx, InterpResult, MemoryKind, OpTy, Scalar, StackPopCleanup, +}; +use rustc_const_eval::ReportErrorExt; +use rustc_hir::def::DefKind; +use rustc_hir::HirId; +use rustc_index::bit_set::BitSet; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::ty::layout::{LayoutError, LayoutOf, LayoutOfHelpers, TyAndLayout}; +use rustc_middle::ty::GenericArgs; +use rustc_middle::ty::{ + self, ConstInt, Instance, ParamEnv, ScalarInt, Ty, TyCtxt, TypeVisitableExt, +}; +use rustc_span::Span; +use rustc_target::abi::{HasDataLayout, Size, TargetDataLayout}; + +use crate::const_prop::CanConstProp; +use crate::const_prop::ConstPropMachine; +use crate::const_prop::ConstPropMode; +use crate::errors::AssertLint; +use crate::MirLint; + +/// The maximum number of bytes that we'll allocate space for a local or the return value. +/// Needed for #66397, because otherwise we eval into large places and that can cause OOM or just +/// Severely regress performance. +const MAX_ALLOC_LIMIT: u64 = 1024; + +pub struct ConstPropLint; + +impl<'tcx> MirLint<'tcx> for ConstPropLint { + fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) { + if body.tainted_by_errors.is_some() { + return; + } + + // will be evaluated by miri and produce its errors there + if body.source.promoted.is_some() { + return; + } + + let def_id = body.source.def_id().expect_local(); + let def_kind = tcx.def_kind(def_id); + let is_fn_like = def_kind.is_fn_like(); + let is_assoc_const = def_kind == DefKind::AssocConst; + + // Only run const prop on functions, methods, closures and associated constants + if !is_fn_like && !is_assoc_const { + // skip anon_const/statics/consts because they'll be evaluated by miri anyway + trace!("ConstPropLint skipped for {:?}", def_id); + return; + } + + // FIXME(welseywiser) const prop doesn't work on coroutines because of query cycles + // computing their layout. + if tcx.is_coroutine(def_id.to_def_id()) { + trace!("ConstPropLint skipped for coroutine {:?}", def_id); + return; + } + + trace!("ConstPropLint starting for {:?}", def_id); + + // FIXME(oli-obk, eddyb) Optimize locals (or even local paths) to hold + // constants, instead of just checking for const-folding succeeding. + // That would require a uniform one-def no-mutation analysis + // and RPO (or recursing when needing the value of a local). + let mut linter = ConstPropagator::new(body, tcx); + linter.visit_body(body); + + trace!("ConstPropLint done for {:?}", def_id); + } +} + +/// Finds optimization opportunities on the MIR. +struct ConstPropagator<'mir, 'tcx> { + ecx: InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + worklist: Vec<BasicBlock>, + visited_blocks: BitSet<BasicBlock>, +} + +impl<'tcx> LayoutOfHelpers<'tcx> for ConstPropagator<'_, 'tcx> { + type LayoutOfResult = Result<TyAndLayout<'tcx>, LayoutError<'tcx>>; + + #[inline] + fn handle_layout_err(&self, err: LayoutError<'tcx>, _: Span, _: Ty<'tcx>) -> LayoutError<'tcx> { + err + } +} + +impl HasDataLayout for ConstPropagator<'_, '_> { + #[inline] + fn data_layout(&self) -> &TargetDataLayout { + &self.tcx.data_layout + } +} + +impl<'tcx> ty::layout::HasTyCtxt<'tcx> for ConstPropagator<'_, 'tcx> { + #[inline] + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } +} + +impl<'tcx> ty::layout::HasParamEnv<'tcx> for ConstPropagator<'_, 'tcx> { + #[inline] + fn param_env(&self) -> ty::ParamEnv<'tcx> { + self.param_env + } +} + +impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> { + fn new(body: &'mir Body<'tcx>, tcx: TyCtxt<'tcx>) -> ConstPropagator<'mir, 'tcx> { + let def_id = body.source.def_id(); + let args = &GenericArgs::identity_for_item(tcx, def_id); + let param_env = tcx.param_env_reveal_all_normalized(def_id); + + let can_const_prop = CanConstProp::check(tcx, param_env, body); + let mut ecx = InterpCx::new( + tcx, + tcx.def_span(def_id), + param_env, + ConstPropMachine::new(can_const_prop), + ); + + let ret_layout = ecx + .layout_of(body.bound_return_ty().instantiate(tcx, args)) + .ok() + // Don't bother allocating memory for large values. + // I don't know how return types can seem to be unsized but this happens in the + // `type/type-unsatisfiable.rs` test. + .filter(|ret_layout| { + ret_layout.is_sized() && ret_layout.size < Size::from_bytes(MAX_ALLOC_LIMIT) + }) + .unwrap_or_else(|| ecx.layout_of(tcx.types.unit).unwrap()); + + let ret = ecx + .allocate(ret_layout, MemoryKind::Stack) + .expect("couldn't perform small allocation") + .into(); + + ecx.push_stack_frame( + Instance::new(def_id, args), + body, + &ret, + StackPopCleanup::Root { cleanup: false }, + ) + .expect("failed to push initial stack frame"); + + for local in body.local_decls.indices() { + // Mark everything initially live. + // This is somewhat dicey since some of them might be unsized and it is incoherent to + // mark those as live... We rely on `local_to_place`/`local_to_op` in the interpreter + // stopping us before those unsized immediates can cause issues deeper in the + // interpreter. + ecx.frame_mut().locals[local].make_live_uninit(); + } + + ConstPropagator { + ecx, + tcx, + param_env, + worklist: vec![START_BLOCK], + visited_blocks: BitSet::new_empty(body.basic_blocks.len()), + } + } + + fn body(&self) -> &'mir Body<'tcx> { + self.ecx.frame().body + } + + fn local_decls(&self) -> &'mir LocalDecls<'tcx> { + &self.body().local_decls + } + + fn get_const(&self, place: Place<'tcx>) -> Option<OpTy<'tcx>> { + let op = match self.ecx.eval_place_to_op(place, None) { + Ok(op) => { + if op + .as_mplace_or_imm() + .right() + .is_some_and(|imm| matches!(*imm, Immediate::Uninit)) + { + // Make sure nobody accidentally uses this value. + return None; + } + op + } + Err(e) => { + trace!("get_const failed: {:?}", e.into_kind().debug()); + return None; + } + }; + + // Try to read the local as an immediate so that if it is representable as a scalar, we can + // handle it as such, but otherwise, just return the value as is. + Some(match self.ecx.read_immediate_raw(&op) { + Ok(Left(imm)) => imm.into(), + _ => op, + }) + } + + /// Remove `local` from the pool of `Locals`. Allows writing to them, + /// but not reading from them anymore. + fn remove_const(ecx: &mut InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, local: Local) { + ecx.frame_mut().locals[local].make_live_uninit(); + ecx.machine.written_only_inside_own_block_locals.remove(&local); + } + + fn lint_root(&self, source_info: SourceInfo) -> Option<HirId> { + source_info.scope.lint_root(&self.body().source_scopes) + } + + fn use_ecx<F, T>(&mut self, location: Location, f: F) -> Option<T> + where + F: FnOnce(&mut Self) -> InterpResult<'tcx, T>, + { + // Overwrite the PC -- whatever the interpreter does to it does not make any sense anyway. + self.ecx.frame_mut().loc = Left(location); + match f(self) { + Ok(val) => Some(val), + Err(error) => { + trace!("InterpCx operation failed: {:?}", error); + // Some errors shouldn't come up because creating them causes + // an allocation, which we should avoid. When that happens, + // dedicated error variants should be introduced instead. + assert!( + !error.kind().formatted_string(), + "const-prop encountered formatting error: {}", + self.ecx.format_error(error), + ); + None + } + } + } + + /// Returns the value, if any, of evaluating `c`. + fn eval_constant(&mut self, c: &ConstOperand<'tcx>, location: Location) -> Option<OpTy<'tcx>> { + // FIXME we need to revisit this for #67176 + if c.has_param() { + return None; + } + + // Normalization needed b/c const prop lint runs in + // `mir_drops_elaborated_and_const_checked`, which happens before + // optimized MIR. Only after optimizing the MIR can we guarantee + // that the `RevealAll` pass has happened and that the body's consts + // are normalized, so any call to resolve before that needs to be + // manually normalized. + let val = self.tcx.try_normalize_erasing_regions(self.param_env, c.const_).ok()?; + + self.use_ecx(location, |this| this.ecx.eval_mir_constant(&val, Some(c.span), None)) + } + + /// Returns the value, if any, of evaluating `place`. + fn eval_place(&mut self, place: Place<'tcx>, location: Location) -> Option<OpTy<'tcx>> { + trace!("eval_place(place={:?})", place); + self.use_ecx(location, |this| this.ecx.eval_place_to_op(place, None)) + } + + /// Returns the value, if any, of evaluating `op`. Calls upon `eval_constant` + /// or `eval_place`, depending on the variant of `Operand` used. + fn eval_operand(&mut self, op: &Operand<'tcx>, location: Location) -> Option<OpTy<'tcx>> { + match *op { + Operand::Constant(ref c) => self.eval_constant(c, location), + Operand::Move(place) | Operand::Copy(place) => self.eval_place(place, location), + } + } + + fn report_assert_as_lint(&self, source_info: &SourceInfo, lint: AssertLint<impl Debug>) { + if let Some(lint_root) = self.lint_root(*source_info) { + self.tcx.emit_spanned_lint(lint.lint(), lint_root, source_info.span, lint); + } + } + + fn check_unary_op(&mut self, op: UnOp, arg: &Operand<'tcx>, location: Location) -> Option<()> { + if let (val, true) = self.use_ecx(location, |this| { + let val = this.ecx.read_immediate(&this.ecx.eval_operand(arg, None)?)?; + let (_res, overflow) = this.ecx.overflowing_unary_op(op, &val)?; + Ok((val, overflow)) + })? { + // `AssertKind` only has an `OverflowNeg` variant, so make sure that is + // appropriate to use. + assert_eq!(op, UnOp::Neg, "Neg is the only UnOp that can overflow"); + let source_info = self.body().source_info(location); + self.report_assert_as_lint( + source_info, + AssertLint::ArithmeticOverflow( + source_info.span, + AssertKind::OverflowNeg(val.to_const_int()), + ), + ); + return None; + } + + Some(()) + } + + fn check_binary_op( + &mut self, + op: BinOp, + left: &Operand<'tcx>, + right: &Operand<'tcx>, + location: Location, + ) -> Option<()> { + let r = self.use_ecx(location, |this| { + this.ecx.read_immediate(&this.ecx.eval_operand(right, None)?) + }); + let l = self + .use_ecx(location, |this| this.ecx.read_immediate(&this.ecx.eval_operand(left, None)?)); + // Check for exceeding shifts *even if* we cannot evaluate the LHS. + if matches!(op, BinOp::Shr | BinOp::Shl) { + let r = r.clone()?; + // We need the type of the LHS. We cannot use `place_layout` as that is the type + // of the result, which for checked binops is not the same! + let left_ty = left.ty(self.local_decls(), self.tcx); + let left_size = self.ecx.layout_of(left_ty).ok()?.size; + let right_size = r.layout.size; + let r_bits = r.to_scalar().to_bits(right_size).ok(); + if r_bits.is_some_and(|b| b >= left_size.bits() as u128) { + debug!("check_binary_op: reporting assert for {:?}", location); + let source_info = self.body().source_info(location); + let panic = AssertKind::Overflow( + op, + match l { + Some(l) => l.to_const_int(), + // Invent a dummy value, the diagnostic ignores it anyway + None => ConstInt::new( + ScalarInt::try_from_uint(1_u8, left_size).unwrap(), + left_ty.is_signed(), + left_ty.is_ptr_sized_integral(), + ), + }, + r.to_const_int(), + ); + self.report_assert_as_lint( + source_info, + AssertLint::ArithmeticOverflow(source_info.span, panic), + ); + return None; + } + } + + if let (Some(l), Some(r)) = (l, r) { + // The remaining operators are handled through `overflowing_binary_op`. + if self.use_ecx(location, |this| { + let (_res, overflow) = this.ecx.overflowing_binary_op(op, &l, &r)?; + Ok(overflow) + })? { + let source_info = self.body().source_info(location); + self.report_assert_as_lint( + source_info, + AssertLint::ArithmeticOverflow( + source_info.span, + AssertKind::Overflow(op, l.to_const_int(), r.to_const_int()), + ), + ); + return None; + } + } + Some(()) + } + + fn check_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) -> Option<()> { + // Perform any special handling for specific Rvalue types. + // Generally, checks here fall into one of two categories: + // 1. Additional checking to provide useful lints to the user + // - In this case, we will do some validation and then fall through to the + // end of the function which evals the assignment. + // 2. Working around bugs in other parts of the compiler + // - In this case, we'll return `None` from this function to stop evaluation. + match rvalue { + // Additional checking: give lints to the user if an overflow would occur. + // We do this here and not in the `Assert` terminator as that terminator is + // only sometimes emitted (overflow checks can be disabled), but we want to always + // lint. + Rvalue::UnaryOp(op, arg) => { + trace!("checking UnaryOp(op = {:?}, arg = {:?})", op, arg); + self.check_unary_op(*op, arg, location)?; + } + Rvalue::BinaryOp(op, box (left, right)) => { + trace!("checking BinaryOp(op = {:?}, left = {:?}, right = {:?})", op, left, right); + self.check_binary_op(*op, left, right, location)?; + } + Rvalue::CheckedBinaryOp(op, box (left, right)) => { + trace!( + "checking CheckedBinaryOp(op = {:?}, left = {:?}, right = {:?})", + op, + left, + right + ); + self.check_binary_op(*op, left, right, location)?; + } + + // Do not try creating references (#67862) + Rvalue::AddressOf(_, place) | Rvalue::Ref(_, _, place) => { + trace!("skipping AddressOf | Ref for {:?}", place); + + // This may be creating mutable references or immutable references to cells. + // If that happens, the pointed to value could be mutated via that reference. + // Since we aren't tracking references, the const propagator loses track of what + // value the local has right now. + // Thus, all locals that have their reference taken + // must not take part in propagation. + Self::remove_const(&mut self.ecx, place.local); + + return None; + } + Rvalue::ThreadLocalRef(def_id) => { + trace!("skipping ThreadLocalRef({:?})", def_id); + + return None; + } + + // There's no other checking to do at this time. + Rvalue::Aggregate(..) + | Rvalue::Use(..) + | Rvalue::CopyForDeref(..) + | Rvalue::Repeat(..) + | Rvalue::Len(..) + | Rvalue::Cast(..) + | Rvalue::ShallowInitBox(..) + | Rvalue::Discriminant(..) + | Rvalue::NullaryOp(..) => {} + } + + // FIXME we need to revisit this for #67176 + if rvalue.has_param() { + return None; + } + if !rvalue.ty(self.local_decls(), self.tcx).is_sized(self.tcx, self.param_env) { + // the interpreter doesn't support unsized locals (only unsized arguments), + // but rustc does (in a kinda broken way), so we have to skip them here + return None; + } + + Some(()) + } + + fn check_assertion( + &mut self, + expected: bool, + msg: &AssertKind<Operand<'tcx>>, + cond: &Operand<'tcx>, + location: Location, + ) -> Option<!> { + let value = &self.eval_operand(cond, location)?; + trace!("assertion on {:?} should be {:?}", value, expected); + + let expected = Scalar::from_bool(expected); + let value_const = self.use_ecx(location, |this| this.ecx.read_scalar(value))?; + + if expected != value_const { + // Poison all places this operand references so that further code + // doesn't use the invalid value + if let Some(place) = cond.place() { + Self::remove_const(&mut self.ecx, place.local); + } + + enum DbgVal<T> { + Val(T), + Underscore, + } + impl<T: std::fmt::Debug> std::fmt::Debug for DbgVal<T> { + fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Val(val) => val.fmt(fmt), + Self::Underscore => fmt.write_str("_"), + } + } + } + let mut eval_to_int = |op| { + // This can be `None` if the lhs wasn't const propagated and we just + // triggered the assert on the value of the rhs. + self.eval_operand(op, location) + .and_then(|op| self.ecx.read_immediate(&op).ok()) + .map_or(DbgVal::Underscore, |op| DbgVal::Val(op.to_const_int())) + }; + let msg = match msg { + AssertKind::DivisionByZero(op) => AssertKind::DivisionByZero(eval_to_int(op)), + AssertKind::RemainderByZero(op) => AssertKind::RemainderByZero(eval_to_int(op)), + AssertKind::Overflow(bin_op @ (BinOp::Div | BinOp::Rem), op1, op2) => { + // Division overflow is *UB* in the MIR, and different than the + // other overflow checks. + AssertKind::Overflow(*bin_op, eval_to_int(op1), eval_to_int(op2)) + } + AssertKind::BoundsCheck { ref len, ref index } => { + let len = eval_to_int(len); + let index = eval_to_int(index); + AssertKind::BoundsCheck { len, index } + } + // Remaining overflow errors are already covered by checks on the binary operators. + AssertKind::Overflow(..) | AssertKind::OverflowNeg(_) => return None, + // Need proper const propagator for these. + _ => return None, + }; + let source_info = self.body().source_info(location); + self.report_assert_as_lint( + source_info, + AssertLint::UnconditionalPanic(source_info.span, msg), + ); + } + + None + } + + fn ensure_not_propagated(&self, local: Local) { + if cfg!(debug_assertions) { + assert!( + self.get_const(local.into()).is_none() + || self + .layout_of(self.local_decls()[local].ty) + .map_or(true, |layout| layout.is_zst()), + "failed to remove values for `{local:?}`, value={:?}", + self.get_const(local.into()), + ) + } + } +} + +impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> { + fn visit_body(&mut self, body: &Body<'tcx>) { + while let Some(bb) = self.worklist.pop() { + if !self.visited_blocks.insert(bb) { + continue; + } + + let data = &body.basic_blocks[bb]; + self.visit_basic_block_data(bb, data); + } + } + + fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) { + self.super_operand(operand, location); + } + + fn visit_constant(&mut self, constant: &ConstOperand<'tcx>, location: Location) { + trace!("visit_constant: {:?}", constant); + self.super_constant(constant, location); + self.eval_constant(constant, location); + } + + fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, location: Location) { + self.super_assign(place, rvalue, location); + + let Some(()) = self.check_rvalue(rvalue, location) else { return }; + + match self.ecx.machine.can_const_prop[place.local] { + // Do nothing if the place is indirect. + _ if place.is_indirect() => {} + ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local), + ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => { + if self + .use_ecx(location, |this| this.ecx.eval_rvalue_into_place(rvalue, *place)) + .is_none() + { + // Const prop failed, so erase the destination, ensuring that whatever happens + // from here on, does not know about the previous value. + // This is important in case we have + // ```rust + // let mut x = 42; + // x = SOME_MUTABLE_STATIC; + // // x must now be uninit + // ``` + // FIXME: we overzealously erase the entire local, because that's easier to + // implement. + trace!( + "propagation into {:?} failed. + Nuking the entire site from orbit, it's the only way to be sure", + place, + ); + Self::remove_const(&mut self.ecx, place.local); + } + } + } + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + trace!("visit_statement: {:?}", statement); + + // We want to evaluate operands before any change to the assigned-to value, + // so we recurse first. + self.super_statement(statement, location); + + match statement.kind { + StatementKind::SetDiscriminant { ref place, .. } => { + match self.ecx.machine.can_const_prop[place.local] { + // Do nothing if the place is indirect. + _ if place.is_indirect() => {} + ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local), + ConstPropMode::FullConstProp | ConstPropMode::OnlyInsideOwnBlock => { + if self.use_ecx(location, |this| this.ecx.statement(statement)).is_some() { + trace!("propped discriminant into {:?}", place); + } else { + Self::remove_const(&mut self.ecx, place.local); + } + } + } + } + StatementKind::StorageLive(local) => { + let frame = self.ecx.frame_mut(); + frame.locals[local].make_live_uninit(); + } + StatementKind::StorageDead(local) => { + let frame = self.ecx.frame_mut(); + // We don't actually track liveness, so the local remains live. But forget its value. + frame.locals[local].make_live_uninit(); + } + _ => {} + } + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + self.super_terminator(terminator, location); + match &terminator.kind { + TerminatorKind::Assert { expected, ref msg, ref cond, .. } => { + self.check_assertion(*expected, msg, cond, location); + } + TerminatorKind::SwitchInt { ref discr, ref targets } => { + if let Some(ref value) = self.eval_operand(discr, location) + && let Some(value_const) = + self.use_ecx(location, |this| this.ecx.read_scalar(value)) + && let Ok(constant) = value_const.try_to_int() + && let Ok(constant) = constant.to_bits(constant.size()) + { + // We managed to evaluate the discriminant, so we know we only need to visit + // one target. + let target = targets.target_for_value(constant); + self.worklist.push(target); + return; + } + // We failed to evaluate the discriminant, fallback to visiting all successors. + } + // None of these have Operands to const-propagate. + TerminatorKind::Goto { .. } + | TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::Drop { .. } + | TerminatorKind::Yield { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::Call { .. } + | TerminatorKind::InlineAsm { .. } => {} + } + + self.worklist.extend(terminator.successors()); + } + + fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) { + self.super_basic_block_data(block, data); + + // We remove all Locals which are restricted in propagation to their containing blocks and + // which were modified in the current block. + // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`. + let mut written_only_inside_own_block_locals = + std::mem::take(&mut self.ecx.machine.written_only_inside_own_block_locals); + + // This loop can get very hot for some bodies: it check each local in each bb. + // To avoid this quadratic behaviour, we only clear the locals that were modified inside + // the current block. + for local in written_only_inside_own_block_locals.drain() { + debug_assert_eq!( + self.ecx.machine.can_const_prop[local], + ConstPropMode::OnlyInsideOwnBlock + ); + Self::remove_const(&mut self.ecx, local); + } + self.ecx.machine.written_only_inside_own_block_locals = + written_only_inside_own_block_locals; + + if cfg!(debug_assertions) { + for (local, &mode) in self.ecx.machine.can_const_prop.iter_enumerated() { + match mode { + ConstPropMode::FullConstProp => {} + ConstPropMode::NoPropagation | ConstPropMode::OnlyInsideOwnBlock => { + self.ensure_not_propagated(local); + } + } + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/copy_prop.rs b/compiler/rustc_mir_transform/src/copy_prop.rs new file mode 100644 index 00000000000..0119b95cced --- /dev/null +++ b/compiler/rustc_mir_transform/src/copy_prop.rs @@ -0,0 +1,184 @@ +use rustc_index::bit_set::BitSet; +use rustc_index::IndexSlice; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_mir_dataflow::impls::borrowed_locals; + +use crate::ssa::SsaLocals; + +/// Unify locals that copy each other. +/// +/// We consider patterns of the form +/// _a = rvalue +/// _b = move? _a +/// _c = move? _a +/// _d = move? _c +/// where each of the locals is only assigned once. +/// +/// We want to replace all those locals by `_a`, either copied or moved. +pub struct CopyProp; + +impl<'tcx> MirPass<'tcx> for CopyProp { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 1 + } + + #[instrument(level = "trace", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + propagate_ssa(tcx, body); + } +} + +fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let borrowed_locals = borrowed_locals(body); + let ssa = SsaLocals::new(body); + + let fully_moved = fully_moved_locals(&ssa, body); + debug!(?fully_moved); + + let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size()); + for (local, &head) in ssa.copy_classes().iter_enumerated() { + if local != head { + storage_to_remove.insert(head); + } + } + + let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h); + + Replacer { + tcx, + copy_classes: ssa.copy_classes(), + fully_moved, + borrowed_locals, + storage_to_remove, + } + .visit_body_preserves_cfg(body); + + if any_replacement { + crate::simplify::remove_unused_definitions(body); + } +} + +/// `SsaLocals` computed equivalence classes between locals considering copy/move assignments. +/// +/// This function also returns whether all the `move?` in the pattern are `move` and not copies. +/// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be +/// replaced by `copy _a`, as we cannot move multiple times from `_a`. +/// +/// If an operand copies `_c`, it must happen before the assignment `_d = _c`, otherwise it is UB. +/// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is +/// moved, and therefore that `_d` is moved. +#[instrument(level = "trace", skip(ssa, body))] +fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> { + let mut fully_moved = BitSet::new_filled(body.local_decls.len()); + + for (_, rvalue, _) in ssa.assignments(body) { + let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) + | Rvalue::CopyForDeref(place)) = rvalue + else { + continue; + }; + + let Some(rhs) = place.as_local() else { continue }; + if !ssa.is_ssa(rhs) { + continue; + } + + if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue { + fully_moved.remove(rhs); + } + } + + ssa.meet_copy_equivalence(&mut fully_moved); + + fully_moved +} + +/// Utility to help performing substitution of `*pattern` by `target`. +struct Replacer<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + fully_moved: BitSet<Local>, + storage_to_remove: BitSet<Local>, + borrowed_locals: BitSet<Local>, + copy_classes: &'a IndexSlice<Local, Local>, +} + +impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) { + let new_local = self.copy_classes[*local]; + match ctxt { + // Do not modify the local in storage statements. + PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {} + // The local should have been marked as non-SSA. + PlaceContext::MutatingUse(_) => assert_eq!(*local, new_local), + // We access the value. + _ => *local = new_local, + } + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) { + if let Some(new_projection) = self.process_projection(place.projection, loc) { + place.projection = self.tcx().mk_place_elems(&new_projection); + } + + let observes_address = match ctxt { + PlaceContext::NonMutatingUse( + NonMutatingUseContext::SharedBorrow + | NonMutatingUseContext::FakeBorrow + | NonMutatingUseContext::AddressOf, + ) => true, + // For debuginfo, merging locals is ok. + PlaceContext::NonUse(NonUseContext::VarDebugInfo) => { + self.borrowed_locals.contains(place.local) + } + _ => false, + }; + if observes_address && !place.is_indirect() { + // We observe the address of `place.local`. Do not replace it. + } else { + self.visit_local( + &mut place.local, + PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy), + loc, + ) + } + } + + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) { + if let Operand::Move(place) = *operand + // A move out of a projection of a copy is equivalent to a copy of the original projection. + && !place.is_indirect_first_projection() + && !self.fully_moved.contains(place.local) + { + *operand = Operand::Copy(place); + } + self.super_operand(operand, loc); + } + + fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) { + // When removing storage statements, we need to remove both (#107511). + if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind + && self.storage_to_remove.contains(l) + { + stmt.make_nop(); + return; + } + + self.super_statement(stmt, loc); + + // Do not leave tautological assignments around. + if let StatementKind::Assign(box (lhs, ref rhs)) = stmt.kind + && let Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)) | Rvalue::CopyForDeref(rhs) = + *rhs + && lhs == rhs + { + stmt.make_nop(); + } + } +} diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs new file mode 100644 index 00000000000..ed8b4ef3ba4 --- /dev/null +++ b/compiler/rustc_mir_transform/src/coroutine.rs @@ -0,0 +1,2083 @@ +//! This is the implementation of the pass which transforms coroutines into state machines. +//! +//! MIR generation for coroutines creates a function which has a self argument which +//! passes by value. This argument is effectively a coroutine type which only contains upvars and +//! is only used for this argument inside the MIR for the coroutine. +//! It is passed by value to enable upvars to be moved out of it. Drop elaboration runs on that +//! MIR before this pass and creates drop flags for MIR locals. +//! It will also drop the coroutine argument (which only consists of upvars) if any of the upvars +//! are moved out of. This pass elaborates the drops of upvars / coroutine argument in the case +//! that none of the upvars were moved out of. This is because we cannot have any drops of this +//! coroutine in the MIR, since it is used to create the drop glue for the coroutine. We'd get +//! infinite recursion otherwise. +//! +//! This pass creates the implementation for either the `Coroutine::resume` or `Future::poll` +//! function and the drop shim for the coroutine based on the MIR input. +//! It converts the coroutine argument from Self to &mut Self adding derefs in the MIR as needed. +//! It computes the final layout of the coroutine struct which looks like this: +//! First upvars are stored +//! It is followed by the coroutine state field. +//! Then finally the MIR locals which are live across a suspension point are stored. +//! ```ignore (illustrative) +//! struct Coroutine { +//! upvars..., +//! state: u32, +//! mir_locals..., +//! } +//! ``` +//! This pass computes the meaning of the state field and the MIR locals which are live +//! across a suspension point. There are however three hardcoded coroutine states: +//! 0 - Coroutine have not been resumed yet +//! 1 - Coroutine has returned / is completed +//! 2 - Coroutine has been poisoned +//! +//! It also rewrites `return x` and `yield y` as setting a new coroutine state and returning +//! `CoroutineState::Complete(x)` and `CoroutineState::Yielded(y)`, +//! or `Poll::Ready(x)` and `Poll::Pending` respectively. +//! MIR locals which are live across a suspension point are moved to the coroutine struct +//! with references to them being updated with references to the coroutine struct. +//! +//! The pass creates two functions which have a switch on the coroutine state giving +//! the action to take. +//! +//! One of them is the implementation of `Coroutine::resume` / `Future::poll`. +//! For coroutines with state 0 (unresumed) it starts the execution of the coroutine. +//! For coroutines with state 1 (returned) and state 2 (poisoned) it panics. +//! Otherwise it continues the execution from the last suspension point. +//! +//! The other function is the drop glue for the coroutine. +//! For coroutines with state 0 (unresumed) it drops the upvars of the coroutine. +//! For coroutines with state 1 (returned) and state 2 (poisoned) it does nothing. +//! Otherwise it drops all the values in scope at the last suspension point. + +use crate::abort_unwinding_calls; +use crate::deref_separator::deref_finder; +use crate::errors; +use crate::pass_manager as pm; +use crate::simplify; +use rustc_data_structures::fx::{FxHashMap, FxHashSet}; +use rustc_errors::pluralize; +use rustc_hir as hir; +use rustc_hir::lang_items::LangItem; +use rustc_hir::CoroutineKind; +use rustc_index::bit_set::{BitMatrix, BitSet, GrowableBitSet}; +use rustc_index::{Idx, IndexVec}; +use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::CoroutineArgs; +use rustc_middle::ty::InstanceDef; +use rustc_middle::ty::{self, Ty, TyCtxt}; +use rustc_mir_dataflow::impls::{ + MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive, +}; +use rustc_mir_dataflow::storage::always_storage_live_locals; +use rustc_mir_dataflow::Analysis; +use rustc_span::def_id::{DefId, LocalDefId}; +use rustc_span::symbol::sym; +use rustc_span::Span; +use rustc_target::abi::{FieldIdx, VariantIdx}; +use rustc_target::spec::PanicStrategy; +use std::{iter, ops}; + +pub struct StateTransform; + +struct RenameLocalVisitor<'tcx> { + from: Local, + to: Local, + tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { + if *local == self.from { + *local = self.to; + } + } + + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) { + match terminator.kind { + TerminatorKind::Return => { + // Do not replace the implicit `_0` access here, as that's not possible. The + // transform already handles `return` correctly. + } + _ => self.super_terminator(terminator, location), + } + } +} + +struct DerefArgVisitor<'tcx> { + tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { + assert_ne!(*local, SELF_ARG); + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) { + if place.local == SELF_ARG { + replace_base( + place, + Place { + local: SELF_ARG, + projection: self.tcx().mk_place_elems(&[ProjectionElem::Deref]), + }, + self.tcx, + ); + } else { + self.visit_local(&mut place.local, context, location); + + for elem in place.projection.iter() { + if let PlaceElem::Index(local) = elem { + assert_ne!(local, SELF_ARG); + } + } + } + } +} + +struct PinArgVisitor<'tcx> { + ref_coroutine_ty: Ty<'tcx>, + tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { + assert_ne!(*local, SELF_ARG); + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) { + if place.local == SELF_ARG { + replace_base( + place, + Place { + local: SELF_ARG, + projection: self.tcx().mk_place_elems(&[ProjectionElem::Field( + FieldIdx::new(0), + self.ref_coroutine_ty, + )]), + }, + self.tcx, + ); + } else { + self.visit_local(&mut place.local, context, location); + + for elem in place.projection.iter() { + if let PlaceElem::Index(local) = elem { + assert_ne!(local, SELF_ARG); + } + } + } + } +} + +fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) { + place.local = new_base.local; + + let mut new_projection = new_base.projection.to_vec(); + new_projection.append(&mut place.projection.to_vec()); + + place.projection = tcx.mk_place_elems(&new_projection); +} + +const SELF_ARG: Local = Local::from_u32(1); + +/// Coroutine has not been resumed yet. +const UNRESUMED: usize = CoroutineArgs::UNRESUMED; +/// Coroutine has returned / is completed. +const RETURNED: usize = CoroutineArgs::RETURNED; +/// Coroutine has panicked and is poisoned. +const POISONED: usize = CoroutineArgs::POISONED; + +/// Number of variants to reserve in coroutine state. Corresponds to +/// `UNRESUMED` (beginning of a coroutine) and `RETURNED`/`POISONED` +/// (end of a coroutine) states. +const RESERVED_VARIANTS: usize = 3; + +/// A `yield` point in the coroutine. +struct SuspensionPoint<'tcx> { + /// State discriminant used when suspending or resuming at this point. + state: usize, + /// The block to jump to after resumption. + resume: BasicBlock, + /// Where to move the resume argument after resumption. + resume_arg: Place<'tcx>, + /// Which block to jump to if the coroutine is dropped in this state. + drop: Option<BasicBlock>, + /// Set of locals that have live storage while at this suspension point. + storage_liveness: GrowableBitSet<Local>, +} + +struct TransformVisitor<'tcx> { + tcx: TyCtxt<'tcx>, + coroutine_kind: hir::CoroutineKind, + + // The type of the discriminant in the coroutine struct + discr_ty: Ty<'tcx>, + + // Mapping from Local to (type of local, coroutine struct index) + // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`. + remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, FieldIdx)>, + + // A map from a suspension point in a block to the locals which have live storage at that point + storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>, + + // A list of suspension points, generated during the transform + suspension_points: Vec<SuspensionPoint<'tcx>>, + + // The set of locals that have no `StorageLive`/`StorageDead` annotations. + always_live_locals: BitSet<Local>, + + // The original RETURN_PLACE local + old_ret_local: Local, + + old_yield_ty: Ty<'tcx>, + + old_ret_ty: Ty<'tcx>, +} + +impl<'tcx> TransformVisitor<'tcx> { + fn insert_none_ret_block(&self, body: &mut Body<'tcx>) -> BasicBlock { + let block = BasicBlock::new(body.basic_blocks.len()); + let source_info = SourceInfo::outermost(body.span); + + let none_value = match self.coroutine_kind { + CoroutineKind::Async(_) => span_bug!(body.span, "`Future`s are not fused inherently"), + CoroutineKind::Coroutine => span_bug!(body.span, "`Coroutine`s cannot be fused"), + // `gen` continues return `None` + CoroutineKind::Gen(_) => { + let option_def_id = self.tcx.require_lang_item(LangItem::Option, None); + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + option_def_id, + VariantIdx::from_usize(0), + self.tcx.mk_args(&[self.old_yield_ty.into()]), + None, + None, + )), + IndexVec::new(), + ) + } + // `async gen` continues to return `Poll::Ready(None)` + CoroutineKind::AsyncGen(_) => { + let ty::Adt(_poll_adt, args) = *self.old_yield_ty.kind() else { bug!() }; + let ty::Adt(_option_adt, args) = *args.type_at(0).kind() else { bug!() }; + let yield_ty = args.type_at(0); + Rvalue::Use(Operand::Constant(Box::new(ConstOperand { + span: source_info.span, + const_: Const::Unevaluated( + UnevaluatedConst::new( + self.tcx.require_lang_item(LangItem::AsyncGenFinished, None), + self.tcx.mk_args(&[yield_ty.into()]), + ), + self.old_yield_ty, + ), + user_ty: None, + }))) + } + }; + + let statements = vec![Statement { + kind: StatementKind::Assign(Box::new((Place::return_place(), none_value))), + source_info, + }]; + + body.basic_blocks_mut().push(BasicBlockData { + statements, + terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }), + is_cleanup: false, + }); + + block + } + + // Make a `CoroutineState` or `Poll` variant assignment. + // + // `core::ops::CoroutineState` only has single element tuple variants, + // so we can just write to the downcasted first field and then set the + // discriminant to the appropriate variant. + fn make_state( + &self, + val: Operand<'tcx>, + source_info: SourceInfo, + is_return: bool, + statements: &mut Vec<Statement<'tcx>>, + ) { + let rvalue = match self.coroutine_kind { + CoroutineKind::Async(_) => { + let poll_def_id = self.tcx.require_lang_item(LangItem::Poll, None); + let args = self.tcx.mk_args(&[self.old_ret_ty.into()]); + if is_return { + // Poll::Ready(val) + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + poll_def_id, + VariantIdx::from_usize(0), + args, + None, + None, + )), + IndexVec::from_raw(vec![val]), + ) + } else { + // Poll::Pending + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + poll_def_id, + VariantIdx::from_usize(1), + args, + None, + None, + )), + IndexVec::new(), + ) + } + } + CoroutineKind::Gen(_) => { + let option_def_id = self.tcx.require_lang_item(LangItem::Option, None); + let args = self.tcx.mk_args(&[self.old_yield_ty.into()]); + if is_return { + // None + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + option_def_id, + VariantIdx::from_usize(0), + args, + None, + None, + )), + IndexVec::new(), + ) + } else { + // Some(val) + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + option_def_id, + VariantIdx::from_usize(1), + args, + None, + None, + )), + IndexVec::from_raw(vec![val]), + ) + } + } + CoroutineKind::AsyncGen(_) => { + if is_return { + let ty::Adt(_poll_adt, args) = *self.old_yield_ty.kind() else { bug!() }; + let ty::Adt(_option_adt, args) = *args.type_at(0).kind() else { bug!() }; + let yield_ty = args.type_at(0); + Rvalue::Use(Operand::Constant(Box::new(ConstOperand { + span: source_info.span, + const_: Const::Unevaluated( + UnevaluatedConst::new( + self.tcx.require_lang_item(LangItem::AsyncGenFinished, None), + self.tcx.mk_args(&[yield_ty.into()]), + ), + self.old_yield_ty, + ), + user_ty: None, + }))) + } else { + Rvalue::Use(val) + } + } + CoroutineKind::Coroutine => { + let coroutine_state_def_id = + self.tcx.require_lang_item(LangItem::CoroutineState, None); + let args = self.tcx.mk_args(&[self.old_yield_ty.into(), self.old_ret_ty.into()]); + if is_return { + // CoroutineState::Complete(val) + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + coroutine_state_def_id, + VariantIdx::from_usize(1), + args, + None, + None, + )), + IndexVec::from_raw(vec![val]), + ) + } else { + // CoroutineState::Yielded(val) + Rvalue::Aggregate( + Box::new(AggregateKind::Adt( + coroutine_state_def_id, + VariantIdx::from_usize(0), + args, + None, + None, + )), + IndexVec::from_raw(vec![val]), + ) + } + } + }; + + statements.push(Statement { + kind: StatementKind::Assign(Box::new((Place::return_place(), rvalue))), + source_info, + }); + } + + // Create a Place referencing a coroutine struct field + fn make_field(&self, variant_index: VariantIdx, idx: FieldIdx, ty: Ty<'tcx>) -> Place<'tcx> { + let self_place = Place::from(SELF_ARG); + let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index); + let mut projection = base.projection.to_vec(); + projection.push(ProjectionElem::Field(idx, ty)); + + Place { local: base.local, projection: self.tcx.mk_place_elems(&projection) } + } + + // Create a statement which changes the discriminant + fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> { + let self_place = Place::from(SELF_ARG); + Statement { + source_info, + kind: StatementKind::SetDiscriminant { + place: Box::new(self_place), + variant_index: state_disc, + }, + } + } + + // Create a statement which reads the discriminant into a temporary + fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) { + let temp_decl = LocalDecl::new(self.discr_ty, body.span); + let local_decls_len = body.local_decls.push(temp_decl); + let temp = Place::from(local_decls_len); + + let self_place = Place::from(SELF_ARG); + let assign = Statement { + source_info: SourceInfo::outermost(body.span), + kind: StatementKind::Assign(Box::new((temp, Rvalue::Discriminant(self_place)))), + }; + (assign, temp) + } +} + +impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { + assert_eq!(self.remap.get(local), None); + } + + fn visit_place( + &mut self, + place: &mut Place<'tcx>, + _context: PlaceContext, + _location: Location, + ) { + // Replace an Local in the remap with a coroutine struct access + if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) { + replace_base(place, self.make_field(variant_index, idx, ty), self.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_key(&l) + } + _ => true, + }); + + let ret_val = match data.terminator().kind { + TerminatorKind::Return => { + Some((true, None, Operand::Move(Place::from(self.old_ret_local)), None)) + } + TerminatorKind::Yield { ref value, resume, resume_arg, drop } => { + Some((false, Some((resume, resume_arg)), value.clone(), drop)) + } + _ => None, + }; + + if let Some((is_return, resume, v, drop)) = ret_val { + let source_info = data.terminator().source_info; + // We must assign the value first in case it gets declared dead below + self.make_state(v, source_info, is_return, &mut data.statements); + let state = if let Some((resume, mut resume_arg)) = resume { + // Yield + let state = RESERVED_VARIANTS + self.suspension_points.len(); + + // The resume arg target location might itself be remapped if its base local is + // live across a yield. + let resume_arg = + if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) { + replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx); + resume_arg + } else { + resume_arg + }; + + self.suspension_points.push(SuspensionPoint { + state, + resume, + resume_arg, + drop, + storage_liveness: self.storage_liveness[block].clone().unwrap().into(), + }); + + VariantIdx::new(state) + } else { + // Return + VariantIdx::new(RETURNED) // state for returned + }; + data.statements.push(self.set_discr(state, source_info)); + data.terminator_mut().kind = TerminatorKind::Return; + } + + self.super_basic_block_data(block, data); + } +} + +fn make_coroutine_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let coroutine_ty = body.local_decls.raw[1].ty; + + let ref_coroutine_ty = Ty::new_ref( + tcx, + tcx.lifetimes.re_erased, + ty::TypeAndMut { ty: coroutine_ty, mutbl: Mutability::Mut }, + ); + + // Replace the by value coroutine argument + body.local_decls.raw[1].ty = ref_coroutine_ty; + + // Add a deref to accesses of the coroutine state + DerefArgVisitor { tcx }.visit_body(body); +} + +fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let ref_coroutine_ty = body.local_decls.raw[1].ty; + + let pin_did = tcx.require_lang_item(LangItem::Pin, Some(body.span)); + let pin_adt_ref = tcx.adt_def(pin_did); + let args = tcx.mk_args(&[ref_coroutine_ty.into()]); + let pin_ref_coroutine_ty = Ty::new_adt(tcx, pin_adt_ref, args); + + // Replace the by ref coroutine argument + body.local_decls.raw[1].ty = pin_ref_coroutine_ty; + + // Add the Pin field access to accesses of the coroutine state + PinArgVisitor { ref_coroutine_ty, tcx }.visit_body(body); +} + +/// Allocates a new local and replaces all references of `local` with it. Returns the new local. +/// +/// `local` will be changed to a new local decl with type `ty`. +/// +/// Note that the new local will be uninitialized. It is the caller's responsibility to assign some +/// valid value to it before its first use. +fn replace_local<'tcx>( + local: Local, + ty: Ty<'tcx>, + body: &mut Body<'tcx>, + tcx: TyCtxt<'tcx>, +) -> Local { + let new_decl = LocalDecl::new(ty, body.span); + let new_local = body.local_decls.push(new_decl); + body.local_decls.swap(local, new_local); + + RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body); + + new_local +} + +/// Transforms the `body` of the coroutine applying the following transforms: +/// +/// - Eliminates all the `get_context` calls that async lowering created. +/// - Replace all `Local` `ResumeTy` types with `&mut Context<'_>` (`context_mut_ref`). +/// +/// The `Local`s that have their types replaced are: +/// - The `resume` argument itself. +/// - The argument to `get_context`. +/// - The yielded value of a `yield`. +/// +/// The `ResumeTy` hides a `&mut Context<'_>` behind an unsafe raw pointer, and the +/// `get_context` function is being used to convert that back to a `&mut Context<'_>`. +/// +/// Ideally the async lowering would not use the `ResumeTy`/`get_context` indirection, +/// but rather directly use `&mut Context<'_>`, however that would currently +/// lead to higher-kinded lifetime errors. +/// See <https://github.com/rust-lang/rust/issues/105501>. +/// +/// The async lowering step and the type / lifetime inference / checking are +/// still using the `ResumeTy` indirection for the time being, and that indirection +/// is removed here. After this transform, the coroutine body only knows about `&mut Context<'_>`. +fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let context_mut_ref = Ty::new_task_context(tcx); + + // replace the type of the `resume` argument + replace_resume_ty_local(tcx, body, Local::new(2), context_mut_ref); + + let get_context_def_id = tcx.require_lang_item(LangItem::GetContext, None); + + for bb in START_BLOCK..body.basic_blocks.next_index() { + let bb_data = &body[bb]; + if bb_data.is_cleanup { + continue; + } + + match &bb_data.terminator().kind { + TerminatorKind::Call { func, .. } => { + let func_ty = func.ty(body, tcx); + if let ty::FnDef(def_id, _) = *func_ty.kind() { + if def_id == get_context_def_id { + let local = eliminate_get_context_call(&mut body[bb]); + replace_resume_ty_local(tcx, body, local, context_mut_ref); + } + } else { + continue; + } + } + TerminatorKind::Yield { resume_arg, .. } => { + replace_resume_ty_local(tcx, body, resume_arg.local, context_mut_ref); + } + _ => {} + } + } +} + +fn eliminate_get_context_call<'tcx>(bb_data: &mut BasicBlockData<'tcx>) -> Local { + let terminator = bb_data.terminator.take().unwrap(); + if let TerminatorKind::Call { mut args, destination, target, .. } = terminator.kind { + let arg = args.pop().unwrap(); + let local = arg.place().unwrap().local; + + let arg = Rvalue::Use(arg); + let assign = Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new((destination, arg))), + }; + bb_data.statements.push(assign); + bb_data.terminator = Some(Terminator { + source_info: terminator.source_info, + kind: TerminatorKind::Goto { target: target.unwrap() }, + }); + local + } else { + bug!(); + } +} + +#[cfg_attr(not(debug_assertions), allow(unused))] +fn replace_resume_ty_local<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + local: Local, + context_mut_ref: Ty<'tcx>, +) { + let local_ty = std::mem::replace(&mut body.local_decls[local].ty, context_mut_ref); + // We have to replace the `ResumeTy` that is used for type and borrow checking + // with `&mut Context<'_>` in MIR. + #[cfg(debug_assertions)] + { + if let ty::Adt(resume_ty_adt, _) = local_ty.kind() { + let expected_adt = tcx.adt_def(tcx.require_lang_item(LangItem::ResumeTy, None)); + assert_eq!(*resume_ty_adt, expected_adt); + } else { + panic!("expected `ResumeTy`, found `{:?}`", local_ty); + }; + } +} + +/// Transforms the `body` of the coroutine applying the following transform: +/// +/// - Remove the `resume` argument. +/// +/// Ideally the async lowering would not add the `resume` argument. +/// +/// The async lowering step and the type / lifetime inference / checking are +/// still using the `resume` argument for the time being. After this transform, +/// the coroutine body doesn't have the `resume` argument. +fn transform_gen_context<'tcx>(_tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // This leaves the local representing the `resume` argument in place, + // but turns it into a regular local variable. This is cheaper than + // adjusting all local references in the body after removing it. + body.arg_count = 1; +} + +struct LivenessInfo { + /// Which locals are live across any suspension point. + saved_locals: CoroutineSavedLocals, + + /// The set of saved locals live at each suspension point. + live_locals_at_suspension_points: Vec<BitSet<CoroutineSavedLocal>>, + + /// Parallel vec to the above with SourceInfo for each yield terminator. + source_info_at_suspension_points: Vec<SourceInfo>, + + /// For every saved local, the set of other saved locals that are + /// storage-live at the same time as this local. We cannot overlap locals in + /// the layout which have conflicting storage. + storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>, + + /// For every suspending block, the locals which are storage-live across + /// that suspension point. + storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>, +} + +/// Computes which locals have to be stored in the state-machine for the +/// given coroutine. +/// +/// The basic idea is as follows: +/// - a local is live until we encounter a `StorageDead` statement. In +/// case none exist, the local is considered to be always live. +/// - a local has to be stored if it is either directly used after the +/// the suspend point, or if it is live and has been previously borrowed. +fn locals_live_across_suspend_points<'tcx>( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + always_live_locals: &BitSet<Local>, + movable: bool, +) -> LivenessInfo { + // Calculate when MIR locals have live storage. This gives us an upper bound of their + // lifetimes. + let mut storage_live = MaybeStorageLive::new(std::borrow::Cow::Borrowed(always_live_locals)) + .into_engine(tcx, body) + .iterate_to_fixpoint() + .into_results_cursor(body); + + // Calculate the MIR locals which have been previously + // borrowed (even if they are still active). + let borrowed_locals_results = + MaybeBorrowedLocals.into_engine(tcx, body).pass_name("coroutine").iterate_to_fixpoint(); + + let mut borrowed_locals_cursor = borrowed_locals_results.clone().into_results_cursor(body); + + // Calculate the MIR locals that we actually need to keep storage around + // for. + let mut requires_storage_cursor = + MaybeRequiresStorage::new(borrowed_locals_results.into_results_cursor(body)) + .into_engine(tcx, body) + .iterate_to_fixpoint() + .into_results_cursor(body); + + // Calculate the liveness of MIR locals ignoring borrows. + let mut liveness = MaybeLiveLocals + .into_engine(tcx, body) + .pass_name("coroutine") + .iterate_to_fixpoint() + .into_results_cursor(body); + + let mut storage_liveness_map = IndexVec::from_elem(None, &body.basic_blocks); + let mut live_locals_at_suspension_points = Vec::new(); + let mut source_info_at_suspension_points = Vec::new(); + let mut live_locals_at_any_suspension_point = BitSet::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: BitSet<_> = BitSet::new_empty(body.local_decls.len()); + live_locals.union(liveness.get()); + + 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_cursor.seek_before_primary_effect(loc); + live_locals.union(borrowed_locals_cursor.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()); + + // 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); + + debug!("loc = {:?}, live_locals = {:?}", 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); + + 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); + let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point); + + // Renumber our liveness_map bitsets to include only the locals we are + // saving. + let live_locals_at_suspension_points = live_locals_at_suspension_points + .iter() + .map(|live_here| saved_locals.renumber_bitset(live_here)) + .collect(); + + let storage_conflicts = compute_storage_conflicts( + body, + &saved_locals, + always_live_locals.clone(), + requires_storage_cursor.into_results(), + ); + + LivenessInfo { + saved_locals, + live_locals_at_suspension_points, + source_info_at_suspension_points, + storage_conflicts, + storage_liveness: storage_liveness_map, + } +} + +/// The set of `Local`s that must be saved across yield points. +/// +/// `CoroutineSavedLocal` is indexed in terms of the elements in this set; +/// i.e. `CoroutineSavedLocal::new(1)` corresponds to the second local +/// included in this set. +struct CoroutineSavedLocals(BitSet<Local>); + +impl CoroutineSavedLocals { + /// Returns an iterator over each `CoroutineSavedLocal` along with the `Local` it corresponds + /// to. + fn iter_enumerated(&self) -> impl '_ + Iterator<Item = (CoroutineSavedLocal, Local)> { + self.iter().enumerate().map(|(i, l)| (CoroutineSavedLocal::from(i), l)) + } + + /// Transforms a `BitSet<Local>` that contains only locals saved across yield points to the + /// equivalent `BitSet<CoroutineSavedLocal>`. + fn renumber_bitset(&self, input: &BitSet<Local>) -> BitSet<CoroutineSavedLocal> { + assert!(self.superset(input), "{:?} not a superset of {:?}", self.0, input); + let mut out = BitSet::new_empty(self.count()); + for (saved_local, local) in self.iter_enumerated() { + if input.contains(local) { + out.insert(saved_local); + } + } + out + } + + fn get(&self, local: Local) -> Option<CoroutineSavedLocal> { + if !self.contains(local) { + return None; + } + + let idx = self.iter().take_while(|&l| l < local).count(); + Some(CoroutineSavedLocal::new(idx)) + } +} + +impl ops::Deref for CoroutineSavedLocals { + type Target = BitSet<Local>; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +/// For every saved local, looks for which locals are StorageLive at the same +/// time. Generates a bitset for every local of all the other locals that may be +/// StorageLive simultaneously with that local. This is used in the layout +/// computation; see `CoroutineLayout` for more. +fn compute_storage_conflicts<'mir, 'tcx>( + body: &'mir Body<'tcx>, + saved_locals: &CoroutineSavedLocals, + always_live_locals: BitSet<Local>, + mut requires_storage: rustc_mir_dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>, +) -> BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal> { + assert_eq!(body.local_decls.len(), saved_locals.domain_size()); + + debug!("compute_storage_conflicts({:?})", body.span); + debug!("always_live = {:?}", always_live_locals); + + // Locals that are always live or ones that need to be stored across + // suspension points are not eligible for overlap. + let mut ineligible_locals = always_live_locals; + ineligible_locals.intersect(&**saved_locals); + + // Compute the storage conflicts for all eligible locals. + let mut visitor = StorageConflictVisitor { + body, + saved_locals: saved_locals, + local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()), + }; + + requires_storage.visit_reachable_with(body, &mut visitor); + + let local_conflicts = visitor.local_conflicts; + + // Compress the matrix using only stored locals (Local -> CoroutineSavedLocal). + // + // NOTE: Today we store a full conflict bitset for every local. Technically + // this is twice as many bits as we need, since the relation is symmetric. + // However, in practice these bitsets are not usually large. The layout code + // also needs to keep track of how many conflicts each local has, so it's + // simpler to keep it this way for now. + let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count()); + for (saved_local_a, local_a) in saved_locals.iter_enumerated() { + if ineligible_locals.contains(local_a) { + // Conflicts with everything. + storage_conflicts.insert_all_into_row(saved_local_a); + } else { + // Keep overlap information only for stored locals. + for (saved_local_b, local_b) in saved_locals.iter_enumerated() { + if local_conflicts.contains(local_a, local_b) { + storage_conflicts.insert(saved_local_a, saved_local_b); + } + } + } + } + storage_conflicts +} + +struct StorageConflictVisitor<'mir, 'tcx, 's> { + body: &'mir Body<'tcx>, + saved_locals: &'s CoroutineSavedLocals, + // FIXME(tmandry): Consider using sparse bitsets here once we have good + // benchmarks for coroutines. + local_conflicts: BitMatrix<Local, Local>, +} + +impl<'mir, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'mir, 'tcx, R> + for StorageConflictVisitor<'mir, 'tcx, '_> +{ + type FlowState = BitSet<Local>; + + fn visit_statement_before_primary_effect( + &mut self, + _results: &mut R, + state: &Self::FlowState, + _statement: &'mir Statement<'tcx>, + loc: Location, + ) { + self.apply_state(state, loc); + } + + fn visit_terminator_before_primary_effect( + &mut self, + _results: &mut R, + state: &Self::FlowState, + _terminator: &'mir Terminator<'tcx>, + loc: Location, + ) { + self.apply_state(state, loc); + } +} + +impl StorageConflictVisitor<'_, '_, '_> { + fn apply_state(&mut self, flow_state: &BitSet<Local>, loc: Location) { + // Ignore unreachable blocks. + if self.body.basic_blocks[loc.block].terminator().kind == TerminatorKind::Unreachable { + return; + } + + let mut eligible_storage_live = flow_state.clone(); + eligible_storage_live.intersect(&**self.saved_locals); + + for local in eligible_storage_live.iter() { + self.local_conflicts.union_row_with(&eligible_storage_live, local); + } + + if eligible_storage_live.count() > 1 { + trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live); + } + } +} + +fn compute_layout<'tcx>( + liveness: LivenessInfo, + body: &Body<'tcx>, +) -> ( + FxHashMap<Local, (Ty<'tcx>, VariantIdx, FieldIdx)>, + CoroutineLayout<'tcx>, + IndexVec<BasicBlock, Option<BitSet<Local>>>, +) { + let LivenessInfo { + saved_locals, + live_locals_at_suspension_points, + source_info_at_suspension_points, + storage_conflicts, + storage_liveness, + } = liveness; + + // Gather live local types and their indices. + let mut locals = IndexVec::<CoroutineSavedLocal, _>::new(); + let mut tys = IndexVec::<CoroutineSavedLocal, _>::new(); + for (saved_local, local) in saved_locals.iter_enumerated() { + debug!("coroutine saved local {:?} => {:?}", saved_local, local); + + locals.push(local); + let decl = &body.local_decls[local]; + debug!(?decl); + + // Do not `assert_crate_local` here, as post-borrowck cleanup may have already cleared + // the information. This is alright, since `ignore_for_traits` is only relevant when + // this code runs on pre-cleanup MIR, and `ignore_for_traits = false` is the safer + // default. + let ignore_for_traits = match decl.local_info { + // Do not include raw pointers created from accessing `static` items, as those could + // well be re-created by another access to the same static. + ClearCrossCrate::Set(box LocalInfo::StaticRef { is_thread_local, .. }) => { + !is_thread_local + } + // Fake borrows are only read by fake reads, so do not have any reality in + // post-analysis MIR. + ClearCrossCrate::Set(box LocalInfo::FakeBorrow) => true, + _ => false, + }; + let decl = + CoroutineSavedTy { ty: decl.ty, source_info: decl.source_info, ignore_for_traits }; + debug!(?decl); + + tys.push(decl); + } + + // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states. + // 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> = [ + 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(RESERVED_VARIANTS).collect(); + let mut remap = FxHashMap::default(); + for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() { + let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx); + let mut fields = IndexVec::new(); + for (idx, saved_local) in live_locals.iter().enumerate() { + fields.push(saved_local); + // 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.entry(locals[saved_local]).or_insert((tys[saved_local].ty, variant_index, idx)); + } + variant_fields.push(fields); + variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]); + } + debug!("coroutine variant_fields = {:?}", variant_fields); + debug!("coroutine storage_conflicts = {:#?}", storage_conflicts); + + let mut field_names = IndexVec::from_elem(None, &tys); + for var in &body.var_debug_info { + let VarDebugInfoContents::Place(place) = &var.value else { continue }; + let Some(local) = place.as_local() else { continue }; + let Some(&(_, variant, field)) = remap.get(&local) else { continue }; + + let saved_local = variant_fields[variant][field]; + field_names.get_or_insert_with(saved_local, || var.name); + } + + let layout = CoroutineLayout { + field_tys: tys, + field_names, + variant_fields, + variant_source_info, + storage_conflicts, + }; + debug!(?layout); + + (remap, layout, storage_liveness) +} + +/// Replaces the entry point of `body` with a block that switches on the coroutine discriminant and +/// dispatches to blocks according to `cases`. +/// +/// After this function, the former entry point of the function will be bb1. +fn insert_switch<'tcx>( + body: &mut Body<'tcx>, + cases: Vec<(usize, BasicBlock)>, + transform: &TransformVisitor<'tcx>, + default: TerminatorKind<'tcx>, +) { + let default_block = insert_term_block(body, default); + let (assign, discr) = transform.get_discr(body); + let switch_targets = + SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block); + let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets }; + + let source_info = SourceInfo::outermost(body.span); + body.basic_blocks_mut().raw.insert( + 0, + BasicBlockData { + statements: vec![assign], + terminator: Some(Terminator { source_info, kind: switch }), + is_cleanup: false, + }, + ); + + let blocks = body.basic_blocks_mut().iter_mut(); + + for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) { + *target = BasicBlock::new(target.index() + 1); + } +} + +fn elaborate_coroutine_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + use crate::shim::DropShimElaborator; + use rustc_middle::mir::patch::MirPatch; + use rustc_mir_dataflow::elaborate_drops::{elaborate_drop, Unwind}; + + // Note that `elaborate_drops` only drops the upvars of a coroutine, and + // this is ok because `open_drop` can only be reached within that own + // coroutine's resume function. + + let def_id = body.source.def_id(); + let param_env = tcx.param_env(def_id); + + let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env }; + + for (block, block_data) in body.basic_blocks.iter_enumerated() { + let (target, unwind, source_info) = match block_data.terminator() { + Terminator { + source_info, + kind: TerminatorKind::Drop { place, target, unwind, replace: _ }, + } => { + if let Some(local) = place.as_local() { + if local == SELF_ARG { + (target, unwind, source_info) + } else { + continue; + } + } else { + continue; + } + } + _ => continue, + }; + let unwind = if block_data.is_cleanup { + Unwind::InCleanup + } else { + Unwind::To(match *unwind { + UnwindAction::Cleanup(tgt) => tgt, + UnwindAction::Continue => elaborator.patch.resume_block(), + UnwindAction::Unreachable => elaborator.patch.unreachable_cleanup_block(), + UnwindAction::Terminate(reason) => elaborator.patch.terminate_block(reason), + }) + }; + elaborate_drop( + &mut elaborator, + *source_info, + Place::from(SELF_ARG), + (), + *target, + unwind, + block, + ); + } + elaborator.patch.apply(body); +} + +fn create_coroutine_drop_shim<'tcx>( + tcx: TyCtxt<'tcx>, + transform: &TransformVisitor<'tcx>, + coroutine_ty: Ty<'tcx>, + body: &mut Body<'tcx>, + drop_clean: BasicBlock, +) -> Body<'tcx> { + let mut body = body.clone(); + body.arg_count = 1; // make sure the resume argument is not included here + + let source_info = SourceInfo::outermost(body.span); + + let mut cases = create_cases(&mut body, transform, Operation::Drop); + + cases.insert(0, (UNRESUMED, drop_clean)); + + // The returned state and the poisoned state fall through to the default + // case which is just to return + + insert_switch(&mut body, cases, transform, TerminatorKind::Return); + + for block in body.basic_blocks_mut() { + let kind = &mut block.terminator_mut().kind; + if let TerminatorKind::CoroutineDrop = *kind { + *kind = TerminatorKind::Return; + } + } + + // Replace the return variable + body.local_decls[RETURN_PLACE] = LocalDecl::with_source_info(Ty::new_unit(tcx), source_info); + + make_coroutine_state_argument_indirect(tcx, &mut body); + + // Change the coroutine argument from &mut to *mut + body.local_decls[SELF_ARG] = LocalDecl::with_source_info( + Ty::new_ptr(tcx, ty::TypeAndMut { ty: coroutine_ty, mutbl: hir::Mutability::Mut }), + source_info, + ); + + // Make sure we remove dead blocks to remove + // unrelated code from the resume part of the function + simplify::remove_dead_blocks(&mut body); + + // Update the body's def to become the drop glue. + let coroutine_instance = body.source.instance; + let drop_in_place = tcx.require_lang_item(LangItem::DropInPlace, None); + let drop_instance = InstanceDef::DropGlue(drop_in_place, Some(coroutine_ty)); + + // Temporary change MirSource to coroutine's instance so that dump_mir produces more sensible + // filename. + body.source.instance = coroutine_instance; + dump_mir(tcx, false, "coroutine_drop", &0, &body, |_, _| Ok(())); + body.source.instance = drop_instance; + + body +} + +fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock { + let source_info = SourceInfo::outermost(body.span); + body.basic_blocks_mut().push(BasicBlockData { + statements: Vec::new(), + terminator: Some(Terminator { source_info, kind }), + is_cleanup: false, + }) +} + +fn insert_panic_block<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + message: AssertMessage<'tcx>, +) -> BasicBlock { + let assert_block = BasicBlock::new(body.basic_blocks.len()); + let term = TerminatorKind::Assert { + cond: Operand::Constant(Box::new(ConstOperand { + span: body.span, + user_ty: None, + const_: Const::from_bool(tcx, false), + })), + expected: true, + msg: Box::new(message), + target: assert_block, + unwind: UnwindAction::Continue, + }; + + let source_info = SourceInfo::outermost(body.span); + body.basic_blocks_mut().push(BasicBlockData { + statements: Vec::new(), + terminator: Some(Terminator { source_info, kind: term }), + is_cleanup: false, + }); + + assert_block +} + +fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, param_env: ty::ParamEnv<'tcx>) -> bool { + // Returning from a function with an uninhabited return type is undefined behavior. + if body.return_ty().is_privately_uninhabited(tcx, param_env) { + return false; + } + + // If there's a return terminator the function may return. + for block in body.basic_blocks.iter() { + if let TerminatorKind::Return = block.terminator().kind { + return true; + } + } + + // Otherwise the function can't return. + false +} + +fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool { + // Nothing can unwind when landing pads are off. + if tcx.sess.panic_strategy() == PanicStrategy::Abort { + return false; + } + + // Unwinds can only start at certain terminators. + for block in body.basic_blocks.iter() { + match block.terminator().kind { + // These never unwind. + TerminatorKind::Goto { .. } + | TerminatorKind::SwitchInt { .. } + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::CoroutineDrop + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } => {} + + // Resume will *continue* unwinding, but if there's no other unwinding terminator it + // will never be reached. + TerminatorKind::UnwindResume => {} + + TerminatorKind::Yield { .. } => { + unreachable!("`can_unwind` called before coroutine transform") + } + + // These may unwind. + TerminatorKind::Drop { .. } + | TerminatorKind::Call { .. } + | TerminatorKind::InlineAsm { .. } + | TerminatorKind::Assert { .. } => return true, + } + } + + // If we didn't find an unwinding terminator, the function cannot unwind. + false +} + +fn create_coroutine_resume_function<'tcx>( + tcx: TyCtxt<'tcx>, + transform: TransformVisitor<'tcx>, + body: &mut Body<'tcx>, + can_return: bool, +) { + let can_unwind = can_unwind(tcx, body); + + // Poison the coroutine when it unwinds + if can_unwind { + let source_info = SourceInfo::outermost(body.span); + let poison_block = body.basic_blocks_mut().push(BasicBlockData { + statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)], + terminator: Some(Terminator { source_info, kind: TerminatorKind::UnwindResume }), + is_cleanup: true, + }); + + for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() { + let source_info = block.terminator().source_info; + + if let TerminatorKind::UnwindResume = block.terminator().kind { + // An existing `Resume` terminator is redirected to jump to our dedicated + // "poisoning block" above. + if idx != poison_block { + *block.terminator_mut() = Terminator { + source_info, + kind: TerminatorKind::Goto { target: poison_block }, + }; + } + } else if !block.is_cleanup { + // Any terminators that *can* unwind but don't have an unwind target set are also + // pointed at our poisoning block (unless they're part of the cleanup path). + if let Some(unwind @ UnwindAction::Continue) = block.terminator_mut().unwind_mut() { + *unwind = UnwindAction::Cleanup(poison_block); + } + } + } + } + + let mut cases = create_cases(body, &transform, Operation::Resume); + + use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn}; + + // Jump to the entry point on the unresumed + cases.insert(0, (UNRESUMED, START_BLOCK)); + + // Panic when resumed on the returned or poisoned state + let coroutine_kind = body.coroutine_kind().unwrap(); + + if can_unwind { + cases.insert( + 1, + (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(coroutine_kind))), + ); + } + + if can_return { + let block = match coroutine_kind { + CoroutineKind::Async(_) | CoroutineKind::Coroutine => { + insert_panic_block(tcx, body, ResumedAfterReturn(coroutine_kind)) + } + CoroutineKind::AsyncGen(_) | CoroutineKind::Gen(_) => { + transform.insert_none_ret_block(body) + } + }; + cases.insert(1, (RETURNED, block)); + } + + insert_switch(body, cases, &transform, TerminatorKind::Unreachable); + + make_coroutine_state_argument_indirect(tcx, body); + + match coroutine_kind { + // Iterator::next doesn't accept a pinned argument, + // unlike for all other coroutine kinds. + CoroutineKind::Gen(_) => {} + _ => { + make_coroutine_state_argument_pinned(tcx, body); + } + } + + // Make sure we remove dead blocks to remove + // unrelated code from the drop part of the function + simplify::remove_dead_blocks(body); + + pm::run_passes_no_validate(tcx, body, &[&abort_unwinding_calls::AbortUnwindingCalls], None); + + dump_mir(tcx, false, "coroutine_resume", &0, body, |_, _| Ok(())); +} + +fn insert_clean_drop(body: &mut Body<'_>) -> BasicBlock { + let return_block = insert_term_block(body, TerminatorKind::Return); + + let term = TerminatorKind::Drop { + place: Place::from(SELF_ARG), + target: return_block, + unwind: UnwindAction::Continue, + replace: false, + }; + let source_info = SourceInfo::outermost(body.span); + + // Create a block to destroy an unresumed coroutines. This can only destroy upvars. + body.basic_blocks_mut().push(BasicBlockData { + statements: Vec::new(), + terminator: Some(Terminator { source_info, kind: term }), + is_cleanup: false, + }) +} + +/// An operation that can be performed on a coroutine. +#[derive(PartialEq, Copy, Clone)] +enum Operation { + Resume, + Drop, +} + +impl Operation { + fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> { + match self { + Operation::Resume => Some(point.resume), + Operation::Drop => point.drop, + } + } +} + +fn create_cases<'tcx>( + body: &mut Body<'tcx>, + transform: &TransformVisitor<'tcx>, + operation: Operation, +) -> Vec<(usize, BasicBlock)> { + let source_info = SourceInfo::outermost(body.span); + + transform + .suspension_points + .iter() + .filter_map(|point| { + // Find the target for this suspension point, if applicable + operation.target_block(point).map(|target| { + let mut statements = Vec::new(); + + // Create StorageLive instructions for locals with live storage + for i in 0..(body.local_decls.len()) { + if i == 2 { + // The resume argument is live on function entry. Don't insert a + // `StorageLive`, or the following `Assign` will read from uninitialized + // memory. + continue; + } + + let l = Local::new(i); + let needs_storage_live = point.storage_liveness.contains(l) + && !transform.remap.contains_key(&l) + && !transform.always_live_locals.contains(l); + if needs_storage_live { + statements + .push(Statement { source_info, kind: StatementKind::StorageLive(l) }); + } + } + + if operation == Operation::Resume { + // Move the resume argument to the destination place of the `Yield` terminator + let resume_arg = Local::new(2); // 0 = return, 1 = self + statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + point.resume_arg, + Rvalue::Use(Operand::Move(resume_arg.into())), + ))), + }); + } + + // Then jump to the real target + let block = body.basic_blocks_mut().push(BasicBlockData { + statements, + terminator: Some(Terminator { + source_info, + kind: TerminatorKind::Goto { target }, + }), + is_cleanup: false, + }); + + (point.state, block) + }) + }) + .collect() +} + +#[instrument(level = "debug", skip(tcx), ret)] +pub(crate) fn mir_coroutine_witnesses<'tcx>( + tcx: TyCtxt<'tcx>, + def_id: LocalDefId, +) -> Option<CoroutineLayout<'tcx>> { + let (body, _) = tcx.mir_promoted(def_id); + let body = body.borrow(); + let body = &*body; + + // The first argument is the coroutine type passed by value + let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty; + + let movable = match *coroutine_ty.kind() { + ty::Coroutine(_, _, movability) => movability == hir::Movability::Movable, + ty::Error(_) => return None, + _ => span_bug!(body.span, "unexpected coroutine type {}", coroutine_ty), + }; + + // The witness simply contains all locals live across suspend points. + + let always_live_locals = always_storage_live_locals(body); + let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable); + + // Extract locals which are live across suspension point into `layout` + // `remap` gives a mapping from local indices onto coroutine struct indices + // `storage_liveness` tells us which locals have live storage at suspension points + let (_, coroutine_layout, _) = compute_layout(liveness_info, body); + + check_suspend_tys(tcx, &coroutine_layout, body); + + Some(coroutine_layout) +} + +impl<'tcx> MirPass<'tcx> for StateTransform { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let Some(old_yield_ty) = body.yield_ty() else { + // This only applies to coroutines + return; + }; + let old_ret_ty = body.return_ty(); + + assert!(body.coroutine_drop().is_none()); + + // The first argument is the coroutine type passed by value + let coroutine_ty = body.local_decls.raw[1].ty; + + // Get the discriminant type and args which typeck computed + let (discr_ty, movable) = match *coroutine_ty.kind() { + ty::Coroutine(_, args, movability) => { + let args = args.as_coroutine(); + (args.discr_ty(tcx), movability == hir::Movability::Movable) + } + _ => { + tcx.sess.span_delayed_bug( + body.span, + format!("unexpected coroutine type {coroutine_ty}"), + ); + return; + } + }; + + let is_async_kind = matches!(body.coroutine_kind(), Some(CoroutineKind::Async(_))); + let is_async_gen_kind = matches!(body.coroutine_kind(), Some(CoroutineKind::AsyncGen(_))); + let is_gen_kind = matches!(body.coroutine_kind(), Some(CoroutineKind::Gen(_))); + let new_ret_ty = match body.coroutine_kind().unwrap() { + CoroutineKind::Async(_) => { + // Compute Poll<return_ty> + let poll_did = tcx.require_lang_item(LangItem::Poll, None); + let poll_adt_ref = tcx.adt_def(poll_did); + let poll_args = tcx.mk_args(&[old_ret_ty.into()]); + Ty::new_adt(tcx, poll_adt_ref, poll_args) + } + CoroutineKind::Gen(_) => { + // Compute Option<yield_ty> + let option_did = tcx.require_lang_item(LangItem::Option, None); + let option_adt_ref = tcx.adt_def(option_did); + let option_args = tcx.mk_args(&[old_yield_ty.into()]); + Ty::new_adt(tcx, option_adt_ref, option_args) + } + CoroutineKind::AsyncGen(_) => { + // The yield ty is already `Poll<Option<yield_ty>>` + old_yield_ty + } + CoroutineKind::Coroutine => { + // Compute CoroutineState<yield_ty, return_ty> + let state_did = tcx.require_lang_item(LangItem::CoroutineState, None); + let state_adt_ref = tcx.adt_def(state_did); + let state_args = tcx.mk_args(&[old_yield_ty.into(), old_ret_ty.into()]); + Ty::new_adt(tcx, state_adt_ref, state_args) + } + }; + + // We rename RETURN_PLACE which has type mir.return_ty to old_ret_local + // RETURN_PLACE then is a fresh unused local with type ret_ty. + let old_ret_local = replace_local(RETURN_PLACE, new_ret_ty, body, tcx); + + // Replace all occurrences of `ResumeTy` with `&mut Context<'_>` within async bodies. + if is_async_kind || is_async_gen_kind { + transform_async_context(tcx, body); + } + + // We also replace the resume argument and insert an `Assign`. + // This is needed because the resume argument `_2` might be live across a `yield`, in which + // case there is no `Assign` to it that the transform can turn into a store to the coroutine + // state. After the yield the slot in the coroutine state would then be uninitialized. + let resume_local = Local::new(2); + let resume_ty = if is_async_kind { + Ty::new_task_context(tcx) + } else { + body.local_decls[resume_local].ty + }; + let old_resume_local = replace_local(resume_local, resume_ty, body, tcx); + + // When first entering the coroutine, move the resume argument into its old local + // (which is now a generator interior). + let source_info = SourceInfo::outermost(body.span); + let stmts = &mut body.basic_blocks_mut()[START_BLOCK].statements; + stmts.insert( + 0, + Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + old_resume_local.into(), + Rvalue::Use(Operand::Move(resume_local.into())), + ))), + }, + ); + + let always_live_locals = always_storage_live_locals(body); + + let liveness_info = + locals_live_across_suspend_points(tcx, body, &always_live_locals, movable); + + if tcx.sess.opts.unstable_opts.validate_mir { + let mut vis = EnsureCoroutineFieldAssignmentsNeverAlias { + assigned_local: None, + saved_locals: &liveness_info.saved_locals, + storage_conflicts: &liveness_info.storage_conflicts, + }; + + vis.visit_body(body); + } + + // Extract locals which are live across suspension point into `layout` + // `remap` gives a mapping from local indices onto coroutine struct indices + // `storage_liveness` tells us which locals have live storage at suspension points + let (remap, layout, storage_liveness) = compute_layout(liveness_info, body); + + let can_return = can_return(tcx, body, tcx.param_env(body.source.def_id())); + + // Run the transformation which converts Places from Local to coroutine struct + // accesses for locals in `remap`. + // It also rewrites `return x` and `yield y` as writing a new coroutine state and returning + // either CoroutineState::Complete(x) and CoroutineState::Yielded(y), + // or Poll::Ready(x) and Poll::Pending respectively depending on `is_async_kind`. + let mut transform = TransformVisitor { + tcx, + coroutine_kind: body.coroutine_kind().unwrap(), + remap, + storage_liveness, + always_live_locals, + suspension_points: Vec::new(), + old_ret_local, + discr_ty, + old_ret_ty, + old_yield_ty, + }; + transform.visit_body(body); + + // Update our MIR struct to reflect the changes we've made + body.arg_count = 2; // self, resume arg + body.spread_arg = None; + + // Remove the context argument within generator bodies. + if is_gen_kind { + transform_gen_context(tcx, body); + } + + // The original arguments to the function are no longer arguments, mark them as such. + // Otherwise they'll conflict with our new arguments, which although they don't have + // argument_index set, will get emitted as unnamed arguments. + for var in &mut body.var_debug_info { + var.argument_index = None; + } + + body.coroutine.as_mut().unwrap().yield_ty = None; + body.coroutine.as_mut().unwrap().coroutine_layout = Some(layout); + + // Insert `drop(coroutine_struct)` which is used to drop upvars for coroutines in + // the unresumed state. + // This is expanded to a drop ladder in `elaborate_coroutine_drops`. + let drop_clean = insert_clean_drop(body); + + dump_mir(tcx, false, "coroutine_pre-elab", &0, body, |_, _| Ok(())); + + // Expand `drop(coroutine_struct)` to a drop ladder which destroys upvars. + // If any upvars are moved out of, drop elaboration will handle upvar destruction. + // However we need to also elaborate the code generated by `insert_clean_drop`. + elaborate_coroutine_drops(tcx, body); + + dump_mir(tcx, false, "coroutine_post-transform", &0, body, |_, _| Ok(())); + + // Create a copy of our MIR and use it to create the drop shim for the coroutine + let drop_shim = create_coroutine_drop_shim(tcx, &transform, coroutine_ty, body, drop_clean); + + body.coroutine.as_mut().unwrap().coroutine_drop = Some(drop_shim); + + // Create the Coroutine::resume / Future::poll function + create_coroutine_resume_function(tcx, transform, body, can_return); + + // Run derefer to fix Derefs that are not in the first place + deref_finder(tcx, body); + } +} + +/// Looks for any assignments between locals (e.g., `_4 = _5`) that will both be converted to fields +/// in the coroutine state machine but whose storage is not marked as conflicting +/// +/// Validation needs to happen immediately *before* `TransformVisitor` is invoked, not after. +/// +/// This condition would arise when the assignment is the last use of `_5` but the initial +/// definition of `_4` if we weren't extra careful to mark all locals used inside a statement as +/// conflicting. Non-conflicting coroutine saved locals may be stored at the same location within +/// the coroutine state machine, which would result in ill-formed MIR: the left-hand and right-hand +/// sides of an assignment may not alias. This caused a miscompilation in [#73137]. +/// +/// [#73137]: https://github.com/rust-lang/rust/issues/73137 +struct EnsureCoroutineFieldAssignmentsNeverAlias<'a> { + saved_locals: &'a CoroutineSavedLocals, + storage_conflicts: &'a BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>, + assigned_local: Option<CoroutineSavedLocal>, +} + +impl EnsureCoroutineFieldAssignmentsNeverAlias<'_> { + fn saved_local_for_direct_place(&self, place: Place<'_>) -> Option<CoroutineSavedLocal> { + if place.is_indirect() { + return None; + } + + self.saved_locals.get(place.local) + } + + fn check_assigned_place(&mut self, place: Place<'_>, f: impl FnOnce(&mut Self)) { + if let Some(assigned_local) = self.saved_local_for_direct_place(place) { + assert!(self.assigned_local.is_none(), "`check_assigned_place` must not recurse"); + + self.assigned_local = Some(assigned_local); + f(self); + self.assigned_local = None; + } + } +} + +impl<'tcx> Visitor<'tcx> for EnsureCoroutineFieldAssignmentsNeverAlias<'_> { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) { + let Some(lhs) = self.assigned_local else { + // This visitor only invokes `visit_place` for the right-hand side of an assignment + // and only after setting `self.assigned_local`. However, the default impl of + // `Visitor::super_body` may call `visit_place` with a `NonUseContext` for places + // with debuginfo. Ignore them here. + assert!(!context.is_use()); + return; + }; + + let Some(rhs) = self.saved_local_for_direct_place(*place) else { return }; + + if !self.storage_conflicts.contains(lhs, rhs) { + bug!( + "Assignment between coroutine saved locals whose storage is not \ + marked as conflicting: {:?}: {:?} = {:?}", + location, + lhs, + rhs, + ); + } + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + match &statement.kind { + StatementKind::Assign(box (lhs, rhs)) => { + self.check_assigned_place(*lhs, |this| this.visit_rvalue(rhs, location)); + } + + StatementKind::FakeRead(..) + | StatementKind::SetDiscriminant { .. } + | StatementKind::Deinit(..) + | StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::Retag(..) + | StatementKind::AscribeUserType(..) + | StatementKind::PlaceMention(..) + | StatementKind::Coverage(..) + | StatementKind::Intrinsic(..) + | StatementKind::ConstEvalCounter + | StatementKind::Nop => {} + } + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + // Checking for aliasing in terminators is probably overkill, but until we have actual + // semantics, we should be conservative here. + match &terminator.kind { + TerminatorKind::Call { + func, + args, + destination, + target: Some(_), + unwind: _, + call_source: _, + fn_span: _, + } => { + self.check_assigned_place(*destination, |this| { + this.visit_operand(func, location); + for arg in args { + this.visit_operand(arg, location); + } + }); + } + + TerminatorKind::Yield { value, resume: _, resume_arg, drop: _ } => { + self.check_assigned_place(*resume_arg, |this| this.visit_operand(value, location)); + } + + // FIXME: Does `asm!` have any aliasing requirements? + TerminatorKind::InlineAsm { .. } => {} + + TerminatorKind::Call { .. } + | TerminatorKind::Goto { .. } + | TerminatorKind::SwitchInt { .. } + | TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::Drop { .. } + | TerminatorKind::Assert { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } => {} + } + } +} + +fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &CoroutineLayout<'tcx>, body: &Body<'tcx>) { + let mut linted_tys = FxHashSet::default(); + + // We want a user-facing param-env. + let param_env = tcx.param_env(body.source.def_id()); + + for (variant, yield_source_info) in + layout.variant_fields.iter().zip(&layout.variant_source_info) + { + debug!(?variant); + for &local in variant { + let decl = &layout.field_tys[local]; + debug!(?decl); + + if !decl.ignore_for_traits && linted_tys.insert(decl.ty) { + let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else { + continue; + }; + + check_must_not_suspend_ty( + tcx, + decl.ty, + hir_id, + param_env, + SuspendCheckData { + source_span: decl.source_info.span, + yield_span: yield_source_info.span, + plural_len: 1, + ..Default::default() + }, + ); + } + } + } +} + +#[derive(Default)] +struct SuspendCheckData<'a> { + source_span: Span, + yield_span: Span, + descr_pre: &'a str, + descr_post: &'a str, + plural_len: usize, +} + +// Returns whether it emitted a diagnostic or not +// Note that this fn and the proceeding one are based on the code +// for creating must_use diagnostics +// +// Note that this technique was chosen over things like a `Suspend` marker trait +// as it is simpler and has precedent in the compiler +fn check_must_not_suspend_ty<'tcx>( + tcx: TyCtxt<'tcx>, + ty: Ty<'tcx>, + hir_id: hir::HirId, + param_env: ty::ParamEnv<'tcx>, + data: SuspendCheckData<'_>, +) -> bool { + if ty.is_unit() { + return false; + } + + let plural_suffix = pluralize!(data.plural_len); + + debug!("Checking must_not_suspend for {}", ty); + + match *ty.kind() { + ty::Adt(..) if ty.is_box() => { + let boxed_ty = ty.boxed_ty(); + let descr_pre = &format!("{}boxed ", data.descr_pre); + check_must_not_suspend_ty( + tcx, + boxed_ty, + hir_id, + param_env, + SuspendCheckData { descr_pre, ..data }, + ) + } + ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data), + // FIXME: support adding the attribute to TAITs + ty::Alias(ty::Opaque, ty::AliasTy { def_id: def, .. }) => { + let mut has_emitted = false; + for &(predicate, _) in tcx.explicit_item_bounds(def).skip_binder() { + // We only look at the `DefId`, so it is safe to skip the binder here. + if let ty::ClauseKind::Trait(ref poly_trait_predicate) = + predicate.kind().skip_binder() + { + let def_id = poly_trait_predicate.trait_ref.def_id; + let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix); + if check_must_not_suspend_def( + tcx, + def_id, + hir_id, + SuspendCheckData { descr_pre, ..data }, + ) { + has_emitted = true; + break; + } + } + } + has_emitted + } + ty::Dynamic(binder, _, _) => { + let mut has_emitted = false; + for predicate in binder.iter() { + if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() { + let def_id = trait_ref.def_id; + let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post); + if check_must_not_suspend_def( + tcx, + def_id, + hir_id, + SuspendCheckData { descr_post, ..data }, + ) { + has_emitted = true; + break; + } + } + } + has_emitted + } + ty::Tuple(fields) => { + let mut has_emitted = false; + for (i, ty) in fields.iter().enumerate() { + let descr_post = &format!(" in tuple element {i}"); + if check_must_not_suspend_ty( + tcx, + ty, + hir_id, + param_env, + SuspendCheckData { descr_post, ..data }, + ) { + has_emitted = true; + } + } + has_emitted + } + ty::Array(ty, len) => { + let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix); + check_must_not_suspend_ty( + tcx, + ty, + hir_id, + param_env, + SuspendCheckData { + descr_pre, + plural_len: len.try_eval_target_usize(tcx, param_env).unwrap_or(0) as usize + 1, + ..data + }, + ) + } + // If drop tracking is enabled, we want to look through references, since the referent + // may not be considered live across the await point. + ty::Ref(_region, ty, _mutability) => { + let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix); + check_must_not_suspend_ty( + tcx, + ty, + hir_id, + param_env, + SuspendCheckData { descr_pre, ..data }, + ) + } + _ => false, + } +} + +fn check_must_not_suspend_def( + tcx: TyCtxt<'_>, + def_id: DefId, + hir_id: hir::HirId, + data: SuspendCheckData<'_>, +) -> bool { + if let Some(attr) = tcx.get_attr(def_id, sym::must_not_suspend) { + let reason = attr.value_str().map(|s| errors::MustNotSuspendReason { + span: data.source_span, + reason: s.as_str().to_string(), + }); + tcx.emit_spanned_lint( + rustc_session::lint::builtin::MUST_NOT_SUSPEND, + hir_id, + data.source_span, + errors::MustNotSupend { + tcx, + yield_sp: data.yield_span, + reason, + src_sp: data.source_span, + pre: data.descr_pre, + def_id, + post: data.descr_post, + }, + ); + + true + } else { + false + } +} diff --git a/compiler/rustc_mir_transform/src/cost_checker.rs b/compiler/rustc_mir_transform/src/cost_checker.rs new file mode 100644 index 00000000000..79bed960b95 --- /dev/null +++ b/compiler/rustc_mir_transform/src/cost_checker.rs @@ -0,0 +1,100 @@ +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, ParamEnv, Ty, TyCtxt}; + +const INSTR_COST: usize = 5; +const CALL_PENALTY: usize = 25; +const LANDINGPAD_PENALTY: usize = 50; +const RESUME_PENALTY: usize = 45; + +/// Verify that the callee body is compatible with the caller. +#[derive(Clone)] +pub(crate) struct CostChecker<'b, 'tcx> { + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + cost: usize, + callee_body: &'b Body<'tcx>, + instance: Option<ty::Instance<'tcx>>, +} + +impl<'b, 'tcx> CostChecker<'b, 'tcx> { + pub fn new( + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + instance: Option<ty::Instance<'tcx>>, + callee_body: &'b Body<'tcx>, + ) -> CostChecker<'b, 'tcx> { + CostChecker { tcx, param_env, callee_body, instance, cost: 0 } + } + + pub fn cost(&self) -> usize { + self.cost + } + + fn instantiate_ty(&self, v: Ty<'tcx>) -> Ty<'tcx> { + if let Some(instance) = self.instance { + instance.instantiate_mir(self.tcx, ty::EarlyBinder::bind(&v)) + } else { + v + } + } +} + +impl<'tcx> Visitor<'tcx> for CostChecker<'_, 'tcx> { + fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) { + // Don't count StorageLive/StorageDead in the inlining cost. + match statement.kind { + StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::Deinit(_) + | StatementKind::Nop => {} + _ => self.cost += INSTR_COST, + } + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, _: Location) { + let tcx = self.tcx; + match terminator.kind { + TerminatorKind::Drop { ref place, unwind, .. } => { + // If the place doesn't actually need dropping, treat it like a regular goto. + let ty = self.instantiate_ty(place.ty(self.callee_body, tcx).ty); + if ty.needs_drop(tcx, self.param_env) { + self.cost += CALL_PENALTY; + if let UnwindAction::Cleanup(_) = unwind { + self.cost += LANDINGPAD_PENALTY; + } + } else { + self.cost += INSTR_COST; + } + } + TerminatorKind::Call { func: Operand::Constant(ref f), unwind, .. } => { + let fn_ty = self.instantiate_ty(f.const_.ty()); + self.cost += if let ty::FnDef(def_id, _) = *fn_ty.kind() + && tcx.is_intrinsic(def_id) + { + // Don't give intrinsics the extra penalty for calls + INSTR_COST + } else { + CALL_PENALTY + }; + if let UnwindAction::Cleanup(_) = unwind { + self.cost += LANDINGPAD_PENALTY; + } + } + TerminatorKind::Assert { unwind, .. } => { + self.cost += CALL_PENALTY; + if let UnwindAction::Cleanup(_) = unwind { + self.cost += LANDINGPAD_PENALTY; + } + } + TerminatorKind::UnwindResume => self.cost += RESUME_PENALTY, + TerminatorKind::InlineAsm { unwind, .. } => { + self.cost += INSTR_COST; + if let UnwindAction::Cleanup(_) = unwind { + self.cost += LANDINGPAD_PENALTY; + } + } + _ => self.cost += INSTR_COST, + } + } +} diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs new file mode 100644 index 00000000000..604589e5b96 --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/counters.rs @@ -0,0 +1,517 @@ +use rustc_data_structures::fx::FxHashMap; +use rustc_data_structures::graph::WithNumNodes; +use rustc_index::bit_set::BitSet; +use rustc_index::IndexVec; +use rustc_middle::mir::coverage::*; + +use super::graph::{BasicCoverageBlock, CoverageGraph, TraverseCoverageGraphWithLoops}; + +use std::fmt::{self, Debug}; + +/// The coverage counter or counter expression associated with a particular +/// BCB node or BCB edge. +#[derive(Clone, Copy)] +pub(super) enum BcbCounter { + Counter { id: CounterId }, + Expression { id: ExpressionId }, +} + +impl BcbCounter { + fn is_expression(&self) -> bool { + matches!(self, Self::Expression { .. }) + } + + pub(super) fn as_term(&self) -> CovTerm { + match *self { + BcbCounter::Counter { id, .. } => CovTerm::Counter(id), + BcbCounter::Expression { id, .. } => CovTerm::Expression(id), + } + } +} + +impl Debug for BcbCounter { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Counter { id, .. } => write!(fmt, "Counter({:?})", id.index()), + Self::Expression { id } => write!(fmt, "Expression({:?})", id.index()), + } + } +} + +/// Generates and stores coverage counter and coverage expression information +/// associated with nodes/edges in the BCB graph. +pub(super) struct CoverageCounters { + next_counter_id: CounterId, + + /// Coverage counters/expressions that are associated with individual BCBs. + bcb_counters: IndexVec<BasicCoverageBlock, Option<BcbCounter>>, + /// Coverage counters/expressions that are associated with the control-flow + /// edge between two BCBs. + bcb_edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), BcbCounter>, + /// Tracks which BCBs have a counter associated with some incoming edge. + /// Only used by assertions, to verify that BCBs with incoming edge + /// counters do not have their own physical counters (expressions are allowed). + bcb_has_incoming_edge_counters: BitSet<BasicCoverageBlock>, + /// Table of expression data, associating each expression ID with its + /// corresponding operator (+ or -) and its LHS/RHS operands. + expressions: IndexVec<ExpressionId, Expression>, +} + +impl CoverageCounters { + pub(super) fn new(basic_coverage_blocks: &CoverageGraph) -> Self { + let num_bcbs = basic_coverage_blocks.num_nodes(); + + Self { + next_counter_id: CounterId::START, + bcb_counters: IndexVec::from_elem_n(None, num_bcbs), + bcb_edge_counters: FxHashMap::default(), + bcb_has_incoming_edge_counters: BitSet::new_empty(num_bcbs), + expressions: IndexVec::new(), + } + } + + /// Makes [`BcbCounter`] `Counter`s and `Expressions` for the `BasicCoverageBlock`s directly or + /// indirectly associated with coverage spans, and accumulates additional `Expression`s + /// representing intermediate values. + pub fn make_bcb_counters( + &mut self, + basic_coverage_blocks: &CoverageGraph, + bcb_has_coverage_spans: impl Fn(BasicCoverageBlock) -> bool, + ) { + MakeBcbCounters::new(self, basic_coverage_blocks).make_bcb_counters(bcb_has_coverage_spans) + } + + fn make_counter(&mut self) -> BcbCounter { + let id = self.next_counter(); + BcbCounter::Counter { id } + } + + fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter { + let expression = Expression { lhs: lhs.as_term(), op, rhs: rhs.as_term() }; + let id = self.expressions.push(expression); + BcbCounter::Expression { id } + } + + /// Variant of `make_expression` that makes `lhs` optional and assumes [`Op::Add`]. + /// + /// This is useful when using [`Iterator::fold`] to build an arbitrary-length sum. + fn make_sum_expression(&mut self, lhs: Option<BcbCounter>, rhs: BcbCounter) -> BcbCounter { + let Some(lhs) = lhs else { return rhs }; + self.make_expression(lhs, Op::Add, rhs) + } + + /// Counter IDs start from one and go up. + fn next_counter(&mut self) -> CounterId { + let next = self.next_counter_id; + self.next_counter_id = self.next_counter_id + 1; + next + } + + pub(super) fn num_counters(&self) -> usize { + self.next_counter_id.as_usize() + } + + #[cfg(test)] + pub(super) fn num_expressions(&self) -> usize { + self.expressions.len() + } + + fn set_bcb_counter(&mut self, bcb: BasicCoverageBlock, counter_kind: BcbCounter) -> BcbCounter { + assert!( + // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also + // have an expression (to be injected into an existing `BasicBlock` represented by this + // `BasicCoverageBlock`). + counter_kind.is_expression() || !self.bcb_has_incoming_edge_counters.contains(bcb), + "attempt to add a `Counter` to a BCB target with existing incoming edge counters" + ); + + if let Some(replaced) = self.bcb_counters[bcb].replace(counter_kind) { + bug!( + "attempt to set a BasicCoverageBlock coverage counter more than once; \ + {bcb:?} already had counter {replaced:?}", + ); + } else { + counter_kind + } + } + + fn set_bcb_edge_counter( + &mut self, + from_bcb: BasicCoverageBlock, + to_bcb: BasicCoverageBlock, + counter_kind: BcbCounter, + ) -> BcbCounter { + // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also + // have an expression (to be injected into an existing `BasicBlock` represented by this + // `BasicCoverageBlock`). + if let Some(node_counter) = self.bcb_counter(to_bcb) + && !node_counter.is_expression() + { + bug!( + "attempt to add an incoming edge counter from {from_bcb:?} \ + when the target BCB already has {node_counter:?}" + ); + } + + self.bcb_has_incoming_edge_counters.insert(to_bcb); + if let Some(replaced) = self.bcb_edge_counters.insert((from_bcb, to_bcb), counter_kind) { + bug!( + "attempt to set an edge counter more than once; from_bcb: \ + {from_bcb:?} already had counter {replaced:?}", + ); + } else { + counter_kind + } + } + + pub(super) fn bcb_counter(&self, bcb: BasicCoverageBlock) -> Option<BcbCounter> { + self.bcb_counters[bcb] + } + + pub(super) fn bcb_node_counters( + &self, + ) -> impl Iterator<Item = (BasicCoverageBlock, &BcbCounter)> { + self.bcb_counters + .iter_enumerated() + .filter_map(|(bcb, counter_kind)| Some((bcb, counter_kind.as_ref()?))) + } + + /// For each edge in the BCB graph that has an associated counter, yields + /// that edge's *from* and *to* nodes, and its counter. + pub(super) fn bcb_edge_counters( + &self, + ) -> impl Iterator<Item = (BasicCoverageBlock, BasicCoverageBlock, &BcbCounter)> { + self.bcb_edge_counters + .iter() + .map(|(&(from_bcb, to_bcb), counter_kind)| (from_bcb, to_bcb, counter_kind)) + } + + pub(super) fn take_expressions(&mut self) -> IndexVec<ExpressionId, Expression> { + std::mem::take(&mut self.expressions) + } +} + +/// Traverse the `CoverageGraph` and add either a `Counter` or `Expression` to every BCB, to be +/// injected with coverage spans. `Expressions` have no runtime overhead, so if a viable expression +/// (adding or subtracting two other counters or expressions) can compute the same result as an +/// embedded counter, an `Expression` should be used. +struct MakeBcbCounters<'a> { + coverage_counters: &'a mut CoverageCounters, + basic_coverage_blocks: &'a CoverageGraph, +} + +impl<'a> MakeBcbCounters<'a> { + fn new( + coverage_counters: &'a mut CoverageCounters, + basic_coverage_blocks: &'a CoverageGraph, + ) -> Self { + Self { coverage_counters, basic_coverage_blocks } + } + + /// If two `BasicCoverageBlock`s branch from another `BasicCoverageBlock`, one of the branches + /// can be counted by `Expression` by subtracting the other branch from the branching + /// block. Otherwise, the `BasicCoverageBlock` executed the least should have the `Counter`. + /// One way to predict which branch executes the least is by considering loops. A loop is exited + /// at a branch, so the branch that jumps to a `BasicCoverageBlock` outside the loop is almost + /// always executed less than the branch that does not exit the loop. + fn make_bcb_counters(&mut self, bcb_has_coverage_spans: impl Fn(BasicCoverageBlock) -> bool) { + debug!("make_bcb_counters(): adding a counter or expression to each BasicCoverageBlock"); + + // Walk the `CoverageGraph`. For each `BasicCoverageBlock` node with an associated + // coverage span, add a counter. If the `BasicCoverageBlock` branches, add a counter or + // expression to each branch `BasicCoverageBlock` (if the branch BCB has only one incoming + // edge) or edge from the branching BCB to the branch BCB (if the branch BCB has multiple + // incoming edges). + // + // The `TraverseCoverageGraphWithLoops` traversal ensures that, when a loop is encountered, + // all `BasicCoverageBlock` nodes in the loop are visited before visiting any node outside + // the loop. The `traversal` state includes a `context_stack`, providing a way to know if + // the current BCB is in one or more nested loops or not. + let mut traversal = TraverseCoverageGraphWithLoops::new(self.basic_coverage_blocks); + while let Some(bcb) = traversal.next() { + if bcb_has_coverage_spans(bcb) { + debug!("{:?} has at least one coverage span. Get or make its counter", bcb); + self.make_node_and_branch_counters(&traversal, bcb); + } else { + debug!( + "{:?} does not have any coverage spans. A counter will only be added if \ + and when a covered BCB has an expression dependency.", + bcb, + ); + } + } + + assert!( + traversal.is_complete(), + "`TraverseCoverageGraphWithLoops` missed some `BasicCoverageBlock`s: {:?}", + traversal.unvisited(), + ); + } + + fn make_node_and_branch_counters( + &mut self, + traversal: &TraverseCoverageGraphWithLoops<'_>, + from_bcb: BasicCoverageBlock, + ) { + // First, ensure that this node has a counter of some kind. + // We might also use its term later to compute one of the branch counters. + let from_bcb_operand = self.get_or_make_counter_operand(from_bcb); + + let branch_target_bcbs = self.basic_coverage_blocks.successors[from_bcb].as_slice(); + + // If this node doesn't have multiple out-edges, or all of its out-edges + // already have counters, then we don't need to create edge counters. + let needs_branch_counters = branch_target_bcbs.len() > 1 + && branch_target_bcbs + .iter() + .any(|&to_bcb| self.branch_has_no_counter(from_bcb, to_bcb)); + if !needs_branch_counters { + return; + } + + debug!( + "{from_bcb:?} has some branch(es) without counters:\n {}", + branch_target_bcbs + .iter() + .map(|&to_bcb| { + format!("{from_bcb:?}->{to_bcb:?}: {:?}", self.branch_counter(from_bcb, to_bcb)) + }) + .collect::<Vec<_>>() + .join("\n "), + ); + + // Of the branch edges that don't have counters yet, one can be given an expression + // (computed from the other edges) instead of a dedicated counter. + let expression_to_bcb = self.choose_preferred_expression_branch(traversal, from_bcb); + + // For each branch arm other than the one that was chosen to get an expression, + // ensure that it has a counter (existing counter/expression or a new counter), + // and accumulate the corresponding terms into a single sum term. + let sum_of_all_other_branches: BcbCounter = { + let _span = debug_span!("sum_of_all_other_branches", ?expression_to_bcb).entered(); + branch_target_bcbs + .iter() + .copied() + // Skip the chosen branch, since we'll calculate it from the other branches. + .filter(|&to_bcb| to_bcb != expression_to_bcb) + .fold(None, |accum, to_bcb| { + let _span = debug_span!("to_bcb", ?accum, ?to_bcb).entered(); + let branch_counter = self.get_or_make_edge_counter_operand(from_bcb, to_bcb); + Some(self.coverage_counters.make_sum_expression(accum, branch_counter)) + }) + .expect("there must be at least one other branch") + }; + + // For the branch that was chosen to get an expression, create that expression + // by taking the count of the node we're branching from, and subtracting the + // sum of all the other branches. + debug!( + "Making an expression for the selected expression_branch: \ + {expression_to_bcb:?} (expression_branch predecessors: {:?})", + self.bcb_predecessors(expression_to_bcb), + ); + let expression = self.coverage_counters.make_expression( + from_bcb_operand, + Op::Subtract, + sum_of_all_other_branches, + ); + debug!("{expression_to_bcb:?} gets an expression: {expression:?}"); + if self.basic_coverage_blocks.bcb_has_multiple_in_edges(expression_to_bcb) { + self.coverage_counters.set_bcb_edge_counter(from_bcb, expression_to_bcb, expression); + } else { + self.coverage_counters.set_bcb_counter(expression_to_bcb, expression); + } + } + + #[instrument(level = "debug", skip(self))] + fn get_or_make_counter_operand(&mut self, bcb: BasicCoverageBlock) -> BcbCounter { + // If the BCB already has a counter, return it. + if let Some(counter_kind) = self.coverage_counters.bcb_counters[bcb] { + debug!("{bcb:?} already has a counter: {counter_kind:?}"); + return counter_kind; + } + + // A BCB with only one incoming edge gets a simple `Counter` (via `make_counter()`). + // Also, a BCB that loops back to itself gets a simple `Counter`. This may indicate the + // program results in a tight infinite loop, but it should still compile. + let one_path_to_target = !self.basic_coverage_blocks.bcb_has_multiple_in_edges(bcb); + if one_path_to_target || self.bcb_predecessors(bcb).contains(&bcb) { + let counter_kind = self.coverage_counters.make_counter(); + if one_path_to_target { + debug!("{bcb:?} gets a new counter: {counter_kind:?}"); + } else { + debug!( + "{bcb:?} has itself as its own predecessor. It can't be part of its own \ + Expression sum, so it will get its own new counter: {counter_kind:?}. \ + (Note, the compiled code will generate an infinite loop.)", + ); + } + return self.coverage_counters.set_bcb_counter(bcb, counter_kind); + } + + // A BCB with multiple incoming edges can compute its count by ensuring that counters + // exist for each of those edges, and then adding them up to get a total count. + let sum_of_in_edges: BcbCounter = { + let _span = debug_span!("sum_of_in_edges", ?bcb).entered(); + // We avoid calling `self.bcb_predecessors` here so that we can + // call methods on `&mut self` inside the fold. + self.basic_coverage_blocks.predecessors[bcb] + .iter() + .copied() + .fold(None, |accum, from_bcb| { + let _span = debug_span!("from_bcb", ?accum, ?from_bcb).entered(); + let edge_counter = self.get_or_make_edge_counter_operand(from_bcb, bcb); + Some(self.coverage_counters.make_sum_expression(accum, edge_counter)) + }) + .expect("there must be at least one in-edge") + }; + + debug!("{bcb:?} gets a new counter (sum of predecessor counters): {sum_of_in_edges:?}"); + self.coverage_counters.set_bcb_counter(bcb, sum_of_in_edges) + } + + #[instrument(level = "debug", skip(self))] + fn get_or_make_edge_counter_operand( + &mut self, + from_bcb: BasicCoverageBlock, + to_bcb: BasicCoverageBlock, + ) -> BcbCounter { + // If the target BCB has only one in-edge (i.e. this one), then create + // a node counter instead, since it will have the same value. + if !self.basic_coverage_blocks.bcb_has_multiple_in_edges(to_bcb) { + assert_eq!([from_bcb].as_slice(), self.basic_coverage_blocks.predecessors[to_bcb]); + return self.get_or_make_counter_operand(to_bcb); + } + + // If the source BCB has only one successor (assumed to be the given target), an edge + // counter is unnecessary. Just get or make a counter for the source BCB. + if self.bcb_successors(from_bcb).len() == 1 { + return self.get_or_make_counter_operand(from_bcb); + } + + // If the edge already has a counter, return it. + if let Some(&counter_kind) = + self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb)) + { + debug!("Edge {from_bcb:?}->{to_bcb:?} already has a counter: {counter_kind:?}"); + return counter_kind; + } + + // Make a new counter to count this edge. + let counter_kind = self.coverage_counters.make_counter(); + debug!("Edge {from_bcb:?}->{to_bcb:?} gets a new counter: {counter_kind:?}"); + self.coverage_counters.set_bcb_edge_counter(from_bcb, to_bcb, counter_kind) + } + + /// Select a branch for the expression, either the recommended `reloop_branch`, or if none was + /// found, select any branch. + fn choose_preferred_expression_branch( + &self, + traversal: &TraverseCoverageGraphWithLoops<'_>, + from_bcb: BasicCoverageBlock, + ) -> BasicCoverageBlock { + let good_reloop_branch = self.find_good_reloop_branch(traversal, from_bcb); + if let Some(reloop_target) = good_reloop_branch { + assert!(self.branch_has_no_counter(from_bcb, reloop_target)); + debug!("Selecting reloop target {reloop_target:?} to get an expression"); + reloop_target + } else { + let &branch_without_counter = self + .bcb_successors(from_bcb) + .iter() + .find(|&&to_bcb| self.branch_has_no_counter(from_bcb, to_bcb)) + .expect( + "needs_branch_counters was `true` so there should be at least one \ + branch", + ); + debug!( + "Selecting any branch={:?} that still needs a counter, to get the \ + `Expression` because there was no `reloop_branch`, or it already had a \ + counter", + branch_without_counter + ); + branch_without_counter + } + } + + /// Tries to find a branch that leads back to the top of a loop, and that + /// doesn't already have a counter. Such branches are good candidates to + /// be given an expression (instead of a physical counter), because they + /// will tend to be executed more times than a loop-exit branch. + fn find_good_reloop_branch( + &self, + traversal: &TraverseCoverageGraphWithLoops<'_>, + from_bcb: BasicCoverageBlock, + ) -> Option<BasicCoverageBlock> { + let branch_target_bcbs = self.bcb_successors(from_bcb); + + // Consider each loop on the current traversal context stack, top-down. + for reloop_bcbs in traversal.reloop_bcbs_per_loop() { + let mut all_branches_exit_this_loop = true; + + // Try to find a branch that doesn't exit this loop and doesn't + // already have a counter. + for &branch_target_bcb in branch_target_bcbs { + // A branch is a reloop branch if it dominates any BCB that has + // an edge back to the loop header. (Other branches are exits.) + let is_reloop_branch = reloop_bcbs.iter().any(|&reloop_bcb| { + self.basic_coverage_blocks.dominates(branch_target_bcb, reloop_bcb) + }); + + if is_reloop_branch { + all_branches_exit_this_loop = false; + if self.branch_has_no_counter(from_bcb, branch_target_bcb) { + // We found a good branch to be given an expression. + return Some(branch_target_bcb); + } + // Keep looking for another reloop branch without a counter. + } else { + // This branch exits the loop. + } + } + + if !all_branches_exit_this_loop { + // We found one or more reloop branches, but all of them already + // have counters. Let the caller choose one of the exit branches. + debug!("All reloop branches had counters; skip checking the other loops"); + return None; + } + + // All of the branches exit this loop, so keep looking for a good + // reloop branch for one of the outer loops. + } + + None + } + + #[inline] + fn bcb_predecessors(&self, bcb: BasicCoverageBlock) -> &[BasicCoverageBlock] { + &self.basic_coverage_blocks.predecessors[bcb] + } + + #[inline] + fn bcb_successors(&self, bcb: BasicCoverageBlock) -> &[BasicCoverageBlock] { + &self.basic_coverage_blocks.successors[bcb] + } + + #[inline] + fn branch_has_no_counter( + &self, + from_bcb: BasicCoverageBlock, + to_bcb: BasicCoverageBlock, + ) -> bool { + self.branch_counter(from_bcb, to_bcb).is_none() + } + + fn branch_counter( + &self, + from_bcb: BasicCoverageBlock, + to_bcb: BasicCoverageBlock, + ) -> Option<&BcbCounter> { + if self.basic_coverage_blocks.bcb_has_multiple_in_edges(to_bcb) { + self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb)) + } else { + self.coverage_counters.bcb_counters[to_bcb].as_ref() + } + } +} diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs new file mode 100644 index 00000000000..263bfdaaaba --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/graph.rs @@ -0,0 +1,566 @@ +use rustc_data_structures::captures::Captures; +use rustc_data_structures::graph::dominators::{self, Dominators}; +use rustc_data_structures::graph::{self, GraphSuccessors, WithNumNodes, WithStartNode}; +use rustc_index::bit_set::BitSet; +use rustc_index::{IndexSlice, IndexVec}; +use rustc_middle::mir::{self, BasicBlock, TerminatorKind}; + +use std::cmp::Ordering; +use std::collections::VecDeque; +use std::ops::{Index, IndexMut}; + +/// A coverage-specific simplification of the MIR control flow graph (CFG). The `CoverageGraph`s +/// nodes are `BasicCoverageBlock`s, which encompass one or more MIR `BasicBlock`s. +#[derive(Debug)] +pub(super) struct CoverageGraph { + bcbs: IndexVec<BasicCoverageBlock, BasicCoverageBlockData>, + bb_to_bcb: IndexVec<BasicBlock, Option<BasicCoverageBlock>>, + pub successors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>, + pub predecessors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>, + dominators: Option<Dominators<BasicCoverageBlock>>, +} + +impl CoverageGraph { + pub fn from_mir(mir_body: &mir::Body<'_>) -> Self { + let (bcbs, bb_to_bcb) = Self::compute_basic_coverage_blocks(mir_body); + + // Pre-transform MIR `BasicBlock` successors and predecessors into the BasicCoverageBlock + // equivalents. Note that since the BasicCoverageBlock graph has been fully simplified, the + // each predecessor of a BCB leader_bb should be in a unique BCB. It is possible for a + // `SwitchInt` to have multiple targets to the same destination `BasicBlock`, so + // de-duplication is required. This is done without reordering the successors. + + let mut seen = IndexVec::from_elem(false, &bcbs); + let successors = IndexVec::from_fn_n( + |bcb| { + for b in seen.iter_mut() { + *b = false; + } + let bcb_data = &bcbs[bcb]; + let mut bcb_successors = Vec::new(); + for successor in bcb_filtered_successors(mir_body, bcb_data.last_bb()) + .filter_map(|successor_bb| bb_to_bcb[successor_bb]) + { + if !seen[successor] { + seen[successor] = true; + bcb_successors.push(successor); + } + } + bcb_successors + }, + bcbs.len(), + ); + + let mut predecessors = IndexVec::from_elem(Vec::new(), &bcbs); + for (bcb, bcb_successors) in successors.iter_enumerated() { + for &successor in bcb_successors { + predecessors[successor].push(bcb); + } + } + + let mut basic_coverage_blocks = + Self { bcbs, bb_to_bcb, successors, predecessors, dominators: None }; + let dominators = dominators::dominators(&basic_coverage_blocks); + basic_coverage_blocks.dominators = Some(dominators); + + // The coverage graph's entry-point node (bcb0) always starts with bb0, + // which never has predecessors. Any other blocks merged into bcb0 can't + // have multiple (coverage-relevant) predecessors, so bcb0 always has + // zero in-edges. + assert!(basic_coverage_blocks[START_BCB].leader_bb() == mir::START_BLOCK); + assert!(basic_coverage_blocks.predecessors[START_BCB].is_empty()); + + basic_coverage_blocks + } + + fn compute_basic_coverage_blocks( + mir_body: &mir::Body<'_>, + ) -> ( + IndexVec<BasicCoverageBlock, BasicCoverageBlockData>, + IndexVec<BasicBlock, Option<BasicCoverageBlock>>, + ) { + let num_basic_blocks = mir_body.basic_blocks.len(); + let mut bcbs = IndexVec::with_capacity(num_basic_blocks); + let mut bb_to_bcb = IndexVec::from_elem_n(None, num_basic_blocks); + + // Walk the MIR CFG using a Preorder traversal, which starts from `START_BLOCK` and follows + // each block terminator's `successors()`. Coverage spans must map to actual source code, + // so compiler generated blocks and paths can be ignored. To that end, the CFG traversal + // intentionally omits unwind paths. + // FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and + // `catch_unwind()` handlers. + + let mut basic_blocks = Vec::new(); + for bb in short_circuit_preorder(mir_body, bcb_filtered_successors) { + if let Some(last) = basic_blocks.last() { + let predecessors = &mir_body.basic_blocks.predecessors()[bb]; + if predecessors.len() > 1 || !predecessors.contains(last) { + // The `bb` has more than one _incoming_ edge, and should start its own + // `BasicCoverageBlockData`. (Note, the `basic_blocks` vector does not yet + // include `bb`; it contains a sequence of one or more sequential basic_blocks + // with no intermediate branches in or out. Save these as a new + // `BasicCoverageBlockData` before starting the new one.) + Self::add_basic_coverage_block( + &mut bcbs, + &mut bb_to_bcb, + basic_blocks.split_off(0), + ); + debug!( + " because {}", + if predecessors.len() > 1 { + "predecessors.len() > 1".to_owned() + } else { + format!("bb {} is not in predecessors: {:?}", bb.index(), predecessors) + } + ); + } + } + basic_blocks.push(bb); + + let term = mir_body[bb].terminator(); + + match term.kind { + TerminatorKind::Return { .. } + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Yield { .. } + | TerminatorKind::SwitchInt { .. } => { + // The `bb` has more than one _outgoing_ edge, or exits the function. Save the + // current sequence of `basic_blocks` gathered to this point, as a new + // `BasicCoverageBlockData`. + Self::add_basic_coverage_block( + &mut bcbs, + &mut bb_to_bcb, + basic_blocks.split_off(0), + ); + debug!(" because term.kind = {:?}", term.kind); + // Note that this condition is based on `TerminatorKind`, even though it + // theoretically boils down to `successors().len() != 1`; that is, either zero + // (e.g., `Return`, `Terminate`) or multiple successors (e.g., `SwitchInt`), but + // since the BCB CFG ignores things like unwind branches (which exist in the + // `Terminator`s `successors()` list) checking the number of successors won't + // work. + } + + // The following `TerminatorKind`s are either not expected outside an unwind branch, + // or they should not (under normal circumstances) branch. Coverage graphs are + // simplified by assuring coverage results are accurate for program executions that + // don't panic. + // + // Programs that panic and unwind may record slightly inaccurate coverage results + // for a coverage region containing the `Terminator` that began the panic. This + // is as intended. (See Issue #78544 for a possible future option to support + // coverage in test programs that panic.) + TerminatorKind::Goto { .. } + | TerminatorKind::UnwindResume + | TerminatorKind::Unreachable + | TerminatorKind::Drop { .. } + | TerminatorKind::Call { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::Assert { .. } + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::InlineAsm { .. } => {} + } + } + + if !basic_blocks.is_empty() { + // process any remaining basic_blocks into a final `BasicCoverageBlockData` + Self::add_basic_coverage_block(&mut bcbs, &mut bb_to_bcb, basic_blocks.split_off(0)); + debug!(" because the end of the MIR CFG was reached while traversing"); + } + + (bcbs, bb_to_bcb) + } + + fn add_basic_coverage_block( + bcbs: &mut IndexVec<BasicCoverageBlock, BasicCoverageBlockData>, + bb_to_bcb: &mut IndexSlice<BasicBlock, Option<BasicCoverageBlock>>, + basic_blocks: Vec<BasicBlock>, + ) { + let bcb = bcbs.next_index(); + for &bb in basic_blocks.iter() { + bb_to_bcb[bb] = Some(bcb); + } + let bcb_data = BasicCoverageBlockData::from(basic_blocks); + debug!("adding bcb{}: {:?}", bcb.index(), bcb_data); + bcbs.push(bcb_data); + } + + #[inline(always)] + pub fn iter_enumerated( + &self, + ) -> impl Iterator<Item = (BasicCoverageBlock, &BasicCoverageBlockData)> { + self.bcbs.iter_enumerated() + } + + #[inline(always)] + pub fn bcb_from_bb(&self, bb: BasicBlock) -> Option<BasicCoverageBlock> { + if bb.index() < self.bb_to_bcb.len() { self.bb_to_bcb[bb] } else { None } + } + + #[inline(always)] + pub fn dominates(&self, dom: BasicCoverageBlock, node: BasicCoverageBlock) -> bool { + self.dominators.as_ref().unwrap().dominates(dom, node) + } + + #[inline(always)] + pub fn cmp_in_dominator_order(&self, a: BasicCoverageBlock, b: BasicCoverageBlock) -> Ordering { + self.dominators.as_ref().unwrap().cmp_in_dominator_order(a, b) + } + + /// Returns true if the given node has 2 or more in-edges, i.e. 2 or more + /// predecessors. + /// + /// This property is interesting to code that assigns counters to nodes and + /// edges, because if a node _doesn't_ have multiple in-edges, then there's + /// no benefit in having a separate counter for its in-edge, because it + /// would have the same value as the node's own counter. + /// + /// FIXME: That assumption might not be true for [`TerminatorKind::Yield`]? + #[inline(always)] + pub(super) fn bcb_has_multiple_in_edges(&self, bcb: BasicCoverageBlock) -> bool { + // Even though bcb0 conceptually has an extra virtual in-edge due to + // being the entry point, we've already asserted that it has no _other_ + // in-edges, so there's no possibility of it having _multiple_ in-edges. + // (And since its virtual in-edge doesn't exist in the graph, that edge + // can't have a separate counter anyway.) + self.predecessors[bcb].len() > 1 + } +} + +impl Index<BasicCoverageBlock> for CoverageGraph { + type Output = BasicCoverageBlockData; + + #[inline] + fn index(&self, index: BasicCoverageBlock) -> &BasicCoverageBlockData { + &self.bcbs[index] + } +} + +impl IndexMut<BasicCoverageBlock> for CoverageGraph { + #[inline] + fn index_mut(&mut self, index: BasicCoverageBlock) -> &mut BasicCoverageBlockData { + &mut self.bcbs[index] + } +} + +impl graph::DirectedGraph for CoverageGraph { + type Node = BasicCoverageBlock; +} + +impl graph::WithNumNodes for CoverageGraph { + #[inline] + fn num_nodes(&self) -> usize { + self.bcbs.len() + } +} + +impl graph::WithStartNode for CoverageGraph { + #[inline] + fn start_node(&self) -> Self::Node { + self.bcb_from_bb(mir::START_BLOCK) + .expect("mir::START_BLOCK should be in a BasicCoverageBlock") + } +} + +type BcbSuccessors<'graph> = std::slice::Iter<'graph, BasicCoverageBlock>; + +impl<'graph> graph::GraphSuccessors<'graph> for CoverageGraph { + type Item = BasicCoverageBlock; + type Iter = std::iter::Cloned<BcbSuccessors<'graph>>; +} + +impl graph::WithSuccessors for CoverageGraph { + #[inline] + fn successors(&self, node: Self::Node) -> <Self as GraphSuccessors<'_>>::Iter { + self.successors[node].iter().cloned() + } +} + +impl<'graph> graph::GraphPredecessors<'graph> for CoverageGraph { + type Item = BasicCoverageBlock; + type Iter = std::iter::Copied<std::slice::Iter<'graph, BasicCoverageBlock>>; +} + +impl graph::WithPredecessors for CoverageGraph { + #[inline] + fn predecessors(&self, node: Self::Node) -> <Self as graph::GraphPredecessors<'_>>::Iter { + self.predecessors[node].iter().copied() + } +} + +rustc_index::newtype_index! { + /// A node in the control-flow graph of CoverageGraph. + #[orderable] + #[debug_format = "bcb{}"] + pub(super) struct BasicCoverageBlock { + const START_BCB = 0; + } +} + +/// `BasicCoverageBlockData` holds the data indexed by a `BasicCoverageBlock`. +/// +/// A `BasicCoverageBlock` (BCB) represents the maximal-length sequence of MIR `BasicBlock`s without +/// conditional branches, and form a new, simplified, coverage-specific Control Flow Graph, without +/// altering the original MIR CFG. +/// +/// Note that running the MIR `SimplifyCfg` transform is not sufficient (and therefore not +/// necessary). The BCB-based CFG is a more aggressive simplification. For example: +/// +/// * The BCB CFG ignores (trims) branches not relevant to coverage, such as unwind-related code, +/// that is injected by the Rust compiler but has no physical source code to count. This also +/// means a BasicBlock with a `Call` terminator can be merged into its primary successor target +/// block, in the same BCB. (But, note: Issue #78544: "MIR InstrumentCoverage: Improve coverage +/// of `#[should_panic]` tests and `catch_unwind()` handlers") +/// * Some BasicBlock terminators support Rust-specific concerns--like borrow-checking--that are +/// not relevant to coverage analysis. `FalseUnwind`, for example, can be treated the same as +/// a `Goto`, and merged with its successor into the same BCB. +/// +/// Each BCB with at least one computed coverage span will have no more than one `Counter`. +/// In some cases, a BCB's execution count can be computed by `Expression`. Additional +/// disjoint coverage spans in a BCB can also be counted by `Expression` (by adding `ZERO` +/// to the BCB's primary counter or expression). +/// +/// The BCB CFG is critical to simplifying the coverage analysis by ensuring graph path-based +/// queries (`dominates()`, `predecessors`, `successors`, etc.) have branch (control flow) +/// significance. +#[derive(Debug, Clone)] +pub(super) struct BasicCoverageBlockData { + pub basic_blocks: Vec<BasicBlock>, +} + +impl BasicCoverageBlockData { + pub fn from(basic_blocks: Vec<BasicBlock>) -> Self { + assert!(basic_blocks.len() > 0); + Self { basic_blocks } + } + + #[inline(always)] + pub fn leader_bb(&self) -> BasicBlock { + self.basic_blocks[0] + } + + #[inline(always)] + pub fn last_bb(&self) -> BasicBlock { + *self.basic_blocks.last().unwrap() + } +} + +// Returns the subset of a block's successors that are relevant to the coverage +// graph, i.e. those that do not represent unwinds or unreachable branches. +// FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and +// `catch_unwind()` handlers. +fn bcb_filtered_successors<'a, 'tcx>( + body: &'a mir::Body<'tcx>, + bb: BasicBlock, +) -> impl Iterator<Item = BasicBlock> + Captures<'a> + Captures<'tcx> { + let terminator = body[bb].terminator(); + + let take_n_successors = match terminator.kind { + // SwitchInt successors are never unwinds, so all of them should be traversed. + TerminatorKind::SwitchInt { .. } => usize::MAX, + // For all other kinds, return only the first successor (if any), ignoring any + // unwind successors. + _ => 1, + }; + + terminator + .successors() + .take(take_n_successors) + .filter(move |&successor| body[successor].terminator().kind != TerminatorKind::Unreachable) +} + +/// Maintains separate worklists for each loop in the BasicCoverageBlock CFG, plus one for the +/// CoverageGraph outside all loops. This supports traversing the BCB CFG in a way that +/// ensures a loop is completely traversed before processing Blocks after the end of the loop. +#[derive(Debug)] +pub(super) struct TraversalContext { + /// BCB with one or more incoming loop backedges, indicating which loop + /// this context is for. + /// + /// If `None`, this is the non-loop context for the function as a whole. + loop_header: Option<BasicCoverageBlock>, + + /// Worklist of BCBs to be processed in this context. + worklist: VecDeque<BasicCoverageBlock>, +} + +pub(super) struct TraverseCoverageGraphWithLoops<'a> { + basic_coverage_blocks: &'a CoverageGraph, + + backedges: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>, + context_stack: Vec<TraversalContext>, + visited: BitSet<BasicCoverageBlock>, +} + +impl<'a> TraverseCoverageGraphWithLoops<'a> { + pub(super) fn new(basic_coverage_blocks: &'a CoverageGraph) -> Self { + let backedges = find_loop_backedges(basic_coverage_blocks); + + let worklist = VecDeque::from([basic_coverage_blocks.start_node()]); + let context_stack = vec![TraversalContext { loop_header: None, worklist }]; + + // `context_stack` starts with a `TraversalContext` for the main function context (beginning + // with the `start` BasicCoverageBlock of the function). New worklists are pushed to the top + // of the stack as loops are entered, and popped off of the stack when a loop's worklist is + // exhausted. + let visited = BitSet::new_empty(basic_coverage_blocks.num_nodes()); + Self { basic_coverage_blocks, backedges, context_stack, visited } + } + + /// For each loop on the loop context stack (top-down), yields a list of BCBs + /// within that loop that have an outgoing edge back to the loop header. + pub(super) fn reloop_bcbs_per_loop(&self) -> impl Iterator<Item = &[BasicCoverageBlock]> { + self.context_stack + .iter() + .rev() + .filter_map(|context| context.loop_header) + .map(|header_bcb| self.backedges[header_bcb].as_slice()) + } + + pub(super) fn next(&mut self) -> Option<BasicCoverageBlock> { + debug!( + "TraverseCoverageGraphWithLoops::next - context_stack: {:?}", + self.context_stack.iter().rev().collect::<Vec<_>>() + ); + + while let Some(context) = self.context_stack.last_mut() { + if let Some(bcb) = context.worklist.pop_front() { + if !self.visited.insert(bcb) { + debug!("Already visited: {bcb:?}"); + continue; + } + debug!("Visiting {bcb:?}"); + + if self.backedges[bcb].len() > 0 { + debug!("{bcb:?} is a loop header! Start a new TraversalContext..."); + self.context_stack.push(TraversalContext { + loop_header: Some(bcb), + worklist: VecDeque::new(), + }); + } + self.add_successors_to_worklists(bcb); + return Some(bcb); + } else { + // Strip contexts with empty worklists from the top of the stack + self.context_stack.pop(); + } + } + + None + } + + pub fn add_successors_to_worklists(&mut self, bcb: BasicCoverageBlock) { + let successors = &self.basic_coverage_blocks.successors[bcb]; + debug!("{:?} has {} successors:", bcb, successors.len()); + + for &successor in successors { + if successor == bcb { + debug!( + "{:?} has itself as its own successor. (Note, the compiled code will \ + generate an infinite loop.)", + bcb + ); + // Don't re-add this successor to the worklist. We are already processing it. + // FIXME: This claims to skip just the self-successor, but it actually skips + // all other successors as well. Does that matter? + break; + } + + // Add successors of the current BCB to the appropriate context. Successors that + // stay within a loop are added to the BCBs context worklist. Successors that + // exit the loop (they are not dominated by the loop header) must be reachable + // from other BCBs outside the loop, and they will be added to a different + // worklist. + // + // Branching blocks (with more than one successor) must be processed before + // blocks with only one successor, to prevent unnecessarily complicating + // `Expression`s by creating a Counter in a `BasicCoverageBlock` that the + // branching block would have given an `Expression` (or vice versa). + + let context = self + .context_stack + .iter_mut() + .rev() + .find(|context| match context.loop_header { + Some(loop_header) => { + self.basic_coverage_blocks.dominates(loop_header, successor) + } + None => true, + }) + .unwrap_or_else(|| bug!("should always fall back to the root non-loop context")); + debug!("adding to worklist for {:?}", context.loop_header); + + // FIXME: The code below had debug messages claiming to add items to a + // particular end of the worklist, but was confused about which end was + // which. The existing behaviour has been preserved for now, but it's + // unclear what the intended behaviour was. + + if self.basic_coverage_blocks.successors[successor].len() > 1 { + context.worklist.push_back(successor); + } else { + context.worklist.push_front(successor); + } + } + } + + pub fn is_complete(&self) -> bool { + self.visited.count() == self.visited.domain_size() + } + + pub fn unvisited(&self) -> Vec<BasicCoverageBlock> { + let mut unvisited_set: BitSet<BasicCoverageBlock> = + BitSet::new_filled(self.visited.domain_size()); + unvisited_set.subtract(&self.visited); + unvisited_set.iter().collect::<Vec<_>>() + } +} + +pub(super) fn find_loop_backedges( + basic_coverage_blocks: &CoverageGraph, +) -> IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>> { + let num_bcbs = basic_coverage_blocks.num_nodes(); + let mut backedges = IndexVec::from_elem_n(Vec::<BasicCoverageBlock>::new(), num_bcbs); + + // Identify loops by their backedges. + for (bcb, _) in basic_coverage_blocks.iter_enumerated() { + for &successor in &basic_coverage_blocks.successors[bcb] { + if basic_coverage_blocks.dominates(successor, bcb) { + let loop_header = successor; + let backedge_from_bcb = bcb; + debug!( + "Found BCB backedge: {:?} -> loop_header: {:?}", + backedge_from_bcb, loop_header + ); + backedges[loop_header].push(backedge_from_bcb); + } + } + } + backedges +} + +fn short_circuit_preorder<'a, 'tcx, F, Iter>( + body: &'a mir::Body<'tcx>, + filtered_successors: F, +) -> impl Iterator<Item = BasicBlock> + Captures<'a> + Captures<'tcx> +where + F: Fn(&'a mir::Body<'tcx>, BasicBlock) -> Iter, + Iter: Iterator<Item = BasicBlock>, +{ + let mut visited = BitSet::new_empty(body.basic_blocks.len()); + let mut worklist = vec![mir::START_BLOCK]; + + std::iter::from_fn(move || { + while let Some(bb) = worklist.pop() { + if !visited.insert(bb) { + continue; + } + + worklist.extend(filtered_successors(body, bb)); + + return Some(bb); + } + + None + }) +} diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs new file mode 100644 index 00000000000..796150f9315 --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/mod.rs @@ -0,0 +1,378 @@ +pub mod query; + +mod counters; +mod graph; +mod spans; + +#[cfg(test)] +mod tests; + +use self::counters::{BcbCounter, CoverageCounters}; +use self::graph::CoverageGraph; +use self::spans::CoverageSpans; + +use crate::MirPass; + +use rustc_data_structures::sync::Lrc; +use rustc_middle::hir; +use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags; +use rustc_middle::mir::coverage::*; +use rustc_middle::mir::{ + self, BasicBlock, BasicBlockData, Coverage, SourceInfo, Statement, StatementKind, Terminator, + TerminatorKind, +}; +use rustc_middle::ty::TyCtxt; +use rustc_span::def_id::DefId; +use rustc_span::source_map::SourceMap; +use rustc_span::{ExpnKind, SourceFile, Span, Symbol}; + +/// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected +/// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen +/// to construct the coverage map. +pub struct InstrumentCoverage; + +impl<'tcx> MirPass<'tcx> for InstrumentCoverage { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.instrument_coverage() + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, mir_body: &mut mir::Body<'tcx>) { + let mir_source = mir_body.source; + + // If the InstrumentCoverage pass is called on promoted MIRs, skip them. + // See: https://github.com/rust-lang/rust/pull/73011#discussion_r438317601 + if mir_source.promoted.is_some() { + trace!( + "InstrumentCoverage skipped for {:?} (already promoted for Miri evaluation)", + mir_source.def_id() + ); + return; + } + + let is_fn_like = + tcx.hir().get_by_def_id(mir_source.def_id().expect_local()).fn_kind().is_some(); + + // Only instrument functions, methods, and closures (not constants since they are evaluated + // at compile time by Miri). + // FIXME(#73156): Handle source code coverage in const eval, but note, if and when const + // expressions get coverage spans, we will probably have to "carve out" space for const + // expressions from coverage spans in enclosing MIR's, like we do for closures. (That might + // be tricky if const expressions have no corresponding statements in the enclosing MIR. + // Closures are carved out by their initial `Assign` statement.) + if !is_fn_like { + trace!("InstrumentCoverage skipped for {:?} (not an fn-like)", mir_source.def_id()); + return; + } + + match mir_body.basic_blocks[mir::START_BLOCK].terminator().kind { + TerminatorKind::Unreachable => { + trace!("InstrumentCoverage skipped for unreachable `START_BLOCK`"); + return; + } + _ => {} + } + + let codegen_fn_attrs = tcx.codegen_fn_attrs(mir_source.def_id()); + if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_COVERAGE) { + return; + } + + trace!("InstrumentCoverage starting for {:?}", mir_source.def_id()); + Instrumentor::new(tcx, mir_body).inject_counters(); + trace!("InstrumentCoverage done for {:?}", mir_source.def_id()); + } +} + +struct Instrumentor<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + mir_body: &'a mut mir::Body<'tcx>, + source_file: Lrc<SourceFile>, + fn_sig_span: Span, + body_span: Span, + function_source_hash: u64, + basic_coverage_blocks: CoverageGraph, + coverage_counters: CoverageCounters, +} + +impl<'a, 'tcx> Instrumentor<'a, 'tcx> { + fn new(tcx: TyCtxt<'tcx>, mir_body: &'a mut mir::Body<'tcx>) -> Self { + let source_map = tcx.sess.source_map(); + let def_id = mir_body.source.def_id(); + let (some_fn_sig, hir_body) = fn_sig_and_body(tcx, def_id); + + let body_span = get_body_span(tcx, hir_body, mir_body); + + let source_file = source_map.lookup_source_file(body_span.lo()); + let fn_sig_span = match some_fn_sig.filter(|fn_sig| { + fn_sig.span.eq_ctxt(body_span) + && Lrc::ptr_eq(&source_file, &source_map.lookup_source_file(fn_sig.span.lo())) + }) { + Some(fn_sig) => fn_sig.span.with_hi(body_span.lo()), + None => body_span.shrink_to_lo(), + }; + + debug!( + "instrumenting {}: {:?}, fn sig span: {:?}, body span: {:?}", + if tcx.is_closure(def_id) { "closure" } else { "function" }, + def_id, + fn_sig_span, + body_span + ); + + let function_source_hash = hash_mir_source(tcx, hir_body); + let basic_coverage_blocks = CoverageGraph::from_mir(mir_body); + let coverage_counters = CoverageCounters::new(&basic_coverage_blocks); + + Self { + tcx, + mir_body, + source_file, + fn_sig_span, + body_span, + function_source_hash, + basic_coverage_blocks, + coverage_counters, + } + } + + fn inject_counters(&'a mut self) { + let fn_sig_span = self.fn_sig_span; + let body_span = self.body_span; + + //////////////////////////////////////////////////// + // Compute coverage spans from the `CoverageGraph`. + let coverage_spans = CoverageSpans::generate_coverage_spans( + self.mir_body, + fn_sig_span, + body_span, + &self.basic_coverage_blocks, + ); + + //////////////////////////////////////////////////// + // Create an optimized mix of `Counter`s and `Expression`s for the `CoverageGraph`. Ensure + // every coverage span has a `Counter` or `Expression` assigned to its `BasicCoverageBlock` + // and all `Expression` dependencies (operands) are also generated, for any other + // `BasicCoverageBlock`s not already associated with a coverage span. + let bcb_has_coverage_spans = |bcb| coverage_spans.bcb_has_coverage_spans(bcb); + self.coverage_counters + .make_bcb_counters(&self.basic_coverage_blocks, bcb_has_coverage_spans); + + let mappings = self.create_mappings_and_inject_coverage_statements(&coverage_spans); + + self.mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo { + function_source_hash: self.function_source_hash, + num_counters: self.coverage_counters.num_counters(), + expressions: self.coverage_counters.take_expressions(), + mappings, + })); + } + + /// For each [`BcbCounter`] associated with a BCB node or BCB edge, create + /// any corresponding mappings (for BCB nodes only), and inject any necessary + /// coverage statements into MIR. + fn create_mappings_and_inject_coverage_statements( + &mut self, + coverage_spans: &CoverageSpans, + ) -> Vec<Mapping> { + let source_map = self.tcx.sess.source_map(); + let body_span = self.body_span; + + use rustc_session::RemapFileNameExt; + let file_name = + Symbol::intern(&self.source_file.name.for_codegen(self.tcx.sess).to_string_lossy()); + + let mut mappings = Vec::new(); + + // Process the counters and spans associated with BCB nodes. + for (bcb, counter_kind) in self.coverage_counters.bcb_node_counters() { + let spans = coverage_spans.spans_for_bcb(bcb); + let has_mappings = !spans.is_empty(); + + // If this BCB has any coverage spans, add corresponding mappings to + // the mappings table. + if has_mappings { + let term = counter_kind.as_term(); + mappings.extend(spans.iter().map(|&span| { + let code_region = make_code_region(source_map, file_name, span, body_span); + Mapping { code_region, term } + })); + } + + let do_inject = match counter_kind { + // Counter-increment statements always need to be injected. + BcbCounter::Counter { .. } => true, + // The only purpose of expression-used statements is to detect + // when a mapping is unreachable, so we only inject them for + // expressions with one or more mappings. + BcbCounter::Expression { .. } => has_mappings, + }; + if do_inject { + inject_statement( + self.mir_body, + self.make_mir_coverage_kind(counter_kind), + self.basic_coverage_blocks[bcb].leader_bb(), + ); + } + } + + // Process the counters associated with BCB edges. + for (from_bcb, to_bcb, counter_kind) in self.coverage_counters.bcb_edge_counters() { + let do_inject = match counter_kind { + // Counter-increment statements always need to be injected. + BcbCounter::Counter { .. } => true, + // BCB-edge expressions never have mappings, so they never need + // a corresponding statement. + BcbCounter::Expression { .. } => false, + }; + if !do_inject { + continue; + } + + // We need to inject a coverage statement into a new BB between the + // last BB of `from_bcb` and the first BB of `to_bcb`. + let from_bb = self.basic_coverage_blocks[from_bcb].last_bb(); + let to_bb = self.basic_coverage_blocks[to_bcb].leader_bb(); + + let new_bb = inject_edge_counter_basic_block(self.mir_body, from_bb, to_bb); + debug!( + "Edge {from_bcb:?} (last {from_bb:?}) -> {to_bcb:?} (leader {to_bb:?}) \ + requires a new MIR BasicBlock {new_bb:?} for edge counter {counter_kind:?}", + ); + + // Inject a counter into the newly-created BB. + inject_statement(self.mir_body, self.make_mir_coverage_kind(counter_kind), new_bb); + } + + mappings + } + + fn make_mir_coverage_kind(&self, counter_kind: &BcbCounter) -> CoverageKind { + match *counter_kind { + BcbCounter::Counter { id } => CoverageKind::CounterIncrement { id }, + BcbCounter::Expression { id } => CoverageKind::ExpressionUsed { id }, + } + } +} + +fn inject_edge_counter_basic_block( + mir_body: &mut mir::Body<'_>, + from_bb: BasicBlock, + to_bb: BasicBlock, +) -> BasicBlock { + let span = mir_body[from_bb].terminator().source_info.span.shrink_to_hi(); + let new_bb = mir_body.basic_blocks_mut().push(BasicBlockData { + statements: vec![], // counter will be injected here + terminator: Some(Terminator { + source_info: SourceInfo::outermost(span), + kind: TerminatorKind::Goto { target: to_bb }, + }), + is_cleanup: false, + }); + let edge_ref = mir_body[from_bb] + .terminator_mut() + .successors_mut() + .find(|successor| **successor == to_bb) + .expect("from_bb should have a successor for to_bb"); + *edge_ref = new_bb; + new_bb +} + +fn inject_statement(mir_body: &mut mir::Body<'_>, counter_kind: CoverageKind, bb: BasicBlock) { + debug!(" injecting statement {counter_kind:?} for {bb:?}"); + let data = &mut mir_body[bb]; + let source_info = data.terminator().source_info; + let statement = Statement { + source_info, + kind: StatementKind::Coverage(Box::new(Coverage { kind: counter_kind })), + }; + data.statements.insert(0, statement); +} + +/// Convert the Span into its file name, start line and column, and end line and column +fn make_code_region( + source_map: &SourceMap, + file_name: Symbol, + span: Span, + body_span: Span, +) -> CodeRegion { + debug!( + "Called make_code_region(file_name={}, span={}, body_span={})", + file_name, + source_map.span_to_diagnostic_string(span), + source_map.span_to_diagnostic_string(body_span) + ); + + let (file, mut start_line, mut start_col, mut end_line, mut end_col) = + source_map.span_to_location_info(span); + if span.hi() == span.lo() { + // Extend an empty span by one character so the region will be counted. + if span.hi() == body_span.hi() { + start_col = start_col.saturating_sub(1); + } else { + end_col = start_col + 1; + } + }; + if let Some(file) = file { + start_line = source_map.doctest_offset_line(&file.name, start_line); + end_line = source_map.doctest_offset_line(&file.name, end_line); + } + CodeRegion { + file_name, + start_line: start_line as u32, + start_col: start_col as u32, + end_line: end_line as u32, + end_col: end_col as u32, + } +} + +fn fn_sig_and_body( + tcx: TyCtxt<'_>, + def_id: DefId, +) -> (Option<&rustc_hir::FnSig<'_>>, &rustc_hir::Body<'_>) { + // FIXME(#79625): Consider improving MIR to provide the information needed, to avoid going back + // to HIR for it. + let hir_node = tcx.hir().get_if_local(def_id).expect("expected DefId is local"); + let (_, fn_body_id) = + hir::map::associated_body(hir_node).expect("HIR node is a function with body"); + (hir_node.fn_sig(), tcx.hir().body(fn_body_id)) +} + +fn get_body_span<'tcx>( + tcx: TyCtxt<'tcx>, + hir_body: &rustc_hir::Body<'tcx>, + mir_body: &mut mir::Body<'tcx>, +) -> Span { + let mut body_span = hir_body.value.span; + let def_id = mir_body.source.def_id(); + + if tcx.is_closure(def_id) { + // If the MIR function is a closure, and if the closure body span + // starts from a macro, but it's content is not in that macro, try + // to find a non-macro callsite, and instrument the spans there + // instead. + loop { + let expn_data = body_span.ctxt().outer_expn_data(); + if expn_data.is_root() { + break; + } + if let ExpnKind::Macro { .. } = expn_data.kind { + body_span = expn_data.call_site; + } else { + break; + } + } + } + + body_span +} + +fn hash_mir_source<'tcx>(tcx: TyCtxt<'tcx>, hir_body: &'tcx rustc_hir::Body<'tcx>) -> u64 { + // FIXME(cjgillot) Stop hashing HIR manually here. + let owner = hir_body.id().hir_id.owner; + tcx.hir_owner_nodes(owner) + .unwrap() + .opt_hash_including_bodies + .unwrap() + .to_smaller_hash() + .as_u64() +} diff --git a/compiler/rustc_mir_transform/src/coverage/query.rs b/compiler/rustc_mir_transform/src/coverage/query.rs new file mode 100644 index 00000000000..dfc7c3a713b --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/query.rs @@ -0,0 +1,48 @@ +use super::*; + +use rustc_data_structures::captures::Captures; +use rustc_middle::mir::coverage::*; +use rustc_middle::mir::{Body, CoverageIdsInfo}; +use rustc_middle::query::Providers; +use rustc_middle::ty::{self}; + +/// A `query` provider for retrieving coverage information injected into MIR. +pub(crate) fn provide(providers: &mut Providers) { + providers.coverage_ids_info = |tcx, def_id| coverage_ids_info(tcx, def_id); +} + +/// Query implementation for `coverage_ids_info`. +fn coverage_ids_info<'tcx>( + tcx: TyCtxt<'tcx>, + instance_def: ty::InstanceDef<'tcx>, +) -> CoverageIdsInfo { + let mir_body = tcx.instance_mir(instance_def); + + let max_counter_id = all_coverage_in_mir_body(mir_body) + .filter_map(|coverage| match coverage.kind { + CoverageKind::CounterIncrement { id } => Some(id), + _ => None, + }) + .max() + .unwrap_or(CounterId::START); + + CoverageIdsInfo { max_counter_id } +} + +fn all_coverage_in_mir_body<'a, 'tcx>( + body: &'a Body<'tcx>, +) -> impl Iterator<Item = &'a Coverage> + Captures<'tcx> { + body.basic_blocks.iter().flat_map(|bb_data| &bb_data.statements).filter_map(|statement| { + match statement.kind { + StatementKind::Coverage(box ref coverage) if !is_inlined(body, statement) => { + Some(coverage) + } + _ => None, + } + }) +} + +fn is_inlined(body: &Body<'_>, statement: &Statement<'_>) -> bool { + let scope_data = &body.source_scopes[statement.source_info.scope]; + scope_data.inlined.is_some() || scope_data.inlined_parent_scope.is_some() +} diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs new file mode 100644 index 00000000000..05ad14f1525 --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/spans.rs @@ -0,0 +1,627 @@ +use std::cell::OnceCell; + +use rustc_data_structures::graph::WithNumNodes; +use rustc_index::IndexVec; +use rustc_middle::mir; +use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP}; + +use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB}; + +mod from_mir; + +pub(super) struct CoverageSpans { + /// Map from BCBs to their list of coverage spans. + bcb_to_spans: IndexVec<BasicCoverageBlock, Vec<Span>>, +} + +impl CoverageSpans { + pub(super) fn generate_coverage_spans( + mir_body: &mir::Body<'_>, + fn_sig_span: Span, + body_span: Span, + basic_coverage_blocks: &CoverageGraph, + ) -> Self { + let coverage_spans = CoverageSpansGenerator::generate_coverage_spans( + mir_body, + fn_sig_span, + body_span, + basic_coverage_blocks, + ); + + // Group the coverage spans by BCB, with the BCBs in sorted order. + let mut bcb_to_spans = IndexVec::from_elem_n(Vec::new(), basic_coverage_blocks.num_nodes()); + for CoverageSpan { bcb, span, .. } in coverage_spans { + bcb_to_spans[bcb].push(span); + } + + Self { bcb_to_spans } + } + + pub(super) fn bcb_has_coverage_spans(&self, bcb: BasicCoverageBlock) -> bool { + !self.bcb_to_spans[bcb].is_empty() + } + + pub(super) fn spans_for_bcb(&self, bcb: BasicCoverageBlock) -> &[Span] { + &self.bcb_to_spans[bcb] + } +} + +/// A BCB is deconstructed into one or more `Span`s. Each `Span` maps to a `CoverageSpan` that +/// references the originating BCB and one or more MIR `Statement`s and/or `Terminator`s. +/// Initially, the `Span`s come from the `Statement`s and `Terminator`s, but subsequent +/// transforms can combine adjacent `Span`s and `CoverageSpan` from the same BCB, merging the +/// `merged_spans` vectors, and the `Span`s to cover the extent of the combined `Span`s. +/// +/// Note: A span merged into another CoverageSpan may come from a `BasicBlock` that +/// is not part of the `CoverageSpan` bcb if the statement was included because it's `Span` matches +/// or is subsumed by the `Span` associated with this `CoverageSpan`, and it's `BasicBlock` +/// `dominates()` the `BasicBlock`s in this `CoverageSpan`. +#[derive(Debug, Clone)] +struct CoverageSpan { + pub span: Span, + pub expn_span: Span, + pub current_macro_or_none: OnceCell<Option<Symbol>>, + pub bcb: BasicCoverageBlock, + /// List of all the original spans from MIR that have been merged into this + /// span. Mainly used to precisely skip over gaps when truncating a span. + pub merged_spans: Vec<Span>, + pub is_closure: bool, +} + +impl CoverageSpan { + pub fn for_fn_sig(fn_sig_span: Span) -> Self { + Self::new(fn_sig_span, fn_sig_span, START_BCB, false) + } + + pub(super) fn new( + span: Span, + expn_span: Span, + bcb: BasicCoverageBlock, + is_closure: bool, + ) -> Self { + Self { + span, + expn_span, + current_macro_or_none: Default::default(), + bcb, + merged_spans: vec![span], + is_closure, + } + } + + pub fn merge_from(&mut self, other: &Self) { + debug_assert!(self.is_mergeable(other)); + self.span = self.span.to(other.span); + self.merged_spans.extend_from_slice(&other.merged_spans); + } + + pub fn cutoff_statements_at(&mut self, cutoff_pos: BytePos) { + self.merged_spans.retain(|span| span.hi() <= cutoff_pos); + if let Some(max_hi) = self.merged_spans.iter().map(|span| span.hi()).max() { + self.span = self.span.with_hi(max_hi); + } + } + + #[inline] + pub fn is_mergeable(&self, other: &Self) -> bool { + self.is_in_same_bcb(other) && !(self.is_closure || other.is_closure) + } + + #[inline] + pub fn is_in_same_bcb(&self, other: &Self) -> bool { + self.bcb == other.bcb + } + + /// If the span is part of a macro, returns the macro name symbol. + pub fn current_macro(&self) -> Option<Symbol> { + self.current_macro_or_none + .get_or_init(|| { + if let ExpnKind::Macro(MacroKind::Bang, current_macro) = + self.expn_span.ctxt().outer_expn_data().kind + { + return Some(current_macro); + } + None + }) + .map(|symbol| symbol) + } + + /// If the span is part of a macro, and the macro is visible (expands directly to the given + /// body_span), returns the macro name symbol. + pub fn visible_macro(&self, body_span: Span) -> Option<Symbol> { + let current_macro = self.current_macro()?; + let parent_callsite = self.expn_span.parent_callsite()?; + + // In addition to matching the context of the body span, the parent callsite + // must also be the source callsite, i.e. the parent must have no parent. + let is_visible_macro = + parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span); + is_visible_macro.then_some(current_macro) + } + + pub fn is_macro_expansion(&self) -> bool { + self.current_macro().is_some() + } +} + +/// Converts the initial set of `CoverageSpan`s (one per MIR `Statement` or `Terminator`) into a +/// minimal set of `CoverageSpan`s, using the BCB CFG to determine where it is safe and useful to: +/// +/// * Remove duplicate source code coverage regions +/// * Merge spans that represent continuous (both in source code and control flow), non-branching +/// execution +/// * Carve out (leave uncovered) any span that will be counted by another MIR (notably, closures) +struct CoverageSpansGenerator<'a> { + /// A `Span` covering the function body of the MIR (typically from left curly brace to right + /// curly brace). + body_span: Span, + + /// The BasicCoverageBlock Control Flow Graph (BCB CFG). + basic_coverage_blocks: &'a CoverageGraph, + + /// The initial set of `CoverageSpan`s, sorted by `Span` (`lo` and `hi`) and by relative + /// dominance between the `BasicCoverageBlock`s of equal `Span`s. + sorted_spans_iter: std::vec::IntoIter<CoverageSpan>, + + /// The current `CoverageSpan` to compare to its `prev`, to possibly merge, discard, force the + /// discard of the `prev` (and or `pending_dups`), or keep both (with `prev` moved to + /// `pending_dups`). If `curr` is not discarded or merged, it becomes `prev` for the next + /// iteration. + some_curr: Option<CoverageSpan>, + + /// The original `span` for `curr`, in case `curr.span()` is modified. The `curr_original_span` + /// **must not be mutated** (except when advancing to the next `curr`), even if `curr.span()` + /// is mutated. + curr_original_span: Span, + + /// The CoverageSpan from a prior iteration; typically assigned from that iteration's `curr`. + /// If that `curr` was discarded, `prev` retains its value from the previous iteration. + some_prev: Option<CoverageSpan>, + + /// Assigned from `curr_original_span` from the previous iteration. The `prev_original_span` + /// **must not be mutated** (except when advancing to the next `prev`), even if `prev.span()` + /// is mutated. + prev_original_span: Span, + + /// One or more `CoverageSpan`s with the same `Span` but different `BasicCoverageBlock`s, and + /// no `BasicCoverageBlock` in this list dominates another `BasicCoverageBlock` in the list. + /// If a new `curr` span also fits this criteria (compared to an existing list of + /// `pending_dups`), that `curr` `CoverageSpan` moves to `prev` before possibly being added to + /// the `pending_dups` list, on the next iteration. As a result, if `prev` and `pending_dups` + /// have the same `Span`, the criteria for `pending_dups` holds for `prev` as well: a `prev` + /// with a matching `Span` does not dominate any `pending_dup` and no `pending_dup` dominates a + /// `prev` with a matching `Span`) + pending_dups: Vec<CoverageSpan>, + + /// The final `CoverageSpan`s to add to the coverage map. A `Counter` or `Expression` + /// will also be injected into the MIR for each `CoverageSpan`. + refined_spans: Vec<CoverageSpan>, +} + +impl<'a> CoverageSpansGenerator<'a> { + /// Generate a minimal set of `CoverageSpan`s, each representing a contiguous code region to be + /// counted. + /// + /// The basic steps are: + /// + /// 1. Extract an initial set of spans from the `Statement`s and `Terminator`s of each + /// `BasicCoverageBlockData`. + /// 2. Sort the spans by span.lo() (starting position). Spans that start at the same position + /// are sorted with longer spans before shorter spans; and equal spans are sorted + /// (deterministically) based on "dominator" relationship (if any). + /// 3. Traverse the spans in sorted order to identify spans that can be dropped (for instance, + /// if another span or spans are already counting the same code region), or should be merged + /// into a broader combined span (because it represents a contiguous, non-branching, and + /// uninterrupted region of source code). + /// + /// Closures are exposed in their enclosing functions as `Assign` `Rvalue`s, and since + /// closures have their own MIR, their `Span` in their enclosing function should be left + /// "uncovered". + /// + /// Note the resulting vector of `CoverageSpan`s may not be fully sorted (and does not need + /// to be). + pub(super) fn generate_coverage_spans( + mir_body: &mir::Body<'_>, + fn_sig_span: Span, // Ensured to be same SourceFile and SyntaxContext as `body_span` + body_span: Span, + basic_coverage_blocks: &'a CoverageGraph, + ) -> Vec<CoverageSpan> { + let sorted_spans = from_mir::mir_to_initial_sorted_coverage_spans( + mir_body, + fn_sig_span, + body_span, + basic_coverage_blocks, + ); + + let coverage_spans = Self { + body_span, + basic_coverage_blocks, + sorted_spans_iter: sorted_spans.into_iter(), + some_curr: None, + curr_original_span: DUMMY_SP, + some_prev: None, + prev_original_span: DUMMY_SP, + pending_dups: Vec::new(), + refined_spans: Vec::with_capacity(basic_coverage_blocks.num_nodes() * 2), + }; + + coverage_spans.to_refined_spans() + } + + /// Iterate through the sorted `CoverageSpan`s, and return the refined list of merged and + /// de-duplicated `CoverageSpan`s. + fn to_refined_spans(mut self) -> Vec<CoverageSpan> { + while self.next_coverage_span() { + // For the first span we don't have `prev` set, so most of the + // span-processing steps don't make sense yet. + if self.some_prev.is_none() { + debug!(" initial span"); + self.maybe_push_macro_name_span(); + continue; + } + + // The remaining cases assume that `prev` and `curr` are set. + let prev = self.prev(); + let curr = self.curr(); + + if curr.is_mergeable(prev) { + debug!(" same bcb (and neither is a closure), merge with prev={prev:?}"); + let prev = self.take_prev(); + self.curr_mut().merge_from(&prev); + self.maybe_push_macro_name_span(); + // Note that curr.span may now differ from curr_original_span + } else if prev.span.hi() <= curr.span.lo() { + debug!( + " different bcbs and disjoint spans, so keep curr for next iter, and add prev={prev:?}", + ); + let prev = self.take_prev(); + self.refined_spans.push(prev); + self.maybe_push_macro_name_span(); + } else if prev.is_closure { + // drop any equal or overlapping span (`curr`) and keep `prev` to test again in the + // next iter + debug!( + " curr overlaps a closure (prev). Drop curr and keep prev for next iter. prev={prev:?}", + ); + self.take_curr(); // Discards curr. + } else if curr.is_closure { + self.carve_out_span_for_closure(); + } else if self.prev_original_span == curr.span { + // Note that this compares the new (`curr`) span to `prev_original_span`. + // In this branch, the actual span byte range of `prev_original_span` is not + // important. What is important is knowing whether the new `curr` span was + // **originally** the same as the original span of `prev()`. The original spans + // reflect their original sort order, and for equal spans, conveys a partial + // ordering based on CFG dominator priority. + if prev.is_macro_expansion() && curr.is_macro_expansion() { + // Macros that expand to include branching (such as + // `assert_eq!()`, `assert_ne!()`, `info!()`, `debug!()`, or + // `trace!()`) typically generate callee spans with identical + // ranges (typically the full span of the macro) for all + // `BasicBlocks`. This makes it impossible to distinguish + // the condition (`if val1 != val2`) from the optional + // branched statements (such as the call to `panic!()` on + // assert failure). In this case it is better (or less + // worse) to drop the optional branch bcbs and keep the + // non-conditional statements, to count when reached. + debug!( + " curr and prev are part of a macro expansion, and curr has the same span \ + as prev, but is in a different bcb. Drop curr and keep prev for next iter. \ + prev={prev:?}", + ); + self.take_curr(); // Discards curr. + } else { + self.update_pending_dups(); + } + } else { + self.cutoff_prev_at_overlapping_curr(); + self.maybe_push_macro_name_span(); + } + } + + // Drain any remaining dups into the output. + for dup in self.pending_dups.drain(..) { + debug!(" ...adding at least one pending dup={:?}", dup); + self.refined_spans.push(dup); + } + + // There is usually a final span remaining in `prev` after the loop ends, + // so add it to the output as well. + if let Some(prev) = self.some_prev.take() { + debug!(" AT END, adding last prev={prev:?}"); + self.refined_spans.push(prev); + } + + // Do one last merge pass, to simplify the output. + self.refined_spans.dedup_by(|b, a| { + if a.is_mergeable(b) { + debug!(?a, ?b, "merging list-adjacent refined spans"); + a.merge_from(b); + true + } else { + false + } + }); + + // Remove `CoverageSpan`s derived from closures, originally added to ensure the coverage + // regions for the current function leave room for the closure's own coverage regions + // (injected separately, from the closure's own MIR). + self.refined_spans.retain(|covspan| !covspan.is_closure); + self.refined_spans + } + + /// If `curr` is part of a new macro expansion, carve out and push a separate + /// span that ends just after the macro name and its subsequent `!`. + fn maybe_push_macro_name_span(&mut self) { + let curr = self.curr(); + + let Some(visible_macro) = curr.visible_macro(self.body_span) else { return }; + if let Some(prev) = &self.some_prev + && prev.expn_span.eq_ctxt(curr.expn_span) + { + return; + } + + // The split point is relative to `curr_original_span`, + // because `curr.span` may have been merged with preceding spans. + let split_point_after_macro_bang = self.curr_original_span.lo() + + BytePos(visible_macro.as_str().len() as u32) + + BytePos(1); // add 1 for the `!` + debug_assert!(split_point_after_macro_bang <= curr.span.hi()); + if split_point_after_macro_bang > curr.span.hi() { + // Something is wrong with the macro name span; + // return now to avoid emitting malformed mappings (e.g. #117788). + return; + } + + let mut macro_name_cov = curr.clone(); + macro_name_cov.span = macro_name_cov.span.with_hi(split_point_after_macro_bang); + self.curr_mut().span = curr.span.with_lo(split_point_after_macro_bang); + + debug!( + " and curr starts a new macro expansion, so add a new span just for \ + the macro `{visible_macro}!`, new span={macro_name_cov:?}", + ); + self.refined_spans.push(macro_name_cov); + } + + #[track_caller] + fn curr(&self) -> &CoverageSpan { + self.some_curr.as_ref().unwrap_or_else(|| bug!("some_curr is None (curr)")) + } + + #[track_caller] + fn curr_mut(&mut self) -> &mut CoverageSpan { + self.some_curr.as_mut().unwrap_or_else(|| bug!("some_curr is None (curr_mut)")) + } + + /// If called, then the next call to `next_coverage_span()` will *not* update `prev` with the + /// `curr` coverage span. + #[track_caller] + fn take_curr(&mut self) -> CoverageSpan { + self.some_curr.take().unwrap_or_else(|| bug!("some_curr is None (take_curr)")) + } + + #[track_caller] + fn prev(&self) -> &CoverageSpan { + self.some_prev.as_ref().unwrap_or_else(|| bug!("some_prev is None (prev)")) + } + + #[track_caller] + fn prev_mut(&mut self) -> &mut CoverageSpan { + self.some_prev.as_mut().unwrap_or_else(|| bug!("some_prev is None (prev_mut)")) + } + + #[track_caller] + fn take_prev(&mut self) -> CoverageSpan { + self.some_prev.take().unwrap_or_else(|| bug!("some_prev is None (take_prev)")) + } + + /// If there are `pending_dups` but `prev` is not a matching dup (`prev.span` doesn't match the + /// `pending_dups` spans), then one of the following two things happened during the previous + /// iteration: + /// * the previous `curr` span (which is now `prev`) was not a duplicate of the pending_dups + /// (in which case there should be at least two spans in `pending_dups`); or + /// * the `span` of `prev` was modified by `curr_mut().merge_from(prev)` (in which case + /// `pending_dups` could have as few as one span) + /// In either case, no more spans will match the span of `pending_dups`, so + /// add the `pending_dups` if they don't overlap `curr`, and clear the list. + fn maybe_flush_pending_dups(&mut self) { + let Some(last_dup) = self.pending_dups.last() else { return }; + if last_dup.span == self.prev().span { + return; + } + + debug!( + " SAME spans, but pending_dups are NOT THE SAME, so BCBs matched on \ + previous iteration, or prev started a new disjoint span" + ); + if last_dup.span.hi() <= self.curr().span.lo() { + for dup in self.pending_dups.drain(..) { + debug!(" ...adding at least one pending={:?}", dup); + self.refined_spans.push(dup); + } + } else { + self.pending_dups.clear(); + } + assert!(self.pending_dups.is_empty()); + } + + /// Advance `prev` to `curr` (if any), and `curr` to the next `CoverageSpan` in sorted order. + fn next_coverage_span(&mut self) -> bool { + if let Some(curr) = self.some_curr.take() { + self.some_prev = Some(curr); + self.prev_original_span = self.curr_original_span; + } + while let Some(curr) = self.sorted_spans_iter.next() { + debug!("FOR curr={:?}", curr); + if let Some(prev) = &self.some_prev + && prev.span.lo() > curr.span.lo() + { + // Skip curr because prev has already advanced beyond the end of curr. + // This can only happen if a prior iteration updated `prev` to skip past + // a region of code, such as skipping past a closure. + debug!( + " prev.span starts after curr.span, so curr will be dropped (skipping past \ + closure?); prev={prev:?}", + ); + } else { + // Save a copy of the original span for `curr` in case the `CoverageSpan` is changed + // by `self.curr_mut().merge_from(prev)`. + self.curr_original_span = curr.span; + self.some_curr.replace(curr); + self.maybe_flush_pending_dups(); + return true; + } + } + false + } + + /// If `prev`s span extends left of the closure (`curr`), carve out the closure's span from + /// `prev`'s span. (The closure's coverage counters will be injected when processing the + /// closure's own MIR.) Add the portion of the span to the left of the closure; and if the span + /// extends to the right of the closure, update `prev` to that portion of the span. For any + /// `pending_dups`, repeat the same process. + fn carve_out_span_for_closure(&mut self) { + let prev = self.prev(); + let curr = self.curr(); + + let left_cutoff = curr.span.lo(); + let right_cutoff = curr.span.hi(); + let has_pre_closure_span = prev.span.lo() < right_cutoff; + let has_post_closure_span = prev.span.hi() > right_cutoff; + + if has_pre_closure_span { + let mut pre_closure = self.prev().clone(); + pre_closure.span = pre_closure.span.with_hi(left_cutoff); + debug!(" prev overlaps a closure. Adding span for pre_closure={:?}", pre_closure); + + for mut dup in self.pending_dups.iter().cloned() { + dup.span = dup.span.with_hi(left_cutoff); + debug!(" ...and at least one pre_closure dup={:?}", dup); + self.refined_spans.push(dup); + } + + self.refined_spans.push(pre_closure); + } + + if has_post_closure_span { + // Mutate `prev.span()` to start after the closure (and discard curr). + // (**NEVER** update `prev_original_span` because it affects the assumptions + // about how the `CoverageSpan`s are ordered.) + self.prev_mut().span = self.prev().span.with_lo(right_cutoff); + debug!(" Mutated prev.span to start after the closure. prev={:?}", self.prev()); + + for dup in &mut self.pending_dups { + debug!(" ...and at least one overlapping dup={:?}", dup); + dup.span = dup.span.with_lo(right_cutoff); + } + + let closure_covspan = self.take_curr(); // Prevent this curr from becoming prev. + self.refined_spans.push(closure_covspan); // since self.prev() was already updated + } else { + self.pending_dups.clear(); + } + } + + /// Called if `curr.span` equals `prev_original_span` (and potentially equal to all + /// `pending_dups` spans, if any). Keep in mind, `prev.span()` may have been changed. + /// If prev.span() was merged into other spans (with matching BCB, for instance), + /// `prev.span.hi()` will be greater than (further right of) `prev_original_span.hi()`. + /// If prev.span() was split off to the right of a closure, prev.span().lo() will be + /// greater than prev_original_span.lo(). The actual span of `prev_original_span` is + /// not as important as knowing that `prev()` **used to have the same span** as `curr()`, + /// which means their sort order is still meaningful for determining the dominator + /// relationship. + /// + /// When two `CoverageSpan`s have the same `Span`, dominated spans can be discarded; but if + /// neither `CoverageSpan` dominates the other, both (or possibly more than two) are held, + /// until their disposition is determined. In this latter case, the `prev` dup is moved into + /// `pending_dups` so the new `curr` dup can be moved to `prev` for the next iteration. + fn update_pending_dups(&mut self) { + let prev_bcb = self.prev().bcb; + let curr_bcb = self.curr().bcb; + + // Equal coverage spans are ordered by dominators before dominated (if any), so it should be + // impossible for `curr` to dominate any previous `CoverageSpan`. + debug_assert!(!self.basic_coverage_blocks.dominates(curr_bcb, prev_bcb)); + + let initial_pending_count = self.pending_dups.len(); + if initial_pending_count > 0 { + self.pending_dups + .retain(|dup| !self.basic_coverage_blocks.dominates(dup.bcb, curr_bcb)); + + let n_discarded = initial_pending_count - self.pending_dups.len(); + if n_discarded > 0 { + debug!( + " discarded {n_discarded} of {initial_pending_count} pending_dups that dominated curr", + ); + } + } + + if self.basic_coverage_blocks.dominates(prev_bcb, curr_bcb) { + debug!( + " different bcbs but SAME spans, and prev dominates curr. Discard prev={:?}", + self.prev() + ); + self.cutoff_prev_at_overlapping_curr(); + // If one span dominates the other, associate the span with the code from the dominated + // block only (`curr`), and discard the overlapping portion of the `prev` span. (Note + // that if `prev.span` is wider than `prev_original_span`, a `CoverageSpan` will still + // be created for `prev`s block, for the non-overlapping portion, left of `curr.span`.) + // + // For example: + // match somenum { + // x if x < 1 => { ... } + // }... + // + // The span for the first `x` is referenced by both the pattern block (every time it is + // evaluated) and the arm code (only when matched). The counter will be applied only to + // the dominated block. This allows coverage to track and highlight things like the + // assignment of `x` above, if the branch is matched, making `x` available to the arm + // code; and to track and highlight the question mark `?` "try" operator at the end of + // a function call returning a `Result`, so the `?` is covered when the function returns + // an `Err`, and not counted as covered if the function always returns `Ok`. + } else { + // Save `prev` in `pending_dups`. (`curr` will become `prev` in the next iteration.) + // If the `curr` CoverageSpan is later discarded, `pending_dups` can be discarded as + // well; but if `curr` is added to refined_spans, the `pending_dups` will also be added. + debug!( + " different bcbs but SAME spans, and neither dominates, so keep curr for \ + next iter, and, pending upcoming spans (unless overlapping) add prev={:?}", + self.prev() + ); + let prev = self.take_prev(); + self.pending_dups.push(prev); + } + } + + /// `curr` overlaps `prev`. If `prev`s span extends left of `curr`s span, keep _only_ + /// statements that end before `curr.lo()` (if any), and add the portion of the + /// combined span for those statements. Any other statements have overlapping spans + /// that can be ignored because `curr` and/or other upcoming statements/spans inside + /// the overlap area will produce their own counters. This disambiguation process + /// avoids injecting multiple counters for overlapping spans, and the potential for + /// double-counting. + fn cutoff_prev_at_overlapping_curr(&mut self) { + debug!( + " different bcbs, overlapping spans, so ignore/drop pending and only add prev \ + if it has statements that end before curr; prev={:?}", + self.prev() + ); + if self.pending_dups.is_empty() { + let curr_span = self.curr().span; + self.prev_mut().cutoff_statements_at(curr_span.lo()); + if self.prev().merged_spans.is_empty() { + debug!(" ... no non-overlapping statements to add"); + } else { + debug!(" ... adding modified prev={:?}", self.prev()); + let prev = self.take_prev(); + self.refined_spans.push(prev); + } + } else { + // with `pending_dups`, `prev` cannot have any statements that don't overlap + self.pending_dups.clear(); + } + } +} diff --git a/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs new file mode 100644 index 00000000000..eab9a9c98f8 --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs @@ -0,0 +1,214 @@ +use rustc_data_structures::captures::Captures; +use rustc_middle::mir::{ + self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator, + TerminatorKind, +}; +use rustc_span::Span; + +use crate::coverage::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph}; +use crate::coverage::spans::CoverageSpan; + +pub(super) fn mir_to_initial_sorted_coverage_spans( + mir_body: &mir::Body<'_>, + fn_sig_span: Span, + body_span: Span, + basic_coverage_blocks: &CoverageGraph, +) -> Vec<CoverageSpan> { + let mut initial_spans = Vec::with_capacity(mir_body.basic_blocks.len() * 2); + for (bcb, bcb_data) in basic_coverage_blocks.iter_enumerated() { + initial_spans.extend(bcb_to_initial_coverage_spans(mir_body, body_span, bcb, bcb_data)); + } + + if initial_spans.is_empty() { + // This can happen if, for example, the function is unreachable (contains only a + // `BasicBlock`(s) with an `Unreachable` terminator). + return initial_spans; + } + + initial_spans.push(CoverageSpan::for_fn_sig(fn_sig_span)); + + initial_spans.sort_by(|a, b| { + // First sort by span start. + Ord::cmp(&a.span.lo(), &b.span.lo()) + // If span starts are the same, sort by span end in reverse order. + // This ensures that if spans A and B are adjacent in the list, + // and they overlap but are not equal, then either: + // - Span A extends further left, or + // - Both have the same start and span A extends further right + .then_with(|| Ord::cmp(&a.span.hi(), &b.span.hi()).reverse()) + // If both spans are equal, sort the BCBs in dominator order, + // so that dominating BCBs come before other BCBs they dominate. + .then_with(|| basic_coverage_blocks.cmp_in_dominator_order(a.bcb, b.bcb)) + // If two spans are otherwise identical, put closure spans first, + // as this seems to be what the refinement step expects. + .then_with(|| Ord::cmp(&a.is_closure, &b.is_closure).reverse()) + }); + + // The desugaring of an async function includes a closure containing the + // original function body, and a terminator that returns the `impl Future`. + // That terminator will cause a confusing coverage count for the function's + // closing brace, so discard everything after the body closure span. + if let Some(body_closure_index) = + initial_spans.iter().rposition(|covspan| covspan.is_closure && covspan.span == body_span) + { + initial_spans.truncate(body_closure_index + 1); + } + + initial_spans +} + +// Generate a set of `CoverageSpan`s from the filtered set of `Statement`s and `Terminator`s of +// the `BasicBlock`(s) in the given `BasicCoverageBlockData`. One `CoverageSpan` is generated +// for each `Statement` and `Terminator`. (Note that subsequent stages of coverage analysis will +// merge some `CoverageSpan`s, at which point a `CoverageSpan` may represent multiple +// `Statement`s and/or `Terminator`s.) +fn bcb_to_initial_coverage_spans<'a, 'tcx>( + mir_body: &'a mir::Body<'tcx>, + body_span: Span, + bcb: BasicCoverageBlock, + bcb_data: &'a BasicCoverageBlockData, +) -> impl Iterator<Item = CoverageSpan> + Captures<'a> + Captures<'tcx> { + bcb_data.basic_blocks.iter().flat_map(move |&bb| { + let data = &mir_body[bb]; + + let statement_spans = data.statements.iter().filter_map(move |statement| { + let expn_span = filtered_statement_span(statement)?; + let span = unexpand_into_body_span(expn_span, body_span)?; + + Some(CoverageSpan::new(span, expn_span, bcb, is_closure(statement))) + }); + + let terminator_span = Some(data.terminator()).into_iter().filter_map(move |terminator| { + let expn_span = filtered_terminator_span(terminator)?; + let span = unexpand_into_body_span(expn_span, body_span)?; + + Some(CoverageSpan::new(span, expn_span, bcb, false)) + }); + + statement_spans.chain(terminator_span) + }) +} + +fn is_closure(statement: &Statement<'_>) -> bool { + match statement.kind { + StatementKind::Assign(box (_, Rvalue::Aggregate(box ref agg_kind, _))) => match agg_kind { + AggregateKind::Closure(_, _) | AggregateKind::Coroutine(_, _, _) => true, + _ => false, + }, + _ => false, + } +} + +/// If the MIR `Statement` has a span contributive to computing coverage spans, +/// return it; otherwise return `None`. +fn filtered_statement_span(statement: &Statement<'_>) -> Option<Span> { + use mir::coverage::CoverageKind; + + match statement.kind { + // These statements have spans that are often outside the scope of the executed source code + // for their parent `BasicBlock`. + StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + // Ignore `ConstEvalCounter`s + | StatementKind::ConstEvalCounter + // Ignore `Nop`s + | StatementKind::Nop => None, + + // FIXME(#78546): MIR InstrumentCoverage - Can the source_info.span for `FakeRead` + // statements be more consistent? + // + // FakeReadCause::ForGuardBinding, in this example: + // match somenum { + // x if x < 1 => { ... } + // }... + // The BasicBlock within the match arm code included one of these statements, but the span + // for it covered the `1` in this source. The actual statements have nothing to do with that + // source span: + // FakeRead(ForGuardBinding, _4); + // where `_4` is: + // _4 = &_1; (at the span for the first `x`) + // and `_1` is the `Place` for `somenum`. + // + // If and when the Issue is resolved, remove this special case match pattern: + StatementKind::FakeRead(box (FakeReadCause::ForGuardBinding, _)) => None, + + // Retain spans from most other statements. + StatementKind::FakeRead(box (_, _)) // Not including `ForGuardBinding` + | StatementKind::Intrinsic(..) + | StatementKind::Coverage(box mir::Coverage { + // The purpose of `SpanMarker` is to be matched and accepted here. + kind: CoverageKind::SpanMarker + }) + | StatementKind::Assign(_) + | StatementKind::SetDiscriminant { .. } + | StatementKind::Deinit(..) + | StatementKind::Retag(_, _) + | StatementKind::PlaceMention(..) + | StatementKind::AscribeUserType(_, _) => { + Some(statement.source_info.span) + } + + StatementKind::Coverage(box mir::Coverage { + // These coverage statements should not exist prior to coverage instrumentation. + kind: CoverageKind::CounterIncrement { .. } | CoverageKind::ExpressionUsed { .. } + }) => bug!("Unexpected coverage statement found during coverage instrumentation: {statement:?}"), + } +} + +/// If the MIR `Terminator` has a span contributive to computing coverage spans, +/// return it; otherwise return `None`. +fn filtered_terminator_span(terminator: &Terminator<'_>) -> Option<Span> { + match terminator.kind { + // These terminators have spans that don't positively contribute to computing a reasonable + // span of actually executed source code. (For example, SwitchInt terminators extracted from + // an `if condition { block }` has a span that includes the executed block, if true, + // but for coverage, the code region executed, up to *and* through the SwitchInt, + // actually stops before the if's block.) + TerminatorKind::Unreachable // Unreachable blocks are not connected to the MIR CFG + | TerminatorKind::Assert { .. } + | TerminatorKind::Drop { .. } + | TerminatorKind::SwitchInt { .. } + // For `FalseEdge`, only the `real` branch is taken, so it is similar to a `Goto`. + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::Goto { .. } => None, + + // Call `func` operand can have a more specific span when part of a chain of calls + | TerminatorKind::Call { ref func, .. } => { + let mut span = terminator.source_info.span; + if let mir::Operand::Constant(box constant) = func { + if constant.span.lo() > span.lo() { + span = span.with_lo(constant.span.lo()); + } + } + Some(span) + } + + // Retain spans from all other terminators + TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Yield { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::InlineAsm { .. } => { + Some(terminator.source_info.span) + } + } +} + +/// Returns an extrapolated span (pre-expansion[^1]) corresponding to a range +/// within the function's body source. This span is guaranteed to be contained +/// within, or equal to, the `body_span`. If the extrapolated span is not +/// contained within the `body_span`, `None` is returned. +/// +/// [^1]Expansions result from Rust syntax including macros, syntactic sugar, +/// etc.). +#[inline] +fn unexpand_into_body_span(span: Span, body_span: Span) -> Option<Span> { + use rustc_span::source_map::original_sp; + + // FIXME(#118525): Consider switching from `original_sp` to `Span::find_ancestor_inside`, + // which is similar but gives slightly different results in some edge cases. + let original_span = original_sp(span, body_span).with_ctxt(body_span.ctxt()); + body_span.contains(original_span).then_some(original_span) +} diff --git a/compiler/rustc_mir_transform/src/coverage/tests.rs b/compiler/rustc_mir_transform/src/coverage/tests.rs new file mode 100644 index 00000000000..931bc8e58ff --- /dev/null +++ b/compiler/rustc_mir_transform/src/coverage/tests.rs @@ -0,0 +1,655 @@ +//! This crate hosts a selection of "unit tests" for components of the `InstrumentCoverage` MIR +//! pass. +//! +//! ```shell +//! ./x.py test --keep-stage 1 compiler/rustc_mir --test-args '--show-output coverage' +//! ``` +//! +//! The tests construct a few "mock" objects, as needed, to support the `InstrumentCoverage` +//! functions and algorithms. Mocked objects include instances of `mir::Body`; including +//! `Terminator`s of various `kind`s, and `Span` objects. Some functions used by or used on +//! real, runtime versions of these mocked-up objects have constraints (such as cross-thread +//! limitations) and deep dependencies on other elements of the full Rust compiler (which is +//! *not* constructed or mocked for these tests). +//! +//! Of particular note, attempting to simply print elements of the `mir::Body` with default +//! `Debug` formatting can fail because some `Debug` format implementations require the +//! `TyCtxt`, obtained via a static global variable that is *not* set for these tests. +//! Initializing the global type context is prohibitively complex for the scope and scale of these +//! tests (essentially requiring initializing the entire compiler). +//! +//! Also note, some basic features of `Span` also rely on the `Span`s own "session globals", which +//! are unrelated to the `TyCtxt` global. Without initializing the `Span` session globals, some +//! basic, coverage-specific features would be impossible to test, but thankfully initializing these +//! globals is comparatively simpler. The easiest way is to wrap the test in a closure argument +//! to: `rustc_span::create_default_session_globals_then(|| { test_here(); })`. + +use super::counters; +use super::graph::{self, BasicCoverageBlock}; + +use itertools::Itertools; +use rustc_data_structures::graph::WithNumNodes; +use rustc_data_structures::graph::WithSuccessors; +use rustc_index::{Idx, IndexVec}; +use rustc_middle::mir::*; +use rustc_middle::ty; +use rustc_span::{BytePos, Pos, Span, DUMMY_SP}; + +fn bcb(index: u32) -> BasicCoverageBlock { + BasicCoverageBlock::from_u32(index) +} + +// All `TEMP_BLOCK` targets should be replaced before calling `to_body() -> mir::Body`. +const TEMP_BLOCK: BasicBlock = BasicBlock::MAX; + +struct MockBlocks<'tcx> { + blocks: IndexVec<BasicBlock, BasicBlockData<'tcx>>, + dummy_place: Place<'tcx>, + next_local: usize, +} + +impl<'tcx> MockBlocks<'tcx> { + fn new() -> Self { + Self { + blocks: IndexVec::new(), + dummy_place: Place { local: RETURN_PLACE, projection: ty::List::empty() }, + next_local: 0, + } + } + + fn new_temp(&mut self) -> Local { + let index = self.next_local; + self.next_local += 1; + Local::new(index) + } + + fn push(&mut self, kind: TerminatorKind<'tcx>) -> BasicBlock { + let next_lo = if let Some(last) = self.blocks.last_index() { + self.blocks[last].terminator().source_info.span.hi() + } else { + BytePos(1) + }; + let next_hi = next_lo + BytePos(1); + self.blocks.push(BasicBlockData { + statements: vec![], + terminator: Some(Terminator { + source_info: SourceInfo::outermost(Span::with_root_ctxt(next_lo, next_hi)), + kind, + }), + is_cleanup: false, + }) + } + + fn link(&mut self, from_block: BasicBlock, to_block: BasicBlock) { + match self.blocks[from_block].terminator_mut().kind { + TerminatorKind::Assert { ref mut target, .. } + | TerminatorKind::Call { target: Some(ref mut target), .. } + | TerminatorKind::Drop { ref mut target, .. } + | TerminatorKind::FalseEdge { real_target: ref mut target, .. } + | TerminatorKind::FalseUnwind { real_target: ref mut target, .. } + | TerminatorKind::Goto { ref mut target } + | TerminatorKind::InlineAsm { destination: Some(ref mut target), .. } + | TerminatorKind::Yield { resume: ref mut target, .. } => *target = to_block, + ref invalid => bug!("Invalid from_block: {:?}", invalid), + } + } + + fn add_block_from( + &mut self, + some_from_block: Option<BasicBlock>, + to_kind: TerminatorKind<'tcx>, + ) -> BasicBlock { + let new_block = self.push(to_kind); + if let Some(from_block) = some_from_block { + self.link(from_block, new_block); + } + new_block + } + + fn set_branch(&mut self, switchint: BasicBlock, branch_index: usize, to_block: BasicBlock) { + match self.blocks[switchint].terminator_mut().kind { + TerminatorKind::SwitchInt { ref mut targets, .. } => { + let mut branches = targets.iter().collect::<Vec<_>>(); + let otherwise = if branch_index == branches.len() { + to_block + } else { + let old_otherwise = targets.otherwise(); + if branch_index > branches.len() { + branches.push((branches.len() as u128, old_otherwise)); + while branches.len() < branch_index { + branches.push((branches.len() as u128, TEMP_BLOCK)); + } + to_block + } else { + branches[branch_index] = (branch_index as u128, to_block); + old_otherwise + } + }; + *targets = SwitchTargets::new(branches.into_iter(), otherwise); + } + ref invalid => bug!("Invalid BasicBlock kind or no to_block: {:?}", invalid), + } + } + + fn call(&mut self, some_from_block: Option<BasicBlock>) -> BasicBlock { + self.add_block_from( + some_from_block, + TerminatorKind::Call { + func: Operand::Copy(self.dummy_place.clone()), + args: vec![], + destination: self.dummy_place.clone(), + target: Some(TEMP_BLOCK), + unwind: UnwindAction::Continue, + call_source: CallSource::Misc, + fn_span: DUMMY_SP, + }, + ) + } + + fn goto(&mut self, some_from_block: Option<BasicBlock>) -> BasicBlock { + self.add_block_from(some_from_block, TerminatorKind::Goto { target: TEMP_BLOCK }) + } + + fn switchint(&mut self, some_from_block: Option<BasicBlock>) -> BasicBlock { + let switchint_kind = TerminatorKind::SwitchInt { + discr: Operand::Move(Place::from(self.new_temp())), + targets: SwitchTargets::static_if(0, TEMP_BLOCK, TEMP_BLOCK), + }; + self.add_block_from(some_from_block, switchint_kind) + } + + fn return_(&mut self, some_from_block: Option<BasicBlock>) -> BasicBlock { + self.add_block_from(some_from_block, TerminatorKind::Return) + } + + fn to_body(self) -> Body<'tcx> { + Body::new_cfg_only(self.blocks) + } +} + +fn debug_basic_blocks(mir_body: &Body<'_>) -> String { + format!( + "{:?}", + mir_body + .basic_blocks + .iter_enumerated() + .map(|(bb, data)| { + let term = &data.terminator(); + let kind = &term.kind; + let span = term.source_info.span; + let sp = format!("(span:{},{})", span.lo().to_u32(), span.hi().to_u32()); + match kind { + TerminatorKind::Assert { target, .. } + | TerminatorKind::Call { target: Some(target), .. } + | TerminatorKind::Drop { target, .. } + | TerminatorKind::FalseEdge { real_target: target, .. } + | TerminatorKind::FalseUnwind { real_target: target, .. } + | TerminatorKind::Goto { target } + | TerminatorKind::InlineAsm { destination: Some(target), .. } + | TerminatorKind::Yield { resume: target, .. } => { + format!("{}{:?}:{} -> {:?}", sp, bb, kind.name(), target) + } + TerminatorKind::SwitchInt { targets, .. } => { + format!("{}{:?}:{} -> {:?}", sp, bb, kind.name(), targets) + } + _ => format!("{}{:?}:{}", sp, bb, kind.name()), + } + }) + .collect::<Vec<_>>() + ) +} + +static PRINT_GRAPHS: bool = false; + +fn print_mir_graphviz(name: &str, mir_body: &Body<'_>) { + if PRINT_GRAPHS { + println!( + "digraph {} {{\n{}\n}}", + name, + mir_body + .basic_blocks + .iter_enumerated() + .map(|(bb, data)| { + format!( + " {:?} [label=\"{:?}: {}\"];\n{}", + bb, + bb, + data.terminator().kind.name(), + mir_body + .basic_blocks + .successors(bb) + .map(|successor| { format!(" {:?} -> {:?};", bb, successor) }) + .join("\n") + ) + }) + .join("\n") + ); + } +} + +fn print_coverage_graphviz( + name: &str, + mir_body: &Body<'_>, + basic_coverage_blocks: &graph::CoverageGraph, +) { + if PRINT_GRAPHS { + println!( + "digraph {} {{\n{}\n}}", + name, + basic_coverage_blocks + .iter_enumerated() + .map(|(bcb, bcb_data)| { + format!( + " {:?} [label=\"{:?}: {}\"];\n{}", + bcb, + bcb, + mir_body[bcb_data.last_bb()].terminator().kind.name(), + basic_coverage_blocks + .successors(bcb) + .map(|successor| { format!(" {:?} -> {:?};", bcb, successor) }) + .join("\n") + ) + }) + .join("\n") + ); + } +} + +/// Create a mock `Body` with a simple flow. +fn goto_switchint<'a>() -> Body<'a> { + let mut blocks = MockBlocks::new(); + let start = blocks.call(None); + let goto = blocks.goto(Some(start)); + let switchint = blocks.switchint(Some(goto)); + let then_call = blocks.call(None); + let else_call = blocks.call(None); + blocks.set_branch(switchint, 0, then_call); + blocks.set_branch(switchint, 1, else_call); + blocks.return_(Some(then_call)); + blocks.return_(Some(else_call)); + + let mir_body = blocks.to_body(); + print_mir_graphviz("mir_goto_switchint", &mir_body); + /* Graphviz character plots created using: `graph-easy --as=boxart`: + ┌────────────────┐ + │ bb0: Call │ + └────────────────┘ + │ + │ + ▼ + ┌────────────────┐ + │ bb1: Goto │ + └────────────────┘ + │ + │ + ▼ + ┌─────────────┐ ┌────────────────┐ + │ bb4: Call │ ◀── │ bb2: SwitchInt │ + └─────────────┘ └────────────────┘ + │ │ + │ │ + ▼ ▼ + ┌─────────────┐ ┌────────────────┐ + │ bb6: Return │ │ bb3: Call │ + └─────────────┘ └────────────────┘ + │ + │ + ▼ + ┌────────────────┐ + │ bb5: Return │ + └────────────────┘ + */ + mir_body +} + +#[track_caller] +fn assert_successors( + basic_coverage_blocks: &graph::CoverageGraph, + bcb: BasicCoverageBlock, + expected_successors: &[BasicCoverageBlock], +) { + let mut successors = basic_coverage_blocks.successors[bcb].clone(); + successors.sort_unstable(); + assert_eq!(successors, expected_successors); +} + +#[test] +fn test_covgraph_goto_switchint() { + let mir_body = goto_switchint(); + if false { + eprintln!("basic_blocks = {}", debug_basic_blocks(&mir_body)); + } + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + print_coverage_graphviz("covgraph_goto_switchint ", &mir_body, &basic_coverage_blocks); + /* + ┌──────────────┐ ┌─────────────────┐ + │ bcb2: Return │ ◀── │ bcb0: SwitchInt │ + └──────────────┘ └─────────────────┘ + │ + │ + ▼ + ┌─────────────────┐ + │ bcb1: Return │ + └─────────────────┘ + */ + assert_eq!( + basic_coverage_blocks.num_nodes(), + 3, + "basic_coverage_blocks: {:?}", + basic_coverage_blocks.iter_enumerated().collect::<Vec<_>>() + ); + + assert_successors(&basic_coverage_blocks, bcb(0), &[bcb(1), bcb(2)]); + assert_successors(&basic_coverage_blocks, bcb(1), &[]); + assert_successors(&basic_coverage_blocks, bcb(2), &[]); +} + +/// Create a mock `Body` with a loop. +fn switchint_then_loop_else_return<'a>() -> Body<'a> { + let mut blocks = MockBlocks::new(); + let start = blocks.call(None); + let switchint = blocks.switchint(Some(start)); + let then_call = blocks.call(None); + blocks.set_branch(switchint, 0, then_call); + let backedge_goto = blocks.goto(Some(then_call)); + blocks.link(backedge_goto, switchint); + let else_return = blocks.return_(None); + blocks.set_branch(switchint, 1, else_return); + + let mir_body = blocks.to_body(); + print_mir_graphviz("mir_switchint_then_loop_else_return", &mir_body); + /* + ┌────────────────┐ + │ bb0: Call │ + └────────────────┘ + │ + │ + ▼ + ┌─────────────┐ ┌────────────────┐ + │ bb4: Return │ ◀── │ bb1: SwitchInt │ ◀┐ + └─────────────┘ └────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌────────────────┐ │ + │ bb2: Call │ │ + └────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌────────────────┐ │ + │ bb3: Goto │ ─┘ + └────────────────┘ + */ + mir_body +} + +#[test] +fn test_covgraph_switchint_then_loop_else_return() { + let mir_body = switchint_then_loop_else_return(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + print_coverage_graphviz( + "covgraph_switchint_then_loop_else_return", + &mir_body, + &basic_coverage_blocks, + ); + /* + ┌─────────────────┐ + │ bcb0: Call │ + └─────────────────┘ + │ + │ + ▼ + ┌────────────┐ ┌─────────────────┐ + │ bcb3: Goto │ ◀── │ bcb1: SwitchInt │ ◀┐ + └────────────┘ └─────────────────┘ │ + │ │ │ + │ │ │ + │ ▼ │ + │ ┌─────────────────┐ │ + │ │ bcb2: Return │ │ + │ └─────────────────┘ │ + │ │ + └─────────────────────────────────────┘ + */ + assert_eq!( + basic_coverage_blocks.num_nodes(), + 4, + "basic_coverage_blocks: {:?}", + basic_coverage_blocks.iter_enumerated().collect::<Vec<_>>() + ); + + assert_successors(&basic_coverage_blocks, bcb(0), &[bcb(1)]); + assert_successors(&basic_coverage_blocks, bcb(1), &[bcb(2), bcb(3)]); + assert_successors(&basic_coverage_blocks, bcb(2), &[]); + assert_successors(&basic_coverage_blocks, bcb(3), &[bcb(1)]); +} + +/// Create a mock `Body` with nested loops. +fn switchint_loop_then_inner_loop_else_break<'a>() -> Body<'a> { + let mut blocks = MockBlocks::new(); + let start = blocks.call(None); + let switchint = blocks.switchint(Some(start)); + let then_call = blocks.call(None); + blocks.set_branch(switchint, 0, then_call); + let else_return = blocks.return_(None); + blocks.set_branch(switchint, 1, else_return); + + let inner_start = blocks.call(Some(then_call)); + let inner_switchint = blocks.switchint(Some(inner_start)); + let inner_then_call = blocks.call(None); + blocks.set_branch(inner_switchint, 0, inner_then_call); + let inner_backedge_goto = blocks.goto(Some(inner_then_call)); + blocks.link(inner_backedge_goto, inner_switchint); + let inner_else_break_goto = blocks.goto(None); + blocks.set_branch(inner_switchint, 1, inner_else_break_goto); + + let backedge_goto = blocks.goto(Some(inner_else_break_goto)); + blocks.link(backedge_goto, switchint); + + let mir_body = blocks.to_body(); + print_mir_graphviz("mir_switchint_loop_then_inner_loop_else_break", &mir_body); + /* + ┌────────────────┐ + │ bb0: Call │ + └────────────────┘ + │ + │ + ▼ + ┌─────────────┐ ┌────────────────┐ + │ bb3: Return │ ◀── │ bb1: SwitchInt │ ◀─────┐ + └─────────────┘ └────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌────────────────┐ │ + │ bb2: Call │ │ + └────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌────────────────┐ │ + │ bb4: Call │ │ + └────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌─────────────┐ ┌────────────────┐ │ + │ bb8: Goto │ ◀── │ bb5: SwitchInt │ ◀┐ │ + └─────────────┘ └────────────────┘ │ │ + │ │ │ │ + │ │ │ │ + ▼ ▼ │ │ + ┌─────────────┐ ┌────────────────┐ │ │ + │ bb9: Goto │ ─┐ │ bb6: Call │ │ │ + └─────────────┘ │ └────────────────┘ │ │ + │ │ │ │ + │ │ │ │ + │ ▼ │ │ + │ ┌────────────────┐ │ │ + │ │ bb7: Goto │ ─┘ │ + │ └────────────────┘ │ + │ │ + └───────────────────────────┘ + */ + mir_body +} + +#[test] +fn test_covgraph_switchint_loop_then_inner_loop_else_break() { + let mir_body = switchint_loop_then_inner_loop_else_break(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + print_coverage_graphviz( + "covgraph_switchint_loop_then_inner_loop_else_break", + &mir_body, + &basic_coverage_blocks, + ); + /* + ┌─────────────────┐ + │ bcb0: Call │ + └─────────────────┘ + │ + │ + ▼ + ┌──────────────┐ ┌─────────────────┐ + │ bcb2: Return │ ◀── │ bcb1: SwitchInt │ ◀┐ + └──────────────┘ └─────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌─────────────────┐ │ + │ bcb3: Call │ │ + └─────────────────┘ │ + │ │ + │ │ + ▼ │ + ┌──────────────┐ ┌─────────────────┐ │ + │ bcb6: Goto │ ◀── │ bcb4: SwitchInt │ ◀┼────┐ + └──────────────┘ └─────────────────┘ │ │ + │ │ │ │ + │ │ │ │ + │ ▼ │ │ + │ ┌─────────────────┐ │ │ + │ │ bcb5: Goto │ ─┘ │ + │ └─────────────────┘ │ + │ │ + └────────────────────────────────────────────┘ + */ + assert_eq!( + basic_coverage_blocks.num_nodes(), + 7, + "basic_coverage_blocks: {:?}", + basic_coverage_blocks.iter_enumerated().collect::<Vec<_>>() + ); + + assert_successors(&basic_coverage_blocks, bcb(0), &[bcb(1)]); + assert_successors(&basic_coverage_blocks, bcb(1), &[bcb(2), bcb(3)]); + assert_successors(&basic_coverage_blocks, bcb(2), &[]); + assert_successors(&basic_coverage_blocks, bcb(3), &[bcb(4)]); + assert_successors(&basic_coverage_blocks, bcb(4), &[bcb(5), bcb(6)]); + assert_successors(&basic_coverage_blocks, bcb(5), &[bcb(1)]); + assert_successors(&basic_coverage_blocks, bcb(6), &[bcb(4)]); +} + +#[test] +fn test_find_loop_backedges_none() { + let mir_body = goto_switchint(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + if false { + eprintln!( + "basic_coverage_blocks = {:?}", + basic_coverage_blocks.iter_enumerated().collect::<Vec<_>>() + ); + eprintln!("successors = {:?}", basic_coverage_blocks.successors); + } + let backedges = graph::find_loop_backedges(&basic_coverage_blocks); + assert_eq!( + backedges.iter_enumerated().map(|(_bcb, backedges)| backedges.len()).sum::<usize>(), + 0, + "backedges: {:?}", + backedges + ); +} + +#[test] +fn test_find_loop_backedges_one() { + let mir_body = switchint_then_loop_else_return(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + let backedges = graph::find_loop_backedges(&basic_coverage_blocks); + assert_eq!( + backedges.iter_enumerated().map(|(_bcb, backedges)| backedges.len()).sum::<usize>(), + 1, + "backedges: {:?}", + backedges + ); + + assert_eq!(backedges[bcb(1)], &[bcb(3)]); +} + +#[test] +fn test_find_loop_backedges_two() { + let mir_body = switchint_loop_then_inner_loop_else_break(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + let backedges = graph::find_loop_backedges(&basic_coverage_blocks); + assert_eq!( + backedges.iter_enumerated().map(|(_bcb, backedges)| backedges.len()).sum::<usize>(), + 2, + "backedges: {:?}", + backedges + ); + + assert_eq!(backedges[bcb(1)], &[bcb(5)]); + assert_eq!(backedges[bcb(4)], &[bcb(6)]); +} + +#[test] +fn test_traverse_coverage_with_loops() { + let mir_body = switchint_loop_then_inner_loop_else_break(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + let mut traversed_in_order = Vec::new(); + let mut traversal = graph::TraverseCoverageGraphWithLoops::new(&basic_coverage_blocks); + while let Some(bcb) = traversal.next() { + traversed_in_order.push(bcb); + } + + // bcb0 is visited first. Then bcb1 starts the first loop, and all remaining nodes, *except* + // bcb6 are inside the first loop. + assert_eq!( + *traversed_in_order.last().expect("should have elements"), + bcb(6), + "bcb6 should not be visited until all nodes inside the first loop have been visited" + ); +} + +#[test] +fn test_make_bcb_counters() { + rustc_span::create_default_session_globals_then(|| { + let mir_body = goto_switchint(); + let basic_coverage_blocks = graph::CoverageGraph::from_mir(&mir_body); + // Historically this test would use `spans` internals to set up fake + // coverage spans for BCBs 1 and 2. Now we skip that step and just tell + // BCB counter construction that those BCBs have spans. + let bcb_has_coverage_spans = |bcb: BasicCoverageBlock| (1..=2).contains(&bcb.as_usize()); + let mut coverage_counters = counters::CoverageCounters::new(&basic_coverage_blocks); + coverage_counters.make_bcb_counters(&basic_coverage_blocks, bcb_has_coverage_spans); + assert_eq!(coverage_counters.num_expressions(), 0); + + assert_eq!( + 0, // bcb1 has a `Counter` with id = 0 + match coverage_counters.bcb_counter(bcb(1)).expect("should have a counter") { + counters::BcbCounter::Counter { id, .. } => id, + _ => panic!("expected a Counter"), + } + .as_u32() + ); + + assert_eq!( + 1, // bcb2 has a `Counter` with id = 1 + match coverage_counters.bcb_counter(bcb(2)).expect("should have a counter") { + counters::BcbCounter::Counter { id, .. } => id, + _ => panic!("expected a Counter"), + } + .as_u32() + ); + }); +} diff --git a/compiler/rustc_mir_transform/src/cross_crate_inline.rs b/compiler/rustc_mir_transform/src/cross_crate_inline.rs new file mode 100644 index 00000000000..261d9dd448d --- /dev/null +++ b/compiler/rustc_mir_transform/src/cross_crate_inline.rs @@ -0,0 +1,130 @@ +use crate::inline; +use crate::pass_manager as pm; +use rustc_attr::InlineAttr; +use rustc_hir::def::DefKind; +use rustc_hir::def_id::LocalDefId; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::query::Providers; +use rustc_middle::ty::TyCtxt; +use rustc_session::config::InliningThreshold; +use rustc_session::config::OptLevel; + +pub fn provide(providers: &mut Providers) { + providers.cross_crate_inlinable = cross_crate_inlinable; +} + +fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool { + let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id); + // If this has an extern indicator, then this function is globally shared and thus will not + // generate cgu-internal copies which would make it cross-crate inlinable. + if codegen_fn_attrs.contains_extern_indicator() { + return false; + } + + // Obey source annotations first; this is important because it means we can use + // #[inline(never)] to force code generation. + match codegen_fn_attrs.inline { + InlineAttr::Never => return false, + InlineAttr::Hint | InlineAttr::Always => return true, + _ => {} + } + + // This just reproduces the logic from Instance::requires_inline. + match tcx.def_kind(def_id) { + DefKind::Ctor(..) | DefKind::Closure => return true, + DefKind::Fn | DefKind::AssocFn => {} + _ => return false, + } + + // Don't do any inference when incremental compilation is enabled; the additional inlining that + // inference permits also creates more work for small edits. + if tcx.sess.opts.incremental.is_some() { + return false; + } + + // Don't do any inference if codegen optimizations are disabled and also MIR inlining is not + // enabled. This ensures that we do inference even if someone only passes -Zinline-mir, + // which is less confusing than having to also enable -Copt-level=1. + if matches!(tcx.sess.opts.optimize, OptLevel::No) && !pm::should_run_pass(tcx, &inline::Inline) + { + return false; + } + + if !tcx.is_mir_available(def_id) { + return false; + } + + let threshold = match tcx.sess.opts.unstable_opts.cross_crate_inline_threshold { + InliningThreshold::Always => return true, + InliningThreshold::Sometimes(threshold) => threshold, + InliningThreshold::Never => return false, + }; + + let mir = tcx.optimized_mir(def_id); + let mut checker = + CostChecker { tcx, callee_body: mir, calls: 0, statements: 0, landing_pads: 0, resumes: 0 }; + checker.visit_body(mir); + checker.calls == 0 + && checker.resumes == 0 + && checker.landing_pads == 0 + && checker.statements <= threshold +} + +struct CostChecker<'b, 'tcx> { + tcx: TyCtxt<'tcx>, + callee_body: &'b Body<'tcx>, + calls: usize, + statements: usize, + landing_pads: usize, + resumes: usize, +} + +impl<'tcx> Visitor<'tcx> for CostChecker<'_, 'tcx> { + fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) { + // Don't count StorageLive/StorageDead in the inlining cost. + match statement.kind { + StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::Deinit(_) + | StatementKind::Nop => {} + _ => self.statements += 1, + } + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, _: Location) { + let tcx = self.tcx; + match terminator.kind { + TerminatorKind::Drop { ref place, unwind, .. } => { + let ty = place.ty(self.callee_body, tcx).ty; + if !ty.is_trivially_pure_clone_copy() { + self.calls += 1; + if let UnwindAction::Cleanup(_) = unwind { + self.landing_pads += 1; + } + } + } + TerminatorKind::Call { unwind, .. } => { + self.calls += 1; + if let UnwindAction::Cleanup(_) = unwind { + self.landing_pads += 1; + } + } + TerminatorKind::Assert { unwind, .. } => { + self.calls += 1; + if let UnwindAction::Cleanup(_) = unwind { + self.landing_pads += 1; + } + } + TerminatorKind::UnwindResume => self.resumes += 1, + TerminatorKind::InlineAsm { unwind, .. } => { + self.statements += 1; + if let UnwindAction::Cleanup(_) = unwind { + self.landing_pads += 1; + } + } + TerminatorKind::Return => {} + _ => self.statements += 1, + } + } +} diff --git a/compiler/rustc_mir_transform/src/ctfe_limit.rs b/compiler/rustc_mir_transform/src/ctfe_limit.rs new file mode 100644 index 00000000000..dcc960e1e02 --- /dev/null +++ b/compiler/rustc_mir_transform/src/ctfe_limit.rs @@ -0,0 +1,58 @@ +//! A pass that inserts the `ConstEvalCounter` instruction into any blocks that have a back edge +//! (thus indicating there is a loop in the CFG), or whose terminator is a function call. +use crate::MirPass; + +use rustc_data_structures::graph::dominators::Dominators; +use rustc_middle::mir::{ + BasicBlock, BasicBlockData, Body, Statement, StatementKind, TerminatorKind, +}; +use rustc_middle::ty::TyCtxt; + +pub struct CtfeLimit; + +impl<'tcx> MirPass<'tcx> for CtfeLimit { + #[instrument(skip(self, _tcx, body))] + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let doms = body.basic_blocks.dominators(); + let indices: Vec<BasicBlock> = body + .basic_blocks + .iter_enumerated() + .filter_map(|(node, node_data)| { + if matches!(node_data.terminator().kind, TerminatorKind::Call { .. }) + // Back edges in a CFG indicate loops + || has_back_edge(doms, node, node_data) + { + Some(node) + } else { + None + } + }) + .collect(); + for index in indices { + insert_counter( + body.basic_blocks_mut() + .get_mut(index) + .expect("basic_blocks index {index} should exist"), + ); + } + } +} + +fn has_back_edge( + doms: &Dominators<BasicBlock>, + node: BasicBlock, + node_data: &BasicBlockData<'_>, +) -> bool { + if !doms.is_reachable(node) { + return false; + } + // Check if any of the dominators of the node are also the node's successor. + node_data.terminator().successors().any(|succ| doms.dominates(succ, node)) +} + +fn insert_counter(basic_block_data: &mut BasicBlockData<'_>) { + basic_block_data.statements.push(Statement { + source_info: basic_block_data.terminator().source_info, + kind: StatementKind::ConstEvalCounter, + }); +} diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs new file mode 100644 index 00000000000..146cc33e8c6 --- /dev/null +++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs @@ -0,0 +1,978 @@ +//! A constant propagation optimization pass based on dataflow analysis. +//! +//! Currently, this pass only propagates scalar values. + +use rustc_const_eval::interpret::{ImmTy, Immediate, InterpCx, OpTy, PlaceTy, Projectable}; +use rustc_data_structures::fx::FxHashMap; +use rustc_hir::def::DefKind; +use rustc_middle::mir::interpret::{AllocId, ConstAllocation, InterpResult, Scalar}; +use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::query::TyCtxtAt; +use rustc_middle::ty::layout::{LayoutOf, TyAndLayout}; +use rustc_middle::ty::{self, Ty, TyCtxt}; +use rustc_mir_dataflow::value_analysis::{ + Map, PlaceIndex, State, TrackElem, ValueAnalysis, ValueAnalysisWrapper, ValueOrPlace, +}; +use rustc_mir_dataflow::{lattice::FlatSet, Analysis, Results, ResultsVisitor}; +use rustc_span::def_id::DefId; +use rustc_span::DUMMY_SP; +use rustc_target::abi::{Abi, FieldIdx, Size, VariantIdx, FIRST_VARIANT}; + +use crate::const_prop::throw_machine_stop_str; + +// These constants are somewhat random guesses and have not been optimized. +// If `tcx.sess.mir_opt_level() >= 4`, we ignore the limits (this can become very expensive). +const BLOCK_LIMIT: usize = 100; +const PLACE_LIMIT: usize = 100; + +pub struct DataflowConstProp; + +impl<'tcx> MirPass<'tcx> for DataflowConstProp { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 3 + } + + #[instrument(skip_all level = "debug")] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + if tcx.sess.mir_opt_level() < 4 && body.basic_blocks.len() > BLOCK_LIMIT { + debug!("aborted dataflow const prop due too many basic blocks"); + return; + } + + // We want to have a somewhat linear runtime w.r.t. the number of statements/terminators. + // Let's call this number `n`. Dataflow analysis has `O(h*n)` transfer function + // applications, where `h` is the height of the lattice. Because the height of our lattice + // is linear w.r.t. the number of tracked places, this is `O(tracked_places * n)`. However, + // because every transfer function application could traverse the whole map, this becomes + // `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of + // map nodes is strongly correlated to the number of tracked places, this becomes more or + // less `O(n)` if we place a constant limit on the number of tracked places. + let place_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None }; + + // Decide which places to track during the analysis. + let map = Map::new(tcx, body, place_limit); + + // Perform the actual dataflow analysis. + let analysis = ConstAnalysis::new(tcx, body, map); + let mut results = debug_span!("analyze") + .in_scope(|| analysis.wrap().into_engine(tcx, body).iterate_to_fixpoint()); + + // Collect results and patch the body afterwards. + let mut visitor = Collector::new(tcx, &body.local_decls); + debug_span!("collect").in_scope(|| results.visit_reachable_with(body, &mut visitor)); + let mut patch = visitor.patch; + debug_span!("patch").in_scope(|| patch.visit_body_preserves_cfg(body)); + } +} + +struct ConstAnalysis<'a, 'tcx> { + map: Map, + tcx: TyCtxt<'tcx>, + local_decls: &'a LocalDecls<'tcx>, + ecx: InterpCx<'tcx, 'tcx, DummyMachine>, + param_env: ty::ParamEnv<'tcx>, +} + +impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> { + type Value = FlatSet<Scalar>; + + const NAME: &'static str = "ConstAnalysis"; + + fn map(&self) -> &Map { + &self.map + } + + fn handle_set_discriminant( + &self, + place: Place<'tcx>, + variant_index: VariantIdx, + state: &mut State<Self::Value>, + ) { + state.flood_discr(place.as_ref(), &self.map); + if self.map.find_discr(place.as_ref()).is_some() { + let enum_ty = place.ty(self.local_decls, self.tcx).ty; + if let Some(discr) = self.eval_discriminant(enum_ty, variant_index) { + state.assign_discr( + place.as_ref(), + ValueOrPlace::Value(FlatSet::Elem(discr)), + &self.map, + ); + } + } + } + + fn handle_assign( + &self, + target: Place<'tcx>, + rvalue: &Rvalue<'tcx>, + state: &mut State<Self::Value>, + ) { + match rvalue { + Rvalue::Use(operand) => { + state.flood(target.as_ref(), self.map()); + if let Some(target) = self.map.find(target.as_ref()) { + self.assign_operand(state, target, operand); + } + } + Rvalue::CopyForDeref(rhs) => { + state.flood(target.as_ref(), self.map()); + if let Some(target) = self.map.find(target.as_ref()) { + self.assign_operand(state, target, &Operand::Copy(*rhs)); + } + } + Rvalue::Aggregate(kind, operands) => { + // If we assign `target = Enum::Variant#0(operand)`, + // we must make sure that all `target as Variant#i` are `Top`. + state.flood(target.as_ref(), self.map()); + + let Some(target_idx) = self.map().find(target.as_ref()) else { return }; + + let (variant_target, variant_index) = match **kind { + AggregateKind::Tuple | AggregateKind::Closure(..) => (Some(target_idx), None), + AggregateKind::Adt(def_id, variant_index, ..) => { + match self.tcx.def_kind(def_id) { + DefKind::Struct => (Some(target_idx), None), + DefKind::Enum => ( + self.map.apply(target_idx, TrackElem::Variant(variant_index)), + Some(variant_index), + ), + _ => return, + } + } + _ => return, + }; + if let Some(variant_target_idx) = variant_target { + for (field_index, operand) in operands.iter().enumerate() { + if let Some(field) = self.map().apply( + variant_target_idx, + TrackElem::Field(FieldIdx::from_usize(field_index)), + ) { + self.assign_operand(state, field, operand); + } + } + } + if let Some(variant_index) = variant_index + && let Some(discr_idx) = self.map().apply(target_idx, TrackElem::Discriminant) + { + // We are assigning the discriminant as part of an aggregate. + // This discriminant can only alias a variant field's value if the operand + // had an invalid value for that type. + // Using invalid values is UB, so we are allowed to perform the assignment + // without extra flooding. + let enum_ty = target.ty(self.local_decls, self.tcx).ty; + if let Some(discr_val) = self.eval_discriminant(enum_ty, variant_index) { + state.insert_value_idx(discr_idx, FlatSet::Elem(discr_val), &self.map); + } + } + } + Rvalue::CheckedBinaryOp(op, box (left, right)) => { + // Flood everything now, so we can use `insert_value_idx` directly later. + state.flood(target.as_ref(), self.map()); + + let Some(target) = self.map().find(target.as_ref()) else { return }; + + let value_target = self.map().apply(target, TrackElem::Field(0_u32.into())); + let overflow_target = self.map().apply(target, TrackElem::Field(1_u32.into())); + + if value_target.is_some() || overflow_target.is_some() { + let (val, overflow) = self.binary_op(state, *op, left, right); + + if let Some(value_target) = value_target { + // We have flooded `target` earlier. + state.insert_value_idx(value_target, val, self.map()); + } + if let Some(overflow_target) = overflow_target { + let overflow = match overflow { + FlatSet::Top => FlatSet::Top, + FlatSet::Elem(overflow) => FlatSet::Elem(Scalar::from_bool(overflow)), + FlatSet::Bottom => FlatSet::Bottom, + }; + // We have flooded `target` earlier. + state.insert_value_idx(overflow_target, overflow, self.map()); + } + } + } + Rvalue::Cast( + CastKind::PointerCoercion(ty::adjustment::PointerCoercion::Unsize), + operand, + _, + ) => { + let pointer = self.handle_operand(operand, state); + state.assign(target.as_ref(), pointer, self.map()); + + if let Some(target_len) = self.map().find_len(target.as_ref()) + && let operand_ty = operand.ty(self.local_decls, self.tcx) + && let Some(operand_ty) = operand_ty.builtin_deref(true) + && let ty::Array(_, len) = operand_ty.ty.kind() + && let Some(len) = Const::Ty(*len).try_eval_scalar_int(self.tcx, self.param_env) + { + state.insert_value_idx(target_len, FlatSet::Elem(len.into()), self.map()); + } + } + _ => self.super_assign(target, rvalue, state), + } + } + + fn handle_rvalue( + &self, + rvalue: &Rvalue<'tcx>, + state: &mut State<Self::Value>, + ) -> ValueOrPlace<Self::Value> { + let val = match rvalue { + Rvalue::Len(place) => { + let place_ty = place.ty(self.local_decls, self.tcx); + if let ty::Array(_, len) = place_ty.ty.kind() { + Const::Ty(*len) + .try_eval_scalar(self.tcx, self.param_env) + .map_or(FlatSet::Top, FlatSet::Elem) + } else if let [ProjectionElem::Deref] = place.projection[..] { + state.get_len(place.local.into(), self.map()) + } else { + FlatSet::Top + } + } + Rvalue::Cast(CastKind::IntToInt | CastKind::IntToFloat, operand, ty) => { + let Ok(layout) = self.tcx.layout_of(self.param_env.and(*ty)) else { + return ValueOrPlace::Value(FlatSet::Top); + }; + match self.eval_operand(operand, state) { + FlatSet::Elem(op) => self + .ecx + .int_to_int_or_float(&op, layout) + .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)), + FlatSet::Bottom => FlatSet::Bottom, + FlatSet::Top => FlatSet::Top, + } + } + Rvalue::Cast(CastKind::FloatToInt | CastKind::FloatToFloat, operand, ty) => { + let Ok(layout) = self.tcx.layout_of(self.param_env.and(*ty)) else { + return ValueOrPlace::Value(FlatSet::Top); + }; + match self.eval_operand(operand, state) { + FlatSet::Elem(op) => self + .ecx + .float_to_float_or_int(&op, layout) + .map_or(FlatSet::Top, |result| self.wrap_immediate(*result)), + FlatSet::Bottom => FlatSet::Bottom, + FlatSet::Top => FlatSet::Top, + } + } + Rvalue::Cast(CastKind::Transmute, operand, _) => { + match self.eval_operand(operand, state) { + FlatSet::Elem(op) => self.wrap_immediate(*op), + FlatSet::Bottom => FlatSet::Bottom, + FlatSet::Top => FlatSet::Top, + } + } + Rvalue::BinaryOp(op, box (left, right)) => { + // Overflows must be ignored here. + let (val, _overflow) = self.binary_op(state, *op, left, right); + val + } + Rvalue::UnaryOp(op, operand) => match self.eval_operand(operand, state) { + FlatSet::Elem(value) => self + .ecx + .wrapping_unary_op(*op, &value) + .map_or(FlatSet::Top, |val| self.wrap_immediate(*val)), + FlatSet::Bottom => FlatSet::Bottom, + FlatSet::Top => FlatSet::Top, + }, + Rvalue::NullaryOp(null_op, ty) => { + let Ok(layout) = self.tcx.layout_of(self.param_env.and(*ty)) else { + return ValueOrPlace::Value(FlatSet::Top); + }; + let val = match null_op { + NullOp::SizeOf if layout.is_sized() => layout.size.bytes(), + NullOp::AlignOf if layout.is_sized() => layout.align.abi.bytes(), + NullOp::OffsetOf(fields) => { + layout.offset_of_subfield(&self.ecx, fields.iter()).bytes() + } + _ => return ValueOrPlace::Value(FlatSet::Top), + }; + FlatSet::Elem(Scalar::from_target_usize(val, &self.tcx)) + } + Rvalue::Discriminant(place) => state.get_discr(place.as_ref(), self.map()), + _ => return self.super_rvalue(rvalue, state), + }; + ValueOrPlace::Value(val) + } + + fn handle_constant( + &self, + constant: &ConstOperand<'tcx>, + _state: &mut State<Self::Value>, + ) -> Self::Value { + constant + .const_ + .try_eval_scalar(self.tcx, self.param_env) + .map_or(FlatSet::Top, FlatSet::Elem) + } + + fn handle_switch_int<'mir>( + &self, + discr: &'mir Operand<'tcx>, + targets: &'mir SwitchTargets, + state: &mut State<Self::Value>, + ) -> TerminatorEdges<'mir, 'tcx> { + let value = match self.handle_operand(discr, state) { + ValueOrPlace::Value(value) => value, + ValueOrPlace::Place(place) => state.get_idx(place, self.map()), + }; + match value { + // We are branching on uninitialized data, this is UB, treat it as unreachable. + // This allows the set of visited edges to grow monotonically with the lattice. + FlatSet::Bottom => TerminatorEdges::None, + FlatSet::Elem(scalar) => { + let choice = scalar.assert_bits(scalar.size()); + TerminatorEdges::Single(targets.target_for_value(choice)) + } + FlatSet::Top => TerminatorEdges::SwitchInt { discr, targets }, + } + } +} + +impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> { + pub fn new(tcx: TyCtxt<'tcx>, body: &'a Body<'tcx>, map: Map) -> Self { + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + Self { + map, + tcx, + local_decls: &body.local_decls, + ecx: InterpCx::new(tcx, DUMMY_SP, param_env, DummyMachine), + param_env: param_env, + } + } + + /// The caller must have flooded `place`. + fn assign_operand( + &self, + state: &mut State<FlatSet<Scalar>>, + place: PlaceIndex, + operand: &Operand<'tcx>, + ) { + match operand { + Operand::Copy(rhs) | Operand::Move(rhs) => { + if let Some(rhs) = self.map.find(rhs.as_ref()) { + state.insert_place_idx(place, rhs, &self.map); + } else if rhs.projection.first() == Some(&PlaceElem::Deref) + && let FlatSet::Elem(pointer) = state.get(rhs.local.into(), &self.map) + && let rhs_ty = self.local_decls[rhs.local].ty + && let Ok(rhs_layout) = self.tcx.layout_of(self.param_env.and(rhs_ty)) + { + let op = ImmTy::from_scalar(pointer, rhs_layout).into(); + self.assign_constant(state, place, op, rhs.projection); + } + } + Operand::Constant(box constant) => { + if let Ok(constant) = self.ecx.eval_mir_constant(&constant.const_, None, None) { + self.assign_constant(state, place, constant, &[]); + } + } + } + } + + /// The caller must have flooded `place`. + /// + /// Perform: `place = operand.projection`. + #[instrument(level = "trace", skip(self, state))] + fn assign_constant( + &self, + state: &mut State<FlatSet<Scalar>>, + place: PlaceIndex, + mut operand: OpTy<'tcx>, + projection: &[PlaceElem<'tcx>], + ) -> Option<!> { + for &(mut proj_elem) in projection { + if let PlaceElem::Index(index) = proj_elem { + if let FlatSet::Elem(index) = state.get(index.into(), &self.map) + && let Ok(offset) = index.to_target_usize(&self.tcx) + && let Some(min_length) = offset.checked_add(1) + { + proj_elem = PlaceElem::ConstantIndex { offset, min_length, from_end: false }; + } else { + return None; + } + } + operand = self.ecx.project(&operand, proj_elem).ok()?; + } + + self.map.for_each_projection_value( + place, + operand, + &mut |elem, op| match elem { + TrackElem::Field(idx) => self.ecx.project_field(op, idx.as_usize()).ok(), + TrackElem::Variant(idx) => self.ecx.project_downcast(op, idx).ok(), + TrackElem::Discriminant => { + let variant = self.ecx.read_discriminant(op).ok()?; + let discr_value = + self.ecx.discriminant_for_variant(op.layout.ty, variant).ok()?; + Some(discr_value.into()) + } + TrackElem::DerefLen => { + let op: OpTy<'_> = self.ecx.deref_pointer(op).ok()?.into(); + let len_usize = op.len(&self.ecx).ok()?; + let layout = + self.tcx.layout_of(self.param_env.and(self.tcx.types.usize)).unwrap(); + Some(ImmTy::from_uint(len_usize, layout).into()) + } + }, + &mut |place, op| { + if let Ok(imm) = self.ecx.read_immediate_raw(op) + && let Some(imm) = imm.right() + { + let elem = self.wrap_immediate(*imm); + state.insert_value_idx(place, elem, &self.map); + } + }, + ); + + None + } + + fn binary_op( + &self, + state: &mut State<FlatSet<Scalar>>, + op: BinOp, + left: &Operand<'tcx>, + right: &Operand<'tcx>, + ) -> (FlatSet<Scalar>, FlatSet<bool>) { + let left = self.eval_operand(left, state); + let right = self.eval_operand(right, state); + + match (left, right) { + (FlatSet::Bottom, _) | (_, FlatSet::Bottom) => (FlatSet::Bottom, FlatSet::Bottom), + // Both sides are known, do the actual computation. + (FlatSet::Elem(left), FlatSet::Elem(right)) => { + match self.ecx.overflowing_binary_op(op, &left, &right) { + Ok((val, overflow)) => { + (FlatSet::Elem(val.to_scalar()), FlatSet::Elem(overflow)) + } + _ => (FlatSet::Top, FlatSet::Top), + } + } + // Exactly one side is known, attempt some algebraic simplifications. + (FlatSet::Elem(const_arg), _) | (_, FlatSet::Elem(const_arg)) => { + let layout = const_arg.layout; + if !matches!(layout.abi, rustc_target::abi::Abi::Scalar(..)) { + return (FlatSet::Top, FlatSet::Top); + } + + let arg_scalar = const_arg.to_scalar(); + let Ok(arg_value) = arg_scalar.to_bits(layout.size) else { + return (FlatSet::Top, FlatSet::Top); + }; + + match op { + BinOp::BitAnd if arg_value == 0 => (FlatSet::Elem(arg_scalar), FlatSet::Bottom), + BinOp::BitOr + if arg_value == layout.size.truncate(u128::MAX) + || (layout.ty.is_bool() && arg_value == 1) => + { + (FlatSet::Elem(arg_scalar), FlatSet::Bottom) + } + BinOp::Mul if layout.ty.is_integral() && arg_value == 0 => { + (FlatSet::Elem(arg_scalar), FlatSet::Elem(false)) + } + _ => (FlatSet::Top, FlatSet::Top), + } + } + (FlatSet::Top, FlatSet::Top) => (FlatSet::Top, FlatSet::Top), + } + } + + fn eval_operand( + &self, + op: &Operand<'tcx>, + state: &mut State<FlatSet<Scalar>>, + ) -> FlatSet<ImmTy<'tcx>> { + let value = match self.handle_operand(op, state) { + ValueOrPlace::Value(value) => value, + ValueOrPlace::Place(place) => state.get_idx(place, &self.map), + }; + match value { + FlatSet::Top => FlatSet::Top, + FlatSet::Elem(scalar) => { + let ty = op.ty(self.local_decls, self.tcx); + self.tcx.layout_of(self.param_env.and(ty)).map_or(FlatSet::Top, |layout| { + FlatSet::Elem(ImmTy::from_scalar(scalar.into(), layout)) + }) + } + FlatSet::Bottom => FlatSet::Bottom, + } + } + + fn eval_discriminant(&self, enum_ty: Ty<'tcx>, variant_index: VariantIdx) -> Option<Scalar> { + if !enum_ty.is_enum() { + return None; + } + let enum_ty_layout = self.tcx.layout_of(self.param_env.and(enum_ty)).ok()?; + let discr_value = + self.ecx.discriminant_for_variant(enum_ty_layout.ty, variant_index).ok()?; + Some(discr_value.to_scalar()) + } + + fn wrap_immediate(&self, imm: Immediate) -> FlatSet<Scalar> { + match imm { + Immediate::Scalar(scalar) => FlatSet::Elem(scalar), + Immediate::Uninit => FlatSet::Bottom, + _ => FlatSet::Top, + } + } +} + +pub(crate) struct Patch<'tcx> { + tcx: TyCtxt<'tcx>, + + /// For a given MIR location, this stores the values of the operands used by that location. In + /// particular, this is before the effect, such that the operands of `_1 = _1 + _2` are + /// properly captured. (This may become UB soon, but it is currently emitted even by safe code.) + pub(crate) before_effect: FxHashMap<(Location, Place<'tcx>), Const<'tcx>>, + + /// Stores the assigned values for assignments where the Rvalue is constant. + pub(crate) assignments: FxHashMap<Location, Const<'tcx>>, +} + +impl<'tcx> Patch<'tcx> { + pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self { + Self { tcx, before_effect: FxHashMap::default(), assignments: FxHashMap::default() } + } + + fn make_operand(&self, const_: Const<'tcx>) -> Operand<'tcx> { + Operand::Constant(Box::new(ConstOperand { span: DUMMY_SP, user_ty: None, const_ })) + } +} + +struct Collector<'tcx, 'locals> { + patch: Patch<'tcx>, + local_decls: &'locals LocalDecls<'tcx>, +} + +impl<'tcx, 'locals> Collector<'tcx, 'locals> { + pub(crate) fn new(tcx: TyCtxt<'tcx>, local_decls: &'locals LocalDecls<'tcx>) -> Self { + Self { patch: Patch::new(tcx), local_decls } + } + + fn try_make_constant( + &self, + ecx: &mut InterpCx<'tcx, 'tcx, DummyMachine>, + place: Place<'tcx>, + state: &State<FlatSet<Scalar>>, + map: &Map, + ) -> Option<Const<'tcx>> { + let ty = place.ty(self.local_decls, self.patch.tcx).ty; + let layout = ecx.layout_of(ty).ok()?; + + if layout.is_zst() { + return Some(Const::zero_sized(ty)); + } + + if layout.is_unsized() { + return None; + } + + let place = map.find(place.as_ref())?; + if layout.abi.is_scalar() + && let Some(value) = propagatable_scalar(place, state, map) + { + return Some(Const::Val(ConstValue::Scalar(value), ty)); + } + + if matches!(layout.abi, Abi::Scalar(..) | Abi::ScalarPair(..)) { + let alloc_id = ecx + .intern_with_temp_alloc(layout, |ecx, dest| { + try_write_constant(ecx, dest, place, ty, state, map) + }) + .ok()?; + return Some(Const::Val(ConstValue::Indirect { alloc_id, offset: Size::ZERO }, ty)); + } + + None + } +} + +fn propagatable_scalar( + place: PlaceIndex, + state: &State<FlatSet<Scalar>>, + map: &Map, +) -> Option<Scalar> { + if let FlatSet::Elem(value) = state.get_idx(place, map) + && value.try_to_int().is_ok() + { + // Do not attempt to propagate pointers, as we may fail to preserve their identity. + Some(value) + } else { + None + } +} + +#[instrument(level = "trace", skip(ecx, state, map))] +fn try_write_constant<'tcx>( + ecx: &mut InterpCx<'_, 'tcx, DummyMachine>, + dest: &PlaceTy<'tcx>, + place: PlaceIndex, + ty: Ty<'tcx>, + state: &State<FlatSet<Scalar>>, + map: &Map, +) -> InterpResult<'tcx> { + let layout = ecx.layout_of(ty)?; + + // Fast path for ZSTs. + if layout.is_zst() { + return Ok(()); + } + + // Fast path for scalars. + if layout.abi.is_scalar() + && let Some(value) = propagatable_scalar(place, state, map) + { + return ecx.write_immediate(Immediate::Scalar(value), dest); + } + + match ty.kind() { + // ZSTs. Nothing to do. + ty::FnDef(..) => {} + + // Those are scalars, must be handled above. + ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char => throw_machine_stop_str!("primitive type with provenance"), + + ty::Tuple(elem_tys) => { + for (i, elem) in elem_tys.iter().enumerate() { + let Some(field) = map.apply(place, TrackElem::Field(FieldIdx::from_usize(i))) else { + throw_machine_stop_str!("missing field in tuple") + }; + let field_dest = ecx.project_field(dest, i)?; + try_write_constant(ecx, &field_dest, field, elem, state, map)?; + } + } + + ty::Adt(def, args) => { + if def.is_union() { + throw_machine_stop_str!("cannot propagate unions") + } + + let (variant_idx, variant_def, variant_place, variant_dest) = if def.is_enum() { + let Some(discr) = map.apply(place, TrackElem::Discriminant) else { + throw_machine_stop_str!("missing discriminant for enum") + }; + let FlatSet::Elem(Scalar::Int(discr)) = state.get_idx(discr, map) else { + throw_machine_stop_str!("discriminant with provenance") + }; + let discr_bits = discr.assert_bits(discr.size()); + let Some((variant, _)) = def.discriminants(*ecx.tcx).find(|(_, var)| discr_bits == var.val) else { + throw_machine_stop_str!("illegal discriminant for enum") + }; + let Some(variant_place) = map.apply(place, TrackElem::Variant(variant)) else { + throw_machine_stop_str!("missing variant for enum") + }; + let variant_dest = ecx.project_downcast(dest, variant)?; + (variant, def.variant(variant), variant_place, variant_dest) + } else { + (FIRST_VARIANT, def.non_enum_variant(), place, dest.clone()) + }; + + for (i, field) in variant_def.fields.iter_enumerated() { + let ty = field.ty(*ecx.tcx, args); + let Some(field) = map.apply(variant_place, TrackElem::Field(i)) else { + throw_machine_stop_str!("missing field in ADT") + }; + let field_dest = ecx.project_field(&variant_dest, i.as_usize())?; + try_write_constant(ecx, &field_dest, field, ty, state, map)?; + } + ecx.write_discriminant(variant_idx, dest)?; + } + + // Unsupported for now. + ty::Array(_, _) + + // Do not attempt to support indirection in constants. + | ty::Ref(..) | ty::RawPtr(..) | ty::FnPtr(..) | ty::Str | ty::Slice(_) + + | ty::Never + | ty::Foreign(..) + | ty::Alias(..) + | ty::Param(_) + | ty::Bound(..) + | ty::Placeholder(..) + | ty::Closure(..) + | ty::Coroutine(..) + | ty::Dynamic(..) => throw_machine_stop_str!("unsupported type"), + + ty::Error(_) | ty::Infer(..) | ty::CoroutineWitness(..) => bug!(), + } + + Ok(()) +} + +impl<'mir, 'tcx> + ResultsVisitor<'mir, 'tcx, Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>> + for Collector<'tcx, '_> +{ + type FlowState = State<FlatSet<Scalar>>; + + fn visit_statement_before_primary_effect( + &mut self, + results: &mut Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>, + state: &Self::FlowState, + statement: &'mir Statement<'tcx>, + location: Location, + ) { + match &statement.kind { + StatementKind::Assign(box (_, rvalue)) => { + OperandCollector { + state, + visitor: self, + ecx: &mut results.analysis.0.ecx, + map: &results.analysis.0.map, + } + .visit_rvalue(rvalue, location); + } + _ => (), + } + } + + fn visit_statement_after_primary_effect( + &mut self, + results: &mut Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>, + state: &Self::FlowState, + statement: &'mir Statement<'tcx>, + location: Location, + ) { + match statement.kind { + StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(_)))) => { + // Don't overwrite the assignment if it already uses a constant (to keep the span). + } + StatementKind::Assign(box (place, _)) => { + if let Some(value) = self.try_make_constant( + &mut results.analysis.0.ecx, + place, + state, + &results.analysis.0.map, + ) { + self.patch.assignments.insert(location, value); + } + } + _ => (), + } + } + + fn visit_terminator_before_primary_effect( + &mut self, + results: &mut Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>, + state: &Self::FlowState, + terminator: &'mir Terminator<'tcx>, + location: Location, + ) { + OperandCollector { + state, + visitor: self, + ecx: &mut results.analysis.0.ecx, + map: &results.analysis.0.map, + } + .visit_terminator(terminator, location); + } +} + +impl<'tcx> MutVisitor<'tcx> for Patch<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) { + if let Some(value) = self.assignments.get(&location) { + match &mut statement.kind { + StatementKind::Assign(box (_, rvalue)) => { + *rvalue = Rvalue::Use(self.make_operand(*value)); + } + _ => bug!("found assignment info for non-assign statement"), + } + } else { + self.super_statement(statement, location); + } + } + + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) { + match operand { + Operand::Copy(place) | Operand::Move(place) => { + if let Some(value) = self.before_effect.get(&(location, *place)) { + *operand = self.make_operand(*value); + } else if !place.projection.is_empty() { + self.super_operand(operand, location) + } + } + Operand::Constant(_) => {} + } + } + + fn process_projection_elem( + &mut self, + elem: PlaceElem<'tcx>, + location: Location, + ) -> Option<PlaceElem<'tcx>> { + if let PlaceElem::Index(local) = elem { + let offset = self.before_effect.get(&(location, local.into()))?; + let offset = offset.try_to_scalar()?; + let offset = offset.to_target_usize(&self.tcx).ok()?; + let min_length = offset.checked_add(1)?; + Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false }) + } else { + None + } + } +} + +struct OperandCollector<'tcx, 'map, 'locals, 'a> { + state: &'a State<FlatSet<Scalar>>, + visitor: &'a mut Collector<'tcx, 'locals>, + ecx: &'map mut InterpCx<'tcx, 'tcx, DummyMachine>, + map: &'map Map, +} + +impl<'tcx> Visitor<'tcx> for OperandCollector<'tcx, '_, '_, '_> { + fn visit_projection_elem( + &mut self, + _: PlaceRef<'tcx>, + elem: PlaceElem<'tcx>, + _: PlaceContext, + location: Location, + ) { + if let PlaceElem::Index(local) = elem + && let Some(value) = + self.visitor.try_make_constant(self.ecx, local.into(), self.state, self.map) + { + self.visitor.patch.before_effect.insert((location, local.into()), value); + } + } + + fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) { + if let Some(place) = operand.place() { + if let Some(value) = + self.visitor.try_make_constant(self.ecx, place, self.state, self.map) + { + self.visitor.patch.before_effect.insert((location, place), value); + } else if !place.projection.is_empty() { + // Try to propagate into `Index` projections. + self.super_operand(operand, location) + } + } + } +} + +pub(crate) struct DummyMachine; + +impl<'mir, 'tcx: 'mir> rustc_const_eval::interpret::Machine<'mir, 'tcx> for DummyMachine { + rustc_const_eval::interpret::compile_time_machine!(<'mir, 'tcx>); + type MemoryKind = !; + const PANIC_ON_ALLOC_FAIL: bool = true; + + #[inline(always)] + fn enforce_alignment(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool { + false // no reason to enforce alignment + } + + fn enforce_validity(_ecx: &InterpCx<'mir, 'tcx, Self>, _layout: TyAndLayout<'tcx>) -> bool { + false + } + + fn before_access_global( + _tcx: TyCtxtAt<'tcx>, + _machine: &Self, + _alloc_id: AllocId, + alloc: ConstAllocation<'tcx>, + _static_def_id: Option<DefId>, + is_write: bool, + ) -> InterpResult<'tcx> { + if is_write { + throw_machine_stop_str!("can't write to global"); + } + + // If the static allocation is mutable, then we can't const prop it as its content + // might be different at runtime. + if alloc.inner().mutability.is_mut() { + throw_machine_stop_str!("can't access mutable globals in ConstProp"); + } + + Ok(()) + } + + fn find_mir_or_eval_fn( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _instance: ty::Instance<'tcx>, + _abi: rustc_target::spec::abi::Abi, + _args: &[rustc_const_eval::interpret::FnArg<'tcx, Self::Provenance>], + _destination: &rustc_const_eval::interpret::PlaceTy<'tcx, Self::Provenance>, + _target: Option<BasicBlock>, + _unwind: UnwindAction, + ) -> interpret::InterpResult<'tcx, Option<(&'mir Body<'tcx>, ty::Instance<'tcx>)>> { + unimplemented!() + } + + fn panic_nounwind( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _msg: &str, + ) -> interpret::InterpResult<'tcx> { + unimplemented!() + } + + fn call_intrinsic( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _instance: ty::Instance<'tcx>, + _args: &[rustc_const_eval::interpret::OpTy<'tcx, Self::Provenance>], + _destination: &rustc_const_eval::interpret::PlaceTy<'tcx, Self::Provenance>, + _target: Option<BasicBlock>, + _unwind: UnwindAction, + ) -> interpret::InterpResult<'tcx> { + unimplemented!() + } + + fn assert_panic( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _msg: &rustc_middle::mir::AssertMessage<'tcx>, + _unwind: UnwindAction, + ) -> interpret::InterpResult<'tcx> { + unimplemented!() + } + + fn binary_ptr_op( + _ecx: &InterpCx<'mir, 'tcx, Self>, + _bin_op: BinOp, + _left: &rustc_const_eval::interpret::ImmTy<'tcx, Self::Provenance>, + _right: &rustc_const_eval::interpret::ImmTy<'tcx, Self::Provenance>, + ) -> interpret::InterpResult<'tcx, (ImmTy<'tcx, Self::Provenance>, bool)> { + throw_machine_stop_str!("can't do pointer arithmetic"); + } + + fn expose_ptr( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _ptr: interpret::Pointer<Self::Provenance>, + ) -> interpret::InterpResult<'tcx> { + unimplemented!() + } + + fn init_frame_extra( + _ecx: &mut InterpCx<'mir, 'tcx, Self>, + _frame: rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance>, + ) -> interpret::InterpResult< + 'tcx, + rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>, + > { + unimplemented!() + } + + fn stack<'a>( + _ecx: &'a InterpCx<'mir, 'tcx, Self>, + ) -> &'a [rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>] + { + // Return an empty stack instead of panicking, as `cur_span` uses it to evaluate constants. + &[] + } + + fn stack_mut<'a>( + _ecx: &'a mut InterpCx<'mir, 'tcx, Self>, + ) -> &'a mut Vec< + rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>, + > { + unimplemented!() + } +} diff --git a/compiler/rustc_mir_transform/src/dead_store_elimination.rs b/compiler/rustc_mir_transform/src/dead_store_elimination.rs new file mode 100644 index 00000000000..3d74ef7e327 --- /dev/null +++ b/compiler/rustc_mir_transform/src/dead_store_elimination.rs @@ -0,0 +1,140 @@ +//! This module implements a dead store elimination (DSE) routine. +//! +//! This transformation was written specifically for the needs of dest prop. Although it is +//! perfectly sound to use it in any context that might need it, its behavior should not be changed +//! without analyzing the interaction this will have with dest prop. Specifically, in addition to +//! the soundness of this pass in general, dest prop needs it to satisfy two additional conditions: +//! +//! 1. It's idempotent, meaning that running this pass a second time immediately after running it a +//! first time will not cause any further changes. +//! 2. This idempotence persists across dest prop's main transform, in other words inserting any +//! number of iterations of dest prop between the first and second application of this transform +//! will still not cause any further changes. +//! + +use crate::util::is_within_packed; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_mir_dataflow::debuginfo::debuginfo_locals; +use rustc_mir_dataflow::impls::{ + borrowed_locals, LivenessTransferFunction, MaybeTransitiveLiveLocals, +}; +use rustc_mir_dataflow::Analysis; + +/// Performs the optimization on the body +/// +/// The `borrowed` set must be a `BitSet` of all the locals that are ever borrowed in this body. It +/// can be generated via the [`borrowed_locals`] function. +pub fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let borrowed_locals = borrowed_locals(body); + + // If the user requests complete debuginfo, mark the locals that appear in it as live, so + // we don't remove assignements to them. + let mut always_live = debuginfo_locals(body); + always_live.union(&borrowed_locals); + + let mut live = MaybeTransitiveLiveLocals::new(&always_live) + .into_engine(tcx, body) + .iterate_to_fixpoint() + .into_results_cursor(body); + + // For blocks with a call terminator, if an argument copy can be turned into a move, + // record it as (block, argument index). + let mut call_operands_to_move = Vec::new(); + let mut patch = Vec::new(); + + for (bb, bb_data) in traversal::preorder(body) { + if let TerminatorKind::Call { ref args, .. } = bb_data.terminator().kind { + let loc = Location { block: bb, statement_index: bb_data.statements.len() }; + + // Position ourselves between the evaluation of `args` and the write to `destination`. + live.seek_to_block_end(bb); + let mut state = live.get().clone(); + + for (index, arg) in args.iter().enumerate().rev() { + if let Operand::Copy(place) = *arg + && !place.is_indirect() + // Do not skip the transformation if the local is in debuginfo, as we do + // not really lose any information for this purpose. + && !borrowed_locals.contains(place.local) + && !state.contains(place.local) + // If `place` is a projection of a disaligned field in a packed ADT, + // the move may be codegened as a pointer to that field. + // Using that disaligned pointer may trigger UB in the callee, + // so do nothing. + && is_within_packed(tcx, body, place).is_none() + { + call_operands_to_move.push((bb, index)); + } + + // Account that `arg` is read from, so we don't promote another argument to a move. + LivenessTransferFunction(&mut state).visit_operand(arg, loc); + } + } + + for (statement_index, statement) in bb_data.statements.iter().enumerate().rev() { + let loc = Location { block: bb, statement_index }; + if let StatementKind::Assign(assign) = &statement.kind { + if !assign.1.is_safe_to_remove() { + continue; + } + } + match &statement.kind { + StatementKind::Assign(box (place, _)) + | StatementKind::SetDiscriminant { place: box place, .. } + | StatementKind::Deinit(box place) => { + if !place.is_indirect() && !always_live.contains(place.local) { + live.seek_before_primary_effect(loc); + if !live.get().contains(place.local) { + patch.push(loc); + } + } + } + StatementKind::Retag(_, _) + | StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::Coverage(_) + | StatementKind::Intrinsic(_) + | StatementKind::ConstEvalCounter + | StatementKind::PlaceMention(_) + | StatementKind::Nop => (), + + StatementKind::FakeRead(_) | StatementKind::AscribeUserType(_, _) => { + bug!("{:?} not found in this MIR phase!", &statement.kind) + } + } + } + } + + if patch.is_empty() && call_operands_to_move.is_empty() { + return; + } + + let bbs = body.basic_blocks.as_mut_preserves_cfg(); + for Location { block, statement_index } in patch { + bbs[block].statements[statement_index].make_nop(); + } + for (block, argument_index) in call_operands_to_move { + let TerminatorKind::Call { ref mut args, .. } = bbs[block].terminator_mut().kind else { + bug!() + }; + let arg = &mut args[argument_index]; + let Operand::Copy(place) = *arg else { bug!() }; + *arg = Operand::Move(place); + } + + crate::simplify::simplify_locals(body, tcx) +} + +pub struct DeadStoreElimination; + +impl<'tcx> MirPass<'tcx> for DeadStoreElimination { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 2 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + eliminate(tcx, body); + } +} diff --git a/compiler/rustc_mir_transform/src/deduce_param_attrs.rs b/compiler/rustc_mir_transform/src/deduce_param_attrs.rs new file mode 100644 index 00000000000..990cfb05e60 --- /dev/null +++ b/compiler/rustc_mir_transform/src/deduce_param_attrs.rs @@ -0,0 +1,226 @@ +//! Deduces supplementary parameter attributes from MIR. +//! +//! Deduced parameter attributes are those that can only be soundly determined by examining the +//! body of the function instead of just the signature. These can be useful for optimization +//! purposes on a best-effort basis. We compute them here and store them into the crate metadata so +//! dependent crates can use them. + +use rustc_hir::def_id::LocalDefId; +use rustc_index::bit_set::BitSet; +use rustc_middle::mir::visit::{NonMutatingUseContext, PlaceContext, Visitor}; +use rustc_middle::mir::{Body, Location, Operand, Place, Terminator, TerminatorKind, RETURN_PLACE}; +use rustc_middle::ty::{self, DeducedParamAttrs, Ty, TyCtxt}; +use rustc_session::config::OptLevel; + +/// A visitor that determines which arguments have been mutated. We can't use the mutability field +/// on LocalDecl for this because it has no meaning post-optimization. +struct DeduceReadOnly { + /// Each bit is indexed by argument number, starting at zero (so 0 corresponds to local decl + /// 1). The bit is true if the argument may have been mutated or false if we know it hasn't + /// been up to the point we're at. + mutable_args: BitSet<usize>, +} + +impl DeduceReadOnly { + /// Returns a new DeduceReadOnly instance. + fn new(arg_count: usize) -> Self { + Self { mutable_args: BitSet::new_empty(arg_count) } + } +} + +impl<'tcx> Visitor<'tcx> for DeduceReadOnly { + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, _location: Location) { + // We're only interested in arguments. + if place.local == RETURN_PLACE || place.local.index() > self.mutable_args.domain_size() { + return; + } + + let mark_as_mutable = match context { + PlaceContext::MutatingUse(..) => { + // This is a mutation, so mark it as such. + true + } + PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf) => { + // Whether mutating though a `&raw const` is allowed is still undecided, so we + // disable any sketchy `readonly` optimizations for now. + // But we only need to do this if the pointer would point into the argument. + // IOW: for indirect places, like `&raw (*local).field`, this surely cannot mutate `local`. + !place.is_indirect() + } + PlaceContext::NonMutatingUse(..) | PlaceContext::NonUse(..) => { + // Not mutating, so it's fine. + false + } + }; + + if mark_as_mutable { + self.mutable_args.insert(place.local.index() - 1); + } + } + + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + // OK, this is subtle. Suppose that we're trying to deduce whether `x` in `f` is read-only + // and we have the following: + // + // fn f(x: BigStruct) { g(x) } + // fn g(mut y: BigStruct) { y.foo = 1 } + // + // If, at the generated MIR level, `f` turned into something like: + // + // fn f(_1: BigStruct) -> () { + // let mut _0: (); + // bb0: { + // _0 = g(move _1) -> bb1; + // } + // ... + // } + // + // then it would be incorrect to mark `x` (i.e. `_1`) as `readonly`, because `g`'s write to + // its copy of the indirect parameter would actually be a write directly to the pointer that + // `f` passes. Note that function arguments are the only situation in which this problem can + // arise: every other use of `move` in MIR doesn't actually write to the value it moves + // from. + // + // Anyway, right now this situation doesn't actually arise in practice. Instead, the MIR for + // that function looks like this: + // + // fn f(_1: BigStruct) -> () { + // let mut _0: (); + // let mut _2: BigStruct; + // bb0: { + // _2 = move _1; + // _0 = g(move _2) -> bb1; + // } + // ... + // } + // + // Because of that extra move that MIR construction inserts, `x` (i.e. `_1`) can *in + // practice* safely be marked `readonly`. + // + // To handle the possibility that other optimizations (for example, destination propagation) + // might someday generate MIR like the first example above, we panic upon seeing an argument + // to *our* function that is directly moved into *another* function as an argument. Having + // eliminated that problematic case, we can safely treat moves as copies in this analysis. + // + // In the future, if MIR optimizations cause arguments of a caller to be directly moved into + // the argument of a callee, we can just add that argument to `mutated_args` instead of + // panicking. + // + // Note that, because the problematic MIR is never actually generated, we can't add a test + // case for this. + + if let TerminatorKind::Call { ref args, .. } = terminator.kind { + for arg in args { + if let Operand::Move(place) = *arg { + let local = place.local; + if place.is_indirect() + || local == RETURN_PLACE + || local.index() > self.mutable_args.domain_size() + { + continue; + } + + self.mutable_args.insert(local.index() - 1); + } + } + }; + + self.super_terminator(terminator, location); + } +} + +/// Returns true if values of a given type will never be passed indirectly, regardless of ABI. +fn type_will_always_be_passed_directly(ty: Ty<'_>) -> bool { + matches!( + ty.kind(), + ty::Bool + | ty::Char + | ty::Float(..) + | ty::Int(..) + | ty::RawPtr(..) + | ty::Ref(..) + | ty::Slice(..) + | ty::Uint(..) + ) +} + +/// Returns the deduced parameter attributes for a function. +/// +/// Deduced parameter attributes are those that can only be soundly determined by examining the +/// body of the function instead of just the signature. These can be useful for optimization +/// purposes on a best-effort basis. We compute them here and store them into the crate metadata so +/// dependent crates can use them. +pub fn deduced_param_attrs<'tcx>( + tcx: TyCtxt<'tcx>, + def_id: LocalDefId, +) -> &'tcx [DeducedParamAttrs] { + // This computation is unfortunately rather expensive, so don't do it unless we're optimizing. + // Also skip it in incremental mode. + if tcx.sess.opts.optimize == OptLevel::No || tcx.sess.opts.incremental.is_some() { + return &[]; + } + + // If the Freeze language item isn't present, then don't bother. + if tcx.lang_items().freeze_trait().is_none() { + return &[]; + } + + // Codegen won't use this information for anything if all the function parameters are passed + // directly. Detect that and bail, for compilation speed. + let fn_ty = tcx.type_of(def_id).instantiate_identity(); + if matches!(fn_ty.kind(), ty::FnDef(..)) { + if fn_ty + .fn_sig(tcx) + .inputs() + .skip_binder() + .iter() + .cloned() + .all(type_will_always_be_passed_directly) + { + return &[]; + } + } + + // Don't deduce any attributes for functions that have no MIR. + if !tcx.is_mir_available(def_id) { + return &[]; + } + + // Grab the optimized MIR. Analyze it to determine which arguments have been mutated. + let body: &Body<'tcx> = tcx.optimized_mir(def_id); + let mut deduce_read_only = DeduceReadOnly::new(body.arg_count); + deduce_read_only.visit_body(body); + + // Set the `readonly` attribute for every argument that we concluded is immutable and that + // contains no UnsafeCells. + // + // FIXME: This is overly conservative around generic parameters: `is_freeze()` will always + // return false for them. For a description of alternatives that could do a better job here, + // see [1]. + // + // [1]: https://github.com/rust-lang/rust/pull/103172#discussion_r999139997 + let param_env = tcx.param_env_reveal_all_normalized(def_id); + let mut deduced_param_attrs = tcx.arena.alloc_from_iter( + body.local_decls.iter().skip(1).take(body.arg_count).enumerate().map( + |(arg_index, local_decl)| DeducedParamAttrs { + read_only: !deduce_read_only.mutable_args.contains(arg_index) + // We must normalize here to reveal opaques and normalize + // their substs, otherwise we'll see exponential blow-up in + // compile times: #113372 + && tcx + .normalize_erasing_regions(param_env, local_decl.ty) + .is_freeze(tcx, param_env), + }, + ), + ); + + // Trailing parameters past the size of the `deduced_param_attrs` array are assumed to have the + // default set of attributes, so we don't have to store them explicitly. Pop them off to save a + // few bytes in metadata. + while deduced_param_attrs.last() == Some(&DeducedParamAttrs::default()) { + let last_index = deduced_param_attrs.len() - 1; + deduced_param_attrs = &mut deduced_param_attrs[0..last_index]; + } + + deduced_param_attrs +} diff --git a/compiler/rustc_mir_transform/src/deduplicate_blocks.rs b/compiler/rustc_mir_transform/src/deduplicate_blocks.rs new file mode 100644 index 00000000000..b40b2ec8bfd --- /dev/null +++ b/compiler/rustc_mir_transform/src/deduplicate_blocks.rs @@ -0,0 +1,187 @@ +//! This pass finds basic blocks that are completely equal, +//! and replaces all uses with just one of them. + +use std::{collections::hash_map::Entry, hash::Hash, hash::Hasher, iter}; + +use rustc_data_structures::fx::FxHashMap; +use rustc_middle::mir::visit::MutVisitor; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +use super::simplify::simplify_cfg; + +pub struct DeduplicateBlocks; + +impl<'tcx> MirPass<'tcx> for DeduplicateBlocks { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 4 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!("Running DeduplicateBlocks on `{:?}`", body.source); + let duplicates = find_duplicates(body); + let has_opts_to_apply = !duplicates.is_empty(); + + if has_opts_to_apply { + let mut opt_applier = OptApplier { tcx, duplicates }; + opt_applier.visit_body(body); + simplify_cfg(tcx, body); + } + } +} + +struct OptApplier<'tcx> { + tcx: TyCtxt<'tcx>, + duplicates: FxHashMap<BasicBlock, BasicBlock>, +} + +impl<'tcx> MutVisitor<'tcx> for OptApplier<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) { + for target in terminator.successors_mut() { + if let Some(replacement) = self.duplicates.get(target) { + debug!("SUCCESS: Replacing: `{:?}` with `{:?}`", target, replacement); + *target = *replacement; + } + } + + self.super_terminator(terminator, location); + } +} + +fn find_duplicates(body: &Body<'_>) -> FxHashMap<BasicBlock, BasicBlock> { + let mut duplicates = FxHashMap::default(); + + let bbs_to_go_through = + body.basic_blocks.iter_enumerated().filter(|(_, bbd)| !bbd.is_cleanup).count(); + + let mut same_hashes = + FxHashMap::with_capacity_and_hasher(bbs_to_go_through, Default::default()); + + // Go through the basic blocks backwards. This means that in case of duplicates, + // we can use the basic block with the highest index as the replacement for all lower ones. + // For example, if bb1, bb2 and bb3 are duplicates, we will first insert bb3 in same_hashes. + // Then we will see that bb2 is a duplicate of bb3, + // and insert bb2 with the replacement bb3 in the duplicates list. + // When we see bb1, we see that it is a duplicate of bb3, and therefore insert it in the duplicates list + // with replacement bb3. + // When the duplicates are removed, we will end up with only bb3. + for (bb, bbd) in body.basic_blocks.iter_enumerated().rev().filter(|(_, bbd)| !bbd.is_cleanup) { + // Basic blocks can get really big, so to avoid checking for duplicates in basic blocks + // that are unlikely to have duplicates, we stop early. The early bail number has been + // found experimentally by eprintln while compiling the crates in the rustc-perf suite. + if bbd.statements.len() > 10 { + continue; + } + + let to_hash = BasicBlockHashable { basic_block_data: bbd }; + let entry = same_hashes.entry(to_hash); + match entry { + Entry::Occupied(occupied) => { + // The basic block was already in the hashmap, which means we have a duplicate + let value = *occupied.get(); + debug!("Inserting {:?} -> {:?}", bb, value); + duplicates.try_insert(bb, value).expect("key was already inserted"); + } + Entry::Vacant(vacant) => { + vacant.insert(bb); + } + } + } + + duplicates +} + +struct BasicBlockHashable<'tcx, 'a> { + basic_block_data: &'a BasicBlockData<'tcx>, +} + +impl Hash for BasicBlockHashable<'_, '_> { + fn hash<H: Hasher>(&self, state: &mut H) { + hash_statements(state, self.basic_block_data.statements.iter()); + // Note that since we only hash the kind, we lose span information if we deduplicate the blocks + self.basic_block_data.terminator().kind.hash(state); + } +} + +impl Eq for BasicBlockHashable<'_, '_> {} + +impl PartialEq for BasicBlockHashable<'_, '_> { + fn eq(&self, other: &Self) -> bool { + self.basic_block_data.statements.len() == other.basic_block_data.statements.len() + && &self.basic_block_data.terminator().kind == &other.basic_block_data.terminator().kind + && iter::zip(&self.basic_block_data.statements, &other.basic_block_data.statements) + .all(|(x, y)| statement_eq(&x.kind, &y.kind)) + } +} + +fn hash_statements<'a, 'tcx, H: Hasher>( + hasher: &mut H, + iter: impl Iterator<Item = &'a Statement<'tcx>>, +) where + 'tcx: 'a, +{ + for stmt in iter { + statement_hash(hasher, &stmt.kind); + } +} + +fn statement_hash<H: Hasher>(hasher: &mut H, stmt: &StatementKind<'_>) { + match stmt { + StatementKind::Assign(box (place, rvalue)) => { + place.hash(hasher); + rvalue_hash(hasher, rvalue) + } + x => x.hash(hasher), + }; +} + +fn rvalue_hash<H: Hasher>(hasher: &mut H, rvalue: &Rvalue<'_>) { + match rvalue { + Rvalue::Use(op) => operand_hash(hasher, op), + x => x.hash(hasher), + }; +} + +fn operand_hash<H: Hasher>(hasher: &mut H, operand: &Operand<'_>) { + match operand { + Operand::Constant(box ConstOperand { user_ty: _, const_, span: _ }) => const_.hash(hasher), + x => x.hash(hasher), + }; +} + +fn statement_eq<'tcx>(lhs: &StatementKind<'tcx>, rhs: &StatementKind<'tcx>) -> bool { + let res = match (lhs, rhs) { + ( + StatementKind::Assign(box (place, rvalue)), + StatementKind::Assign(box (place2, rvalue2)), + ) => place == place2 && rvalue_eq(rvalue, rvalue2), + (x, y) => x == y, + }; + debug!("statement_eq lhs: `{:?}` rhs: `{:?}` result: {:?}", lhs, rhs, res); + res +} + +fn rvalue_eq<'tcx>(lhs: &Rvalue<'tcx>, rhs: &Rvalue<'tcx>) -> bool { + let res = match (lhs, rhs) { + (Rvalue::Use(op1), Rvalue::Use(op2)) => operand_eq(op1, op2), + (x, y) => x == y, + }; + debug!("rvalue_eq lhs: `{:?}` rhs: `{:?}` result: {:?}", lhs, rhs, res); + res +} + +fn operand_eq<'tcx>(lhs: &Operand<'tcx>, rhs: &Operand<'tcx>) -> bool { + let res = match (lhs, rhs) { + ( + Operand::Constant(box ConstOperand { user_ty: _, const_, span: _ }), + Operand::Constant(box ConstOperand { user_ty: _, const_: const2, span: _ }), + ) => const_ == const2, + (x, y) => x == y, + }; + debug!("operand_eq lhs: `{:?}` rhs: `{:?}` result: {:?}", lhs, rhs, res); + res +} diff --git a/compiler/rustc_mir_transform/src/deref_separator.rs b/compiler/rustc_mir_transform/src/deref_separator.rs new file mode 100644 index 00000000000..0e2fccc85da --- /dev/null +++ b/compiler/rustc_mir_transform/src/deref_separator.rs @@ -0,0 +1,85 @@ +use rustc_index::IndexVec; +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; + +pub struct Derefer; + +pub struct DerefChecker<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + patcher: MirPatch<'tcx>, + local_decls: &'a IndexVec<Local, LocalDecl<'tcx>>, +} + +impl<'a, 'tcx> MutVisitor<'tcx> for DerefChecker<'a, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, cntxt: PlaceContext, loc: Location) { + if !place.projection.is_empty() + && cntxt != PlaceContext::NonUse(VarDebugInfo) + && place.projection[1..].contains(&ProjectionElem::Deref) + { + let mut place_local = place.local; + let mut last_len = 0; + let mut last_deref_idx = 0; + + for (idx, elem) in place.projection[0..].iter().enumerate() { + if *elem == ProjectionElem::Deref { + last_deref_idx = idx; + } + } + + for (idx, (p_ref, p_elem)) in place.iter_projections().enumerate() { + if !p_ref.projection.is_empty() && p_elem == ProjectionElem::Deref { + let ty = p_ref.ty(self.local_decls, self.tcx).ty; + let temp = self.patcher.new_local_with_info( + ty, + self.local_decls[p_ref.local].source_info.span, + LocalInfo::DerefTemp, + ); + + // We are adding current p_ref's projections to our + // temp value, excluding projections we already covered. + let deref_place = Place::from(place_local) + .project_deeper(&p_ref.projection[last_len..], self.tcx); + + self.patcher.add_assign( + loc, + Place::from(temp), + Rvalue::CopyForDeref(deref_place), + ); + place_local = temp; + last_len = p_ref.projection.len(); + + // Change `Place` only if we are actually at the Place's last deref + if idx == last_deref_idx { + let temp_place = + Place::from(temp).project_deeper(&place.projection[idx..], self.tcx); + *place = temp_place; + } + } + } + } + } +} + +pub fn deref_finder<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let patch = MirPatch::new(body); + let mut checker = DerefChecker { tcx, patcher: patch, local_decls: &body.local_decls }; + + for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() { + checker.visit_basic_block_data(bb, data); + } + + checker.patcher.apply(body); +} + +impl<'tcx> MirPass<'tcx> for Derefer { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + deref_finder(tcx, body); + } +} diff --git a/compiler/rustc_mir_transform/src/dest_prop.rs b/compiler/rustc_mir_transform/src/dest_prop.rs new file mode 100644 index 00000000000..15502adfb5a --- /dev/null +++ b/compiler/rustc_mir_transform/src/dest_prop.rs @@ -0,0 +1,849 @@ +//! Propagates assignment destinations backwards in the CFG to eliminate redundant assignments. +//! +//! # Motivation +//! +//! MIR building can insert a lot of redundant copies, and Rust code in general often tends to move +//! values around a lot. The result is a lot of assignments of the form `dest = {move} src;` in MIR. +//! MIR building for constants in particular tends to create additional locals that are only used +//! inside a single block to shuffle a value around unnecessarily. +//! +//! LLVM by itself is not good enough at eliminating these redundant copies (eg. see +//! <https://github.com/rust-lang/rust/issues/32966>), so this leaves some performance on the table +//! that we can regain by implementing an optimization for removing these assign statements in rustc +//! itself. When this optimization runs fast enough, it can also speed up the constant evaluation +//! and code generation phases of rustc due to the reduced number of statements and locals. +//! +//! # The Optimization +//! +//! Conceptually, this optimization is "destination propagation". It is similar to the Named Return +//! Value Optimization, or NRVO, known from the C++ world, except that it isn't limited to return +//! values or the return place `_0`. On a very high level, independent of the actual implementation +//! details, it does the following: +//! +//! 1) Identify `dest = src;` statements with values for `dest` and `src` whose storage can soundly +//! be merged. +//! 2) Replace all mentions of `src` with `dest` ("unifying" them and propagating the destination +//! backwards). +//! 3) Delete the `dest = src;` statement (by making it a `nop`). +//! +//! Step 1) is by far the hardest, so it is explained in more detail below. +//! +//! ## Soundness +//! +//! We have a pair of places `p` and `q`, whose memory we would like to merge. In order for this to +//! be sound, we need to check a number of conditions: +//! +//! * `p` and `q` must both be *constant* - it does not make much sense to talk about merging them +//! if they do not consistently refer to the same place in memory. This is satisfied if they do +//! not contain any indirection through a pointer or any indexing projections. +//! +//! * `p` and `q` must have the **same type**. If we replace a local with a subtype or supertype, +//! we may end up with a differnet vtable for that local. See the `subtyping-impacts-selection` +//! tests for an example where that causes issues. +//! +//! * We need to make sure that the goal of "merging the memory" is actually structurally possible +//! in MIR. For example, even if all the other conditions are satisfied, there is no way to +//! "merge" `_5.foo` and `_6.bar`. For now, we ensure this by requiring that both `p` and `q` are +//! locals with no further projections. Future iterations of this pass should improve on this. +//! +//! * Finally, we want `p` and `q` to use the same memory - however, we still need to make sure that +//! each of them has enough "ownership" of that memory to continue "doing its job." More +//! precisely, what we will check is that whenever the program performs a write to `p`, then it +//! does not currently care about what the value in `q` is (and vice versa). We formalize the +//! notion of "does not care what the value in `q` is" by checking the *liveness* of `q`. +//! +//! Because of the difficulty of computing liveness of places that have their address taken, we do +//! not even attempt to do it. Any places that are in a local that has its address taken is +//! excluded from the optimization. +//! +//! The first two conditions are simple structural requirements on the `Assign` statements that can +//! be trivially checked. The third requirement however is more difficult and costly to check. +//! +//! ## Future Improvements +//! +//! There are a number of ways in which this pass could be improved in the future: +//! +//! * Merging storage liveness ranges instead of removing storage statements completely. This may +//! improve stack usage. +//! +//! * Allow merging locals into places with projections, eg `_5` into `_6.foo`. +//! +//! * Liveness analysis with more precision than whole locals at a time. The smaller benefit of this +//! is that it would allow us to dest prop at "sub-local" levels in some cases. The bigger benefit +//! of this is that such liveness analysis can report more accurate results about whole locals at +//! a time. For example, consider: +//! +//! ```ignore (syntax-highlighting-only) +//! _1 = u; +//! // unrelated code +//! _1.f1 = v; +//! _2 = _1.f1; +//! ``` +//! +//! Because the current analysis only thinks in terms of locals, it does not have enough +//! information to report that `_1` is dead in the "unrelated code" section. +//! +//! * Liveness analysis enabled by alias analysis. This would allow us to not just bail on locals +//! that ever have their address taken. Of course that requires actually having alias analysis +//! (and a model to build it on), so this might be a bit of a ways off. +//! +//! * Various perf improvements. There are a bunch of comments in here marked `PERF` with ideas for +//! how to do things more efficiently. However, the complexity of the pass as a whole should be +//! kept in mind. +//! +//! ## Previous Work +//! +//! A [previous attempt][attempt 1] at implementing an optimization like this turned out to be a +//! significant regression in compiler performance. Fixing the regressions introduced a lot of +//! undesirable complexity to the implementation. +//! +//! A [subsequent approach][attempt 2] tried to avoid the costly computation by limiting itself to +//! acyclic CFGs, but still turned out to be far too costly to run due to suboptimal performance +//! within individual basic blocks, requiring a walk across the entire block for every assignment +//! found within the block. For the `tuple-stress` benchmark, which has 458745 statements in a +//! single block, this proved to be far too costly. +//! +//! [Another approach after that][attempt 3] was much closer to correct, but had some soundness +//! issues - it was failing to consider stores outside live ranges, and failed to uphold some of the +//! requirements that MIR has for non-overlapping places within statements. However, it also had +//! performance issues caused by `O(l² * s)` runtime, where `l` is the number of locals and `s` is +//! the number of statements and terminators. +//! +//! Since the first attempt at this, the compiler has improved dramatically, and new analysis +//! frameworks have been added that should make this approach viable without requiring a limited +//! approach that only works for some classes of CFGs: +//! - rustc now has a powerful dataflow analysis framework that can handle forwards and backwards +//! analyses efficiently. +//! - Layout optimizations for coroutines have been added to improve code generation for +//! async/await, which are very similar in spirit to what this optimization does. +//! +//! Also, rustc now has a simple NRVO pass (see `nrvo.rs`), which handles a subset of the cases that +//! this destination propagation pass handles, proving that similar optimizations can be performed +//! on MIR. +//! +//! ## Pre/Post Optimization +//! +//! It is recommended to run `SimplifyCfg` and then `SimplifyLocals` some time after this pass, as +//! it replaces the eliminated assign statements with `nop`s and leaves unused locals behind. +//! +//! [liveness]: https://en.wikipedia.org/wiki/Live_variable_analysis +//! [attempt 1]: https://github.com/rust-lang/rust/pull/47954 +//! [attempt 2]: https://github.com/rust-lang/rust/pull/71003 +//! [attempt 3]: https://github.com/rust-lang/rust/pull/72632 + +use std::collections::hash_map::{Entry, OccupiedEntry}; + +use crate::simplify::remove_dead_blocks; +use crate::MirPass; +use rustc_data_structures::fx::FxHashMap; +use rustc_index::bit_set::BitSet; +use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor}; +use rustc_middle::mir::HasLocalDecls; +use rustc_middle::mir::{dump_mir, PassWhere}; +use rustc_middle::mir::{ + traversal, Body, InlineAsmOperand, Local, LocalKind, Location, Operand, Place, Rvalue, + Statement, StatementKind, TerminatorKind, +}; +use rustc_middle::ty::TyCtxt; +use rustc_mir_dataflow::impls::MaybeLiveLocals; +use rustc_mir_dataflow::{Analysis, ResultsCursor}; + +pub struct DestinationPropagation; + +impl<'tcx> MirPass<'tcx> for DestinationPropagation { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // For now, only run at MIR opt level 3. Two things need to be changed before this can be + // turned on by default: + // 1. Because of the overeager removal of storage statements, this can cause stack space + // regressions. This opt is not the place to fix this though, it's a more general + // problem in MIR. + // 2. Despite being an overall perf improvement, this still causes a 30% regression in + // keccak. We can temporarily fix this by bounding function size, but in the long term + // we should fix this by being smarter about invalidating analysis results. + sess.mir_opt_level() >= 3 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + let mut allocations = Allocations::default(); + trace!(func = ?tcx.def_path_str(def_id)); + + let borrowed = rustc_mir_dataflow::impls::borrowed_locals(body); + + // In order to avoid having to collect data for every single pair of locals in the body, we + // do not allow doing more than one merge for places that are derived from the same local at + // once. To avoid missed opportunities, we instead iterate to a fixed point - we'll refer to + // each of these iterations as a "round." + // + // Reaching a fixed point could in theory take up to `min(l, s)` rounds - however, we do not + // expect to see MIR like that. To verify this, a test was run against `[rust-lang/regex]` - + // the average MIR body saw 1.32 full iterations of this loop. The most that was hit were 30 + // for a single function. Only 80/2801 (2.9%) of functions saw at least 5. + // + // [rust-lang/regex]: + // https://github.com/rust-lang/regex/tree/b5372864e2df6a2f5e543a556a62197f50ca3650 + let mut round_count = 0; + loop { + // PERF: Can we do something smarter than recalculating the candidates and liveness + // results? + let mut candidates = find_candidates( + body, + &borrowed, + &mut allocations.candidates, + &mut allocations.candidates_reverse, + ); + trace!(?candidates); + let mut live = MaybeLiveLocals + .into_engine(tcx, body) + .iterate_to_fixpoint() + .into_results_cursor(body); + dest_prop_mir_dump(tcx, body, &mut live, round_count); + + FilterInformation::filter_liveness( + &mut candidates, + &mut live, + &mut allocations.write_info, + body, + ); + + // Because we do not update liveness information, it is unsound to use a local for more + // than one merge operation within a single round of optimizations. We store here which + // ones we have already used. + let mut merged_locals: BitSet<Local> = BitSet::new_empty(body.local_decls.len()); + + // This is the set of merges we will apply this round. It is a subset of the candidates. + let mut merges = FxHashMap::default(); + + for (src, candidates) in candidates.c.iter() { + if merged_locals.contains(*src) { + continue; + } + let Some(dest) = candidates.iter().find(|dest| !merged_locals.contains(**dest)) + else { + continue; + }; + if !tcx.consider_optimizing(|| { + format!("{} round {}", tcx.def_path_str(def_id), round_count) + }) { + break; + } + merges.insert(*src, *dest); + merged_locals.insert(*src); + merged_locals.insert(*dest); + } + trace!(merging = ?merges); + + if merges.is_empty() { + break; + } + round_count += 1; + + apply_merges(body, tcx, &merges, &merged_locals); + } + + if round_count != 0 { + // Merging can introduce overlap between moved arguments and/or call destination in an + // unreachable code, which validator considers to be ill-formed. + remove_dead_blocks(body); + } + + trace!(round_count); + } +} + +/// Container for the various allocations that we need. +/// +/// We store these here and hand out `&mut` access to them, instead of dropping and recreating them +/// frequently. Everything with a `&'alloc` lifetime points into here. +#[derive(Default)] +struct Allocations { + candidates: FxHashMap<Local, Vec<Local>>, + candidates_reverse: FxHashMap<Local, Vec<Local>>, + write_info: WriteInfo, + // PERF: Do this for `MaybeLiveLocals` allocations too. +} + +#[derive(Debug)] +struct Candidates<'alloc> { + /// The set of candidates we are considering in this optimization. + /// + /// We will always merge the key into at most one of its values. + /// + /// Whether a place ends up in the key or the value does not correspond to whether it appears as + /// the lhs or rhs of any assignment. As a matter of fact, the places in here might never appear + /// in an assignment at all. This happens because if we see an assignment like this: + /// + /// ```ignore (syntax-highlighting-only) + /// _1.0 = _2.0 + /// ``` + /// + /// We will still report that we would like to merge `_1` and `_2` in an attempt to allow us to + /// remove that assignment. + c: &'alloc mut FxHashMap<Local, Vec<Local>>, + /// A reverse index of the `c` set; if the `c` set contains `a => Place { local: b, proj }`, + /// then this contains `b => a`. + // PERF: Possibly these should be `SmallVec`s? + reverse: &'alloc mut FxHashMap<Local, Vec<Local>>, +} + +////////////////////////////////////////////////////////// +// Merging +// +// Applies the actual optimization + +fn apply_merges<'tcx>( + body: &mut Body<'tcx>, + tcx: TyCtxt<'tcx>, + merges: &FxHashMap<Local, Local>, + merged_locals: &BitSet<Local>, +) { + let mut merger = Merger { tcx, merges, merged_locals }; + merger.visit_body_preserves_cfg(body); +} + +struct Merger<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + merges: &'a FxHashMap<Local, Local>, + merged_locals: &'a BitSet<Local>, +} + +impl<'a, 'tcx> MutVisitor<'tcx> for Merger<'a, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _location: Location) { + if let Some(dest) = self.merges.get(local) { + *local = *dest; + } + } + + fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) { + match &statement.kind { + // FIXME: Don't delete storage statements, but "merge" the storage ranges instead. + StatementKind::StorageDead(local) | StatementKind::StorageLive(local) + if self.merged_locals.contains(*local) => + { + statement.make_nop(); + return; + } + _ => (), + }; + self.super_statement(statement, location); + match &statement.kind { + StatementKind::Assign(box (dest, rvalue)) => { + match rvalue { + Rvalue::CopyForDeref(place) + | Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) => { + // These might've been turned into self-assignments by the replacement + // (this includes the original statement we wanted to eliminate). + if dest == place { + debug!("{:?} turned into self-assignment, deleting", location); + statement.make_nop(); + } + } + _ => {} + } + } + + _ => {} + } + } +} + +////////////////////////////////////////////////////////// +// Liveness filtering +// +// This section enforces bullet point 2 + +struct FilterInformation<'a, 'body, 'alloc, 'tcx> { + body: &'body Body<'tcx>, + live: &'a mut ResultsCursor<'body, 'tcx, MaybeLiveLocals>, + candidates: &'a mut Candidates<'alloc>, + write_info: &'alloc mut WriteInfo, + at: Location, +} + +// We first implement some utility functions which we will expose removing candidates according to +// different needs. Throughout the liveness filtering, the `candidates` are only ever accessed +// through these methods, and not directly. +impl<'alloc> Candidates<'alloc> { + /// Just `Vec::retain`, but the condition is inverted and we add debugging output + fn vec_filter_candidates( + src: Local, + v: &mut Vec<Local>, + mut f: impl FnMut(Local) -> CandidateFilter, + at: Location, + ) { + v.retain(|dest| { + let remove = f(*dest); + if remove == CandidateFilter::Remove { + trace!("eliminating {:?} => {:?} due to conflict at {:?}", src, dest, at); + } + remove == CandidateFilter::Keep + }); + } + + /// `vec_filter_candidates` but for an `Entry` + fn entry_filter_candidates( + mut entry: OccupiedEntry<'_, Local, Vec<Local>>, + p: Local, + f: impl FnMut(Local) -> CandidateFilter, + at: Location, + ) { + let candidates = entry.get_mut(); + Self::vec_filter_candidates(p, candidates, f, at); + if candidates.len() == 0 { + entry.remove(); + } + } + + /// For all candidates `(p, q)` or `(q, p)` removes the candidate if `f(q)` says to do so + fn filter_candidates_by( + &mut self, + p: Local, + mut f: impl FnMut(Local) -> CandidateFilter, + at: Location, + ) { + // Cover the cases where `p` appears as a `src` + if let Entry::Occupied(entry) = self.c.entry(p) { + Self::entry_filter_candidates(entry, p, &mut f, at); + } + // And the cases where `p` appears as a `dest` + let Some(srcs) = self.reverse.get_mut(&p) else { + return; + }; + // We use `retain` here to remove the elements from the reverse set if we've removed the + // matching candidate in the forward set. + srcs.retain(|src| { + if f(*src) == CandidateFilter::Keep { + return true; + } + let Entry::Occupied(entry) = self.c.entry(*src) else { + return false; + }; + Self::entry_filter_candidates( + entry, + *src, + |dest| { + if dest == p { CandidateFilter::Remove } else { CandidateFilter::Keep } + }, + at, + ); + false + }); + } +} + +#[derive(Copy, Clone, PartialEq, Eq)] +enum CandidateFilter { + Keep, + Remove, +} + +impl<'a, 'body, 'alloc, 'tcx> FilterInformation<'a, 'body, 'alloc, 'tcx> { + /// Filters the set of candidates to remove those that conflict. + /// + /// The steps we take are exactly those that are outlined at the top of the file. For each + /// statement/terminator, we collect the set of locals that are written to in that + /// statement/terminator, and then we remove all pairs of candidates that contain one such local + /// and another one that is live. + /// + /// We need to be careful about the ordering of operations within each statement/terminator + /// here. Many statements might write and read from more than one place, and we need to consider + /// them all. The strategy for doing this is as follows: We first gather all the places that are + /// written to within the statement/terminator via `WriteInfo`. Then, we use the liveness + /// analysis from *before* the statement/terminator (in the control flow sense) to eliminate + /// candidates - this is because we want to conservatively treat a pair of locals that is both + /// read and written in the statement/terminator to be conflicting, and the liveness analysis + /// before the statement/terminator will correctly report locals that are read in the + /// statement/terminator to be live. We are additionally conservative by treating all written to + /// locals as also being read from. + fn filter_liveness<'b>( + candidates: &mut Candidates<'alloc>, + live: &mut ResultsCursor<'b, 'tcx, MaybeLiveLocals>, + write_info_alloc: &'alloc mut WriteInfo, + body: &'b Body<'tcx>, + ) { + let mut this = FilterInformation { + body, + live, + candidates, + // We don't actually store anything at this scope, we just keep things here to be able + // to reuse the allocation. + write_info: write_info_alloc, + // Doesn't matter what we put here, will be overwritten before being used + at: Location::START, + }; + this.internal_filter_liveness(); + } + + fn internal_filter_liveness(&mut self) { + for (block, data) in traversal::preorder(self.body) { + self.at = Location { block, statement_index: data.statements.len() }; + self.live.seek_after_primary_effect(self.at); + self.write_info.for_terminator(&data.terminator().kind); + self.apply_conflicts(); + + for (i, statement) in data.statements.iter().enumerate().rev() { + self.at = Location { block, statement_index: i }; + self.live.seek_after_primary_effect(self.at); + self.write_info.for_statement(&statement.kind, self.body); + self.apply_conflicts(); + } + } + } + + fn apply_conflicts(&mut self) { + let writes = &self.write_info.writes; + for p in writes { + let other_skip = self.write_info.skip_pair.and_then(|(a, b)| { + if a == *p { + Some(b) + } else if b == *p { + Some(a) + } else { + None + } + }); + self.candidates.filter_candidates_by( + *p, + |q| { + if Some(q) == other_skip { + return CandidateFilter::Keep; + } + // It is possible that a local may be live for less than the + // duration of a statement This happens in the case of function + // calls or inline asm. Because of this, we also mark locals as + // conflicting when both of them are written to in the same + // statement. + if self.live.contains(q) || writes.contains(&q) { + CandidateFilter::Remove + } else { + CandidateFilter::Keep + } + }, + self.at, + ); + } + } +} + +/// Describes where a statement/terminator writes to +#[derive(Default, Debug)] +struct WriteInfo { + writes: Vec<Local>, + /// If this pair of locals is a candidate pair, completely skip processing it during this + /// statement. All other candidates are unaffected. + skip_pair: Option<(Local, Local)>, +} + +impl WriteInfo { + fn for_statement<'tcx>(&mut self, statement: &StatementKind<'tcx>, body: &Body<'tcx>) { + self.reset(); + match statement { + StatementKind::Assign(box (lhs, rhs)) => { + self.add_place(*lhs); + match rhs { + Rvalue::Use(op) => { + self.add_operand(op); + self.consider_skipping_for_assign_use(*lhs, op, body); + } + Rvalue::Repeat(op, _) => { + self.add_operand(op); + } + Rvalue::Cast(_, op, _) + | Rvalue::UnaryOp(_, op) + | Rvalue::ShallowInitBox(op, _) => { + self.add_operand(op); + } + Rvalue::BinaryOp(_, ops) | Rvalue::CheckedBinaryOp(_, ops) => { + for op in [&ops.0, &ops.1] { + self.add_operand(op); + } + } + Rvalue::Aggregate(_, ops) => { + for op in ops { + self.add_operand(op); + } + } + Rvalue::ThreadLocalRef(_) + | Rvalue::NullaryOp(_, _) + | Rvalue::Ref(_, _, _) + | Rvalue::AddressOf(_, _) + | Rvalue::Len(_) + | Rvalue::Discriminant(_) + | Rvalue::CopyForDeref(_) => (), + } + } + // Retags are technically also reads, but reporting them as a write suffices + StatementKind::SetDiscriminant { place, .. } + | StatementKind::Deinit(place) + | StatementKind::Retag(_, place) => { + self.add_place(**place); + } + StatementKind::Intrinsic(_) + | StatementKind::ConstEvalCounter + | StatementKind::Nop + | StatementKind::Coverage(_) + | StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::PlaceMention(_) => (), + StatementKind::FakeRead(_) | StatementKind::AscribeUserType(_, _) => { + bug!("{:?} not found in this MIR phase", statement) + } + } + } + + fn consider_skipping_for_assign_use<'tcx>( + &mut self, + lhs: Place<'tcx>, + rhs: &Operand<'tcx>, + body: &Body<'tcx>, + ) { + let Some(rhs) = rhs.place() else { return }; + if let Some(pair) = places_to_candidate_pair(lhs, rhs, body) { + self.skip_pair = Some(pair); + } + } + + fn for_terminator<'tcx>(&mut self, terminator: &TerminatorKind<'tcx>) { + self.reset(); + match terminator { + TerminatorKind::SwitchInt { discr: op, .. } + | TerminatorKind::Assert { cond: op, .. } => { + self.add_operand(op); + } + TerminatorKind::Call { destination, func, args, .. } => { + self.add_place(*destination); + self.add_operand(func); + for arg in args { + self.add_operand(arg); + } + } + TerminatorKind::InlineAsm { operands, .. } => { + for asm_operand in operands { + match asm_operand { + InlineAsmOperand::In { value, .. } => { + self.add_operand(value); + } + InlineAsmOperand::Out { place, .. } => { + if let Some(place) = place { + self.add_place(*place); + } + } + // Note that the `late` field in `InOut` is about whether the registers used + // for these things overlap, and is of absolutely no interest to us. + InlineAsmOperand::InOut { in_value, out_place, .. } => { + if let Some(place) = out_place { + self.add_place(*place); + } + self.add_operand(in_value); + } + InlineAsmOperand::Const { .. } + | InlineAsmOperand::SymFn { .. } + | InlineAsmOperand::SymStatic { .. } => (), + } + } + } + TerminatorKind::Goto { .. } + | TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable { .. } => (), + TerminatorKind::Drop { .. } => { + // `Drop`s create a `&mut` and so are not considered + } + TerminatorKind::Yield { .. } + | TerminatorKind::CoroutineDrop + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } => { + bug!("{:?} not found in this MIR phase", terminator) + } + } + } + + fn add_place(&mut self, place: Place<'_>) { + self.writes.push(place.local); + } + + fn add_operand<'tcx>(&mut self, op: &Operand<'tcx>) { + match op { + // FIXME(JakobDegen): In a previous version, the `Move` case was incorrectly treated as + // being a read only. This was unsound, however we cannot add a regression test because + // it is not possible to set this off with current MIR. Once we have that ability, a + // regression test should be added. + Operand::Move(p) => self.add_place(*p), + Operand::Copy(_) | Operand::Constant(_) => (), + } + } + + fn reset(&mut self) { + self.writes.clear(); + self.skip_pair = None; + } +} + +///////////////////////////////////////////////////// +// Candidate accumulation + +/// If the pair of places is being considered for merging, returns the candidate which would be +/// merged in order to accomplish this. +/// +/// The contract here is in one direction - there is a guarantee that merging the locals that are +/// outputted by this function would result in an assignment between the inputs becoming a +/// self-assignment. However, there is no guarantee that the returned pair is actually suitable for +/// merging - candidate collection must still check this independently. +/// +/// This output is unique for each unordered pair of input places. +fn places_to_candidate_pair<'tcx>( + a: Place<'tcx>, + b: Place<'tcx>, + body: &Body<'tcx>, +) -> Option<(Local, Local)> { + let (mut a, mut b) = if a.projection.len() == 0 && b.projection.len() == 0 { + (a.local, b.local) + } else { + return None; + }; + + // By sorting, we make sure we're input order independent + if a > b { + std::mem::swap(&mut a, &mut b); + } + + // We could now return `(a, b)`, but then we miss some candidates in the case where `a` can't be + // used as a `src`. + if is_local_required(a, body) { + std::mem::swap(&mut a, &mut b); + } + // We could check `is_local_required` again here, but there's no need - after all, we make no + // promise that the candidate pair is actually valid + Some((a, b)) +} + +/// Collects the candidates for merging +/// +/// This is responsible for enforcing the first and third bullet point. +fn find_candidates<'alloc, 'tcx>( + body: &Body<'tcx>, + borrowed: &BitSet<Local>, + candidates: &'alloc mut FxHashMap<Local, Vec<Local>>, + candidates_reverse: &'alloc mut FxHashMap<Local, Vec<Local>>, +) -> Candidates<'alloc> { + candidates.clear(); + candidates_reverse.clear(); + let mut visitor = FindAssignments { body, candidates, borrowed }; + visitor.visit_body(body); + // Deduplicate candidates + for (_, cands) in candidates.iter_mut() { + cands.sort(); + cands.dedup(); + } + // Generate the reverse map + for (src, cands) in candidates.iter() { + for dest in cands.iter().copied() { + candidates_reverse.entry(dest).or_default().push(*src); + } + } + Candidates { c: candidates, reverse: candidates_reverse } +} + +struct FindAssignments<'a, 'alloc, 'tcx> { + body: &'a Body<'tcx>, + candidates: &'alloc mut FxHashMap<Local, Vec<Local>>, + borrowed: &'a BitSet<Local>, +} + +impl<'tcx> Visitor<'tcx> for FindAssignments<'_, '_, 'tcx> { + fn visit_statement(&mut self, statement: &Statement<'tcx>, _: Location) { + if let StatementKind::Assign(box ( + lhs, + Rvalue::CopyForDeref(rhs) | Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)), + )) = &statement.kind + { + let Some((src, dest)) = places_to_candidate_pair(*lhs, *rhs, self.body) else { + return; + }; + + // As described at the top of the file, we do not go near things that have + // their address taken. + if self.borrowed.contains(src) || self.borrowed.contains(dest) { + return; + } + + // As described at the top of this file, we do not touch locals which have + // different types. + let src_ty = self.body.local_decls()[src].ty; + let dest_ty = self.body.local_decls()[dest].ty; + if src_ty != dest_ty { + // FIXME(#112651): This can be removed afterwards. Also update the module description. + trace!("skipped `{src:?} = {dest:?}` due to subtyping: {src_ty} != {dest_ty}"); + return; + } + + // Also, we need to make sure that MIR actually allows the `src` to be removed + if is_local_required(src, self.body) { + return; + } + + // We may insert duplicates here, but that's fine + self.candidates.entry(src).or_default().push(dest); + } + } +} + +/// Some locals are part of the function's interface and can not be removed. +/// +/// Note that these locals *can* still be merged with non-required locals by removing that other +/// local. +fn is_local_required(local: Local, body: &Body<'_>) -> bool { + match body.local_kind(local) { + LocalKind::Arg | LocalKind::ReturnPointer => true, + LocalKind::Temp => false, + } +} + +///////////////////////////////////////////////////////// +// MIR Dump + +fn dest_prop_mir_dump<'body, 'tcx>( + tcx: TyCtxt<'tcx>, + body: &'body Body<'tcx>, + live: &mut ResultsCursor<'body, 'tcx, MaybeLiveLocals>, + round: usize, +) { + let mut reachable = None; + dump_mir(tcx, false, "DestinationPropagation-dataflow", &round, body, |pass_where, w| { + let reachable = reachable.get_or_insert_with(|| traversal::reachable_as_bitset(body)); + + match pass_where { + PassWhere::BeforeLocation(loc) if reachable.contains(loc.block) => { + live.seek_after_primary_effect(loc); + writeln!(w, " // live: {:?}", live.get())?; + } + PassWhere::AfterTerminator(bb) if reachable.contains(bb) => { + let loc = body.terminator_loc(bb); + live.seek_before_primary_effect(loc); + writeln!(w, " // live: {:?}", live.get())?; + } + + PassWhere::BeforeBlock(bb) if reachable.contains(bb) => { + live.seek_to_block_start(bb); + writeln!(w, " // live: {:?}", live.get())?; + } + + PassWhere::BeforeCFG | PassWhere::AfterCFG | PassWhere::AfterLocation(_) => {} + + PassWhere::BeforeLocation(_) | PassWhere::AfterTerminator(_) => { + writeln!(w, " // live: <unreachable>")?; + } + + PassWhere::BeforeBlock(_) => { + writeln!(w, " // live: <unreachable>")?; + } + } + + Ok(()) + }); +} diff --git a/compiler/rustc_mir_transform/src/dump_mir.rs b/compiler/rustc_mir_transform/src/dump_mir.rs new file mode 100644 index 00000000000..13841be494c --- /dev/null +++ b/compiler/rustc_mir_transform/src/dump_mir.rs @@ -0,0 +1,34 @@ +//! This pass just dumps MIR at a specified point. + +use std::fs::File; +use std::io; + +use crate::MirPass; +use rustc_middle::mir::write_mir_pretty; +use rustc_middle::mir::Body; +use rustc_middle::ty::TyCtxt; +use rustc_session::config::{OutFileName, OutputType}; + +pub struct Marker(pub &'static str); + +impl<'tcx> MirPass<'tcx> for Marker { + fn name(&self) -> &'static str { + self.0 + } + + fn run_pass(&self, _tcx: TyCtxt<'tcx>, _body: &mut Body<'tcx>) {} +} + +pub fn emit_mir(tcx: TyCtxt<'_>) -> io::Result<()> { + match tcx.output_filenames(()).path(OutputType::Mir) { + OutFileName::Stdout => { + let mut f = io::stdout(); + write_mir_pretty(tcx, None, &mut f)?; + } + OutFileName::Real(path) => { + let mut f = io::BufWriter::new(File::create(&path)?); + write_mir_pretty(tcx, None, &mut f)?; + } + } + Ok(()) +} diff --git a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs new file mode 100644 index 00000000000..6eb6cb069fe --- /dev/null +++ b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs @@ -0,0 +1,415 @@ +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, Ty, TyCtxt}; +use std::fmt::Debug; + +use super::simplify::simplify_cfg; + +/// This pass optimizes something like +/// ```ignore (syntax-highlighting-only) +/// let x: Option<()>; +/// let y: Option<()>; +/// match (x,y) { +/// (Some(_), Some(_)) => {0}, +/// _ => {1} +/// } +/// ``` +/// into something like +/// ```ignore (syntax-highlighting-only) +/// let x: Option<()>; +/// let y: Option<()>; +/// let discriminant_x = std::mem::discriminant(x); +/// let discriminant_y = std::mem::discriminant(y); +/// if discriminant_x == discriminant_y { +/// match x { +/// Some(_) => 0, +/// _ => 1, // <---- +/// } // | Actually the same bb +/// } else { // | +/// 1 // <-------------- +/// } +/// ``` +/// +/// Specifically, it looks for instances of control flow like this: +/// ```text +/// +/// ================= +/// | BB1 | +/// |---------------| ============================ +/// | ... | /------> | BBC | +/// |---------------| | |--------------------------| +/// | switchInt(Q) | | | _cl = discriminant(P) | +/// | c | --------/ |--------------------------| +/// | d | -------\ | switchInt(_cl) | +/// | ... | | | c | ---> BBC.2 +/// | otherwise | --\ | /--- | otherwise | +/// ================= | | | ============================ +/// | | | +/// ================= | | | +/// | BBU | <-| | | ============================ +/// |---------------| | \-------> | BBD | +/// |---------------| | | |--------------------------| +/// | unreachable | | | | _dl = discriminant(P) | +/// ================= | | |--------------------------| +/// | | | switchInt(_dl) | +/// ================= | | | d | ---> BBD.2 +/// | BB9 | <--------------- | otherwise | +/// |---------------| ============================ +/// | ... | +/// ================= +/// ``` +/// Where the `otherwise` branch on `BB1` is permitted to either go to `BBU` or to `BB9`. In the +/// code: +/// - `BB1` is `parent` and `BBC, BBD` are children +/// - `P` is `child_place` +/// - `child_ty` is the type of `_cl`. +/// - `Q` is `parent_op`. +/// - `parent_ty` is the type of `Q`. +/// - `BB9` is `destination` +/// All this is then transformed into: +/// ```text +/// +/// ======================= +/// | BB1 | +/// |---------------------| ============================ +/// | ... | /------> | BBEq | +/// | _s = discriminant(P)| | |--------------------------| +/// | _t = Ne(Q, _s) | | |--------------------------| +/// |---------------------| | | switchInt(Q) | +/// | switchInt(_t) | | | c | ---> BBC.2 +/// | false | --------/ | d | ---> BBD.2 +/// | otherwise | ---------------- | otherwise | +/// ======================= | ============================ +/// | +/// ================= | +/// | BB9 | <-----------/ +/// |---------------| +/// | ... | +/// ================= +/// ``` +/// +/// This is only correct for some `P`, since `P` is now computed outside the original `switchInt`. +/// The filter on which `P` are allowed (together with discussion of its correctness) is found in +/// `may_hoist`. +pub struct EarlyOtherwiseBranch; + +impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // unsound: https://github.com/rust-lang/rust/issues/95162 + sess.mir_opt_level() >= 3 && sess.opts.unstable_opts.unsound_mir_opts + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("running EarlyOtherwiseBranch on {:?}", body.source); + + let mut should_cleanup = false; + + // Also consider newly generated bbs in the same pass + for i in 0..body.basic_blocks.len() { + let bbs = &*body.basic_blocks; + let parent = BasicBlock::from_usize(i); + let Some(opt_data) = evaluate_candidate(tcx, body, parent) else { continue }; + + if !tcx.consider_optimizing(|| format!("EarlyOtherwiseBranch {:?}", &opt_data)) { + break; + } + + trace!("SUCCESS: found optimization possibility to apply: {:?}", &opt_data); + + should_cleanup = true; + + let TerminatorKind::SwitchInt { discr: parent_op, targets: parent_targets } = + &bbs[parent].terminator().kind + else { + unreachable!() + }; + // Always correct since we can only switch on `Copy` types + let parent_op = match parent_op { + Operand::Move(x) => Operand::Copy(*x), + Operand::Copy(x) => Operand::Copy(*x), + Operand::Constant(x) => Operand::Constant(x.clone()), + }; + let parent_ty = parent_op.ty(body.local_decls(), tcx); + let statements_before = bbs[parent].statements.len(); + let parent_end = Location { block: parent, statement_index: statements_before }; + + let mut patch = MirPatch::new(body); + + // create temp to store second discriminant in, `_s` in example above + let second_discriminant_temp = + patch.new_temp(opt_data.child_ty, opt_data.child_source.span); + + patch.add_statement(parent_end, StatementKind::StorageLive(second_discriminant_temp)); + + // create assignment of discriminant + patch.add_assign( + parent_end, + Place::from(second_discriminant_temp), + Rvalue::Discriminant(opt_data.child_place), + ); + + // create temp to store inequality comparison between the two discriminants, `_t` in + // example above + let nequal = BinOp::Ne; + let comp_res_type = nequal.ty(tcx, parent_ty, opt_data.child_ty); + let comp_temp = patch.new_temp(comp_res_type, opt_data.child_source.span); + patch.add_statement(parent_end, StatementKind::StorageLive(comp_temp)); + + // create inequality comparison between the two discriminants + let comp_rvalue = Rvalue::BinaryOp( + nequal, + Box::new((parent_op.clone(), Operand::Move(Place::from(second_discriminant_temp)))), + ); + patch.add_statement( + parent_end, + StatementKind::Assign(Box::new((Place::from(comp_temp), comp_rvalue))), + ); + + let eq_new_targets = parent_targets.iter().map(|(value, child)| { + let TerminatorKind::SwitchInt { targets, .. } = &bbs[child].terminator().kind + else { + unreachable!() + }; + (value, targets.target_for_value(value)) + }); + let eq_targets = SwitchTargets::new(eq_new_targets, opt_data.destination); + + // Create `bbEq` in example above + let eq_switch = BasicBlockData::new(Some(Terminator { + source_info: bbs[parent].terminator().source_info, + kind: TerminatorKind::SwitchInt { + // switch on the first discriminant, so we can mark the second one as dead + discr: parent_op, + targets: eq_targets, + }, + })); + + let eq_bb = patch.new_block(eq_switch); + + // Jump to it on the basis of the inequality comparison + let true_case = opt_data.destination; + let false_case = eq_bb; + patch.patch_terminator( + parent, + TerminatorKind::if_(Operand::Move(Place::from(comp_temp)), true_case, false_case), + ); + + // generate StorageDead for the second_discriminant_temp not in use anymore + patch.add_statement(parent_end, StatementKind::StorageDead(second_discriminant_temp)); + + // Generate a StorageDead for comp_temp in each of the targets, since we moved it into + // the switch + for bb in [false_case, true_case].iter() { + patch.add_statement( + Location { block: *bb, statement_index: 0 }, + StatementKind::StorageDead(comp_temp), + ); + } + + patch.apply(body); + } + + // Since this optimization adds new basic blocks and invalidates others, + // clean up the cfg to make it nicer for other passes + if should_cleanup { + simplify_cfg(tcx, body); + } + } +} + +/// Returns true if computing the discriminant of `place` may be hoisted out of the branch +fn may_hoist<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, place: Place<'tcx>) -> bool { + // FIXME(JakobDegen): This is unsound. Someone could write code like this: + // ```rust + // let Q = val; + // if discriminant(P) == otherwise { + // let ptr = &mut Q as *mut _ as *mut u8; + // unsafe { *ptr = 10; } // Any invalid value for the type + // } + // + // match P { + // A => match Q { + // A => { + // // code + // } + // _ => { + // // don't use Q + // } + // } + // _ => { + // // don't use Q + // } + // }; + // ``` + // + // Hoisting the `discriminant(Q)` out of the `A` arm causes us to compute the discriminant of an + // invalid value, which is UB. + // + // In order to fix this, we would either need to show that the discriminant computation of + // `place` is computed in all branches, including the `otherwise` branch, or we would need + // another analysis pass to determine that the place is fully initialized. It might even be best + // to have the hoisting be performed in a different pass and just do the CFG changing in this + // pass. + for (place, proj) in place.iter_projections() { + match proj { + // Dereferencing in the computation of `place` might cause issues from one of two + // categories. First, the referent might be invalid. We protect against this by + // dereferencing references only (not pointers). Second, the use of a reference may + // invalidate other references that are used later (for aliasing reasons). Consider + // where such an invalidated reference may appear: + // - In `Q`: Not possible since `Q` is used as the operand of a `SwitchInt` and so + // cannot contain referenced data. + // - In `BBU`: Not possible since that block contains only the `unreachable` terminator + // - In `BBC.2, BBD.2`: Not possible, since `discriminant(P)` was computed prior to + // reaching that block in the input to our transformation, and so any data + // invalidated by that computation could not have been used there. + // - In `BB9`: Not possible since control flow might have reached `BB9` via the + // `otherwise` branch in `BBC, BBD` in the input to our transformation, which would + // have invalidated the data when computing `discriminant(P)` + // So dereferencing here is correct. + ProjectionElem::Deref => match place.ty(body.local_decls(), tcx).ty.kind() { + ty::Ref(..) => {} + _ => return false, + }, + // Field projections are always valid + ProjectionElem::Field(..) => {} + // We cannot allow + // downcasts either, since the correctness of the downcast may depend on the parent + // branch being taken. An easy example of this is + // ``` + // Q = discriminant(_3) + // P = (_3 as Variant) + // ``` + // However, checking if the child and parent place are the same and only erroring then + // is not sufficient either, since the `discriminant(_3) == 1` (or whatever) check may + // be replaced by another optimization pass with any other condition that can be proven + // equivalent. + ProjectionElem::Downcast(..) => { + return false; + } + // We cannot allow indexing since the index may be out of bounds. + _ => { + return false; + } + } + } + true +} + +#[derive(Debug)] +struct OptimizationData<'tcx> { + destination: BasicBlock, + child_place: Place<'tcx>, + child_ty: Ty<'tcx>, + child_source: SourceInfo, +} + +fn evaluate_candidate<'tcx>( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + parent: BasicBlock, +) -> Option<OptimizationData<'tcx>> { + let bbs = &body.basic_blocks; + let TerminatorKind::SwitchInt { targets, discr: parent_discr } = &bbs[parent].terminator().kind + else { + return None; + }; + let parent_ty = parent_discr.ty(body.local_decls(), tcx); + let parent_dest = { + let poss = targets.otherwise(); + // If the fallthrough on the parent is trivially unreachable, we can let the + // children choose the destination + if bbs[poss].statements.len() == 0 + && bbs[poss].terminator().kind == TerminatorKind::Unreachable + { + None + } else { + Some(poss) + } + }; + let (_, child) = targets.iter().next()?; + let child_terminator = &bbs[child].terminator(); + let TerminatorKind::SwitchInt { targets: child_targets, discr: child_discr } = + &child_terminator.kind + else { + return None; + }; + let child_ty = child_discr.ty(body.local_decls(), tcx); + if child_ty != parent_ty { + return None; + } + let Some(StatementKind::Assign(boxed)) = &bbs[child].statements.first().map(|x| &x.kind) else { + return None; + }; + let (_, Rvalue::Discriminant(child_place)) = &**boxed else { + return None; + }; + let destination = parent_dest.unwrap_or(child_targets.otherwise()); + + // Verify that the optimization is legal in general + // We can hoist evaluating the child discriminant out of the branch + if !may_hoist(tcx, body, *child_place) { + return None; + } + + // Verify that the optimization is legal for each branch + for (value, child) in targets.iter() { + if !verify_candidate_branch(&bbs[child], value, *child_place, destination) { + return None; + } + } + Some(OptimizationData { + destination, + child_place: *child_place, + child_ty, + child_source: child_terminator.source_info, + }) +} + +fn verify_candidate_branch<'tcx>( + branch: &BasicBlockData<'tcx>, + value: u128, + place: Place<'tcx>, + destination: BasicBlock, +) -> bool { + // In order for the optimization to be correct, the branch must... + // ...have exactly one statement + if branch.statements.len() != 1 { + return false; + } + // ...assign the discriminant of `place` in that statement + let StatementKind::Assign(boxed) = &branch.statements[0].kind else { return false }; + let (discr_place, Rvalue::Discriminant(from_place)) = &**boxed else { return false }; + if *from_place != place { + return false; + } + // ...make that assignment to a local + if discr_place.projection.len() != 0 { + return false; + } + // ...terminate on a `SwitchInt` that invalidates that local + let TerminatorKind::SwitchInt { discr: switch_op, targets, .. } = &branch.terminator().kind + else { + return false; + }; + if *switch_op != Operand::Move(*discr_place) { + return false; + } + // ...fall through to `destination` if the switch misses + if destination != targets.otherwise() { + return false; + } + // ...have a branch for value `value` + let mut iter = targets.iter(); + let Some((target_value, _)) = iter.next() else { + return false; + }; + if target_value != value { + return false; + } + // ...and have no more branches + if let Some(_) = iter.next() { + return false; + } + return true; +} diff --git a/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs b/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs new file mode 100644 index 00000000000..96943435bab --- /dev/null +++ b/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs @@ -0,0 +1,151 @@ +//! This pass transforms derefs of Box into a deref of the pointer inside Box. +//! +//! Box is not actually a pointer so it is incorrect to dereference it directly. + +use rustc_hir::def_id::DefId; +use rustc_index::Idx; +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::visit::MutVisitor; +use rustc_middle::mir::*; +use rustc_middle::ty::{Ty, TyCtxt}; +use rustc_target::abi::FieldIdx; + +/// Constructs the types used when accessing a Box's pointer +pub fn build_ptr_tys<'tcx>( + tcx: TyCtxt<'tcx>, + pointee: Ty<'tcx>, + unique_did: DefId, + nonnull_did: DefId, +) -> (Ty<'tcx>, Ty<'tcx>, Ty<'tcx>) { + let args = tcx.mk_args(&[pointee.into()]); + let unique_ty = tcx.type_of(unique_did).instantiate(tcx, args); + let nonnull_ty = tcx.type_of(nonnull_did).instantiate(tcx, args); + let ptr_ty = Ty::new_imm_ptr(tcx, pointee); + + (unique_ty, nonnull_ty, ptr_ty) +} + +/// Constructs the projection needed to access a Box's pointer +pub fn build_projection<'tcx>( + unique_ty: Ty<'tcx>, + nonnull_ty: Ty<'tcx>, + ptr_ty: Ty<'tcx>, +) -> [PlaceElem<'tcx>; 3] { + [ + PlaceElem::Field(FieldIdx::new(0), unique_ty), + PlaceElem::Field(FieldIdx::new(0), nonnull_ty), + PlaceElem::Field(FieldIdx::new(0), ptr_ty), + ] +} + +struct ElaborateBoxDerefVisitor<'tcx, 'a> { + tcx: TyCtxt<'tcx>, + unique_did: DefId, + nonnull_did: DefId, + local_decls: &'a mut LocalDecls<'tcx>, + patch: MirPatch<'tcx>, +} + +impl<'tcx, 'a> MutVisitor<'tcx> for ElaborateBoxDerefVisitor<'tcx, 'a> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_place( + &mut self, + place: &mut Place<'tcx>, + context: visit::PlaceContext, + location: Location, + ) { + let tcx = self.tcx; + + let base_ty = self.local_decls[place.local].ty; + + // Derefer ensures that derefs are always the first projection + if place.projection.first() == Some(&PlaceElem::Deref) && base_ty.is_box() { + let source_info = self.local_decls[place.local].source_info; + + let (unique_ty, nonnull_ty, ptr_ty) = + build_ptr_tys(tcx, base_ty.boxed_ty(), self.unique_did, self.nonnull_did); + + let ptr_local = self.patch.new_temp(ptr_ty, source_info.span); + + self.patch.add_assign( + location, + Place::from(ptr_local), + Rvalue::Use(Operand::Copy( + Place::from(place.local) + .project_deeper(&build_projection(unique_ty, nonnull_ty, ptr_ty), tcx), + )), + ); + + place.local = ptr_local; + } + + self.super_place(place, context, location); + } +} + +pub struct ElaborateBoxDerefs; + +impl<'tcx> MirPass<'tcx> for ElaborateBoxDerefs { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + if let Some(def_id) = tcx.lang_items().owned_box() { + let unique_did = + tcx.adt_def(def_id).non_enum_variant().fields[FieldIdx::from_u32(0)].did; + + let Some(nonnull_def) = tcx.type_of(unique_did).instantiate_identity().ty_adt_def() + else { + span_bug!(tcx.def_span(unique_did), "expected Box to contain Unique") + }; + + let nonnull_did = nonnull_def.non_enum_variant().fields[FieldIdx::from_u32(0)].did; + + let patch = MirPatch::new(body); + + let local_decls = &mut body.local_decls; + + let mut visitor = + ElaborateBoxDerefVisitor { tcx, unique_did, nonnull_did, local_decls, patch }; + + for (block, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() { + visitor.visit_basic_block_data(block, data); + } + + visitor.patch.apply(body); + + for debug_info in body.var_debug_info.iter_mut() { + if let VarDebugInfoContents::Place(place) = &mut debug_info.value { + let mut new_projections: Option<Vec<_>> = None; + let mut last_deref = 0; + + for (i, (base, elem)) in place.iter_projections().enumerate() { + let base_ty = base.ty(&body.local_decls, tcx).ty; + + if elem == PlaceElem::Deref && base_ty.is_box() { + let new_projections = new_projections.get_or_insert_default(); + + let (unique_ty, nonnull_ty, ptr_ty) = + build_ptr_tys(tcx, base_ty.boxed_ty(), unique_did, nonnull_did); + + new_projections.extend_from_slice(&base.projection[last_deref..]); + new_projections.extend_from_slice(&build_projection( + unique_ty, nonnull_ty, ptr_ty, + )); + new_projections.push(PlaceElem::Deref); + + last_deref = i; + } + } + + if let Some(mut new_projections) = new_projections { + new_projections.extend_from_slice(&place.projection[last_deref..]); + place.projection = tcx.mk_place_elems(&new_projections); + } + } + } + } else { + // box is not present, this pass doesn't need to do anything + } + } +} diff --git a/compiler/rustc_mir_transform/src/elaborate_drops.rs b/compiler/rustc_mir_transform/src/elaborate_drops.rs new file mode 100644 index 00000000000..c45badbc559 --- /dev/null +++ b/compiler/rustc_mir_transform/src/elaborate_drops.rs @@ -0,0 +1,506 @@ +use crate::deref_separator::deref_finder; +use rustc_index::bit_set::BitSet; +use rustc_index::IndexVec; +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_mir_dataflow::elaborate_drops::{elaborate_drop, DropFlagState, Unwind}; +use rustc_mir_dataflow::elaborate_drops::{DropElaborator, DropFlagMode, DropStyle}; +use rustc_mir_dataflow::impls::{MaybeInitializedPlaces, MaybeUninitializedPlaces}; +use rustc_mir_dataflow::move_paths::{LookupResult, MoveData, MovePathIndex}; +use rustc_mir_dataflow::on_all_children_bits; +use rustc_mir_dataflow::on_lookup_result_bits; +use rustc_mir_dataflow::MoveDataParamEnv; +use rustc_mir_dataflow::{Analysis, ResultsCursor}; +use rustc_span::Span; +use rustc_target::abi::{FieldIdx, VariantIdx}; +use std::fmt; + +/// 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 will run, +/// as the target of the drop may be uninitialized. +/// In general, the compiler cannot determine at compile time whether a destructor will run or not. +/// +/// At a high level, this pass refines Drop to only run the destructor if the +/// target is initialized. The way this is achieved is by inserting drop flags for every variable +/// that may be dropped, and then using those flags to determine whether a destructor should run. +/// Once this is complete, Drop terminators in the MIR correspond to a call to the "drop glue" or +/// "drop shim" for the type of the dropped place. +/// +/// This pass relies on dropped places having an associated move path, which is then used to determine +/// the initialization status of the place and its descendants. +/// It's worth noting that a MIR containing a Drop without an associated move path is probably ill formed, +/// as it would allow running a destructor on a place behind a reference: +/// +/// ```text +// fn drop_term<T>(t: &mut T) { +// mir!( +// { +// Drop(*t, exit) +// } +// exit = { +// Return() +// } +// ) +// } +/// ``` +pub struct ElaborateDrops; + +impl<'tcx> MirPass<'tcx> for ElaborateDrops { + #[instrument(level = "trace", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!("elaborate_drops({:?} @ {:?})", body.source, body.span); + + let def_id = body.source.def_id(); + let param_env = tcx.param_env_reveal_all_normalized(def_id); + // For types that do not need dropping, the behaviour is trivial. So we only need to track + // init/uninit for types that do need dropping. + let move_data = + MoveData::gather_moves(body, tcx, param_env, |ty| ty.needs_drop(tcx, param_env)); + let elaborate_patch = { + let env = MoveDataParamEnv { move_data, param_env }; + + let mut inits = MaybeInitializedPlaces::new(tcx, body, &env) + .skipping_unreachable_unwind() + .into_engine(tcx, body) + .pass_name("elaborate_drops") + .iterate_to_fixpoint() + .into_results_cursor(body); + let dead_unwinds = compute_dead_unwinds(body, &mut inits); + + let uninits = MaybeUninitializedPlaces::new(tcx, body, &env) + .mark_inactive_variants_as_uninit() + .skipping_unreachable_unwind(dead_unwinds) + .into_engine(tcx, body) + .pass_name("elaborate_drops") + .iterate_to_fixpoint() + .into_results_cursor(body); + + let drop_flags = IndexVec::from_elem(None, &env.move_data.move_paths); + ElaborateDropsCtxt { + tcx, + body, + env: &env, + init_data: InitializationData { inits, uninits }, + drop_flags, + patch: MirPatch::new(body), + } + .elaborate() + }; + elaborate_patch.apply(body); + deref_finder(tcx, body); + } +} + +/// Records unwind edges which are known to be unreachable, because they are in `drop` terminators +/// that can't drop anything. +#[instrument(level = "trace", skip(body, flow_inits), ret)] +fn compute_dead_unwinds<'mir, 'tcx>( + body: &'mir Body<'tcx>, + flow_inits: &mut ResultsCursor<'mir, 'tcx, MaybeInitializedPlaces<'mir, 'tcx>>, +) -> BitSet<BasicBlock> { + // We only need to do this pass once, because unwind edges can only + // reach cleanup blocks, which can't have unwind edges themselves. + let mut dead_unwinds = BitSet::new_empty(body.basic_blocks.len()); + for (bb, bb_data) in body.basic_blocks.iter_enumerated() { + let TerminatorKind::Drop { place, unwind: UnwindAction::Cleanup(_), .. } = + bb_data.terminator().kind + else { + continue; + }; + + flow_inits.seek_before_primary_effect(body.terminator_loc(bb)); + if flow_inits.analysis().is_unwind_dead(place, flow_inits.get()) { + dead_unwinds.insert(bb); + } + } + + dead_unwinds +} + +struct InitializationData<'mir, 'tcx> { + inits: ResultsCursor<'mir, 'tcx, MaybeInitializedPlaces<'mir, 'tcx>>, + uninits: ResultsCursor<'mir, 'tcx, MaybeUninitializedPlaces<'mir, 'tcx>>, +} + +impl InitializationData<'_, '_> { + fn seek_before(&mut self, loc: Location) { + self.inits.seek_before_primary_effect(loc); + self.uninits.seek_before_primary_effect(loc); + } + + fn maybe_live_dead(&self, path: MovePathIndex) -> (bool, bool) { + (self.inits.contains(path), self.uninits.contains(path)) + } +} + +struct Elaborator<'a, 'b, 'tcx> { + ctxt: &'a mut ElaborateDropsCtxt<'b, 'tcx>, +} + +impl fmt::Debug for Elaborator<'_, '_, '_> { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result { + Ok(()) + } +} + +impl<'a, 'tcx> DropElaborator<'a, 'tcx> for Elaborator<'a, '_, 'tcx> { + type Path = MovePathIndex; + + fn patch(&mut self) -> &mut MirPatch<'tcx> { + &mut self.ctxt.patch + } + + fn body(&self) -> &'a Body<'tcx> { + self.ctxt.body + } + + fn tcx(&self) -> TyCtxt<'tcx> { + self.ctxt.tcx + } + + fn param_env(&self) -> ty::ParamEnv<'tcx> { + self.ctxt.param_env() + } + + #[instrument(level = "debug", skip(self), ret)] + fn drop_style(&self, path: Self::Path, mode: DropFlagMode) -> DropStyle { + let ((maybe_live, maybe_dead), multipart) = match mode { + DropFlagMode::Shallow => (self.ctxt.init_data.maybe_live_dead(path), false), + DropFlagMode::Deep => { + let mut some_live = false; + let mut some_dead = false; + let mut children_count = 0; + on_all_children_bits(self.ctxt.move_data(), path, |child| { + let (live, dead) = self.ctxt.init_data.maybe_live_dead(child); + debug!("elaborate_drop: state({:?}) = {:?}", child, (live, dead)); + some_live |= live; + some_dead |= dead; + children_count += 1; + }); + ((some_live, some_dead), children_count != 1) + } + }; + match (maybe_live, maybe_dead, multipart) { + (false, _, _) => DropStyle::Dead, + (true, false, _) => DropStyle::Static, + (true, true, false) => DropStyle::Conditional, + (true, true, true) => DropStyle::Open, + } + } + + fn clear_drop_flag(&mut self, loc: Location, path: Self::Path, mode: DropFlagMode) { + match mode { + DropFlagMode::Shallow => { + self.ctxt.set_drop_flag(loc, path, DropFlagState::Absent); + } + DropFlagMode::Deep => { + on_all_children_bits(self.ctxt.move_data(), path, |child| { + self.ctxt.set_drop_flag(loc, child, DropFlagState::Absent) + }); + } + } + } + + fn field_subpath(&self, path: Self::Path, field: FieldIdx) -> Option<Self::Path> { + rustc_mir_dataflow::move_path_children_matching(self.ctxt.move_data(), path, |e| match e { + ProjectionElem::Field(idx, _) => idx == field, + _ => false, + }) + } + + fn array_subpath(&self, path: Self::Path, index: u64, size: u64) -> Option<Self::Path> { + rustc_mir_dataflow::move_path_children_matching(self.ctxt.move_data(), path, |e| match e { + ProjectionElem::ConstantIndex { offset, min_length, from_end } => { + debug_assert!(size == min_length, "min_length should be exact for arrays"); + assert!(!from_end, "from_end should not be used for array element ConstantIndex"); + offset == index + } + _ => false, + }) + } + + fn deref_subpath(&self, path: Self::Path) -> Option<Self::Path> { + rustc_mir_dataflow::move_path_children_matching(self.ctxt.move_data(), path, |e| { + e == ProjectionElem::Deref + }) + } + + fn downcast_subpath(&self, path: Self::Path, variant: VariantIdx) -> Option<Self::Path> { + rustc_mir_dataflow::move_path_children_matching(self.ctxt.move_data(), path, |e| match e { + ProjectionElem::Downcast(_, idx) => idx == variant, + _ => false, + }) + } + + fn get_drop_flag(&mut self, path: Self::Path) -> Option<Operand<'tcx>> { + self.ctxt.drop_flag(path).map(Operand::Copy) + } +} + +struct ElaborateDropsCtxt<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + body: &'a Body<'tcx>, + env: &'a MoveDataParamEnv<'tcx>, + init_data: InitializationData<'a, 'tcx>, + drop_flags: IndexVec<MovePathIndex, Option<Local>>, + patch: MirPatch<'tcx>, +} + +impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> { + fn move_data(&self) -> &'b MoveData<'tcx> { + &self.env.move_data + } + + fn param_env(&self) -> ty::ParamEnv<'tcx> { + self.env.param_env + } + + fn create_drop_flag(&mut self, index: MovePathIndex, span: Span) { + let patch = &mut self.patch; + debug!("create_drop_flag({:?})", self.body.span); + self.drop_flags[index].get_or_insert_with(|| patch.new_temp(self.tcx.types.bool, span)); + } + + fn drop_flag(&mut self, index: MovePathIndex) -> Option<Place<'tcx>> { + self.drop_flags[index].map(Place::from) + } + + /// create a patch that elaborates all drops in the input + /// MIR. + fn elaborate(mut self) -> MirPatch<'tcx> { + self.collect_drop_flags(); + + self.elaborate_drops(); + + self.drop_flags_on_init(); + self.drop_flags_for_fn_rets(); + self.drop_flags_for_args(); + self.drop_flags_for_locs(); + + self.patch + } + + fn collect_drop_flags(&mut self) { + for (bb, data) in self.body.basic_blocks.iter_enumerated() { + let terminator = data.terminator(); + let TerminatorKind::Drop { ref place, .. } = terminator.kind else { continue }; + + let path = self.move_data().rev_lookup.find(place.as_ref()); + debug!("collect_drop_flags: {:?}, place {:?} ({:?})", bb, place, path); + + match path { + LookupResult::Exact(path) => { + self.init_data.seek_before(self.body.terminator_loc(bb)); + on_all_children_bits(self.move_data(), path, |child| { + let (maybe_live, maybe_dead) = self.init_data.maybe_live_dead(child); + debug!( + "collect_drop_flags: collecting {:?} from {:?}@{:?} - {:?}", + child, + place, + path, + (maybe_live, maybe_dead) + ); + if maybe_live && maybe_dead { + self.create_drop_flag(child, terminator.source_info.span) + } + }); + } + LookupResult::Parent(None) => {} + LookupResult::Parent(Some(parent)) => { + if self.body.local_decls[place.local].is_deref_temp() { + continue; + } + + self.init_data.seek_before(self.body.terminator_loc(bb)); + let (_maybe_live, maybe_dead) = self.init_data.maybe_live_dead(parent); + if maybe_dead { + self.tcx.sess.span_delayed_bug( + terminator.source_info.span, + format!( + "drop of untracked, uninitialized value {bb:?}, place {place:?} ({path:?})" + ), + ); + } + } + }; + } + } + + fn elaborate_drops(&mut self) { + // This function should mirror what `collect_drop_flags` does. + for (bb, data) in self.body.basic_blocks.iter_enumerated() { + let terminator = data.terminator(); + let TerminatorKind::Drop { place, target, unwind, replace } = terminator.kind else { + continue; + }; + + // This place does not need dropping. It does not have an associated move-path, so the + // match below will conservatively keep an unconditional drop. As that drop is useless, + // just remove it here and now. + if !place + .ty(&self.body.local_decls, self.tcx) + .ty + .needs_drop(self.tcx, self.env.param_env) + { + self.patch.patch_terminator(bb, TerminatorKind::Goto { target }); + continue; + } + + let path = self.move_data().rev_lookup.find(place.as_ref()); + match path { + LookupResult::Exact(path) => { + let unwind = match unwind { + _ if data.is_cleanup => Unwind::InCleanup, + UnwindAction::Cleanup(cleanup) => Unwind::To(cleanup), + UnwindAction::Continue => Unwind::To(self.patch.resume_block()), + UnwindAction::Unreachable => { + Unwind::To(self.patch.unreachable_cleanup_block()) + } + UnwindAction::Terminate(reason) => { + debug_assert_ne!( + reason, + UnwindTerminateReason::InCleanup, + "we are not in a cleanup block, InCleanup reason should be impossible" + ); + Unwind::To(self.patch.terminate_block(reason)) + } + }; + self.init_data.seek_before(self.body.terminator_loc(bb)); + elaborate_drop( + &mut Elaborator { ctxt: self }, + terminator.source_info, + place, + path, + target, + unwind, + bb, + ) + } + LookupResult::Parent(None) => {} + LookupResult::Parent(Some(_)) => { + if !replace { + self.tcx.sess.span_delayed_bug( + terminator.source_info.span, + format!("drop of untracked value {bb:?}"), + ); + } + // A drop and replace behind a pointer/array/whatever. + // The borrow checker requires that these locations are initialized before the assignment, + // so we just leave an unconditional drop. + assert!(!data.is_cleanup); + } + } + } + } + + fn constant_bool(&self, span: Span, val: bool) -> Rvalue<'tcx> { + Rvalue::Use(Operand::Constant(Box::new(ConstOperand { + span, + user_ty: None, + const_: Const::from_bool(self.tcx, val), + }))) + } + + fn set_drop_flag(&mut self, loc: Location, path: MovePathIndex, val: DropFlagState) { + if let Some(flag) = self.drop_flags[path] { + let span = self.patch.source_info_for_location(self.body, loc).span; + let val = self.constant_bool(span, val.value()); + self.patch.add_assign(loc, Place::from(flag), val); + } + } + + fn drop_flags_on_init(&mut self) { + let loc = Location::START; + let span = self.patch.source_info_for_location(self.body, loc).span; + let false_ = self.constant_bool(span, false); + for flag in self.drop_flags.iter().flatten() { + self.patch.add_assign(loc, Place::from(*flag), false_.clone()); + } + } + + fn drop_flags_for_fn_rets(&mut self) { + for (bb, data) in self.body.basic_blocks.iter_enumerated() { + if let TerminatorKind::Call { + destination, + target: Some(tgt), + unwind: UnwindAction::Cleanup(_), + .. + } = data.terminator().kind + { + assert!(!self.patch.is_patched(bb)); + + let loc = Location { block: tgt, statement_index: 0 }; + let path = self.move_data().rev_lookup.find(destination.as_ref()); + on_lookup_result_bits(self.move_data(), path, |child| { + self.set_drop_flag(loc, child, DropFlagState::Present) + }); + } + } + } + + fn drop_flags_for_args(&mut self) { + let loc = Location::START; + rustc_mir_dataflow::drop_flag_effects_for_function_entry(self.body, self.env, |path, ds| { + self.set_drop_flag(loc, path, ds); + }) + } + + fn drop_flags_for_locs(&mut self) { + // We intentionally iterate only over the *old* basic blocks. + // + // Basic blocks created by drop elaboration update their + // drop flags by themselves, to avoid the drop flags being + // clobbered before they are read. + + for (bb, data) in self.body.basic_blocks.iter_enumerated() { + debug!("drop_flags_for_locs({:?})", data); + for i in 0..(data.statements.len() + 1) { + debug!("drop_flag_for_locs: stmt {}", i); + if i == data.statements.len() { + match data.terminator().kind { + TerminatorKind::Drop { .. } => { + // drop elaboration should handle that by itself + continue; + } + TerminatorKind::UnwindResume => { + // It is possible for `Resume` to be patched + // (in particular it can be patched to be replaced with + // a Goto; see `MirPatch::new`). + } + _ => { + assert!(!self.patch.is_patched(bb)); + } + } + } + let loc = Location { block: bb, statement_index: i }; + rustc_mir_dataflow::drop_flag_effects_for_location( + self.body, + self.env, + loc, + |path, ds| self.set_drop_flag(loc, path, ds), + ) + } + + // There may be a critical edge after this call, + // so mark the return as initialized *before* the + // call. + if let TerminatorKind::Call { + destination, + target: Some(_), + unwind: + UnwindAction::Continue | UnwindAction::Unreachable | UnwindAction::Terminate(_), + .. + } = data.terminator().kind + { + assert!(!self.patch.is_patched(bb)); + + let loc = Location { block: bb, statement_index: data.statements.len() }; + let path = self.move_data().rev_lookup.find(destination.as_ref()); + on_lookup_result_bits(self.move_data(), path, |child| { + self.set_drop_flag(loc, child, DropFlagState::Present) + }); + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/errors.rs b/compiler/rustc_mir_transform/src/errors.rs new file mode 100644 index 00000000000..2358661738a --- /dev/null +++ b/compiler/rustc_mir_transform/src/errors.rs @@ -0,0 +1,313 @@ +use std::borrow::Cow; + +use rustc_errors::{ + Applicability, DecorateLint, DiagnosticArgValue, DiagnosticBuilder, DiagnosticMessage, + EmissionGuarantee, ErrorGuaranteed, Handler, IntoDiagnostic, +}; +use rustc_macros::{Diagnostic, LintDiagnostic, Subdiagnostic}; +use rustc_middle::mir::{AssertKind, UnsafetyViolationDetails}; +use rustc_middle::ty::TyCtxt; +use rustc_session::lint::{self, Lint}; +use rustc_span::def_id::DefId; +use rustc_span::Span; + +use crate::fluent_generated as fluent; + +#[derive(LintDiagnostic)] +pub(crate) enum ConstMutate { + #[diag(mir_transform_const_modify)] + #[note] + Modify { + #[note(mir_transform_const_defined_here)] + konst: Span, + }, + #[diag(mir_transform_const_mut_borrow)] + #[note] + #[note(mir_transform_note2)] + MutBorrow { + #[note(mir_transform_note3)] + method_call: Option<Span>, + #[note(mir_transform_const_defined_here)] + konst: Span, + }, +} + +#[derive(Diagnostic)] +#[diag(mir_transform_unaligned_packed_ref, code = "E0793")] +#[note] +#[note(mir_transform_note_ub)] +#[help] +pub(crate) struct UnalignedPackedRef { + #[primary_span] + pub span: Span, +} + +#[derive(LintDiagnostic)] +#[diag(mir_transform_unused_unsafe)] +pub(crate) struct UnusedUnsafe { + #[label(mir_transform_unused_unsafe)] + pub span: Span, + #[label] + pub nested_parent: Option<Span>, +} + +pub(crate) struct RequiresUnsafe { + pub span: Span, + pub details: RequiresUnsafeDetail, + pub enclosing: Option<Span>, + pub op_in_unsafe_fn_allowed: bool, +} + +// The primary message for this diagnostic should be '{$label} is unsafe and...', +// so we need to eagerly translate the label here, which isn't supported by the derive API +// We could also exhaustively list out the primary messages for all unsafe violations, +// but this would result in a lot of duplication. +impl<'sess> IntoDiagnostic<'sess> for RequiresUnsafe { + #[track_caller] + fn into_diagnostic(self, handler: &'sess Handler) -> DiagnosticBuilder<'sess, ErrorGuaranteed> { + let mut diag = handler.struct_diagnostic(fluent::mir_transform_requires_unsafe); + diag.code(rustc_errors::DiagnosticId::Error("E0133".to_string())); + diag.set_span(self.span); + diag.span_label(self.span, self.details.label()); + let desc = handler.eagerly_translate_to_string(self.details.label(), [].into_iter()); + diag.set_arg("details", desc); + diag.set_arg("op_in_unsafe_fn_allowed", self.op_in_unsafe_fn_allowed); + self.details.add_subdiagnostics(&mut diag); + if let Some(sp) = self.enclosing { + diag.span_label(sp, fluent::mir_transform_not_inherited); + } + diag + } +} + +#[derive(Clone)] +pub(crate) struct RequiresUnsafeDetail { + pub span: Span, + pub violation: UnsafetyViolationDetails, +} + +impl RequiresUnsafeDetail { + fn add_subdiagnostics<G: EmissionGuarantee>(&self, diag: &mut DiagnosticBuilder<'_, G>) { + use UnsafetyViolationDetails::*; + match self.violation { + CallToUnsafeFunction => { + diag.note(fluent::mir_transform_call_to_unsafe_note); + } + UseOfInlineAssembly => { + diag.note(fluent::mir_transform_use_of_asm_note); + } + InitializingTypeWith => { + diag.note(fluent::mir_transform_initializing_valid_range_note); + } + CastOfPointerToInt => { + diag.note(fluent::mir_transform_const_ptr2int_note); + } + UseOfMutableStatic => { + diag.note(fluent::mir_transform_use_of_static_mut_note); + } + UseOfExternStatic => { + diag.note(fluent::mir_transform_use_of_extern_static_note); + } + DerefOfRawPointer => { + diag.note(fluent::mir_transform_deref_ptr_note); + } + AccessToUnionField => { + diag.note(fluent::mir_transform_union_access_note); + } + MutationOfLayoutConstrainedField => { + diag.note(fluent::mir_transform_mutation_layout_constrained_note); + } + BorrowOfLayoutConstrainedField => { + diag.note(fluent::mir_transform_mutation_layout_constrained_borrow_note); + } + CallToFunctionWith { ref missing, ref build_enabled } => { + diag.help(fluent::mir_transform_target_feature_call_help); + diag.set_arg( + "missing_target_features", + DiagnosticArgValue::StrListSepByAnd( + missing.iter().map(|feature| Cow::from(feature.as_str())).collect(), + ), + ); + diag.set_arg("missing_target_features_count", missing.len()); + if !build_enabled.is_empty() { + diag.note(fluent::mir_transform_target_feature_call_note); + diag.set_arg( + "build_target_features", + DiagnosticArgValue::StrListSepByAnd( + build_enabled + .iter() + .map(|feature| Cow::from(feature.as_str())) + .collect(), + ), + ); + diag.set_arg("build_target_features_count", build_enabled.len()); + } + } + } + } + + fn label(&self) -> DiagnosticMessage { + use UnsafetyViolationDetails::*; + match self.violation { + CallToUnsafeFunction => fluent::mir_transform_call_to_unsafe_label, + UseOfInlineAssembly => fluent::mir_transform_use_of_asm_label, + InitializingTypeWith => fluent::mir_transform_initializing_valid_range_label, + CastOfPointerToInt => fluent::mir_transform_const_ptr2int_label, + UseOfMutableStatic => fluent::mir_transform_use_of_static_mut_label, + UseOfExternStatic => fluent::mir_transform_use_of_extern_static_label, + DerefOfRawPointer => fluent::mir_transform_deref_ptr_label, + AccessToUnionField => fluent::mir_transform_union_access_label, + MutationOfLayoutConstrainedField => { + fluent::mir_transform_mutation_layout_constrained_label + } + BorrowOfLayoutConstrainedField => { + fluent::mir_transform_mutation_layout_constrained_borrow_label + } + CallToFunctionWith { .. } => fluent::mir_transform_target_feature_call_label, + } + } +} + +pub(crate) struct UnsafeOpInUnsafeFn { + pub details: RequiresUnsafeDetail, + + /// These spans point to: + /// 1. the start of the function body + /// 2. the end of the function body + /// 3. the function signature + pub suggest_unsafe_block: Option<(Span, Span, Span)>, +} + +impl<'a> DecorateLint<'a, ()> for UnsafeOpInUnsafeFn { + #[track_caller] + fn decorate_lint<'b>( + self, + diag: &'b mut DiagnosticBuilder<'a, ()>, + ) -> &'b mut DiagnosticBuilder<'a, ()> { + let handler = diag.handler().expect("lint should not yet be emitted"); + let desc = handler.eagerly_translate_to_string(self.details.label(), [].into_iter()); + diag.set_arg("details", desc); + diag.span_label(self.details.span, self.details.label()); + self.details.add_subdiagnostics(diag); + + if let Some((start, end, fn_sig)) = self.suggest_unsafe_block { + diag.span_note(fn_sig, fluent::mir_transform_note); + diag.tool_only_multipart_suggestion( + fluent::mir_transform_suggestion, + vec![(start, " unsafe {".into()), (end, "}".into())], + Applicability::MaybeIncorrect, + ); + } + + diag + } + + fn msg(&self) -> DiagnosticMessage { + fluent::mir_transform_unsafe_op_in_unsafe_fn + } +} + +pub(crate) enum AssertLint<P> { + ArithmeticOverflow(Span, AssertKind<P>), + UnconditionalPanic(Span, AssertKind<P>), +} + +impl<'a, P: std::fmt::Debug> DecorateLint<'a, ()> for AssertLint<P> { + fn decorate_lint<'b>( + self, + diag: &'b mut DiagnosticBuilder<'a, ()>, + ) -> &'b mut DiagnosticBuilder<'a, ()> { + let span = self.span(); + let assert_kind = self.panic(); + let message = assert_kind.diagnostic_message(); + assert_kind.add_args(&mut |name, value| { + diag.set_arg(name, value); + }); + diag.span_label(span, message); + + diag + } + + fn msg(&self) -> DiagnosticMessage { + match self { + AssertLint::ArithmeticOverflow(..) => fluent::mir_transform_arithmetic_overflow, + AssertLint::UnconditionalPanic(..) => fluent::mir_transform_operation_will_panic, + } + } +} + +impl<P> AssertLint<P> { + pub fn lint(&self) -> &'static Lint { + match self { + AssertLint::ArithmeticOverflow(..) => lint::builtin::ARITHMETIC_OVERFLOW, + AssertLint::UnconditionalPanic(..) => lint::builtin::UNCONDITIONAL_PANIC, + } + } + pub fn span(&self) -> Span { + match self { + AssertLint::ArithmeticOverflow(sp, _) | AssertLint::UnconditionalPanic(sp, _) => *sp, + } + } + pub fn panic(self) -> AssertKind<P> { + match self { + AssertLint::ArithmeticOverflow(_, p) | AssertLint::UnconditionalPanic(_, p) => p, + } + } +} + +#[derive(LintDiagnostic)] +#[diag(mir_transform_ffi_unwind_call)] +pub(crate) struct FfiUnwindCall { + #[label(mir_transform_ffi_unwind_call)] + pub span: Span, + pub foreign: bool, +} + +#[derive(LintDiagnostic)] +#[diag(mir_transform_fn_item_ref)] +pub(crate) struct FnItemRef { + #[suggestion(code = "{sugg}", applicability = "unspecified")] + pub span: Span, + pub sugg: String, + pub ident: String, +} + +pub(crate) struct MustNotSupend<'tcx, 'a> { + pub tcx: TyCtxt<'tcx>, + pub yield_sp: Span, + pub reason: Option<MustNotSuspendReason>, + pub src_sp: Span, + pub pre: &'a str, + pub def_id: DefId, + pub post: &'a str, +} + +// Needed for def_path_str +impl<'a> DecorateLint<'a, ()> for MustNotSupend<'_, '_> { + fn decorate_lint<'b>( + self, + diag: &'b mut rustc_errors::DiagnosticBuilder<'a, ()>, + ) -> &'b mut rustc_errors::DiagnosticBuilder<'a, ()> { + diag.span_label(self.yield_sp, fluent::_subdiag::label); + if let Some(reason) = self.reason { + diag.subdiagnostic(reason); + } + diag.span_help(self.src_sp, fluent::_subdiag::help); + diag.set_arg("pre", self.pre); + diag.set_arg("def_path", self.tcx.def_path_str(self.def_id)); + diag.set_arg("post", self.post); + diag + } + + fn msg(&self) -> rustc_errors::DiagnosticMessage { + fluent::mir_transform_must_not_suspend + } +} + +#[derive(Subdiagnostic)] +#[note(mir_transform_note)] +pub(crate) struct MustNotSuspendReason { + #[primary_span] + pub span: Span, + pub reason: String, +} diff --git a/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs b/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs new file mode 100644 index 00000000000..26fcfad8287 --- /dev/null +++ b/compiler/rustc_mir_transform/src/ffi_unwind_calls.rs @@ -0,0 +1,171 @@ +use rustc_hir::def_id::{LocalDefId, LOCAL_CRATE}; +use rustc_middle::mir::*; +use rustc_middle::query::LocalCrate; +use rustc_middle::query::Providers; +use rustc_middle::ty::layout; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_session::lint::builtin::FFI_UNWIND_CALLS; +use rustc_target::spec::abi::Abi; +use rustc_target::spec::PanicStrategy; + +use crate::errors; + +fn abi_can_unwind(abi: Abi) -> bool { + use Abi::*; + match abi { + C { unwind } + | System { unwind } + | Cdecl { unwind } + | Stdcall { unwind } + | Fastcall { unwind } + | Vectorcall { unwind } + | Thiscall { unwind } + | Aapcs { unwind } + | Win64 { unwind } + | SysV64 { unwind } => unwind, + PtxKernel + | Msp430Interrupt + | X86Interrupt + | AmdGpuKernel + | EfiApi + | AvrInterrupt + | AvrNonBlockingInterrupt + | RiscvInterruptM + | RiscvInterruptS + | CCmseNonSecureCall + | Wasm + | RustIntrinsic + | PlatformIntrinsic + | Unadjusted => false, + Rust | RustCall | RustCold => true, + } +} + +// Check if the body of this def_id can possibly leak a foreign unwind into Rust code. +fn has_ffi_unwind_calls(tcx: TyCtxt<'_>, local_def_id: LocalDefId) -> bool { + debug!("has_ffi_unwind_calls({local_def_id:?})"); + + // Only perform check on functions because constants cannot call FFI functions. + let def_id = local_def_id.to_def_id(); + let kind = tcx.def_kind(def_id); + if !kind.is_fn_like() { + return false; + } + + let body = &*tcx.mir_built(local_def_id).borrow(); + + let body_ty = tcx.type_of(def_id).skip_binder(); + let body_abi = match body_ty.kind() { + ty::FnDef(..) => body_ty.fn_sig(tcx).abi(), + ty::Closure(..) => Abi::RustCall, + ty::Coroutine(..) => Abi::Rust, + _ => span_bug!(body.span, "unexpected body ty: {:?}", body_ty), + }; + let body_can_unwind = layout::fn_can_unwind(tcx, Some(def_id), body_abi); + + // Foreign unwinds cannot leak past functions that themselves cannot unwind. + if !body_can_unwind { + return false; + } + + let mut tainted = false; + + for block in body.basic_blocks.iter() { + if block.is_cleanup { + continue; + } + let Some(terminator) = &block.terminator else { continue }; + let TerminatorKind::Call { func, .. } = &terminator.kind else { continue }; + + let ty = func.ty(body, tcx); + let sig = ty.fn_sig(tcx); + + // Rust calls cannot themselves create foreign unwinds. + if let Abi::Rust | Abi::RustCall | Abi::RustCold = sig.abi() { + continue; + }; + + let fn_def_id = match ty.kind() { + ty::FnPtr(_) => None, + &ty::FnDef(def_id, _) => { + // Rust calls cannot themselves create foreign unwinds. + if !tcx.is_foreign_item(def_id) { + continue; + } + Some(def_id) + } + _ => bug!("invalid callee of type {:?}", ty), + }; + + if layout::fn_can_unwind(tcx, fn_def_id, sig.abi()) && abi_can_unwind(sig.abi()) { + // We have detected a call that can possibly leak foreign unwind. + // + // Because the function body itself can unwind, we are not aborting this function call + // upon unwind, so this call can possibly leak foreign unwind into Rust code if the + // panic runtime linked is panic-abort. + + let lint_root = body.source_scopes[terminator.source_info.scope] + .local_data + .as_ref() + .assert_crate_local() + .lint_root; + let span = terminator.source_info.span; + + let foreign = fn_def_id.is_some(); + tcx.emit_spanned_lint( + FFI_UNWIND_CALLS, + lint_root, + span, + errors::FfiUnwindCall { span, foreign }, + ); + + tainted = true; + } + } + + tainted +} + +fn required_panic_strategy(tcx: TyCtxt<'_>, _: LocalCrate) -> Option<PanicStrategy> { + if tcx.is_panic_runtime(LOCAL_CRATE) { + return Some(tcx.sess.panic_strategy()); + } + + if tcx.sess.panic_strategy() == PanicStrategy::Abort { + return Some(PanicStrategy::Abort); + } + + for def_id in tcx.hir().body_owners() { + if tcx.has_ffi_unwind_calls(def_id) { + // Given that this crate is compiled in `-C panic=unwind`, the `AbortUnwindingCalls` + // MIR pass will not be run on FFI-unwind call sites, therefore a foreign exception + // can enter Rust through these sites. + // + // On the other hand, crates compiled with `-C panic=abort` expects that all Rust + // functions cannot unwind (whether it's caused by Rust panic or foreign exception), + // and this expectation mismatch can cause unsoundness (#96926). + // + // To address this issue, we enforce that if FFI-unwind calls are used in a crate + // compiled with `panic=unwind`, then the final panic strategy must be `panic=unwind`. + // This will ensure that no crates will have wrong unwindability assumption. + // + // It should be noted that it is okay to link `panic=unwind` into a `panic=abort` + // program if it contains no FFI-unwind calls. In such case foreign exception can only + // enter Rust in a `panic=abort` crate, which will lead to an abort. There will also + // be no exceptions generated from Rust, so the assumption which `panic=abort` crates + // make, that no Rust function can unwind, indeed holds for crates compiled with + // `panic=unwind` as well. In such case this function returns `None`, indicating that + // the crate does not require a particular final panic strategy, and can be freely + // linked to crates with either strategy (we need such ability for libstd and its + // dependencies). + return Some(PanicStrategy::Unwind); + } + } + + // This crate can be linked with either runtime. + None +} + +pub(crate) fn provide(providers: &mut Providers) { + *providers = Providers { has_ffi_unwind_calls, required_panic_strategy, ..*providers }; +} diff --git a/compiler/rustc_mir_transform/src/function_item_references.rs b/compiler/rustc_mir_transform/src/function_item_references.rs new file mode 100644 index 00000000000..340bb1948eb --- /dev/null +++ b/compiler/rustc_mir_transform/src/function_item_references.rs @@ -0,0 +1,194 @@ +use itertools::Itertools; +use rustc_hir::def_id::DefId; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, EarlyBinder, GenericArgsRef, Ty, TyCtxt}; +use rustc_session::lint::builtin::FUNCTION_ITEM_REFERENCES; +use rustc_span::{symbol::sym, Span}; +use rustc_target::spec::abi::Abi; + +use crate::{errors, MirLint}; + +pub struct FunctionItemReferences; + +impl<'tcx> MirLint<'tcx> for FunctionItemReferences { + fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>) { + let mut checker = FunctionItemRefChecker { tcx, body }; + checker.visit_body(body); + } +} + +struct FunctionItemRefChecker<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + body: &'a Body<'tcx>, +} + +impl<'tcx> Visitor<'tcx> for FunctionItemRefChecker<'_, 'tcx> { + /// Emits a lint for function reference arguments bound by `fmt::Pointer` or passed to + /// `transmute`. This only handles arguments in calls outside macro expansions to avoid double + /// counting function references formatted as pointers by macros. + fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { + if let TerminatorKind::Call { + func, + args, + destination: _, + target: _, + unwind: _, + call_source: _, + fn_span: _, + } = &terminator.kind + { + let source_info = *self.body.source_info(location); + let func_ty = func.ty(self.body, self.tcx); + if let ty::FnDef(def_id, args_ref) = *func_ty.kind() { + // Handle calls to `transmute` + if self.tcx.is_diagnostic_item(sym::transmute, def_id) { + let arg_ty = args[0].ty(self.body, self.tcx); + for inner_ty in arg_ty.walk().filter_map(|arg| arg.as_type()) { + if let Some((fn_id, fn_args)) = FunctionItemRefChecker::is_fn_ref(inner_ty) + { + let span = self.nth_arg_span(args, 0); + self.emit_lint(fn_id, fn_args, source_info, span); + } + } + } else { + self.check_bound_args(def_id, args_ref, args, source_info); + } + } + } + self.super_terminator(terminator, location); + } +} + +impl<'tcx> FunctionItemRefChecker<'_, 'tcx> { + /// Emits a lint for function reference arguments bound by `fmt::Pointer` in calls to the + /// function defined by `def_id` with the substitutions `args_ref`. + fn check_bound_args( + &self, + def_id: DefId, + args_ref: GenericArgsRef<'tcx>, + args: &[Operand<'tcx>], + source_info: SourceInfo, + ) { + let param_env = self.tcx.param_env(def_id); + let bounds = param_env.caller_bounds(); + for bound in bounds { + if let Some(bound_ty) = self.is_pointer_trait(bound) { + // Get the argument types as they appear in the function signature. + let arg_defs = + self.tcx.fn_sig(def_id).instantiate_identity().skip_binder().inputs(); + for (arg_num, arg_def) in arg_defs.iter().enumerate() { + // For all types reachable from the argument type in the fn sig + for inner_ty in arg_def.walk().filter_map(|arg| arg.as_type()) { + // If the inner type matches the type bound by `Pointer` + if inner_ty == bound_ty { + // Do a substitution using the parameters from the callsite + let subst_ty = + EarlyBinder::bind(inner_ty).instantiate(self.tcx, args_ref); + if let Some((fn_id, fn_args)) = + FunctionItemRefChecker::is_fn_ref(subst_ty) + { + let mut span = self.nth_arg_span(args, arg_num); + if span.from_expansion() { + // The operand's ctxt wouldn't display the lint since it's inside a macro so + // we have to use the callsite's ctxt. + let callsite_ctxt = span.source_callsite().ctxt(); + span = span.with_ctxt(callsite_ctxt); + } + self.emit_lint(fn_id, fn_args, source_info, span); + } + } + } + } + } + } + } + + /// If the given predicate is the trait `fmt::Pointer`, returns the bound parameter type. + fn is_pointer_trait(&self, bound: ty::Clause<'tcx>) -> Option<Ty<'tcx>> { + if let ty::ClauseKind::Trait(predicate) = bound.kind().skip_binder() { + self.tcx + .is_diagnostic_item(sym::Pointer, predicate.def_id()) + .then(|| predicate.trait_ref.self_ty()) + } else { + None + } + } + + /// If a type is a reference or raw pointer to the anonymous type of a function definition, + /// returns that function's `DefId` and `GenericArgsRef`. + fn is_fn_ref(ty: Ty<'tcx>) -> Option<(DefId, GenericArgsRef<'tcx>)> { + let referent_ty = match ty.kind() { + ty::Ref(_, referent_ty, _) => Some(referent_ty), + ty::RawPtr(ty_and_mut) => Some(&ty_and_mut.ty), + _ => None, + }; + referent_ty + .map(|ref_ty| { + if let ty::FnDef(def_id, args_ref) = *ref_ty.kind() { + Some((def_id, args_ref)) + } else { + None + } + }) + .unwrap_or(None) + } + + fn nth_arg_span(&self, args: &[Operand<'tcx>], n: usize) -> Span { + match &args[n] { + Operand::Copy(place) | Operand::Move(place) => { + self.body.local_decls[place.local].source_info.span + } + Operand::Constant(constant) => constant.span, + } + } + + fn emit_lint( + &self, + fn_id: DefId, + fn_args: GenericArgsRef<'tcx>, + source_info: SourceInfo, + span: Span, + ) { + let lint_root = self.body.source_scopes[source_info.scope] + .local_data + .as_ref() + .assert_crate_local() + .lint_root; + // FIXME: use existing printing routines to print the function signature + let fn_sig = self.tcx.fn_sig(fn_id).instantiate(self.tcx, fn_args); + let unsafety = fn_sig.unsafety().prefix_str(); + let abi = match fn_sig.abi() { + Abi::Rust => String::from(""), + other_abi => { + let mut s = String::from("extern \""); + s.push_str(other_abi.name()); + s.push_str("\" "); + s + } + }; + let ident = self.tcx.item_name(fn_id).to_ident_string(); + let ty_params = fn_args.types().map(|ty| format!("{ty}")); + let const_params = fn_args.consts().map(|c| format!("{c}")); + let params = ty_params.chain(const_params).join(", "); + let num_args = fn_sig.inputs().map_bound(|inputs| inputs.len()).skip_binder(); + let variadic = if fn_sig.c_variadic() { ", ..." } else { "" }; + let ret = if fn_sig.output().skip_binder().is_unit() { "" } else { " -> _" }; + let sugg = format!( + "{} as {}{}fn({}{}){}", + if params.is_empty() { ident.clone() } else { format!("{ident}::<{params}>") }, + unsafety, + abi, + vec!["_"; num_args].join(", "), + variadic, + ret, + ); + + self.tcx.emit_spanned_lint( + FUNCTION_ITEM_REFERENCES, + lint_root, + span, + errors::FnItemRef { span, sugg, ident }, + ); + } +} diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs new file mode 100644 index 00000000000..b21988e179d --- /dev/null +++ b/compiler/rustc_mir_transform/src/gvn.rs @@ -0,0 +1,1063 @@ +//! Global value numbering. +//! +//! MIR may contain repeated and/or redundant computations. The objective of this pass is to detect +//! such redundancies and re-use the already-computed result when possible. +//! +//! In a first pass, we compute a symbolic representation of values that are assigned to SSA +//! locals. This symbolic representation is defined by the `Value` enum. Each produced instance of +//! `Value` is interned as a `VnIndex`, which allows us to cheaply compute identical values. +//! +//! From those assignments, we construct a mapping `VnIndex -> Vec<(Local, Location)>` of available +//! values, the locals in which they are stored, and a the assignment location. +//! +//! In a second pass, we traverse all (non SSA) assignments `x = rvalue` and operands. For each +//! one, we compute the `VnIndex` of the rvalue. If this `VnIndex` is associated to a constant, we +//! replace the rvalue/operand by that constant. Otherwise, if there is an SSA local `y` +//! associated to this `VnIndex`, and if its definition location strictly dominates the assignment +//! to `x`, we replace the assignment by `x = y`. +//! +//! By opportunity, this pass simplifies some `Rvalue`s based on the accumulated knowledge. +//! +//! # Operational semantic +//! +//! Operationally, this pass attempts to prove bitwise equality between locals. Given this MIR: +//! ```ignore (MIR) +//! _a = some value // has VnIndex i +//! // some MIR +//! _b = some other value // also has VnIndex i +//! ``` +//! +//! We consider it to be replacable by: +//! ```ignore (MIR) +//! _a = some value // has VnIndex i +//! // some MIR +//! _c = some other value // also has VnIndex i +//! assume(_a bitwise equal to _c) // follows from having the same VnIndex +//! _b = _a // follows from the `assume` +//! ``` +//! +//! Which is simplifiable to: +//! ```ignore (MIR) +//! _a = some value // has VnIndex i +//! // some MIR +//! _b = _a +//! ``` +//! +//! # Handling of references +//! +//! We handle references by assigning a different "provenance" index to each Ref/AddressOf rvalue. +//! This ensure that we do not spuriously merge borrows that should not be merged. Meanwhile, we +//! consider all the derefs of an immutable reference to a freeze type to give the same value: +//! ```ignore (MIR) +//! _a = *_b // _b is &Freeze +//! _c = *_b // replaced by _c = _a +//! ``` +//! +//! # Determinism of constant propagation +//! +//! When registering a new `Value`, we attempt to opportunistically evaluate it as a constant. +//! The evaluated form is inserted in `evaluated` as an `OpTy` or `None` if evaluation failed. +//! +//! The difficulty is non-deterministic evaluation of MIR constants. Some `Const` can have +//! different runtime values each time they are evaluated. This is the case with +//! `Const::Slice` which have a new pointer each time they are evaluated, and constants that +//! contain a fn pointer (`AllocId` pointing to a `GlobalAlloc::Function`) pointing to a different +//! symbol in each codegen unit. +//! +//! Meanwhile, we want to be able to read indirect constants. For instance: +//! ``` +//! static A: &'static &'static u8 = &&63; +//! fn foo() -> u8 { +//! **A // We want to replace by 63. +//! } +//! fn bar() -> u8 { +//! b"abc"[1] // We want to replace by 'b'. +//! } +//! ``` +//! +//! The `Value::Constant` variant stores a possibly unevaluated constant. Evaluating that constant +//! may be non-deterministic. When that happens, we assign a disambiguator to ensure that we do not +//! merge the constants. See `duplicate_slice` test in `gvn.rs`. +//! +//! Second, when writing constants in MIR, we do not write `Const::Slice` or `Const` +//! that contain `AllocId`s. + +use rustc_const_eval::interpret::{intern_const_alloc_for_constprop, MemoryKind}; +use rustc_const_eval::interpret::{ImmTy, InterpCx, OpTy, Projectable, Scalar}; +use rustc_data_structures::fx::{FxHashMap, FxIndexSet}; +use rustc_data_structures::graph::dominators::Dominators; +use rustc_hir::def::DefKind; +use rustc_index::bit_set::BitSet; +use rustc_index::newtype_index; +use rustc_index::IndexVec; +use rustc_middle::mir::interpret::GlobalAlloc; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::adjustment::PointerCoercion; +use rustc_middle::ty::layout::LayoutOf; +use rustc_middle::ty::{self, Ty, TyCtxt, TypeAndMut}; +use rustc_span::def_id::DefId; +use rustc_span::DUMMY_SP; +use rustc_target::abi::{self, Abi, Size, VariantIdx, FIRST_VARIANT}; +use std::borrow::Cow; + +use crate::dataflow_const_prop::DummyMachine; +use crate::ssa::{AssignedValue, SsaLocals}; +use either::Either; + +pub struct GVN; + +impl<'tcx> MirPass<'tcx> for GVN { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 4 + } + + #[instrument(level = "trace", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + propagate_ssa(tcx, body); + } +} + +fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + let ssa = SsaLocals::new(body); + // Clone dominators as we need them while mutating the body. + let dominators = body.basic_blocks.dominators().clone(); + + let mut state = VnState::new(tcx, param_env, &ssa, &dominators, &body.local_decls); + ssa.for_each_assignment_mut( + body.basic_blocks.as_mut_preserves_cfg(), + |local, value, location| { + let value = match value { + // We do not know anything of this assigned value. + AssignedValue::Arg | AssignedValue::Terminator(_) => None, + // Try to get some insight. + AssignedValue::Rvalue(rvalue) => { + let value = state.simplify_rvalue(rvalue, location); + // FIXME(#112651) `rvalue` may have a subtype to `local`. We can only mark `local` as + // reusable if we have an exact type match. + if state.local_decls[local].ty != rvalue.ty(state.local_decls, tcx) { + return; + } + value + } + }; + // `next_opaque` is `Some`, so `new_opaque` must return `Some`. + let value = value.or_else(|| state.new_opaque()).unwrap(); + state.assign(local, value); + }, + ); + + // Stop creating opaques during replacement as it is useless. + state.next_opaque = None; + + let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec(); + for bb in reverse_postorder { + let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb]; + state.visit_basic_block_data(bb, data); + } + + // For each local that is reused (`y` above), we remove its storage statements do avoid any + // difficulty. Those locals are SSA, so should be easy to optimize by LLVM without storage + // statements. + StorageRemover { tcx, reused_locals: state.reused_locals }.visit_body_preserves_cfg(body); +} + +newtype_index! { + struct VnIndex {} +} + +/// Computing the aggregate's type can be quite slow, so we only keep the minimal amount of +/// information to reconstruct it when needed. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +enum AggregateTy<'tcx> { + /// Invariant: this must not be used for an empty array. + Array, + Tuple, + Def(DefId, ty::GenericArgsRef<'tcx>), +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +enum AddressKind { + Ref(BorrowKind), + Address(Mutability), +} + +#[derive(Debug, PartialEq, Eq, Hash)] +enum Value<'tcx> { + // Root values. + /// Used to represent values we know nothing about. + /// The `usize` is a counter incremented by `new_opaque`. + Opaque(usize), + /// Evaluated or unevaluated constant value. + Constant { + value: Const<'tcx>, + /// Some constants do not have a deterministic value. To avoid merging two instances of the + /// same `Const`, we assign them an additional integer index. + disambiguator: usize, + }, + /// An aggregate value, either tuple/closure/struct/enum. + /// This does not contain unions, as we cannot reason with the value. + Aggregate(AggregateTy<'tcx>, VariantIdx, Vec<VnIndex>), + /// This corresponds to a `[value; count]` expression. + Repeat(VnIndex, ty::Const<'tcx>), + /// The address of a place. + Address { + place: Place<'tcx>, + kind: AddressKind, + /// Give each borrow and pointer a different provenance, so we don't merge them. + provenance: usize, + }, + + // Extractions. + /// This is the *value* obtained by projecting another value. + Projection(VnIndex, ProjectionElem<VnIndex, Ty<'tcx>>), + /// Discriminant of the given value. + Discriminant(VnIndex), + /// Length of an array or slice. + Len(VnIndex), + + // Operations. + NullaryOp(NullOp<'tcx>, Ty<'tcx>), + UnaryOp(UnOp, VnIndex), + BinaryOp(BinOp, VnIndex, VnIndex), + CheckedBinaryOp(BinOp, VnIndex, VnIndex), + Cast { + kind: CastKind, + value: VnIndex, + from: Ty<'tcx>, + to: Ty<'tcx>, + }, +} + +struct VnState<'body, 'tcx> { + tcx: TyCtxt<'tcx>, + ecx: InterpCx<'tcx, 'tcx, DummyMachine>, + param_env: ty::ParamEnv<'tcx>, + local_decls: &'body LocalDecls<'tcx>, + /// Value stored in each local. + locals: IndexVec<Local, Option<VnIndex>>, + /// First local to be assigned that value. + rev_locals: FxHashMap<VnIndex, Vec<Local>>, + values: FxIndexSet<Value<'tcx>>, + /// Values evaluated as constants if possible. + evaluated: IndexVec<VnIndex, Option<OpTy<'tcx>>>, + /// Counter to generate different values. + /// This is an option to stop creating opaques during replacement. + next_opaque: Option<usize>, + ssa: &'body SsaLocals, + dominators: &'body Dominators<BasicBlock>, + reused_locals: BitSet<Local>, +} + +impl<'body, 'tcx> VnState<'body, 'tcx> { + fn new( + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + ssa: &'body SsaLocals, + dominators: &'body Dominators<BasicBlock>, + local_decls: &'body LocalDecls<'tcx>, + ) -> Self { + VnState { + tcx, + ecx: InterpCx::new(tcx, DUMMY_SP, param_env, DummyMachine), + param_env, + local_decls, + locals: IndexVec::from_elem(None, local_decls), + rev_locals: FxHashMap::default(), + values: FxIndexSet::default(), + evaluated: IndexVec::new(), + next_opaque: Some(0), + ssa, + dominators, + reused_locals: BitSet::new_empty(local_decls.len()), + } + } + + #[instrument(level = "trace", skip(self), ret)] + fn insert(&mut self, value: Value<'tcx>) -> VnIndex { + let (index, new) = self.values.insert_full(value); + let index = VnIndex::from_usize(index); + if new { + let evaluated = self.eval_to_const(index); + let _index = self.evaluated.push(evaluated); + debug_assert_eq!(index, _index); + } + index + } + + /// Create a new `Value` for which we have no information at all, except that it is distinct + /// from all the others. + #[instrument(level = "trace", skip(self), ret)] + fn new_opaque(&mut self) -> Option<VnIndex> { + let next_opaque = self.next_opaque.as_mut()?; + let value = Value::Opaque(*next_opaque); + *next_opaque += 1; + Some(self.insert(value)) + } + + /// Create a new `Value::Address` distinct from all the others. + #[instrument(level = "trace", skip(self), ret)] + fn new_pointer(&mut self, place: Place<'tcx>, kind: AddressKind) -> Option<VnIndex> { + let next_opaque = self.next_opaque.as_mut()?; + let value = Value::Address { place, kind, provenance: *next_opaque }; + *next_opaque += 1; + Some(self.insert(value)) + } + + fn get(&self, index: VnIndex) -> &Value<'tcx> { + self.values.get_index(index.as_usize()).unwrap() + } + + /// Record that `local` is assigned `value`. `local` must be SSA. + #[instrument(level = "trace", skip(self))] + fn assign(&mut self, local: Local, value: VnIndex) { + self.locals[local] = Some(value); + + // Only register the value if its type is `Sized`, as we will emit copies of it. + let is_sized = !self.tcx.features().unsized_locals + || self.local_decls[local].ty.is_sized(self.tcx, self.param_env); + if is_sized { + self.rev_locals.entry(value).or_default().push(local); + } + } + + fn insert_constant(&mut self, value: Const<'tcx>) -> Option<VnIndex> { + let disambiguator = if value.is_deterministic() { + // The constant is deterministic, no need to disambiguate. + 0 + } else { + // Multiple mentions of this constant will yield different values, + // so assign a different `disambiguator` to ensure they do not get the same `VnIndex`. + let next_opaque = self.next_opaque.as_mut()?; + let disambiguator = *next_opaque; + *next_opaque += 1; + disambiguator + }; + Some(self.insert(Value::Constant { value, disambiguator })) + } + + fn insert_scalar(&mut self, scalar: Scalar, ty: Ty<'tcx>) -> VnIndex { + self.insert_constant(Const::from_scalar(self.tcx, scalar, ty)) + .expect("scalars are deterministic") + } + + #[instrument(level = "trace", skip(self), ret)] + fn eval_to_const(&mut self, value: VnIndex) -> Option<OpTy<'tcx>> { + use Value::*; + let op = match *self.get(value) { + Opaque(_) => return None, + // Do not bother evaluating repeat expressions. This would uselessly consume memory. + Repeat(..) => return None, + + Constant { ref value, disambiguator: _ } => { + self.ecx.eval_mir_constant(value, None, None).ok()? + } + Aggregate(kind, variant, ref fields) => { + let fields = fields + .iter() + .map(|&f| self.evaluated[f].as_ref()) + .collect::<Option<Vec<_>>>()?; + let ty = match kind { + AggregateTy::Array => { + assert!(fields.len() > 0); + Ty::new_array(self.tcx, fields[0].layout.ty, fields.len() as u64) + } + AggregateTy::Tuple => { + Ty::new_tup_from_iter(self.tcx, fields.iter().map(|f| f.layout.ty)) + } + AggregateTy::Def(def_id, args) => { + self.tcx.type_of(def_id).instantiate(self.tcx, args) + } + }; + let variant = if ty.is_enum() { Some(variant) } else { None }; + let ty = self.ecx.layout_of(ty).ok()?; + if ty.is_zst() { + ImmTy::uninit(ty).into() + } else if matches!(ty.abi, Abi::Scalar(..) | Abi::ScalarPair(..)) { + let dest = self.ecx.allocate(ty, MemoryKind::Stack).ok()?; + let variant_dest = if let Some(variant) = variant { + self.ecx.project_downcast(&dest, variant).ok()? + } else { + dest.clone() + }; + for (field_index, op) in fields.into_iter().enumerate() { + let field_dest = self.ecx.project_field(&variant_dest, field_index).ok()?; + self.ecx.copy_op(op, &field_dest, /*allow_transmute*/ false).ok()?; + } + self.ecx.write_discriminant(variant.unwrap_or(FIRST_VARIANT), &dest).ok()?; + self.ecx + .alloc_mark_immutable(dest.ptr().provenance.unwrap().alloc_id()) + .ok()?; + dest.into() + } else { + return None; + } + } + + Projection(base, elem) => { + let value = self.evaluated[base].as_ref()?; + let elem = match elem { + ProjectionElem::Deref => ProjectionElem::Deref, + ProjectionElem::Downcast(name, read_variant) => { + ProjectionElem::Downcast(name, read_variant) + } + ProjectionElem::Field(f, ty) => ProjectionElem::Field(f, ty), + ProjectionElem::ConstantIndex { offset, min_length, from_end } => { + ProjectionElem::ConstantIndex { offset, min_length, from_end } + } + ProjectionElem::Subslice { from, to, from_end } => { + ProjectionElem::Subslice { from, to, from_end } + } + ProjectionElem::OpaqueCast(ty) => ProjectionElem::OpaqueCast(ty), + ProjectionElem::Subtype(ty) => ProjectionElem::Subtype(ty), + // This should have been replaced by a `ConstantIndex` earlier. + ProjectionElem::Index(_) => return None, + }; + self.ecx.project(value, elem).ok()? + } + Address { place, kind, provenance: _ } => { + if !place.is_indirect_first_projection() { + return None; + } + let local = self.locals[place.local]?; + let pointer = self.evaluated[local].as_ref()?; + let mut mplace = self.ecx.deref_pointer(pointer).ok()?; + for proj in place.projection.iter().skip(1) { + // We have no call stack to associate a local with a value, so we cannot interpret indexing. + if matches!(proj, ProjectionElem::Index(_)) { + return None; + } + mplace = self.ecx.project(&mplace, proj).ok()?; + } + let pointer = mplace.to_ref(&self.ecx); + let ty = match kind { + AddressKind::Ref(bk) => Ty::new_ref( + self.tcx, + self.tcx.lifetimes.re_erased, + ty::TypeAndMut { ty: mplace.layout.ty, mutbl: bk.to_mutbl_lossy() }, + ), + AddressKind::Address(mutbl) => { + Ty::new_ptr(self.tcx, TypeAndMut { ty: mplace.layout.ty, mutbl }) + } + }; + let layout = self.ecx.layout_of(ty).ok()?; + ImmTy::from_immediate(pointer, layout).into() + } + + Discriminant(base) => { + let base = self.evaluated[base].as_ref()?; + let variant = self.ecx.read_discriminant(base).ok()?; + let discr_value = + self.ecx.discriminant_for_variant(base.layout.ty, variant).ok()?; + discr_value.into() + } + Len(slice) => { + let slice = self.evaluated[slice].as_ref()?; + let usize_layout = self.ecx.layout_of(self.tcx.types.usize).unwrap(); + let len = slice.len(&self.ecx).ok()?; + let imm = ImmTy::try_from_uint(len, usize_layout)?; + imm.into() + } + NullaryOp(null_op, ty) => { + let layout = self.ecx.layout_of(ty).ok()?; + if let NullOp::SizeOf | NullOp::AlignOf = null_op + && layout.is_unsized() + { + return None; + } + let val = match null_op { + NullOp::SizeOf => layout.size.bytes(), + NullOp::AlignOf => layout.align.abi.bytes(), + NullOp::OffsetOf(fields) => { + layout.offset_of_subfield(&self.ecx, fields.iter()).bytes() + } + }; + let usize_layout = self.ecx.layout_of(self.tcx.types.usize).unwrap(); + let imm = ImmTy::try_from_uint(val, usize_layout)?; + imm.into() + } + UnaryOp(un_op, operand) => { + let operand = self.evaluated[operand].as_ref()?; + let operand = self.ecx.read_immediate(operand).ok()?; + let (val, _) = self.ecx.overflowing_unary_op(un_op, &operand).ok()?; + val.into() + } + BinaryOp(bin_op, lhs, rhs) => { + let lhs = self.evaluated[lhs].as_ref()?; + let lhs = self.ecx.read_immediate(lhs).ok()?; + let rhs = self.evaluated[rhs].as_ref()?; + let rhs = self.ecx.read_immediate(rhs).ok()?; + let (val, _) = self.ecx.overflowing_binary_op(bin_op, &lhs, &rhs).ok()?; + val.into() + } + CheckedBinaryOp(bin_op, lhs, rhs) => { + let lhs = self.evaluated[lhs].as_ref()?; + let lhs = self.ecx.read_immediate(lhs).ok()?; + let rhs = self.evaluated[rhs].as_ref()?; + let rhs = self.ecx.read_immediate(rhs).ok()?; + let (val, overflowed) = self.ecx.overflowing_binary_op(bin_op, &lhs, &rhs).ok()?; + let tuple = Ty::new_tup_from_iter( + self.tcx, + [val.layout.ty, self.tcx.types.bool].into_iter(), + ); + let tuple = self.ecx.layout_of(tuple).ok()?; + ImmTy::from_scalar_pair(val.to_scalar(), Scalar::from_bool(overflowed), tuple) + .into() + } + Cast { kind, value, from: _, to } => match kind { + CastKind::IntToInt | CastKind::IntToFloat => { + let value = self.evaluated[value].as_ref()?; + let value = self.ecx.read_immediate(value).ok()?; + let to = self.ecx.layout_of(to).ok()?; + let res = self.ecx.int_to_int_or_float(&value, to).ok()?; + res.into() + } + CastKind::FloatToFloat | CastKind::FloatToInt => { + let value = self.evaluated[value].as_ref()?; + let value = self.ecx.read_immediate(value).ok()?; + let to = self.ecx.layout_of(to).ok()?; + let res = self.ecx.float_to_float_or_int(&value, to).ok()?; + res.into() + } + CastKind::Transmute => { + let value = self.evaluated[value].as_ref()?; + let to = self.ecx.layout_of(to).ok()?; + // `offset` for immediates only supports scalar/scalar-pair ABIs, + // so bail out if the target is not one. + if value.as_mplace_or_imm().is_right() { + match (value.layout.abi, to.abi) { + (Abi::Scalar(..), Abi::Scalar(..)) => {} + (Abi::ScalarPair(..), Abi::ScalarPair(..)) => {} + _ => return None, + } + } + value.offset(Size::ZERO, to, &self.ecx).ok()? + } + _ => return None, + }, + }; + Some(op) + } + + fn project( + &mut self, + place: PlaceRef<'tcx>, + value: VnIndex, + proj: PlaceElem<'tcx>, + ) -> Option<VnIndex> { + let proj = match proj { + ProjectionElem::Deref => { + let ty = place.ty(self.local_decls, self.tcx).ty; + if let Some(Mutability::Not) = ty.ref_mutability() + && let Some(pointee_ty) = ty.builtin_deref(true) + && pointee_ty.ty.is_freeze(self.tcx, self.param_env) + { + // An immutable borrow `_x` always points to the same value for the + // lifetime of the borrow, so we can merge all instances of `*_x`. + ProjectionElem::Deref + } else { + return None; + } + } + ProjectionElem::Downcast(name, index) => ProjectionElem::Downcast(name, index), + ProjectionElem::Field(f, ty) => { + if let Value::Aggregate(_, _, fields) = self.get(value) { + return Some(fields[f.as_usize()]); + } else if let Value::Projection(outer_value, ProjectionElem::Downcast(_, read_variant)) = self.get(value) + && let Value::Aggregate(_, written_variant, fields) = self.get(*outer_value) + // This pass is not aware of control-flow, so we do not know whether the + // replacement we are doing is actually reachable. We could be in any arm of + // ``` + // match Some(x) { + // Some(y) => /* stuff */, + // None => /* other */, + // } + // ``` + // + // In surface rust, the current statement would be unreachable. + // + // However, from the reference chapter on enums and RFC 2195, + // accessing the wrong variant is not UB if the enum has repr. + // So it's not impossible for a series of MIR opts to generate + // a downcast to an inactive variant. + && written_variant == read_variant + { + return Some(fields[f.as_usize()]); + } + ProjectionElem::Field(f, ty) + } + ProjectionElem::Index(idx) => { + if let Value::Repeat(inner, _) = self.get(value) { + return Some(*inner); + } + let idx = self.locals[idx]?; + ProjectionElem::Index(idx) + } + ProjectionElem::ConstantIndex { offset, min_length, from_end } => { + match self.get(value) { + Value::Repeat(inner, _) => { + return Some(*inner); + } + Value::Aggregate(AggregateTy::Array, _, operands) => { + let offset = if from_end { + operands.len() - offset as usize + } else { + offset as usize + }; + return operands.get(offset).copied(); + } + _ => {} + }; + ProjectionElem::ConstantIndex { offset, min_length, from_end } + } + ProjectionElem::Subslice { from, to, from_end } => { + ProjectionElem::Subslice { from, to, from_end } + } + ProjectionElem::OpaqueCast(ty) => ProjectionElem::OpaqueCast(ty), + ProjectionElem::Subtype(ty) => ProjectionElem::Subtype(ty), + }; + + Some(self.insert(Value::Projection(value, proj))) + } + + /// Simplify the projection chain if we know better. + #[instrument(level = "trace", skip(self))] + fn simplify_place_projection(&mut self, place: &mut Place<'tcx>, location: Location) { + // If the projection is indirect, we treat the local as a value, so can replace it with + // another local. + if place.is_indirect() + && let Some(base) = self.locals[place.local] + && let Some(new_local) = self.try_as_local(base, location) + { + place.local = new_local; + self.reused_locals.insert(new_local); + } + + let mut projection = Cow::Borrowed(&place.projection[..]); + + for i in 0..projection.len() { + let elem = projection[i]; + if let ProjectionElem::Index(idx) = elem + && let Some(idx) = self.locals[idx] + { + if let Some(offset) = self.evaluated[idx].as_ref() + && let Ok(offset) = self.ecx.read_target_usize(offset) + { + projection.to_mut()[i] = ProjectionElem::ConstantIndex { + offset, + min_length: offset + 1, + from_end: false, + }; + } else if let Some(new_idx) = self.try_as_local(idx, location) { + projection.to_mut()[i] = ProjectionElem::Index(new_idx); + self.reused_locals.insert(new_idx); + } + } + } + + if projection.is_owned() { + place.projection = self.tcx.mk_place_elems(&projection); + } + + trace!(?place); + } + + /// Represent the *value* which would be read from `place`, and point `place` to a preexisting + /// place with the same value (if that already exists). + #[instrument(level = "trace", skip(self), ret)] + fn simplify_place_value( + &mut self, + place: &mut Place<'tcx>, + location: Location, + ) -> Option<VnIndex> { + self.simplify_place_projection(place, location); + + // Invariant: `place` and `place_ref` point to the same value, even if they point to + // different memory locations. + let mut place_ref = place.as_ref(); + + // Invariant: `value` holds the value up-to the `index`th projection excluded. + let mut value = self.locals[place.local]?; + for (index, proj) in place.projection.iter().enumerate() { + if let Some(local) = self.try_as_local(value, location) { + // Both `local` and `Place { local: place.local, projection: projection[..index] }` + // hold the same value. Therefore, following place holds the value in the original + // `place`. + place_ref = PlaceRef { local, projection: &place.projection[index..] }; + } + + let base = PlaceRef { local: place.local, projection: &place.projection[..index] }; + value = self.project(base, value, proj)?; + } + + if let Some(new_local) = self.try_as_local(value, location) { + place_ref = PlaceRef { local: new_local, projection: &[] }; + } + + if place_ref.local != place.local || place_ref.projection.len() < place.projection.len() { + // By the invariant on `place_ref`. + *place = place_ref.project_deeper(&[], self.tcx); + self.reused_locals.insert(place_ref.local); + } + + Some(value) + } + + #[instrument(level = "trace", skip(self), ret)] + fn simplify_operand( + &mut self, + operand: &mut Operand<'tcx>, + location: Location, + ) -> Option<VnIndex> { + match *operand { + Operand::Constant(ref mut constant) => { + let const_ = constant.const_.normalize(self.tcx, self.param_env); + self.insert_constant(const_) + } + Operand::Copy(ref mut place) | Operand::Move(ref mut place) => { + let value = self.simplify_place_value(place, location)?; + if let Some(const_) = self.try_as_constant(value) { + *operand = Operand::Constant(Box::new(const_)); + } + Some(value) + } + } + } + + #[instrument(level = "trace", skip(self), ret)] + fn simplify_rvalue( + &mut self, + rvalue: &mut Rvalue<'tcx>, + location: Location, + ) -> Option<VnIndex> { + let value = match *rvalue { + // Forward values. + Rvalue::Use(ref mut operand) => return self.simplify_operand(operand, location), + Rvalue::CopyForDeref(place) => { + let mut operand = Operand::Copy(place); + let val = self.simplify_operand(&mut operand, location); + *rvalue = Rvalue::Use(operand); + return val; + } + + // Roots. + Rvalue::Repeat(ref mut op, amount) => { + let op = self.simplify_operand(op, location)?; + Value::Repeat(op, amount) + } + Rvalue::NullaryOp(op, ty) => Value::NullaryOp(op, ty), + Rvalue::Aggregate(..) => return self.simplify_aggregate(rvalue, location), + Rvalue::Ref(_, borrow_kind, ref mut place) => { + self.simplify_place_projection(place, location); + return self.new_pointer(*place, AddressKind::Ref(borrow_kind)); + } + Rvalue::AddressOf(mutbl, ref mut place) => { + self.simplify_place_projection(place, location); + return self.new_pointer(*place, AddressKind::Address(mutbl)); + } + + // Operations. + Rvalue::Len(ref mut place) => { + let place = self.simplify_place_value(place, location)?; + Value::Len(place) + } + Rvalue::Cast(kind, ref mut value, to) => { + let from = value.ty(self.local_decls, self.tcx); + let value = self.simplify_operand(value, location)?; + if let CastKind::PointerCoercion( + PointerCoercion::ReifyFnPointer | PointerCoercion::ClosureFnPointer(_), + ) = kind + { + // Each reification of a generic fn may get a different pointer. + // Do not try to merge them. + return self.new_opaque(); + } + Value::Cast { kind, value, from, to } + } + Rvalue::BinaryOp(op, box (ref mut lhs, ref mut rhs)) => { + let lhs = self.simplify_operand(lhs, location); + let rhs = self.simplify_operand(rhs, location); + Value::BinaryOp(op, lhs?, rhs?) + } + Rvalue::CheckedBinaryOp(op, box (ref mut lhs, ref mut rhs)) => { + let lhs = self.simplify_operand(lhs, location); + let rhs = self.simplify_operand(rhs, location); + Value::CheckedBinaryOp(op, lhs?, rhs?) + } + Rvalue::UnaryOp(op, ref mut arg) => { + let arg = self.simplify_operand(arg, location)?; + Value::UnaryOp(op, arg) + } + Rvalue::Discriminant(ref mut place) => { + let place = self.simplify_place_value(place, location)?; + if let Some(discr) = self.simplify_discriminant(place) { + return Some(discr); + } + Value::Discriminant(place) + } + + // Unsupported values. + Rvalue::ThreadLocalRef(..) | Rvalue::ShallowInitBox(..) => return None, + }; + debug!(?value); + Some(self.insert(value)) + } + + fn simplify_discriminant(&mut self, place: VnIndex) -> Option<VnIndex> { + if let Value::Aggregate(enum_ty, variant, _) = *self.get(place) + && let AggregateTy::Def(enum_did, enum_substs) = enum_ty + && let DefKind::Enum = self.tcx.def_kind(enum_did) + { + let enum_ty = self.tcx.type_of(enum_did).instantiate(self.tcx, enum_substs); + let discr = self.ecx.discriminant_for_variant(enum_ty, variant).ok()?; + return Some(self.insert_scalar(discr.to_scalar(), discr.layout.ty)); + } + + None + } + + fn simplify_aggregate( + &mut self, + rvalue: &mut Rvalue<'tcx>, + location: Location, + ) -> Option<VnIndex> { + let Rvalue::Aggregate(box ref kind, ref mut fields) = *rvalue else { bug!() }; + + let tcx = self.tcx; + if fields.is_empty() { + let is_zst = match *kind { + AggregateKind::Array(..) | AggregateKind::Tuple | AggregateKind::Closure(..) => { + true + } + // Only enums can be non-ZST. + AggregateKind::Adt(did, ..) => tcx.def_kind(did) != DefKind::Enum, + // Coroutines are never ZST, as they at least contain the implicit states. + AggregateKind::Coroutine(..) => false, + }; + + if is_zst { + let ty = rvalue.ty(self.local_decls, tcx); + return self.insert_constant(Const::zero_sized(ty)); + } + } + + let (ty, variant_index) = match *kind { + AggregateKind::Array(..) => { + assert!(!fields.is_empty()); + (AggregateTy::Array, FIRST_VARIANT) + } + AggregateKind::Tuple => { + assert!(!fields.is_empty()); + (AggregateTy::Tuple, FIRST_VARIANT) + } + AggregateKind::Closure(did, substs) | AggregateKind::Coroutine(did, substs, _) => { + (AggregateTy::Def(did, substs), FIRST_VARIANT) + } + AggregateKind::Adt(did, variant_index, substs, _, None) => { + (AggregateTy::Def(did, substs), variant_index) + } + // Do not track unions. + AggregateKind::Adt(_, _, _, _, Some(_)) => return None, + }; + + let fields: Option<Vec<_>> = fields + .iter_mut() + .map(|op| self.simplify_operand(op, location).or_else(|| self.new_opaque())) + .collect(); + let fields = fields?; + + if let AggregateTy::Array = ty + && fields.len() > 4 + { + let first = fields[0]; + if fields.iter().all(|&v| v == first) { + let len = ty::Const::from_target_usize(self.tcx, fields.len().try_into().unwrap()); + if let Some(const_) = self.try_as_constant(first) { + *rvalue = Rvalue::Repeat(Operand::Constant(Box::new(const_)), len); + } else if let Some(local) = self.try_as_local(first, location) { + *rvalue = Rvalue::Repeat(Operand::Copy(local.into()), len); + self.reused_locals.insert(local); + } + return Some(self.insert(Value::Repeat(first, len))); + } + } + + Some(self.insert(Value::Aggregate(ty, variant_index, fields))) + } +} + +fn op_to_prop_const<'tcx>( + ecx: &mut InterpCx<'_, 'tcx, DummyMachine>, + op: &OpTy<'tcx>, +) -> Option<ConstValue<'tcx>> { + // Do not attempt to propagate unsized locals. + if op.layout.is_unsized() { + return None; + } + + // This constant is a ZST, just return an empty value. + if op.layout.is_zst() { + return Some(ConstValue::ZeroSized); + } + + // Do not synthetize too large constants. Codegen will just memcpy them, which we'd like to avoid. + if !matches!(op.layout.abi, Abi::Scalar(..) | Abi::ScalarPair(..)) { + return None; + } + + // If this constant has scalar ABI, return it as a `ConstValue::Scalar`. + if let Abi::Scalar(abi::Scalar::Initialized { .. }) = op.layout.abi + && let Ok(scalar) = ecx.read_scalar(op) + && scalar.try_to_int().is_ok() + { + return Some(ConstValue::Scalar(scalar)); + } + + // If this constant is already represented as an `Allocation`, + // try putting it into global memory to return it. + if let Either::Left(mplace) = op.as_mplace_or_imm() { + let (size, _align) = ecx.size_and_align_of_mplace(&mplace).ok()??; + + // Do not try interning a value that contains provenance. + // Due to https://github.com/rust-lang/rust/issues/79738, doing so could lead to bugs. + // FIXME: remove this hack once that issue is fixed. + let alloc_ref = ecx.get_ptr_alloc(mplace.ptr(), size).ok()??; + if alloc_ref.has_provenance() { + return None; + } + + let pointer = mplace.ptr().into_pointer_or_addr().ok()?; + let (prov, offset) = pointer.into_parts(); + let alloc_id = prov.alloc_id(); + intern_const_alloc_for_constprop(ecx, alloc_id).ok()?; + if matches!(ecx.tcx.global_alloc(alloc_id), GlobalAlloc::Memory(_)) { + // `alloc_id` may point to a static. Codegen will choke on an `Indirect` with anything + // by `GlobalAlloc::Memory`, so do fall through to copying if needed. + // FIXME: find a way to treat this more uniformly + // (probably by fixing codegen) + return Some(ConstValue::Indirect { alloc_id, offset }); + } + } + + // Everything failed: create a new allocation to hold the data. + let alloc_id = + ecx.intern_with_temp_alloc(op.layout, |ecx, dest| ecx.copy_op(op, dest, false)).ok()?; + let value = ConstValue::Indirect { alloc_id, offset: Size::ZERO }; + + // Check that we do not leak a pointer. + // Those pointers may lose part of their identity in codegen. + // FIXME: remove this hack once https://github.com/rust-lang/rust/issues/79738 is fixed. + if ecx.tcx.global_alloc(alloc_id).unwrap_memory().inner().provenance().ptrs().is_empty() { + return Some(value); + } + + None +} + +impl<'tcx> VnState<'_, 'tcx> { + /// If `index` is a `Value::Constant`, return the `Constant` to be put in the MIR. + fn try_as_constant(&mut self, index: VnIndex) -> Option<ConstOperand<'tcx>> { + // This was already constant in MIR, do not change it. + if let Value::Constant { value, disambiguator: _ } = *self.get(index) + // If the constant is not deterministic, adding an additional mention of it in MIR will + // not give the same value as the former mention. + && value.is_deterministic() + { + return Some(ConstOperand { span: rustc_span::DUMMY_SP, user_ty: None, const_: value }); + } + + let op = self.evaluated[index].as_ref()?; + if op.layout.is_unsized() { + // Do not attempt to propagate unsized locals. + return None; + } + + let value = op_to_prop_const(&mut self.ecx, op)?; + + // Check that we do not leak a pointer. + // Those pointers may lose part of their identity in codegen. + // FIXME: remove this hack once https://github.com/rust-lang/rust/issues/79738 is fixed. + assert!(!value.may_have_provenance(self.tcx, op.layout.size)); + + let const_ = Const::Val(value, op.layout.ty); + Some(ConstOperand { span: rustc_span::DUMMY_SP, user_ty: None, const_ }) + } + + /// If there is a local which is assigned `index`, and its assignment strictly dominates `loc`, + /// return it. + fn try_as_local(&mut self, index: VnIndex, loc: Location) -> Option<Local> { + let other = self.rev_locals.get(&index)?; + other + .iter() + .copied() + .find(|&other| self.ssa.assignment_dominates(self.dominators, other, loc)) + } +} + +impl<'tcx> MutVisitor<'tcx> for VnState<'_, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, _: PlaceContext, location: Location) { + self.simplify_place_projection(place, location); + } + + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) { + self.simplify_operand(operand, location); + } + + fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, location: Location) { + if let StatementKind::Assign(box (_, ref mut rvalue)) = stmt.kind + // Do not try to simplify a constant, it's already in canonical shape. + && !matches!(rvalue, Rvalue::Use(Operand::Constant(_))) + { + if let Some(value) = self.simplify_rvalue(rvalue, location) { + if let Some(const_) = self.try_as_constant(value) { + *rvalue = Rvalue::Use(Operand::Constant(Box::new(const_))); + } else if let Some(local) = self.try_as_local(value, location) + && *rvalue != Rvalue::Use(Operand::Move(local.into())) + { + *rvalue = Rvalue::Use(Operand::Copy(local.into())); + self.reused_locals.insert(local); + } + } + } else { + self.super_statement(stmt, location); + } + } +} + +struct StorageRemover<'tcx> { + tcx: TyCtxt<'tcx>, + reused_locals: BitSet<Local>, +} + +impl<'tcx> MutVisitor<'tcx> for StorageRemover<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, _: Location) { + if let Operand::Move(place) = *operand + && let Some(local) = place.as_local() + && self.reused_locals.contains(local) + { + *operand = Operand::Copy(place); + } + } + + fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) { + match stmt.kind { + // When removing storage statements, we need to remove both (#107511). + StatementKind::StorageLive(l) | StatementKind::StorageDead(l) + if self.reused_locals.contains(l) => + { + stmt.make_nop() + } + _ => self.super_statement(stmt, loc), + } + } +} diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs new file mode 100644 index 00000000000..8ad804bf3e7 --- /dev/null +++ b/compiler/rustc_mir_transform/src/inline.rs @@ -0,0 +1,1045 @@ +//! Inlining pass for MIR functions +use crate::deref_separator::deref_finder; +use rustc_attr::InlineAttr; +use rustc_const_eval::transform::validate::validate_types; +use rustc_hir::def_id::DefId; +use rustc_index::bit_set::BitSet; +use rustc_index::Idx; +use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs}; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::TypeVisitableExt; +use rustc_middle::ty::{self, Instance, InstanceDef, ParamEnv, Ty, TyCtxt}; +use rustc_session::config::OptLevel; +use rustc_target::abi::FieldIdx; +use rustc_target::spec::abi::Abi; + +use crate::cost_checker::CostChecker; +use crate::simplify::{remove_dead_blocks, CfgSimplifier}; +use crate::util; +use std::iter; +use std::ops::{Range, RangeFrom}; + +pub(crate) mod cycle; + +const TOP_DOWN_DEPTH_LIMIT: usize = 5; + +pub struct Inline; + +#[derive(Copy, Clone, Debug)] +struct CallSite<'tcx> { + callee: Instance<'tcx>, + fn_sig: ty::PolyFnSig<'tcx>, + block: BasicBlock, + source_info: SourceInfo, +} + +impl<'tcx> MirPass<'tcx> for Inline { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + if let Some(enabled) = sess.opts.unstable_opts.inline_mir { + return enabled; + } + + match sess.mir_opt_level() { + 0 | 1 => false, + 2 => { + (sess.opts.optimize == OptLevel::Default + || sess.opts.optimize == OptLevel::Aggressive) + && sess.opts.incremental == None + } + _ => true, + } + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let span = trace_span!("inline", body = %tcx.def_path_str(body.source.def_id())); + let _guard = span.enter(); + if inline(tcx, body) { + debug!("running simplify cfg on {:?}", body.source); + CfgSimplifier::new(body).simplify(); + remove_dead_blocks(body); + deref_finder(tcx, body); + } + } +} + +fn inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool { + let def_id = body.source.def_id().expect_local(); + + // Only do inlining into fn bodies. + if !tcx.hir().body_owner_kind(def_id).is_fn_or_closure() { + return false; + } + if body.source.promoted.is_some() { + return false; + } + // Avoid inlining into coroutines, since their `optimized_mir` is used for layout computation, + // which can create a cycle, even when no attempt is made to inline the function in the other + // direction. + if body.coroutine.is_some() { + return false; + } + + let param_env = tcx.param_env_reveal_all_normalized(def_id); + + let mut this = Inliner { + tcx, + param_env, + codegen_fn_attrs: tcx.codegen_fn_attrs(def_id), + history: Vec::new(), + changed: false, + }; + let blocks = START_BLOCK..body.basic_blocks.next_index(); + this.process_blocks(body, blocks); + this.changed +} + +struct Inliner<'tcx> { + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + /// Caller codegen attributes. + codegen_fn_attrs: &'tcx CodegenFnAttrs, + /// Stack of inlined instances. + /// We only check the `DefId` and not the args because we want to + /// avoid inlining cases of polymorphic recursion. + /// The number of `DefId`s is finite, so checking history is enough + /// to ensure that we do not loop endlessly while inlining. + history: Vec<DefId>, + /// Indicates that the caller body has been modified. + changed: bool, +} + +impl<'tcx> Inliner<'tcx> { + fn process_blocks(&mut self, caller_body: &mut Body<'tcx>, blocks: Range<BasicBlock>) { + // How many callsites in this body are we allowed to inline? We need to limit this in order + // to prevent super-linear growth in MIR size + let inline_limit = match self.history.len() { + 0 => usize::MAX, + 1..=TOP_DOWN_DEPTH_LIMIT => 1, + _ => return, + }; + let mut inlined_count = 0; + for bb in blocks { + let bb_data = &caller_body[bb]; + if bb_data.is_cleanup { + continue; + } + + let Some(callsite) = self.resolve_callsite(caller_body, bb, bb_data) else { + continue; + }; + + let span = trace_span!("process_blocks", %callsite.callee, ?bb); + let _guard = span.enter(); + + match self.try_inlining(caller_body, &callsite) { + Err(reason) => { + debug!("not-inlined {} [{}]", callsite.callee, reason); + continue; + } + Ok(new_blocks) => { + debug!("inlined {}", callsite.callee); + self.changed = true; + + self.history.push(callsite.callee.def_id()); + self.process_blocks(caller_body, new_blocks); + self.history.pop(); + + inlined_count += 1; + if inlined_count == inline_limit { + debug!("inline count reached"); + return; + } + } + } + } + } + + /// Attempts to inline a callsite into the caller body. When successful returns basic blocks + /// containing the inlined body. Otherwise returns an error describing why inlining didn't take + /// place. + fn try_inlining( + &self, + caller_body: &mut Body<'tcx>, + callsite: &CallSite<'tcx>, + ) -> Result<std::ops::Range<BasicBlock>, &'static str> { + self.check_mir_is_available(caller_body, &callsite.callee)?; + + let callee_attrs = self.tcx.codegen_fn_attrs(callsite.callee.def_id()); + let cross_crate_inlinable = self.tcx.cross_crate_inlinable(callsite.callee.def_id()); + self.check_codegen_attributes(callsite, callee_attrs, cross_crate_inlinable)?; + + let terminator = caller_body[callsite.block].terminator.as_ref().unwrap(); + let TerminatorKind::Call { args, destination, .. } = &terminator.kind else { bug!() }; + let destination_ty = destination.ty(&caller_body.local_decls, self.tcx).ty; + for arg in args { + if !arg.ty(&caller_body.local_decls, self.tcx).is_sized(self.tcx, self.param_env) { + // We do not allow inlining functions with unsized params. Inlining these functions + // could create unsized locals, which are unsound and being phased out. + return Err("Call has unsized argument"); + } + } + + let callee_body = try_instance_mir(self.tcx, callsite.callee.def)?; + self.check_mir_body(callsite, callee_body, callee_attrs, cross_crate_inlinable)?; + + if !self.tcx.consider_optimizing(|| { + format!("Inline {:?} into {:?}", callsite.callee, caller_body.source) + }) { + return Err("optimization fuel exhausted"); + } + + let Ok(callee_body) = callsite.callee.try_instantiate_mir_and_normalize_erasing_regions( + self.tcx, + self.param_env, + ty::EarlyBinder::bind(callee_body.clone()), + ) else { + return Err("failed to normalize callee body"); + }; + + // Normally, this shouldn't be required, but trait normalization failure can create a + // validation ICE. + if !validate_types( + self.tcx, + MirPhase::Runtime(RuntimePhase::Optimized), + self.param_env, + &callee_body, + ) + .is_empty() + { + return Err("failed to validate callee body"); + } + + // Check call signature compatibility. + // Normally, this shouldn't be required, but trait normalization failure can create a + // validation ICE. + let output_type = callee_body.return_ty(); + if !util::relate_types( + self.tcx, + self.param_env, + ty::Variance::Covariant, + output_type, + destination_ty, + ) { + trace!(?output_type, ?destination_ty); + return Err("failed to normalize return type"); + } + if callsite.fn_sig.abi() == Abi::RustCall { + // FIXME: Don't inline user-written `extern "rust-call"` functions, + // since this is generally perf-negative on rustc, and we hope that + // LLVM will inline these functions instead. + if callee_body.spread_arg.is_some() { + return Err("do not inline user-written rust-call functions"); + } + + let (self_arg, arg_tuple) = match &args[..] { + [arg_tuple] => (None, arg_tuple), + [self_arg, arg_tuple] => (Some(self_arg), arg_tuple), + _ => bug!("Expected `rust-call` to have 1 or 2 args"), + }; + + let self_arg_ty = + self_arg.map(|self_arg| self_arg.ty(&caller_body.local_decls, self.tcx)); + + let arg_tuple_ty = arg_tuple.ty(&caller_body.local_decls, self.tcx); + let ty::Tuple(arg_tuple_tys) = *arg_tuple_ty.kind() else { + bug!("Closure arguments are not passed as a tuple"); + }; + + for (arg_ty, input) in + self_arg_ty.into_iter().chain(arg_tuple_tys).zip(callee_body.args_iter()) + { + let input_type = callee_body.local_decls[input].ty; + if !util::relate_types( + self.tcx, + self.param_env, + ty::Variance::Covariant, + input_type, + arg_ty, + ) { + trace!(?arg_ty, ?input_type); + return Err("failed to normalize tuple argument type"); + } + } + } else { + for (arg, input) in args.iter().zip(callee_body.args_iter()) { + let input_type = callee_body.local_decls[input].ty; + let arg_ty = arg.ty(&caller_body.local_decls, self.tcx); + if !util::relate_types( + self.tcx, + self.param_env, + ty::Variance::Covariant, + input_type, + arg_ty, + ) { + trace!(?arg_ty, ?input_type); + return Err("failed to normalize argument type"); + } + } + } + + let old_blocks = caller_body.basic_blocks.next_index(); + self.inline_call(caller_body, callsite, callee_body); + let new_blocks = old_blocks..caller_body.basic_blocks.next_index(); + + Ok(new_blocks) + } + + fn check_mir_is_available( + &self, + caller_body: &Body<'tcx>, + callee: &Instance<'tcx>, + ) -> Result<(), &'static str> { + let caller_def_id = caller_body.source.def_id(); + let callee_def_id = callee.def_id(); + if callee_def_id == caller_def_id { + return Err("self-recursion"); + } + + match callee.def { + InstanceDef::Item(_) => { + // If there is no MIR available (either because it was not in metadata or + // because it has no MIR because it's an extern function), then the inliner + // won't cause cycles on this. + if !self.tcx.is_mir_available(callee_def_id) { + return Err("item MIR unavailable"); + } + } + // These have no own callable MIR. + InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => { + return Err("instance without MIR (intrinsic / virtual)"); + } + // This cannot result in an immediate cycle since the callee MIR is a shim, which does + // not get any optimizations run on it. Any subsequent inlining may cause cycles, but we + // do not need to catch this here, we can wait until the inliner decides to continue + // inlining a second time. + InstanceDef::VTableShim(_) + | InstanceDef::ReifyShim(_) + | InstanceDef::FnPtrShim(..) + | InstanceDef::ClosureOnceShim { .. } + | InstanceDef::DropGlue(..) + | InstanceDef::CloneShim(..) + | InstanceDef::ThreadLocalShim(..) + | InstanceDef::FnPtrAddrShim(..) => return Ok(()), + } + + if self.tcx.is_constructor(callee_def_id) { + trace!("constructors always have MIR"); + // Constructor functions cannot cause a query cycle. + return Ok(()); + } + + if callee_def_id.is_local() { + // Avoid a cycle here by only using `instance_mir` only if we have + // a lower `DefPathHash` than the callee. This ensures that the callee will + // not inline us. This trick even works with incremental compilation, + // since `DefPathHash` is stable. + if self.tcx.def_path_hash(caller_def_id).local_hash() + < self.tcx.def_path_hash(callee_def_id).local_hash() + { + return Ok(()); + } + + // If we know for sure that the function we're calling will itself try to + // call us, then we avoid inlining that function. + if self.tcx.mir_callgraph_reachable((*callee, caller_def_id.expect_local())) { + return Err("caller might be reachable from callee (query cycle avoidance)"); + } + + Ok(()) + } else { + // This cannot result in an immediate cycle since the callee MIR is from another crate + // and is already optimized. Any subsequent inlining may cause cycles, but we do + // not need to catch this here, we can wait until the inliner decides to continue + // inlining a second time. + trace!("functions from other crates always have MIR"); + Ok(()) + } + } + + fn resolve_callsite( + &self, + caller_body: &Body<'tcx>, + bb: BasicBlock, + bb_data: &BasicBlockData<'tcx>, + ) -> Option<CallSite<'tcx>> { + // Only consider direct calls to functions + let terminator = bb_data.terminator(); + if let TerminatorKind::Call { ref func, fn_span, .. } = terminator.kind { + let func_ty = func.ty(caller_body, self.tcx); + if let ty::FnDef(def_id, args) = *func_ty.kind() { + // To resolve an instance its args have to be fully normalized. + let args = self.tcx.try_normalize_erasing_regions(self.param_env, args).ok()?; + let callee = + Instance::resolve(self.tcx, self.param_env, def_id, args).ok().flatten()?; + + if let InstanceDef::Virtual(..) | InstanceDef::Intrinsic(_) = callee.def { + return None; + } + + if self.history.contains(&callee.def_id()) { + return None; + } + + let fn_sig = self.tcx.fn_sig(def_id).instantiate(self.tcx, args); + let source_info = SourceInfo { span: fn_span, ..terminator.source_info }; + + return Some(CallSite { callee, fn_sig, block: bb, source_info }); + } + } + + None + } + + /// Returns an error if inlining is not possible based on codegen attributes alone. A success + /// indicates that inlining decision should be based on other criteria. + fn check_codegen_attributes( + &self, + callsite: &CallSite<'tcx>, + callee_attrs: &CodegenFnAttrs, + cross_crate_inlinable: bool, + ) -> Result<(), &'static str> { + if let InlineAttr::Never = callee_attrs.inline { + return Err("never inline hint"); + } + + // Reachability pass defines which functions are eligible for inlining. Generally inlining + // other functions is incorrect because they could reference symbols that aren't exported. + let is_generic = callsite + .callee + .args + .non_erasable_generics(self.tcx, callsite.callee.def_id()) + .next() + .is_some(); + if !is_generic && !cross_crate_inlinable { + return Err("not exported"); + } + + if callsite.fn_sig.c_variadic() { + return Err("C variadic"); + } + + if callee_attrs.flags.contains(CodegenFnAttrFlags::COLD) { + return Err("cold"); + } + + if callee_attrs.no_sanitize != self.codegen_fn_attrs.no_sanitize { + return Err("incompatible sanitizer set"); + } + + // Two functions are compatible if the callee has no attribute (meaning + // that it's codegen agnostic), or sets an attribute that is identical + // to this function's attribute. + if callee_attrs.instruction_set.is_some() + && callee_attrs.instruction_set != self.codegen_fn_attrs.instruction_set + { + return Err("incompatible instruction set"); + } + + if callee_attrs.target_features != self.codegen_fn_attrs.target_features { + // In general it is not correct to inline a callee with target features that are a + // subset of the caller. This is because the callee might contain calls, and the ABI of + // those calls depends on the target features of the surrounding function. By moving a + // `Call` terminator from one MIR body to another with more target features, we might + // change the ABI of that call! + return Err("incompatible target features"); + } + + Ok(()) + } + + /// Returns inlining decision that is based on the examination of callee MIR body. + /// Assumes that codegen attributes have been checked for compatibility already. + #[instrument(level = "debug", skip(self, callee_body))] + fn check_mir_body( + &self, + callsite: &CallSite<'tcx>, + callee_body: &Body<'tcx>, + callee_attrs: &CodegenFnAttrs, + cross_crate_inlinable: bool, + ) -> Result<(), &'static str> { + let tcx = self.tcx; + + let mut threshold = if cross_crate_inlinable { + self.tcx.sess.opts.unstable_opts.inline_mir_hint_threshold.unwrap_or(100) + } else { + self.tcx.sess.opts.unstable_opts.inline_mir_threshold.unwrap_or(50) + }; + + // Give a bonus functions with a small number of blocks, + // We normally have two or three blocks for even + // very small functions. + if callee_body.basic_blocks.len() <= 3 { + threshold += threshold / 4; + } + debug!(" final inline threshold = {}", threshold); + + // FIXME: Give a bonus to functions with only a single caller + + let mut checker = + CostChecker::new(self.tcx, self.param_env, Some(callsite.callee), callee_body); + + // Traverse the MIR manually so we can account for the effects of inlining on the CFG. + let mut work_list = vec![START_BLOCK]; + let mut visited = BitSet::new_empty(callee_body.basic_blocks.len()); + while let Some(bb) = work_list.pop() { + if !visited.insert(bb.index()) { + continue; + } + + let blk = &callee_body.basic_blocks[bb]; + checker.visit_basic_block_data(bb, blk); + + let term = blk.terminator(); + if let TerminatorKind::Drop { ref place, target, unwind, replace: _ } = term.kind { + work_list.push(target); + + // If the place doesn't actually need dropping, treat it like a regular goto. + let ty = callsite.callee.instantiate_mir( + self.tcx, + ty::EarlyBinder::bind(&place.ty(callee_body, tcx).ty), + ); + if ty.needs_drop(tcx, self.param_env) + && let UnwindAction::Cleanup(unwind) = unwind + { + work_list.push(unwind); + } + } else if callee_attrs.instruction_set != self.codegen_fn_attrs.instruction_set + && matches!(term.kind, TerminatorKind::InlineAsm { .. }) + { + // During the attribute checking stage we allow a callee with no + // instruction_set assigned to count as compatible with a function that does + // assign one. However, during this stage we require an exact match when any + // inline-asm is detected. LLVM will still possibly do an inline later on + // if the no-attribute function ends up with the same instruction set anyway. + return Err("Cannot move inline-asm across instruction sets"); + } else { + work_list.extend(term.successors()) + } + } + + // N.B. We still apply our cost threshold to #[inline(always)] functions. + // That attribute is often applied to very large functions that exceed LLVM's (very + // generous) inlining threshold. Such functions are very poor MIR inlining candidates. + // Always inlining #[inline(always)] functions in MIR, on net, slows down the compiler. + let cost = checker.cost(); + if cost <= threshold { + debug!("INLINING {:?} [cost={} <= threshold={}]", callsite, cost, threshold); + Ok(()) + } else { + debug!("NOT inlining {:?} [cost={} > threshold={}]", callsite, cost, threshold); + Err("cost above threshold") + } + } + + fn inline_call( + &self, + caller_body: &mut Body<'tcx>, + callsite: &CallSite<'tcx>, + mut callee_body: Body<'tcx>, + ) { + let terminator = caller_body[callsite.block].terminator.take().unwrap(); + let TerminatorKind::Call { args, destination, unwind, target, .. } = terminator.kind else { + bug!("unexpected terminator kind {:?}", terminator.kind); + }; + + let return_block = if let Some(block) = target { + // Prepare a new block for code that should execute when call returns. We don't use + // target block directly since it might have other predecessors. + let mut data = BasicBlockData::new(Some(Terminator { + source_info: terminator.source_info, + kind: TerminatorKind::Goto { target: block }, + })); + data.is_cleanup = caller_body[block].is_cleanup; + Some(caller_body.basic_blocks_mut().push(data)) + } else { + None + }; + + // If the call is something like `a[*i] = f(i)`, where + // `i : &mut usize`, then just duplicating the `a[*i]` + // Place could result in two different locations if `f` + // writes to `i`. To prevent this we need to create a temporary + // borrow of the place and pass the destination as `*temp` instead. + fn dest_needs_borrow(place: Place<'_>) -> bool { + for elem in place.projection.iter() { + match elem { + ProjectionElem::Deref | ProjectionElem::Index(_) => return true, + _ => {} + } + } + + false + } + + let dest = if dest_needs_borrow(destination) { + trace!("creating temp for return destination"); + let dest = Rvalue::Ref( + self.tcx.lifetimes.re_erased, + BorrowKind::Mut { kind: MutBorrowKind::Default }, + destination, + ); + let dest_ty = dest.ty(caller_body, self.tcx); + let temp = + Place::from(self.new_call_temp(caller_body, &callsite, dest_ty, return_block)); + caller_body[callsite.block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::Assign(Box::new((temp, dest))), + }); + self.tcx.mk_place_deref(temp) + } else { + destination + }; + + // Always create a local to hold the destination, as `RETURN_PLACE` may appear + // where a full `Place` is not allowed. + let (remap_destination, destination_local) = if let Some(d) = dest.as_local() { + (false, d) + } else { + ( + true, + self.new_call_temp( + caller_body, + &callsite, + destination.ty(caller_body, self.tcx).ty, + return_block, + ), + ) + }; + + // Copy the arguments if needed. + let args: Vec<_> = + self.make_call_args(args, &callsite, caller_body, &callee_body, return_block); + + let mut integrator = Integrator { + args: &args, + new_locals: Local::new(caller_body.local_decls.len()).., + new_scopes: SourceScope::new(caller_body.source_scopes.len()).., + new_blocks: BasicBlock::new(caller_body.basic_blocks.len()).., + destination: destination_local, + callsite_scope: caller_body.source_scopes[callsite.source_info.scope].clone(), + callsite, + cleanup_block: unwind, + in_cleanup_block: false, + return_block, + tcx: self.tcx, + always_live_locals: BitSet::new_filled(callee_body.local_decls.len()), + }; + + // Map all `Local`s, `SourceScope`s and `BasicBlock`s to new ones + // (or existing ones, in a few special cases) in the caller. + integrator.visit_body(&mut callee_body); + + // If there are any locals without storage markers, give them storage only for the + // duration of the call. + for local in callee_body.vars_and_temps_iter() { + if integrator.always_live_locals.contains(local) { + let new_local = integrator.map_local(local); + caller_body[callsite.block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::StorageLive(new_local), + }); + } + } + if let Some(block) = return_block { + // To avoid repeated O(n) insert, push any new statements to the end and rotate + // the slice once. + let mut n = 0; + if remap_destination { + caller_body[block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::Assign(Box::new(( + dest, + Rvalue::Use(Operand::Move(destination_local.into())), + ))), + }); + n += 1; + } + for local in callee_body.vars_and_temps_iter().rev() { + if integrator.always_live_locals.contains(local) { + let new_local = integrator.map_local(local); + caller_body[block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::StorageDead(new_local), + }); + n += 1; + } + } + caller_body[block].statements.rotate_right(n); + } + + // Insert all of the (mapped) parts of the callee body into the caller. + caller_body.local_decls.extend(callee_body.drain_vars_and_temps()); + caller_body.source_scopes.extend(&mut callee_body.source_scopes.drain(..)); + caller_body.var_debug_info.append(&mut callee_body.var_debug_info); + caller_body.basic_blocks_mut().extend(callee_body.basic_blocks_mut().drain(..)); + + caller_body[callsite.block].terminator = Some(Terminator { + source_info: callsite.source_info, + kind: TerminatorKind::Goto { target: integrator.map_block(START_BLOCK) }, + }); + + // Copy only unevaluated constants from the callee_body into the caller_body. + // Although we are only pushing `ConstKind::Unevaluated` consts to + // `required_consts`, here we may not only have `ConstKind::Unevaluated` + // because we are calling `instantiate_and_normalize_erasing_regions`. + caller_body.required_consts.extend(callee_body.required_consts.iter().copied().filter( + |&ct| match ct.const_ { + Const::Ty(_) => { + bug!("should never encounter ty::UnevaluatedConst in `required_consts`") + } + Const::Val(..) | Const::Unevaluated(..) => true, + }, + )); + } + + fn make_call_args( + &self, + args: Vec<Operand<'tcx>>, + callsite: &CallSite<'tcx>, + caller_body: &mut Body<'tcx>, + callee_body: &Body<'tcx>, + return_block: Option<BasicBlock>, + ) -> Vec<Local> { + let tcx = self.tcx; + + // There is a bit of a mismatch between the *caller* of a closure and the *callee*. + // The caller provides the arguments wrapped up in a tuple: + // + // tuple_tmp = (a, b, c) + // Fn::call(closure_ref, tuple_tmp) + // + // meanwhile the closure body expects the arguments (here, `a`, `b`, and `c`) + // as distinct arguments. (This is the "rust-call" ABI hack.) Normally, codegen has + // the job of unpacking this tuple. But here, we are codegen. =) So we want to create + // a vector like + // + // [closure_ref, tuple_tmp.0, tuple_tmp.1, tuple_tmp.2] + // + // Except for one tiny wrinkle: we don't actually want `tuple_tmp.0`. It's more convenient + // if we "spill" that into *another* temporary, so that we can map the argument + // variable in the callee MIR directly to an argument variable on our side. + // So we introduce temporaries like: + // + // tmp0 = tuple_tmp.0 + // tmp1 = tuple_tmp.1 + // tmp2 = tuple_tmp.2 + // + // and the vector is `[closure_ref, tmp0, tmp1, tmp2]`. + if callsite.fn_sig.abi() == Abi::RustCall && callee_body.spread_arg.is_none() { + let mut args = args.into_iter(); + let self_ = self.create_temp_if_necessary( + args.next().unwrap(), + callsite, + caller_body, + return_block, + ); + let tuple = self.create_temp_if_necessary( + args.next().unwrap(), + callsite, + caller_body, + return_block, + ); + assert!(args.next().is_none()); + + let tuple = Place::from(tuple); + let ty::Tuple(tuple_tys) = tuple.ty(caller_body, tcx).ty.kind() else { + bug!("Closure arguments are not passed as a tuple"); + }; + + // The `closure_ref` in our example above. + let closure_ref_arg = iter::once(self_); + + // The `tmp0`, `tmp1`, and `tmp2` in our example above. + let tuple_tmp_args = tuple_tys.iter().enumerate().map(|(i, ty)| { + // This is e.g., `tuple_tmp.0` in our example above. + let tuple_field = Operand::Move(tcx.mk_place_field(tuple, FieldIdx::new(i), ty)); + + // Spill to a local to make e.g., `tmp0`. + self.create_temp_if_necessary(tuple_field, callsite, caller_body, return_block) + }); + + closure_ref_arg.chain(tuple_tmp_args).collect() + } else { + args.into_iter() + .map(|a| self.create_temp_if_necessary(a, callsite, caller_body, return_block)) + .collect() + } + } + + /// If `arg` is already a temporary, returns it. Otherwise, introduces a fresh + /// temporary `T` and an instruction `T = arg`, and returns `T`. + fn create_temp_if_necessary( + &self, + arg: Operand<'tcx>, + callsite: &CallSite<'tcx>, + caller_body: &mut Body<'tcx>, + return_block: Option<BasicBlock>, + ) -> Local { + // Reuse the operand if it is a moved temporary. + if let Operand::Move(place) = &arg + && let Some(local) = place.as_local() + && caller_body.local_kind(local) == LocalKind::Temp + { + return local; + } + + // Otherwise, create a temporary for the argument. + trace!("creating temp for argument {:?}", arg); + let arg_ty = arg.ty(caller_body, self.tcx); + let local = self.new_call_temp(caller_body, callsite, arg_ty, return_block); + caller_body[callsite.block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::Assign(Box::new((Place::from(local), Rvalue::Use(arg)))), + }); + local + } + + /// Introduces a new temporary into the caller body that is live for the duration of the call. + fn new_call_temp( + &self, + caller_body: &mut Body<'tcx>, + callsite: &CallSite<'tcx>, + ty: Ty<'tcx>, + return_block: Option<BasicBlock>, + ) -> Local { + let local = caller_body.local_decls.push(LocalDecl::new(ty, callsite.source_info.span)); + + caller_body[callsite.block].statements.push(Statement { + source_info: callsite.source_info, + kind: StatementKind::StorageLive(local), + }); + + if let Some(block) = return_block { + caller_body[block].statements.insert( + 0, + Statement { + source_info: callsite.source_info, + kind: StatementKind::StorageDead(local), + }, + ); + } + + local + } +} + +/** + * Integrator. + * + * Integrates blocks from the callee function into the calling function. + * Updates block indices, references to locals and other control flow + * stuff. +*/ +struct Integrator<'a, 'tcx> { + args: &'a [Local], + new_locals: RangeFrom<Local>, + new_scopes: RangeFrom<SourceScope>, + new_blocks: RangeFrom<BasicBlock>, + destination: Local, + callsite_scope: SourceScopeData<'tcx>, + callsite: &'a CallSite<'tcx>, + cleanup_block: UnwindAction, + in_cleanup_block: bool, + return_block: Option<BasicBlock>, + tcx: TyCtxt<'tcx>, + always_live_locals: BitSet<Local>, +} + +impl Integrator<'_, '_> { + fn map_local(&self, local: Local) -> Local { + let new = if local == RETURN_PLACE { + self.destination + } else { + let idx = local.index() - 1; + if idx < self.args.len() { + self.args[idx] + } else { + Local::new(self.new_locals.start.index() + (idx - self.args.len())) + } + }; + trace!("mapping local `{:?}` to `{:?}`", local, new); + new + } + + fn map_scope(&self, scope: SourceScope) -> SourceScope { + let new = SourceScope::new(self.new_scopes.start.index() + scope.index()); + trace!("mapping scope `{:?}` to `{:?}`", scope, new); + new + } + + fn map_block(&self, block: BasicBlock) -> BasicBlock { + let new = BasicBlock::new(self.new_blocks.start.index() + block.index()); + trace!("mapping block `{:?}` to `{:?}`", block, new); + new + } + + fn map_unwind(&self, unwind: UnwindAction) -> UnwindAction { + if self.in_cleanup_block { + match unwind { + UnwindAction::Cleanup(_) | UnwindAction::Continue => { + bug!("cleanup on cleanup block"); + } + UnwindAction::Unreachable | UnwindAction::Terminate(_) => return unwind, + } + } + + match unwind { + UnwindAction::Unreachable | UnwindAction::Terminate(_) => unwind, + UnwindAction::Cleanup(target) => UnwindAction::Cleanup(self.map_block(target)), + // Add an unwind edge to the original call's cleanup block + UnwindAction::Continue => self.cleanup_block, + } + } +} + +impl<'tcx> MutVisitor<'tcx> for Integrator<'_, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, local: &mut Local, _ctxt: PlaceContext, _location: Location) { + *local = self.map_local(*local); + } + + fn visit_source_scope_data(&mut self, scope_data: &mut SourceScopeData<'tcx>) { + self.super_source_scope_data(scope_data); + if scope_data.parent_scope.is_none() { + // Attach the outermost callee scope as a child of the callsite + // scope, via the `parent_scope` and `inlined_parent_scope` chains. + scope_data.parent_scope = Some(self.callsite.source_info.scope); + assert_eq!(scope_data.inlined_parent_scope, None); + scope_data.inlined_parent_scope = if self.callsite_scope.inlined.is_some() { + Some(self.callsite.source_info.scope) + } else { + self.callsite_scope.inlined_parent_scope + }; + + // Mark the outermost callee scope as an inlined one. + assert_eq!(scope_data.inlined, None); + scope_data.inlined = Some((self.callsite.callee, self.callsite.source_info.span)); + } else if scope_data.inlined_parent_scope.is_none() { + // Make it easy to find the scope with `inlined` set above. + scope_data.inlined_parent_scope = Some(self.map_scope(OUTERMOST_SOURCE_SCOPE)); + } + } + + fn visit_source_scope(&mut self, scope: &mut SourceScope) { + *scope = self.map_scope(*scope); + } + + fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) { + self.in_cleanup_block = data.is_cleanup; + self.super_basic_block_data(block, data); + self.in_cleanup_block = false; + } + + fn visit_retag(&mut self, kind: &mut RetagKind, place: &mut Place<'tcx>, loc: Location) { + self.super_retag(kind, place, loc); + + // We have to patch all inlined retags to be aware that they are no longer + // happening on function entry. + if *kind == RetagKind::FnEntry { + *kind = RetagKind::Default; + } + } + + fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) { + if let StatementKind::StorageLive(local) | StatementKind::StorageDead(local) = + statement.kind + { + self.always_live_locals.remove(local); + } + self.super_statement(statement, location); + } + + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, loc: Location) { + // Don't try to modify the implicit `_0` access on return (`return` terminators are + // replaced down below anyways). + if !matches!(terminator.kind, TerminatorKind::Return) { + self.super_terminator(terminator, loc); + } + + match terminator.kind { + TerminatorKind::CoroutineDrop | TerminatorKind::Yield { .. } => bug!(), + TerminatorKind::Goto { ref mut target } => { + *target = self.map_block(*target); + } + TerminatorKind::SwitchInt { ref mut targets, .. } => { + for tgt in targets.all_targets_mut() { + *tgt = self.map_block(*tgt); + } + } + TerminatorKind::Drop { ref mut target, ref mut unwind, .. } => { + *target = self.map_block(*target); + *unwind = self.map_unwind(*unwind); + } + TerminatorKind::Call { ref mut target, ref mut unwind, .. } => { + if let Some(ref mut tgt) = *target { + *tgt = self.map_block(*tgt); + } + *unwind = self.map_unwind(*unwind); + } + TerminatorKind::Assert { ref mut target, ref mut unwind, .. } => { + *target = self.map_block(*target); + *unwind = self.map_unwind(*unwind); + } + TerminatorKind::Return => { + terminator.kind = if let Some(tgt) = self.return_block { + TerminatorKind::Goto { target: tgt } + } else { + TerminatorKind::Unreachable + } + } + TerminatorKind::UnwindResume => { + terminator.kind = match self.cleanup_block { + UnwindAction::Cleanup(tgt) => TerminatorKind::Goto { target: tgt }, + UnwindAction::Continue => TerminatorKind::UnwindResume, + UnwindAction::Unreachable => TerminatorKind::Unreachable, + UnwindAction::Terminate(reason) => TerminatorKind::UnwindTerminate(reason), + }; + } + TerminatorKind::UnwindTerminate(_) => {} + TerminatorKind::Unreachable => {} + TerminatorKind::FalseEdge { ref mut real_target, ref mut imaginary_target } => { + *real_target = self.map_block(*real_target); + *imaginary_target = self.map_block(*imaginary_target); + } + TerminatorKind::FalseUnwind { real_target: _, unwind: _ } => + // see the ordering of passes in the optimized_mir query. + { + bug!("False unwinds should have been removed before inlining") + } + TerminatorKind::InlineAsm { ref mut destination, ref mut unwind, .. } => { + if let Some(ref mut tgt) = *destination { + *tgt = self.map_block(*tgt); + } + *unwind = self.map_unwind(*unwind); + } + } + } +} + +#[instrument(skip(tcx), level = "debug")] +fn try_instance_mir<'tcx>( + tcx: TyCtxt<'tcx>, + instance: InstanceDef<'tcx>, +) -> Result<&'tcx Body<'tcx>, &'static str> { + match instance { + ty::InstanceDef::DropGlue(_, Some(ty)) => match ty.kind() { + ty::Adt(def, args) => { + let fields = def.all_fields(); + for field in fields { + let field_ty = field.ty(tcx, args); + if field_ty.has_param() && field_ty.has_projections() { + return Err("cannot build drop shim for polymorphic type"); + } + } + + Ok(tcx.instance_mir(instance)) + } + _ => Ok(tcx.instance_mir(instance)), + }, + _ => Ok(tcx.instance_mir(instance)), + } +} diff --git a/compiler/rustc_mir_transform/src/inline/cycle.rs b/compiler/rustc_mir_transform/src/inline/cycle.rs new file mode 100644 index 00000000000..d30e0bad813 --- /dev/null +++ b/compiler/rustc_mir_transform/src/inline/cycle.rs @@ -0,0 +1,175 @@ +use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet}; +use rustc_data_structures::stack::ensure_sufficient_stack; +use rustc_hir::def_id::{DefId, LocalDefId}; +use rustc_middle::mir::TerminatorKind; +use rustc_middle::ty::TypeVisitableExt; +use rustc_middle::ty::{self, GenericArgsRef, InstanceDef, TyCtxt}; +use rustc_session::Limit; + +// FIXME: check whether it is cheaper to precompute the entire call graph instead of invoking +// this query ridiculously often. +#[instrument(level = "debug", skip(tcx, root, target))] +pub(crate) fn mir_callgraph_reachable<'tcx>( + tcx: TyCtxt<'tcx>, + (root, target): (ty::Instance<'tcx>, LocalDefId), +) -> bool { + trace!(%root, target = %tcx.def_path_str(target)); + let param_env = tcx.param_env_reveal_all_normalized(target); + assert_ne!( + root.def_id().expect_local(), + target, + "you should not call `mir_callgraph_reachable` on immediate self recursion" + ); + assert!( + matches!(root.def, InstanceDef::Item(_)), + "you should not call `mir_callgraph_reachable` on shims" + ); + assert!( + !tcx.is_constructor(root.def_id()), + "you should not call `mir_callgraph_reachable` on enum/struct constructor functions" + ); + #[instrument( + level = "debug", + skip(tcx, param_env, target, stack, seen, recursion_limiter, caller, recursion_limit) + )] + fn process<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + caller: ty::Instance<'tcx>, + target: LocalDefId, + stack: &mut Vec<ty::Instance<'tcx>>, + seen: &mut FxHashSet<ty::Instance<'tcx>>, + recursion_limiter: &mut FxHashMap<DefId, usize>, + recursion_limit: Limit, + ) -> bool { + trace!(%caller); + for &(callee, args) in tcx.mir_inliner_callees(caller.def) { + let Ok(args) = caller.try_instantiate_mir_and_normalize_erasing_regions( + tcx, + param_env, + ty::EarlyBinder::bind(args), + ) else { + trace!(?caller, ?param_env, ?args, "cannot normalize, skipping"); + continue; + }; + let Ok(Some(callee)) = ty::Instance::resolve(tcx, param_env, callee, args) else { + trace!(?callee, "cannot resolve, skipping"); + continue; + }; + + // Found a path. + if callee.def_id() == target.to_def_id() { + return true; + } + + if tcx.is_constructor(callee.def_id()) { + trace!("constructors always have MIR"); + // Constructor functions cannot cause a query cycle. + continue; + } + + match callee.def { + InstanceDef::Item(_) => { + // If there is no MIR available (either because it was not in metadata or + // because it has no MIR because it's an extern function), then the inliner + // won't cause cycles on this. + if !tcx.is_mir_available(callee.def_id()) { + trace!(?callee, "no mir available, skipping"); + continue; + } + } + // These have no own callable MIR. + InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => continue, + // These have MIR and if that MIR is inlined, substituted and then inlining is run + // again, a function item can end up getting inlined. Thus we'll be able to cause + // a cycle that way + InstanceDef::VTableShim(_) + | InstanceDef::ReifyShim(_) + | InstanceDef::FnPtrShim(..) + | InstanceDef::ClosureOnceShim { .. } + | InstanceDef::ThreadLocalShim { .. } + | InstanceDef::CloneShim(..) => {} + + // This shim does not call any other functions, thus there can be no recursion. + InstanceDef::FnPtrAddrShim(..) => continue, + InstanceDef::DropGlue(..) => { + // FIXME: A not fully substituted drop shim can cause ICEs if one attempts to + // have its MIR built. Likely oli-obk just screwed up the `ParamEnv`s, so this + // needs some more analysis. + if callee.has_param() { + continue; + } + } + } + + if seen.insert(callee) { + let recursion = recursion_limiter.entry(callee.def_id()).or_default(); + trace!(?callee, recursion = *recursion); + if recursion_limit.value_within_limit(*recursion) { + *recursion += 1; + stack.push(callee); + let found_recursion = ensure_sufficient_stack(|| { + process( + tcx, + param_env, + callee, + target, + stack, + seen, + recursion_limiter, + recursion_limit, + ) + }); + if found_recursion { + return true; + } + stack.pop(); + } else { + // Pessimistically assume that there could be recursion. + return true; + } + } + } + false + } + process( + tcx, + param_env, + root, + target, + &mut Vec::new(), + &mut FxHashSet::default(), + &mut FxHashMap::default(), + tcx.recursion_limit(), + ) +} + +pub(crate) fn mir_inliner_callees<'tcx>( + tcx: TyCtxt<'tcx>, + instance: ty::InstanceDef<'tcx>, +) -> &'tcx [(DefId, GenericArgsRef<'tcx>)] { + let steal; + let guard; + let body = match (instance, instance.def_id().as_local()) { + (InstanceDef::Item(_), Some(def_id)) => { + steal = tcx.mir_promoted(def_id).0; + guard = steal.borrow(); + &*guard + } + // Functions from other crates and MIR shims + _ => tcx.instance_mir(instance), + }; + let mut calls = FxIndexSet::default(); + for bb_data in body.basic_blocks.iter() { + let terminator = bb_data.terminator(); + if let TerminatorKind::Call { func, .. } = &terminator.kind { + let ty = func.ty(&body.local_decls, tcx); + let call = match ty.kind() { + ty::FnDef(def_id, args) => (*def_id, *args), + _ => continue, + }; + calls.insert(call); + } + } + tcx.arena.alloc_from_iter(calls.iter().copied()) +} diff --git a/compiler/rustc_mir_transform/src/instsimplify.rs b/compiler/rustc_mir_transform/src/instsimplify.rs new file mode 100644 index 00000000000..4f0f63d22a4 --- /dev/null +++ b/compiler/rustc_mir_transform/src/instsimplify.rs @@ -0,0 +1,312 @@ +//! Performs various peephole optimizations. + +use crate::simplify::simplify_duplicate_switch_targets; +use rustc_middle::mir::*; +use rustc_middle::ty::layout::ValidityRequirement; +use rustc_middle::ty::{self, GenericArgsRef, ParamEnv, Ty, TyCtxt}; +use rustc_span::symbol::Symbol; +use rustc_target::abi::FieldIdx; + +pub struct InstSimplify; + +impl<'tcx> MirPass<'tcx> for InstSimplify { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let ctx = InstSimplifyContext { + tcx, + local_decls: &body.local_decls, + param_env: tcx.param_env_reveal_all_normalized(body.source.def_id()), + }; + for block in body.basic_blocks.as_mut() { + for statement in block.statements.iter_mut() { + match statement.kind { + StatementKind::Assign(box (_place, ref mut rvalue)) => { + ctx.simplify_bool_cmp(&statement.source_info, rvalue); + ctx.simplify_ref_deref(&statement.source_info, rvalue); + ctx.simplify_len(&statement.source_info, rvalue); + ctx.simplify_cast(&statement.source_info, rvalue); + } + _ => {} + } + } + + ctx.simplify_primitive_clone(block.terminator.as_mut().unwrap(), &mut block.statements); + ctx.simplify_intrinsic_assert( + block.terminator.as_mut().unwrap(), + &mut block.statements, + ); + simplify_duplicate_switch_targets(block.terminator.as_mut().unwrap()); + } + } +} + +struct InstSimplifyContext<'tcx, 'a> { + tcx: TyCtxt<'tcx>, + local_decls: &'a LocalDecls<'tcx>, + param_env: ParamEnv<'tcx>, +} + +impl<'tcx> InstSimplifyContext<'tcx, '_> { + fn should_simplify(&self, source_info: &SourceInfo, rvalue: &Rvalue<'tcx>) -> bool { + self.tcx.consider_optimizing(|| { + format!("InstSimplify - Rvalue: {rvalue:?} SourceInfo: {source_info:?}") + }) + } + + /// Transform boolean comparisons into logical operations. + fn simplify_bool_cmp(&self, source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) { + match rvalue { + Rvalue::BinaryOp(op @ (BinOp::Eq | BinOp::Ne), box (a, b)) => { + let new = match (op, self.try_eval_bool(a), self.try_eval_bool(b)) { + // Transform "Eq(a, true)" ==> "a" + (BinOp::Eq, _, Some(true)) => Some(Rvalue::Use(a.clone())), + + // Transform "Ne(a, false)" ==> "a" + (BinOp::Ne, _, Some(false)) => Some(Rvalue::Use(a.clone())), + + // Transform "Eq(true, b)" ==> "b" + (BinOp::Eq, Some(true), _) => Some(Rvalue::Use(b.clone())), + + // Transform "Ne(false, b)" ==> "b" + (BinOp::Ne, Some(false), _) => Some(Rvalue::Use(b.clone())), + + // Transform "Eq(false, b)" ==> "Not(b)" + (BinOp::Eq, Some(false), _) => Some(Rvalue::UnaryOp(UnOp::Not, b.clone())), + + // Transform "Ne(true, b)" ==> "Not(b)" + (BinOp::Ne, Some(true), _) => Some(Rvalue::UnaryOp(UnOp::Not, b.clone())), + + // Transform "Eq(a, false)" ==> "Not(a)" + (BinOp::Eq, _, Some(false)) => Some(Rvalue::UnaryOp(UnOp::Not, a.clone())), + + // Transform "Ne(a, true)" ==> "Not(a)" + (BinOp::Ne, _, Some(true)) => Some(Rvalue::UnaryOp(UnOp::Not, a.clone())), + + _ => None, + }; + + if let Some(new) = new + && self.should_simplify(source_info, rvalue) + { + *rvalue = new; + } + } + + _ => {} + } + } + + fn try_eval_bool(&self, a: &Operand<'_>) -> Option<bool> { + let a = a.constant()?; + if a.const_.ty().is_bool() { a.const_.try_to_bool() } else { None } + } + + /// Transform "&(*a)" ==> "a". + fn simplify_ref_deref(&self, source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) { + if let Rvalue::Ref(_, _, place) = rvalue { + if let Some((base, ProjectionElem::Deref)) = place.as_ref().last_projection() { + if rvalue.ty(self.local_decls, self.tcx) != base.ty(self.local_decls, self.tcx).ty { + return; + } + + if !self.should_simplify(source_info, rvalue) { + return; + } + + *rvalue = Rvalue::Use(Operand::Copy(Place { + local: base.local, + projection: self.tcx.mk_place_elems(base.projection), + })); + } + } + } + + /// Transform "Len([_; N])" ==> "N". + fn simplify_len(&self, source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) { + if let Rvalue::Len(ref place) = *rvalue { + let place_ty = place.ty(self.local_decls, self.tcx).ty; + if let ty::Array(_, len) = *place_ty.kind() { + if !self.should_simplify(source_info, rvalue) { + return; + } + + let const_ = Const::from_ty_const(len, self.tcx); + let constant = ConstOperand { span: source_info.span, const_, user_ty: None }; + *rvalue = Rvalue::Use(Operand::Constant(Box::new(constant))); + } + } + } + + fn simplify_cast(&self, _source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) { + if let Rvalue::Cast(kind, operand, cast_ty) = rvalue { + let operand_ty = operand.ty(self.local_decls, self.tcx); + if operand_ty == *cast_ty { + *rvalue = Rvalue::Use(operand.clone()); + } else if *kind == CastKind::Transmute { + // Transmuting an integer to another integer is just a signedness cast + if let (ty::Int(int), ty::Uint(uint)) | (ty::Uint(uint), ty::Int(int)) = + (operand_ty.kind(), cast_ty.kind()) + && int.bit_width() == uint.bit_width() + { + // The width check isn't strictly necessary, as different widths + // are UB and thus we'd be allowed to turn it into a cast anyway. + // But let's keep the UB around for codegen to exploit later. + // (If `CastKind::Transmute` ever becomes *not* UB for mismatched sizes, + // then the width check is necessary for big-endian correctness.) + *kind = CastKind::IntToInt; + return; + } + + // Transmuting a transparent struct/union to a field's type is a projection + if let ty::Adt(adt_def, args) = operand_ty.kind() + && adt_def.repr().transparent() + && (adt_def.is_struct() || adt_def.is_union()) + && let Some(place) = operand.place() + { + let variant = adt_def.non_enum_variant(); + for (i, field) in variant.fields.iter().enumerate() { + let field_ty = field.ty(self.tcx, args); + if field_ty == *cast_ty { + let place = place.project_deeper( + &[ProjectionElem::Field(FieldIdx::from_usize(i), *cast_ty)], + self.tcx, + ); + let operand = if operand.is_move() { + Operand::Move(place) + } else { + Operand::Copy(place) + }; + *rvalue = Rvalue::Use(operand); + return; + } + } + } + } + } + } + + fn simplify_primitive_clone( + &self, + terminator: &mut Terminator<'tcx>, + statements: &mut Vec<Statement<'tcx>>, + ) { + let TerminatorKind::Call { func, args, destination, target, .. } = &mut terminator.kind + else { + return; + }; + + // It's definitely not a clone if there are multiple arguments + if args.len() != 1 { + return; + } + + let Some(destination_block) = *target else { return }; + + // Only bother looking more if it's easy to know what we're calling + let Some((fn_def_id, fn_args)) = func.const_fn_def() else { return }; + + // Clone needs one subst, so we can cheaply rule out other stuff + if fn_args.len() != 1 { + return; + } + + // These types are easily available from locals, so check that before + // doing DefId lookups to figure out what we're actually calling. + let arg_ty = args[0].ty(self.local_decls, self.tcx); + + let ty::Ref(_region, inner_ty, Mutability::Not) = *arg_ty.kind() else { return }; + + if !inner_ty.is_trivially_pure_clone_copy() { + return; + } + + let trait_def_id = self.tcx.trait_of_item(fn_def_id); + if trait_def_id.is_none() || trait_def_id != self.tcx.lang_items().clone_trait() { + return; + } + + if !self.tcx.consider_optimizing(|| { + format!( + "InstSimplify - Call: {:?} SourceInfo: {:?}", + (fn_def_id, fn_args), + terminator.source_info + ) + }) { + return; + } + + let Some(arg_place) = args.pop().unwrap().place() else { return }; + + statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::Use(Operand::Copy( + arg_place.project_deeper(&[ProjectionElem::Deref], self.tcx), + )), + ))), + }); + terminator.kind = TerminatorKind::Goto { target: destination_block }; + } + + fn simplify_intrinsic_assert( + &self, + terminator: &mut Terminator<'tcx>, + _statements: &mut Vec<Statement<'tcx>>, + ) { + let TerminatorKind::Call { func, target, .. } = &mut terminator.kind else { + return; + }; + let Some(target_block) = target else { + return; + }; + let func_ty = func.ty(self.local_decls, self.tcx); + let Some((intrinsic_name, args)) = resolve_rust_intrinsic(self.tcx, func_ty) else { + return; + }; + // The intrinsics we are interested in have one generic parameter + if args.is_empty() { + return; + } + let ty = args.type_at(0); + + let known_is_valid = intrinsic_assert_panics(self.tcx, self.param_env, ty, intrinsic_name); + match known_is_valid { + // We don't know the layout or it's not validity assertion at all, don't touch it + None => {} + Some(true) => { + // If we know the assert panics, indicate to later opts that the call diverges + *target = None; + } + Some(false) => { + // If we know the assert does not panic, turn the call into a Goto + terminator.kind = TerminatorKind::Goto { target: *target_block }; + } + } + } +} + +fn intrinsic_assert_panics<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + ty: Ty<'tcx>, + intrinsic_name: Symbol, +) -> Option<bool> { + let requirement = ValidityRequirement::from_intrinsic(intrinsic_name)?; + Some(!tcx.check_validity_requirement((requirement, param_env.and(ty))).ok()?) +} + +fn resolve_rust_intrinsic<'tcx>( + tcx: TyCtxt<'tcx>, + func_ty: Ty<'tcx>, +) -> Option<(Symbol, GenericArgsRef<'tcx>)> { + if let ty::FnDef(def_id, args) = *func_ty.kind() { + if tcx.is_intrinsic(def_id) { + return Some((tcx.item_name(def_id), args)); + } + } + None +} diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs new file mode 100644 index 00000000000..36a15f47276 --- /dev/null +++ b/compiler/rustc_mir_transform/src/jump_threading.rs @@ -0,0 +1,762 @@ +//! A jump threading optimization. +//! +//! This optimization seeks to replace join-then-switch control flow patterns by straight jumps +//! X = 0 X = 0 +//! ------------\ /-------- ------------ +//! X = 1 X----X SwitchInt(X) => X = 1 +//! ------------/ \-------- ------------ +//! +//! +//! We proceed by walking the cfg backwards starting from each `SwitchInt` terminator, +//! looking for assignments that will turn the `SwitchInt` into a simple `Goto`. +//! +//! The algorithm maintains a set of replacement conditions: +//! - `conditions[place]` contains `Condition { value, polarity: Eq, target }` +//! if assigning `value` to `place` turns the `SwitchInt` into `Goto { target }`. +//! - `conditions[place]` contains `Condition { value, polarity: Ne, target }` +//! if assigning anything different from `value` to `place` turns the `SwitchInt` +//! into `Goto { target }`. +//! +//! In this file, we denote as `place ?= value` the existence of a replacement condition +//! on `place` with given `value`, irrespective of the polarity and target of that +//! replacement condition. +//! +//! We then walk the CFG backwards transforming the set of conditions. +//! When we find a fulfilling assignment, we record a `ThreadingOpportunity`. +//! All `ThreadingOpportunity`s are applied to the body, by duplicating blocks if required. +//! +//! The optimization search can be very heavy, as it performs a DFS on MIR starting from +//! each `SwitchInt` terminator. To manage the complexity, we: +//! - bound the maximum depth by a constant `MAX_BACKTRACK`; +//! - we only traverse `Goto` terminators. +//! +//! We try to avoid creating irreducible control-flow by not threading through a loop header. +//! +//! Likewise, applying the optimisation can create a lot of new MIR, so we bound the instruction +//! cost by `MAX_COST`. + +use rustc_arena::DroplessArena; +use rustc_data_structures::fx::FxHashSet; +use rustc_index::bit_set::BitSet; +use rustc_index::IndexVec; +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt}; +use rustc_mir_dataflow::value_analysis::{Map, PlaceIndex, State, TrackElem}; + +use crate::cost_checker::CostChecker; + +pub struct JumpThreading; + +const MAX_BACKTRACK: usize = 5; +const MAX_COST: usize = 100; +const MAX_PLACES: usize = 100; + +impl<'tcx> MirPass<'tcx> for JumpThreading { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 4 + } + + #[instrument(skip_all level = "debug")] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + debug!(?def_id); + + let param_env = tcx.param_env_reveal_all_normalized(def_id); + let map = Map::new(tcx, body, Some(MAX_PLACES)); + let loop_headers = loop_headers(body); + + let arena = DroplessArena::default(); + let mut finder = TOFinder { + tcx, + param_env, + body, + arena: &arena, + map: &map, + loop_headers: &loop_headers, + opportunities: Vec::new(), + }; + + for (bb, bbdata) in body.basic_blocks.iter_enumerated() { + debug!(?bb, term = ?bbdata.terminator()); + if bbdata.is_cleanup || loop_headers.contains(bb) { + continue; + } + let Some((discr, targets)) = bbdata.terminator().kind.as_switch() else { continue }; + let Some(discr) = discr.place() else { continue }; + debug!(?discr, ?bb); + + let discr_ty = discr.ty(body, tcx).ty; + let Ok(discr_layout) = tcx.layout_of(param_env.and(discr_ty)) else { continue }; + + let Some(discr) = finder.map.find(discr.as_ref()) else { continue }; + debug!(?discr); + + let cost = CostChecker::new(tcx, param_env, None, body); + + let mut state = State::new(ConditionSet::default(), finder.map); + + let conds = if let Some((value, then, else_)) = targets.as_static_if() { + let Some(value) = ScalarInt::try_from_uint(value, discr_layout.size) else { + continue; + }; + arena.alloc_from_iter([ + Condition { value, polarity: Polarity::Eq, target: then }, + Condition { value, polarity: Polarity::Ne, target: else_ }, + ]) + } else { + arena.alloc_from_iter(targets.iter().filter_map(|(value, target)| { + let value = ScalarInt::try_from_uint(value, discr_layout.size)?; + Some(Condition { value, polarity: Polarity::Eq, target }) + })) + }; + let conds = ConditionSet(conds); + state.insert_value_idx(discr, conds, finder.map); + + finder.find_opportunity(bb, state, cost, 0); + } + + let opportunities = finder.opportunities; + debug!(?opportunities); + if opportunities.is_empty() { + return; + } + + // Verify that we do not thread through a loop header. + for to in opportunities.iter() { + assert!(to.chain.iter().all(|&block| !loop_headers.contains(block))); + } + OpportunitySet::new(body, opportunities).apply(body); + } +} + +#[derive(Debug)] +struct ThreadingOpportunity { + /// The list of `BasicBlock`s from the one that found the opportunity to the `SwitchInt`. + chain: Vec<BasicBlock>, + /// The `SwitchInt` will be replaced by `Goto { target }`. + target: BasicBlock, +} + +struct TOFinder<'tcx, 'a> { + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + body: &'a Body<'tcx>, + map: &'a Map, + loop_headers: &'a BitSet<BasicBlock>, + /// We use an arena to avoid cloning the slices when cloning `state`. + arena: &'a DroplessArena, + opportunities: Vec<ThreadingOpportunity>, +} + +/// Represent the following statement. If we can prove that the current local is equal/not-equal +/// to `value`, jump to `target`. +#[derive(Copy, Clone, Debug)] +struct Condition { + value: ScalarInt, + polarity: Polarity, + target: BasicBlock, +} + +#[derive(Copy, Clone, Debug, Eq, PartialEq)] +enum Polarity { + Ne, + Eq, +} + +impl Condition { + fn matches(&self, value: ScalarInt) -> bool { + (self.value == value) == (self.polarity == Polarity::Eq) + } + + fn inv(mut self) -> Self { + self.polarity = match self.polarity { + Polarity::Eq => Polarity::Ne, + Polarity::Ne => Polarity::Eq, + }; + self + } +} + +#[derive(Copy, Clone, Debug, Default)] +struct ConditionSet<'a>(&'a [Condition]); + +impl<'a> ConditionSet<'a> { + fn iter(self) -> impl Iterator<Item = Condition> + 'a { + self.0.iter().copied() + } + + fn iter_matches(self, value: ScalarInt) -> impl Iterator<Item = Condition> + 'a { + self.iter().filter(move |c| c.matches(value)) + } + + fn map(self, arena: &'a DroplessArena, f: impl Fn(Condition) -> Condition) -> ConditionSet<'a> { + ConditionSet(arena.alloc_from_iter(self.iter().map(f))) + } +} + +impl<'tcx, 'a> TOFinder<'tcx, 'a> { + fn is_empty(&self, state: &State<ConditionSet<'a>>) -> bool { + state.all(|cs| cs.0.is_empty()) + } + + /// Recursion entry point to find threading opportunities. + #[instrument(level = "trace", skip(self, cost), ret)] + fn find_opportunity( + &mut self, + bb: BasicBlock, + mut state: State<ConditionSet<'a>>, + mut cost: CostChecker<'_, 'tcx>, + depth: usize, + ) { + // Do not thread through loop headers. + if self.loop_headers.contains(bb) { + return; + } + + debug!(cost = ?cost.cost()); + for (statement_index, stmt) in + self.body.basic_blocks[bb].statements.iter().enumerate().rev() + { + if self.is_empty(&state) { + return; + } + + cost.visit_statement(stmt, Location { block: bb, statement_index }); + if cost.cost() > MAX_COST { + return; + } + + // Attempt to turn the `current_condition` on `lhs` into a condition on another place. + self.process_statement(bb, stmt, &mut state); + + // When a statement mutates a place, assignments to that place that happen + // above the mutation cannot fulfill a condition. + // _1 = 5 // Whatever happens here, it won't change the result of a `SwitchInt`. + // _1 = 6 + if let Some((lhs, tail)) = self.mutated_statement(stmt) { + state.flood_with_tail_elem(lhs.as_ref(), tail, self.map, ConditionSet::default()); + } + } + + if self.is_empty(&state) || depth >= MAX_BACKTRACK { + return; + } + + let last_non_rec = self.opportunities.len(); + + let predecessors = &self.body.basic_blocks.predecessors()[bb]; + if let &[pred] = &predecessors[..] + && bb != START_BLOCK + { + let term = self.body.basic_blocks[pred].terminator(); + match term.kind { + TerminatorKind::SwitchInt { ref discr, ref targets } => { + self.process_switch_int(discr, targets, bb, &mut state); + self.find_opportunity(pred, state, cost, depth + 1); + } + _ => self.recurse_through_terminator(pred, &state, &cost, depth), + } + } else { + for &pred in predecessors { + self.recurse_through_terminator(pred, &state, &cost, depth); + } + } + + let new_tos = &mut self.opportunities[last_non_rec..]; + debug!(?new_tos); + + // Try to deduplicate threading opportunities. + if new_tos.len() > 1 + && new_tos.len() == predecessors.len() + && predecessors + .iter() + .zip(new_tos.iter()) + .all(|(&pred, to)| to.chain == &[pred] && to.target == new_tos[0].target) + { + // All predecessors have a threading opportunity, and they all point to the same block. + debug!(?new_tos, "dedup"); + let first = &mut new_tos[0]; + *first = ThreadingOpportunity { chain: vec![bb], target: first.target }; + self.opportunities.truncate(last_non_rec + 1); + return; + } + + for op in self.opportunities[last_non_rec..].iter_mut() { + op.chain.push(bb); + } + } + + /// Extract the mutated place from a statement. + /// + /// This method returns the `Place` so we can flood the state in case of a partial assignment. + /// (_1 as Ok).0 = _5; + /// (_1 as Err).0 = _6; + /// We want to ensure that a `SwitchInt((_1 as Ok).0)` does not see the first assignment, as + /// the value may have been mangled by the second assignment. + /// + /// In case we assign to a discriminant, we return `Some(TrackElem::Discriminant)`, so we can + /// stop at flooding the discriminant, and preserve the variant fields. + /// (_1 as Some).0 = _6; + /// SetDiscriminant(_1, 1); + /// switchInt((_1 as Some).0) + #[instrument(level = "trace", skip(self), ret)] + fn mutated_statement( + &self, + stmt: &Statement<'tcx>, + ) -> Option<(Place<'tcx>, Option<TrackElem>)> { + match stmt.kind { + StatementKind::Assign(box (place, _)) + | StatementKind::Deinit(box place) => Some((place, None)), + StatementKind::SetDiscriminant { box place, variant_index: _ } => { + Some((place, Some(TrackElem::Discriminant))) + } + StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => { + Some((Place::from(local), None)) + } + StatementKind::Retag(..) + | StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(..)) + // copy_nonoverlapping takes pointers and mutated the pointed-to value. + | StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(..)) + | StatementKind::AscribeUserType(..) + | StatementKind::Coverage(..) + | StatementKind::FakeRead(..) + | StatementKind::ConstEvalCounter + | StatementKind::PlaceMention(..) + | StatementKind::Nop => None, + } + } + + #[instrument(level = "trace", skip(self))] + fn process_operand( + &mut self, + bb: BasicBlock, + lhs: PlaceIndex, + rhs: &Operand<'tcx>, + state: &mut State<ConditionSet<'a>>, + ) -> Option<!> { + let register_opportunity = |c: Condition| { + debug!(?bb, ?c.target, "register"); + self.opportunities.push(ThreadingOpportunity { chain: vec![bb], target: c.target }) + }; + + match rhs { + // If we expect `lhs ?= A`, we have an opportunity if we assume `constant == A`. + Operand::Constant(constant) => { + let conditions = state.try_get_idx(lhs, self.map)?; + let constant = + constant.const_.normalize(self.tcx, self.param_env).try_to_scalar_int()?; + conditions.iter_matches(constant).for_each(register_opportunity); + } + // Transfer the conditions on the copied rhs. + Operand::Move(rhs) | Operand::Copy(rhs) => { + let rhs = self.map.find(rhs.as_ref())?; + state.insert_place_idx(rhs, lhs, self.map); + } + } + + None + } + + #[instrument(level = "trace", skip(self))] + fn process_statement( + &mut self, + bb: BasicBlock, + stmt: &Statement<'tcx>, + state: &mut State<ConditionSet<'a>>, + ) -> Option<!> { + let register_opportunity = |c: Condition| { + debug!(?bb, ?c.target, "register"); + self.opportunities.push(ThreadingOpportunity { chain: vec![bb], target: c.target }) + }; + + // Below, `lhs` is the return value of `mutated_statement`, + // the place to which `conditions` apply. + + let discriminant_for_variant = |enum_ty: Ty<'tcx>, variant_index| { + let discr = enum_ty.discriminant_for_variant(self.tcx, variant_index)?; + let discr_layout = self.tcx.layout_of(self.param_env.and(discr.ty)).ok()?; + let scalar = ScalarInt::try_from_uint(discr.val, discr_layout.size)?; + Some(Operand::const_from_scalar( + self.tcx, + discr.ty, + scalar.into(), + rustc_span::DUMMY_SP, + )) + }; + + match &stmt.kind { + // If we expect `discriminant(place) ?= A`, + // we have an opportunity if `variant_index ?= A`. + StatementKind::SetDiscriminant { box place, variant_index } => { + let discr_target = self.map.find_discr(place.as_ref())?; + let enum_ty = place.ty(self.body, self.tcx).ty; + let discr = discriminant_for_variant(enum_ty, *variant_index)?; + self.process_operand(bb, discr_target, &discr, state)?; + } + // If we expect `lhs ?= true`, we have an opportunity if we assume `lhs == true`. + StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume( + Operand::Copy(place) | Operand::Move(place), + )) => { + let conditions = state.try_get(place.as_ref(), self.map)?; + conditions.iter_matches(ScalarInt::TRUE).for_each(register_opportunity); + } + StatementKind::Assign(box (lhs_place, rhs)) => { + if let Some(lhs) = self.map.find(lhs_place.as_ref()) { + match rhs { + Rvalue::Use(operand) => self.process_operand(bb, lhs, operand, state)?, + // Transfer the conditions on the copy rhs. + Rvalue::CopyForDeref(rhs) => { + self.process_operand(bb, lhs, &Operand::Copy(*rhs), state)? + } + Rvalue::Discriminant(rhs) => { + let rhs = self.map.find_discr(rhs.as_ref())?; + state.insert_place_idx(rhs, lhs, self.map); + } + // If we expect `lhs ?= A`, we have an opportunity if we assume `constant == A`. + Rvalue::Aggregate(box ref kind, ref operands) => { + let agg_ty = lhs_place.ty(self.body, self.tcx).ty; + let lhs = match kind { + // Do not support unions. + AggregateKind::Adt(.., Some(_)) => return None, + AggregateKind::Adt(_, variant_index, ..) if agg_ty.is_enum() => { + if let Some(discr_target) = + self.map.apply(lhs, TrackElem::Discriminant) + && let Some(discr_value) = + discriminant_for_variant(agg_ty, *variant_index) + { + self.process_operand(bb, discr_target, &discr_value, state); + } + self.map.apply(lhs, TrackElem::Variant(*variant_index))? + } + _ => lhs, + }; + for (field_index, operand) in operands.iter_enumerated() { + if let Some(field) = + self.map.apply(lhs, TrackElem::Field(field_index)) + { + self.process_operand(bb, field, operand, state); + } + } + } + // Transfer the conditions on the copy rhs, after inversing polarity. + Rvalue::UnaryOp(UnOp::Not, Operand::Move(place) | Operand::Copy(place)) => { + let conditions = state.try_get_idx(lhs, self.map)?; + let place = self.map.find(place.as_ref())?; + let conds = conditions.map(self.arena, Condition::inv); + state.insert_value_idx(place, conds, self.map); + } + // We expect `lhs ?= A`. We found `lhs = Eq(rhs, B)`. + // Create a condition on `rhs ?= B`. + Rvalue::BinaryOp( + op, + box ( + Operand::Move(place) | Operand::Copy(place), + Operand::Constant(value), + ) + | box ( + Operand::Constant(value), + Operand::Move(place) | Operand::Copy(place), + ), + ) => { + let conditions = state.try_get_idx(lhs, self.map)?; + let place = self.map.find(place.as_ref())?; + let equals = match op { + BinOp::Eq => ScalarInt::TRUE, + BinOp::Ne => ScalarInt::FALSE, + _ => return None, + }; + let value = value + .const_ + .normalize(self.tcx, self.param_env) + .try_to_scalar_int()?; + let conds = conditions.map(self.arena, |c| Condition { + value, + polarity: if c.matches(equals) { + Polarity::Eq + } else { + Polarity::Ne + }, + ..c + }); + state.insert_value_idx(place, conds, self.map); + } + + _ => {} + } + } + } + _ => {} + } + + None + } + + #[instrument(level = "trace", skip(self, cost))] + fn recurse_through_terminator( + &mut self, + bb: BasicBlock, + state: &State<ConditionSet<'a>>, + cost: &CostChecker<'_, 'tcx>, + depth: usize, + ) { + let register_opportunity = |c: Condition| { + debug!(?bb, ?c.target, "register"); + self.opportunities.push(ThreadingOpportunity { chain: vec![bb], target: c.target }) + }; + + let term = self.body.basic_blocks[bb].terminator(); + let place_to_flood = match term.kind { + // We come from a target, so those are not possible. + TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::CoroutineDrop => bug!("{term:?} has no terminators"), + // Disallowed during optimizations. + TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::Yield { .. } => bug!("{term:?} invalid"), + // Cannot reason about inline asm. + TerminatorKind::InlineAsm { .. } => return, + // `SwitchInt` is handled specially. + TerminatorKind::SwitchInt { .. } => return, + // We can recurse, no thing particular to do. + TerminatorKind::Goto { .. } => None, + // Flood the overwritten place, and progress through. + TerminatorKind::Drop { place: destination, .. } + | TerminatorKind::Call { destination, .. } => Some(destination), + // Treat as an `assume(cond == expected)`. + TerminatorKind::Assert { ref cond, expected, .. } => { + if let Some(place) = cond.place() + && let Some(conditions) = state.try_get(place.as_ref(), self.map) + { + let expected = if expected { ScalarInt::TRUE } else { ScalarInt::FALSE }; + conditions.iter_matches(expected).for_each(register_opportunity); + } + None + } + }; + + // We can recurse through this terminator. + let mut state = state.clone(); + if let Some(place_to_flood) = place_to_flood { + state.flood_with(place_to_flood.as_ref(), self.map, ConditionSet::default()); + } + self.find_opportunity(bb, state, cost.clone(), depth + 1); + } + + #[instrument(level = "trace", skip(self))] + fn process_switch_int( + &mut self, + discr: &Operand<'tcx>, + targets: &SwitchTargets, + target_bb: BasicBlock, + state: &mut State<ConditionSet<'a>>, + ) -> Option<!> { + debug_assert_ne!(target_bb, START_BLOCK); + debug_assert_eq!(self.body.basic_blocks.predecessors()[target_bb].len(), 1); + + let discr = discr.place()?; + let discr_ty = discr.ty(self.body, self.tcx).ty; + let discr_layout = self.tcx.layout_of(self.param_env.and(discr_ty)).ok()?; + let conditions = state.try_get(discr.as_ref(), self.map)?; + + if let Some((value, _)) = targets.iter().find(|&(_, target)| target == target_bb) { + let value = ScalarInt::try_from_uint(value, discr_layout.size)?; + debug_assert_eq!(targets.iter().filter(|&(_, target)| target == target_bb).count(), 1); + + // We are inside `target_bb`. Since we have a single predecessor, we know we passed + // through the `SwitchInt` before arriving here. Therefore, we know that + // `discr == value`. If one condition can be fulfilled by `discr == value`, + // that's an opportunity. + for c in conditions.iter_matches(value) { + debug!(?target_bb, ?c.target, "register"); + self.opportunities.push(ThreadingOpportunity { chain: vec![], target: c.target }); + } + } else if let Some((value, _, else_bb)) = targets.as_static_if() + && target_bb == else_bb + { + let value = ScalarInt::try_from_uint(value, discr_layout.size)?; + + // We only know that `discr != value`. That's much weaker information than + // the equality we had in the previous arm. All we can conclude is that + // the replacement condition `discr != value` can be threaded, and nothing else. + for c in conditions.iter() { + if c.value == value && c.polarity == Polarity::Ne { + debug!(?target_bb, ?c.target, "register"); + self.opportunities + .push(ThreadingOpportunity { chain: vec![], target: c.target }); + } + } + } + + None + } +} + +struct OpportunitySet { + opportunities: Vec<ThreadingOpportunity>, + /// For each bb, give the TOs in which it appears. The pair corresponds to the index + /// in `opportunities` and the index in `ThreadingOpportunity::chain`. + involving_tos: IndexVec<BasicBlock, Vec<(usize, usize)>>, + /// Cache the number of predecessors for each block, as we clear the basic block cache.. + predecessors: IndexVec<BasicBlock, usize>, +} + +impl OpportunitySet { + fn new(body: &Body<'_>, opportunities: Vec<ThreadingOpportunity>) -> OpportunitySet { + let mut involving_tos = IndexVec::from_elem(Vec::new(), &body.basic_blocks); + for (index, to) in opportunities.iter().enumerate() { + for (ibb, &bb) in to.chain.iter().enumerate() { + involving_tos[bb].push((index, ibb)); + } + involving_tos[to.target].push((index, to.chain.len())); + } + let predecessors = predecessor_count(body); + OpportunitySet { opportunities, involving_tos, predecessors } + } + + /// Apply the opportunities on the graph. + fn apply(&mut self, body: &mut Body<'_>) { + for i in 0..self.opportunities.len() { + self.apply_once(i, body); + } + } + + #[instrument(level = "trace", skip(self, body))] + fn apply_once(&mut self, index: usize, body: &mut Body<'_>) { + debug!(?self.predecessors); + debug!(?self.involving_tos); + + // Check that `predecessors` satisfies its invariant. + debug_assert_eq!(self.predecessors, predecessor_count(body)); + + // Remove the TO from the vector to allow modifying the other ones later. + let op = &mut self.opportunities[index]; + debug!(?op); + let op_chain = std::mem::take(&mut op.chain); + let op_target = op.target; + debug_assert_eq!(op_chain.len(), op_chain.iter().collect::<FxHashSet<_>>().len()); + + let Some((current, chain)) = op_chain.split_first() else { return }; + let basic_blocks = body.basic_blocks.as_mut(); + + // Invariant: the control-flow is well-formed at the end of each iteration. + let mut current = *current; + for &succ in chain { + debug!(?current, ?succ); + + // `succ` must be a successor of `current`. If it is not, this means this TO is not + // satisfiable and a previous TO erased this edge, so we bail out. + if basic_blocks[current].terminator().successors().find(|s| *s == succ).is_none() { + debug!("impossible"); + return; + } + + // Fast path: `succ` is only used once, so we can reuse it directly. + if self.predecessors[succ] == 1 { + debug!("single"); + current = succ; + continue; + } + + let new_succ = basic_blocks.push(basic_blocks[succ].clone()); + debug!(?new_succ); + + // Replace `succ` by `new_succ` where it appears. + let mut num_edges = 0; + for s in basic_blocks[current].terminator_mut().successors_mut() { + if *s == succ { + *s = new_succ; + num_edges += 1; + } + } + + // Update predecessors with the new block. + let _new_succ = self.predecessors.push(num_edges); + debug_assert_eq!(new_succ, _new_succ); + self.predecessors[succ] -= num_edges; + self.update_predecessor_count(basic_blocks[new_succ].terminator(), Update::Incr); + + // Replace the `current -> succ` edge by `current -> new_succ` in all the following + // TOs. This is necessary to avoid trying to thread through a non-existing edge. We + // use `involving_tos` here to avoid traversing the full set of TOs on each iteration. + let mut new_involved = Vec::new(); + for &(to_index, in_to_index) in &self.involving_tos[current] { + // That TO has already been applied, do nothing. + if to_index <= index { + continue; + } + + let other_to = &mut self.opportunities[to_index]; + if other_to.chain.get(in_to_index) != Some(¤t) { + continue; + } + let s = other_to.chain.get_mut(in_to_index + 1).unwrap_or(&mut other_to.target); + if *s == succ { + // `other_to` references the `current -> succ` edge, so replace `succ`. + *s = new_succ; + new_involved.push((to_index, in_to_index + 1)); + } + } + + // The TOs that we just updated now reference `new_succ`. Update `involving_tos` + // in case we need to duplicate an edge starting at `new_succ` later. + let _new_succ = self.involving_tos.push(new_involved); + debug_assert_eq!(new_succ, _new_succ); + + current = new_succ; + } + + let current = &mut basic_blocks[current]; + self.update_predecessor_count(current.terminator(), Update::Decr); + current.terminator_mut().kind = TerminatorKind::Goto { target: op_target }; + self.predecessors[op_target] += 1; + } + + fn update_predecessor_count(&mut self, terminator: &Terminator<'_>, incr: Update) { + match incr { + Update::Incr => { + for s in terminator.successors() { + self.predecessors[s] += 1; + } + } + Update::Decr => { + for s in terminator.successors() { + self.predecessors[s] -= 1; + } + } + } + } +} + +fn predecessor_count(body: &Body<'_>) -> IndexVec<BasicBlock, usize> { + let mut predecessors: IndexVec<_, _> = + body.basic_blocks.predecessors().iter().map(|ps| ps.len()).collect(); + predecessors[START_BLOCK] += 1; // Account for the implicit entry edge. + predecessors +} + +enum Update { + Incr, + Decr, +} + +/// Compute the set of loop headers in the given body. We define a loop header as a block which has +/// at least a predecessor which it dominates. This definition is only correct for reducible CFGs. +/// But if the CFG is already irreducible, there is no point in trying much harder. +/// is already irreducible. +fn loop_headers(body: &Body<'_>) -> BitSet<BasicBlock> { + let mut loop_headers = BitSet::new_empty(body.basic_blocks.len()); + let dominators = body.basic_blocks.dominators(); + // Only visit reachable blocks. + for (bb, bbdata) in traversal::preorder(body) { + for succ in bbdata.terminator().successors() { + if dominators.dominates(succ, bb) { + loop_headers.insert(succ); + } + } + } + loop_headers +} diff --git a/compiler/rustc_mir_transform/src/large_enums.rs b/compiler/rustc_mir_transform/src/large_enums.rs new file mode 100644 index 00000000000..1d788a55ff8 --- /dev/null +++ b/compiler/rustc_mir_transform/src/large_enums.rs @@ -0,0 +1,299 @@ +use crate::rustc_middle::ty::util::IntTypeExt; +use rustc_data_structures::fx::FxHashMap; +use rustc_middle::mir::interpret::AllocId; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, AdtDef, ParamEnv, Ty, TyCtxt}; +use rustc_session::Session; +use rustc_target::abi::{HasDataLayout, Size, TagEncoding, Variants}; + +/// A pass that seeks to optimize unnecessary moves of large enum types, if there is a large +/// enough discrepancy between them. +/// +/// i.e. If there is are two variants: +/// ``` +/// enum Example { +/// Small, +/// Large([u32; 1024]), +/// } +/// ``` +/// Instead of emitting moves of the large variant, +/// Perform a memcpy instead. +/// Based off of [this HackMD](https://hackmd.io/@ft4bxUsFT5CEUBmRKYHr7w/rJM8BBPzD). +/// +/// In summary, what this does is at runtime determine which enum variant is active, +/// and instead of copying all the bytes of the largest possible variant, +/// copy only the bytes for the currently active variant. +pub struct EnumSizeOpt { + pub(crate) discrepancy: u64, +} + +impl<'tcx> MirPass<'tcx> for EnumSizeOpt { + fn is_enabled(&self, sess: &Session) -> bool { + // There are some differences in behavior on wasm and ARM that are not properly + // understood, so we conservatively treat this optimization as unsound: + // https://github.com/rust-lang/rust/pull/85158#issuecomment-1101836457 + sess.opts.unstable_opts.unsound_mir_opts || sess.mir_opt_level() >= 3 + } + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // NOTE: This pass may produce different MIR based on the alignment of the target + // platform, but it will still be valid. + self.optim(tcx, body); + } +} + +impl EnumSizeOpt { + fn candidate<'tcx>( + &self, + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + ty: Ty<'tcx>, + alloc_cache: &mut FxHashMap<Ty<'tcx>, AllocId>, + ) -> Option<(AdtDef<'tcx>, usize, AllocId)> { + let adt_def = match ty.kind() { + ty::Adt(adt_def, _args) if adt_def.is_enum() => adt_def, + _ => return None, + }; + let layout = tcx.layout_of(param_env.and(ty)).ok()?; + let variants = match &layout.variants { + Variants::Single { .. } => return None, + Variants::Multiple { tag_encoding: TagEncoding::Niche { .. }, .. } => return None, + + Variants::Multiple { variants, .. } if variants.len() <= 1 => return None, + Variants::Multiple { variants, .. } => variants, + }; + let min = variants.iter().map(|v| v.size).min().unwrap(); + let max = variants.iter().map(|v| v.size).max().unwrap(); + if max.bytes() - min.bytes() < self.discrepancy { + return None; + } + + let num_discrs = adt_def.discriminants(tcx).count(); + if variants.iter_enumerated().any(|(var_idx, _)| { + let discr_for_var = adt_def.discriminant_for_variant(tcx, var_idx).val; + (discr_for_var > usize::MAX as u128) || (discr_for_var as usize >= num_discrs) + }) { + return None; + } + if let Some(alloc_id) = alloc_cache.get(&ty) { + return Some((*adt_def, num_discrs, *alloc_id)); + } + + let data_layout = tcx.data_layout(); + let ptr_sized_int = data_layout.ptr_sized_integer(); + let target_bytes = ptr_sized_int.size().bytes() as usize; + let mut data = vec![0; target_bytes * num_discrs]; + macro_rules! encode_store { + ($curr_idx: expr, $endian: expr, $bytes: expr) => { + let bytes = match $endian { + rustc_target::abi::Endian::Little => $bytes.to_le_bytes(), + rustc_target::abi::Endian::Big => $bytes.to_be_bytes(), + }; + for (i, b) in bytes.into_iter().enumerate() { + data[$curr_idx + i] = b; + } + }; + } + + for (var_idx, layout) in variants.iter_enumerated() { + let curr_idx = + target_bytes * adt_def.discriminant_for_variant(tcx, var_idx).val as usize; + let sz = layout.size; + match ptr_sized_int { + rustc_target::abi::Integer::I32 => { + encode_store!(curr_idx, data_layout.endian, sz.bytes() as u32); + } + rustc_target::abi::Integer::I64 => { + encode_store!(curr_idx, data_layout.endian, sz.bytes()); + } + _ => unreachable!(), + }; + } + let alloc = interpret::Allocation::from_bytes( + data, + tcx.data_layout.ptr_sized_integer().align(&tcx.data_layout).abi, + Mutability::Not, + ); + let alloc = tcx.reserve_and_set_memory_alloc(tcx.mk_const_alloc(alloc)); + Some((*adt_def, num_discrs, *alloc_cache.entry(ty).or_insert(alloc))) + } + fn optim<'tcx>(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let mut alloc_cache = FxHashMap::default(); + let body_did = body.source.def_id(); + let param_env = tcx.param_env_reveal_all_normalized(body_did); + + let blocks = body.basic_blocks.as_mut(); + let local_decls = &mut body.local_decls; + + for bb in blocks { + bb.expand_statements(|st| { + if let StatementKind::Assign(box ( + lhs, + Rvalue::Use(Operand::Copy(rhs) | Operand::Move(rhs)), + )) = &st.kind + { + let ty = lhs.ty(local_decls, tcx).ty; + + let source_info = st.source_info; + let span = source_info.span; + + let (adt_def, num_variants, alloc_id) = + self.candidate(tcx, param_env, ty, &mut alloc_cache)?; + + 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 place = Place::from(size_array_local); + let constant_vals = ConstOperand { + span, + user_ty: None, + const_: Const::Val( + ConstValue::Indirect { alloc_id, offset: Size::ZERO }, + tmp_ty, + ), + }; + 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::AddressOf(Mutability::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 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::AddressOf(Mutability::Not, *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 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 store_dead = Statement { + source_info, + kind: StatementKind::StorageDead(size_array_local), + }; + let iter = [ + store_live, + const_assign, + store_discr, + cast_discr, + store_size, + dst_ptr, + dst_cast, + src_ptr, + src_cast, + deinit_old, + copy_bytes, + store_dead, + ] + .into_iter(); + + st.make_nop(); + Some(iter) + } else { + None + } + }); + } + } +} diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs new file mode 100644 index 00000000000..89e897191e8 --- /dev/null +++ b/compiler/rustc_mir_transform/src/lib.rs @@ -0,0 +1,681 @@ +#![allow(rustc::potential_query_instability)] +#![deny(rustc::untranslatable_diagnostic)] +#![deny(rustc::diagnostic_outside_of_impl)] +#![feature(box_patterns)] +#![feature(cow_is_borrowed)] +#![feature(decl_macro)] +#![feature(is_sorted)] +#![feature(let_chains)] +#![feature(map_try_insert)] +#![feature(min_specialization)] +#![feature(never_type)] +#![feature(option_get_or_insert_default)] +#![feature(trusted_step)] +#![feature(try_blocks)] +#![feature(yeet_expr)] +#![feature(if_let_guard)] +#![recursion_limit = "256"] + +#[macro_use] +extern crate tracing; +#[macro_use] +extern crate rustc_middle; + +use hir::ConstContext; +use required_consts::RequiredConstsVisitor; +use rustc_const_eval::util; +use rustc_data_structures::fx::FxIndexSet; +use rustc_data_structures::steal::Steal; +use rustc_hir as hir; +use rustc_hir::def::DefKind; +use rustc_hir::def_id::LocalDefId; +use rustc_hir::intravisit::{self, Visitor}; +use rustc_index::IndexVec; +use rustc_middle::mir::visit::Visitor as _; +use rustc_middle::mir::{ + traversal, AnalysisPhase, Body, CallSource, ClearCrossCrate, ConstOperand, ConstQualifs, + LocalDecl, MirPass, MirPhase, Operand, Place, ProjectionElem, Promoted, RuntimePhase, Rvalue, + SourceInfo, Statement, StatementKind, TerminatorKind, START_BLOCK, +}; +use rustc_middle::query::Providers; +use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt}; +use rustc_span::sym; +use rustc_trait_selection::traits; + +#[macro_use] +mod pass_manager; + +use pass_manager::{self as pm, Lint, MirLint, WithMinOptLevel}; + +mod abort_unwinding_calls; +mod add_call_guards; +mod add_moves_for_packed_drops; +mod add_retag; +mod check_const_item_mutation; +mod check_packed_ref; +pub mod check_unsafety; +mod remove_place_mention; +// This pass is public to allow external drivers to perform MIR cleanup +mod add_subtyping_projections; +pub mod cleanup_post_borrowck; +mod const_debuginfo; +mod const_goto; +mod const_prop; +mod const_prop_lint; +mod copy_prop; +mod coroutine; +mod cost_checker; +mod coverage; +mod cross_crate_inline; +mod ctfe_limit; +mod dataflow_const_prop; +mod dead_store_elimination; +mod deduce_param_attrs; +mod deduplicate_blocks; +mod deref_separator; +mod dest_prop; +pub mod dump_mir; +mod early_otherwise_branch; +mod elaborate_box_derefs; +mod elaborate_drops; +mod errors; +mod ffi_unwind_calls; +mod function_item_references; +mod gvn; +pub mod inline; +mod instsimplify; +mod jump_threading; +mod large_enums; +mod lower_intrinsics; +mod lower_slice_len; +mod match_branches; +mod multiple_return_terminators; +mod normalize_array_len; +mod nrvo; +mod prettify; +mod ref_prop; +mod remove_noop_landing_pads; +mod remove_storage_markers; +mod remove_uninit_drops; +mod remove_unneeded_drops; +mod remove_zsts; +mod required_consts; +mod reveal_all; +mod separate_const_switch; +mod shim; +mod ssa; +// This pass is public to allow external drivers to perform MIR cleanup +mod check_alignment; +pub mod simplify; +mod simplify_branches; +mod simplify_comparison_integral; +mod sroa; +mod uninhabited_enum_branching; +mod unreachable_prop; + +use rustc_const_eval::transform::check_consts::{self, ConstCx}; +use rustc_const_eval::transform::promote_consts; +use rustc_const_eval::transform::validate; +use rustc_mir_dataflow::rustc_peek; + +rustc_fluent_macro::fluent_messages! { "../messages.ftl" } + +pub fn provide(providers: &mut Providers) { + check_unsafety::provide(providers); + coverage::query::provide(providers); + ffi_unwind_calls::provide(providers); + shim::provide(providers); + cross_crate_inline::provide(providers); + *providers = Providers { + mir_keys, + mir_const, + mir_const_qualif, + mir_promoted, + mir_drops_elaborated_and_const_checked, + mir_for_ctfe, + mir_coroutine_witnesses: coroutine::mir_coroutine_witnesses, + optimized_mir, + is_mir_available, + is_ctfe_mir_available: |tcx, did| is_mir_available(tcx, did), + mir_callgraph_reachable: inline::cycle::mir_callgraph_reachable, + mir_inliner_callees: inline::cycle::mir_inliner_callees, + promoted_mir, + deduced_param_attrs: deduce_param_attrs::deduced_param_attrs, + ..*providers + }; +} + +fn remap_mir_for_const_eval_select<'tcx>( + tcx: TyCtxt<'tcx>, + mut body: Body<'tcx>, + context: hir::Constness, +) -> Body<'tcx> { + for bb in body.basic_blocks.as_mut().iter_mut() { + let terminator = bb.terminator.as_mut().expect("invalid terminator"); + match terminator.kind { + TerminatorKind::Call { + func: Operand::Constant(box ConstOperand { ref const_, .. }), + ref mut args, + destination, + target, + unwind, + fn_span, + .. + } if let ty::FnDef(def_id, _) = *const_.ty().kind() + && tcx.item_name(def_id) == sym::const_eval_select + && tcx.is_intrinsic(def_id) => + { + let [tupled_args, called_in_const, called_at_rt]: [_; 3] = + std::mem::take(args).try_into().unwrap(); + let ty = tupled_args.ty(&body.local_decls, tcx); + let fields = ty.tuple_fields(); + let num_args = fields.len(); + let func = + if context == hir::Constness::Const { called_in_const } else { called_at_rt }; + let (method, place): (fn(Place<'tcx>) -> Operand<'tcx>, Place<'tcx>) = + match tupled_args { + Operand::Constant(_) => { + // there is no good way of extracting a tuple arg from a constant (const generic stuff) + // so we just create a temporary and deconstruct that. + let local = body.local_decls.push(LocalDecl::new(ty, fn_span)); + bb.statements.push(Statement { + source_info: SourceInfo::outermost(fn_span), + kind: StatementKind::Assign(Box::new(( + local.into(), + Rvalue::Use(tupled_args.clone()), + ))), + }); + (Operand::Move, local.into()) + } + Operand::Move(place) => (Operand::Move, place), + Operand::Copy(place) => (Operand::Copy, place), + }; + let place_elems = place.projection; + let arguments = (0..num_args) + .map(|x| { + let mut place_elems = place_elems.to_vec(); + place_elems.push(ProjectionElem::Field(x.into(), fields[x])); + let projection = tcx.mk_place_elems(&place_elems); + let place = Place { local: place.local, projection }; + method(place) + }) + .collect(); + terminator.kind = TerminatorKind::Call { + func, + args: arguments, + destination, + target, + unwind, + call_source: CallSource::Misc, + fn_span, + }; + } + _ => {} + } + } + body +} + +fn is_mir_available(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool { + tcx.mir_keys(()).contains(&def_id) +} + +/// Finds the full set of `DefId`s within the current crate that have +/// MIR associated with them. +fn mir_keys(tcx: TyCtxt<'_>, (): ()) -> FxIndexSet<LocalDefId> { + let mut set = FxIndexSet::default(); + + // All body-owners have MIR associated with them. + set.extend(tcx.hir().body_owners()); + + // Additionally, tuple struct/variant constructors have MIR, but + // they don't have a BodyId, so we need to build them separately. + struct GatherCtors<'a> { + set: &'a mut FxIndexSet<LocalDefId>, + } + impl<'tcx> Visitor<'tcx> for GatherCtors<'_> { + fn visit_variant_data(&mut self, v: &'tcx hir::VariantData<'tcx>) { + if let hir::VariantData::Tuple(_, _, def_id) = *v { + self.set.insert(def_id); + } + intravisit::walk_struct_def(self, v) + } + } + tcx.hir().visit_all_item_likes_in_crate(&mut GatherCtors { set: &mut set }); + + set +} + +fn mir_const_qualif(tcx: TyCtxt<'_>, def: LocalDefId) -> ConstQualifs { + let const_kind = tcx.hir().body_const_context(def); + + // No need to const-check a non-const `fn`. + match const_kind { + Some(ConstContext::Const { .. } | ConstContext::Static(_)) + | Some(ConstContext::ConstFn) => {} + None => span_bug!( + tcx.def_span(def), + "`mir_const_qualif` should only be called on const fns and const items" + ), + } + + // N.B., this `borrow()` is guaranteed to be valid (i.e., the value + // cannot yet be stolen), because `mir_promoted()`, which steals + // from `mir_const()`, forces this query to execute before + // performing the steal. + let body = &tcx.mir_const(def).borrow(); + + if body.return_ty().references_error() { + tcx.sess.span_delayed_bug(body.span, "mir_const_qualif: MIR had errors"); + return Default::default(); + } + + let ccx = check_consts::ConstCx { body, tcx, const_kind, param_env: tcx.param_env(def) }; + + let mut validator = check_consts::check::Checker::new(&ccx); + validator.check_body(); + + // We return the qualifs in the return place for every MIR body, even though it is only used + // when deciding to promote a reference to a `const` for now. + validator.qualifs_in_return_place() +} + +/// Make MIR ready for const evaluation. This is run on all MIR, not just on consts! +/// FIXME(oli-obk): it's unclear whether we still need this phase (and its corresponding query). +/// We used to have this for pre-miri MIR based const eval. +fn mir_const(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> { + // Unsafety check uses the raw mir, so make sure it is run. + if !tcx.sess.opts.unstable_opts.thir_unsafeck { + tcx.ensure_with_value().unsafety_check_result(def); + } + + // has_ffi_unwind_calls query uses the raw mir, so make sure it is run. + tcx.ensure_with_value().has_ffi_unwind_calls(def); + + let mut body = tcx.mir_built(def).steal(); + + pass_manager::dump_mir_for_phase_change(tcx, &body); + + pm::run_passes( + tcx, + &mut body, + &[ + // MIR-level lints. + &Lint(check_packed_ref::CheckPackedRef), + &Lint(check_const_item_mutation::CheckConstItemMutation), + &Lint(function_item_references::FunctionItemReferences), + // What we need to do constant evaluation. + &simplify::SimplifyCfg::Initial, + &rustc_peek::SanityCheck, // Just a lint + ], + None, + ); + tcx.alloc_steal_mir(body) +} + +/// Compute the main MIR body and the list of MIR bodies of the promoteds. +fn mir_promoted( + tcx: TyCtxt<'_>, + def: LocalDefId, +) -> (&Steal<Body<'_>>, &Steal<IndexVec<Promoted, Body<'_>>>) { + // Ensure that we compute the `mir_const_qualif` for constants at + // this point, before we steal the mir-const result. + // Also this means promotion can rely on all const checks having been done. + + let const_qualifs = match tcx.def_kind(def) { + DefKind::Fn | DefKind::AssocFn | DefKind::Closure + if tcx.constness(def) == hir::Constness::Const + || tcx.is_const_default_method(def.to_def_id()) => + { + tcx.mir_const_qualif(def) + } + DefKind::AssocConst + | DefKind::Const + | DefKind::Static(_) + | DefKind::InlineConst + | DefKind::AnonConst => tcx.mir_const_qualif(def), + _ => ConstQualifs::default(), + }; + let mut body = tcx.mir_const(def).steal(); + if let Some(error_reported) = const_qualifs.tainted_by_errors { + body.tainted_by_errors = Some(error_reported); + } + + let mut required_consts = Vec::new(); + let mut required_consts_visitor = RequiredConstsVisitor::new(&mut required_consts); + for (bb, bb_data) in traversal::reverse_postorder(&body) { + required_consts_visitor.visit_basic_block_data(bb, bb_data); + } + body.required_consts = required_consts; + + // What we need to run borrowck etc. + let promote_pass = promote_consts::PromoteTemps::default(); + pm::run_passes( + tcx, + &mut body, + &[&promote_pass, &simplify::SimplifyCfg::PromoteConsts, &coverage::InstrumentCoverage], + Some(MirPhase::Analysis(AnalysisPhase::Initial)), + ); + + let promoted = promote_pass.promoted_fragments.into_inner(); + (tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted)) +} + +/// Compute the MIR that is used during CTFE (and thus has no optimizations run on it) +fn mir_for_ctfe(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &Body<'_> { + tcx.arena.alloc(inner_mir_for_ctfe(tcx, def_id)) +} + +fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: LocalDefId) -> Body<'_> { + // FIXME: don't duplicate this between the optimized_mir/mir_for_ctfe queries + if tcx.is_constructor(def.to_def_id()) { + // There's no reason to run all of the MIR passes on constructors when + // we can just output the MIR we want directly. This also saves const + // qualification and borrow checking the trouble of special casing + // constructors. + return shim::build_adt_ctor(tcx, def.to_def_id()); + } + + let body = tcx.mir_drops_elaborated_and_const_checked(def); + let body = match tcx.hir().body_const_context(def) { + // consts and statics do not have `optimized_mir`, so we can steal the body instead of + // cloning it. + Some(hir::ConstContext::Const { .. } | hir::ConstContext::Static(_)) => body.steal(), + Some(hir::ConstContext::ConstFn) => body.borrow().clone(), + None => bug!("`mir_for_ctfe` called on non-const {def:?}"), + }; + + let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::Const); + pm::run_passes(tcx, &mut body, &[&ctfe_limit::CtfeLimit], None); + + body +} + +/// Obtain just the main MIR (no promoteds) and run some cleanups on it. This also runs +/// mir borrowck *before* doing so in order to ensure that borrowck can be run and doesn't +/// end up missing the source MIR due to stealing happening. +fn mir_drops_elaborated_and_const_checked(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> { + if tcx.is_coroutine(def.to_def_id()) { + tcx.ensure_with_value().mir_coroutine_witnesses(def); + } + let mir_borrowck = tcx.mir_borrowck(def); + + let is_fn_like = tcx.def_kind(def).is_fn_like(); + if is_fn_like { + // Do not compute the mir call graph without said call graph actually being used. + if pm::should_run_pass(tcx, &inline::Inline) { + tcx.ensure_with_value().mir_inliner_callees(ty::InstanceDef::Item(def.to_def_id())); + } + } + + let (body, _) = tcx.mir_promoted(def); + let mut body = body.steal(); + if let Some(error_reported) = mir_borrowck.tainted_by_errors { + body.tainted_by_errors = Some(error_reported); + } + + // Check if it's even possible to satisfy the 'where' clauses + // for this item. + // + // This branch will never be taken for any normal function. + // However, it's possible to `#!feature(trivial_bounds)]` to write + // a function with impossible to satisfy clauses, e.g.: + // `fn foo() where String: Copy {}` + // + // We don't usually need to worry about this kind of case, + // since we would get a compilation error if the user tried + // to call it. However, since we optimize even without any + // calls to the function, we need to make sure that it even + // makes sense to try to evaluate the body. + // + // If there are unsatisfiable where clauses, then all bets are + // off, and we just give up. + // + // We manually filter the predicates, skipping anything that's not + // "global". We are in a potentially generic context + // (e.g. we are evaluating a function without substituting generic + // parameters, so this filtering serves two purposes: + // + // 1. We skip evaluating any predicates that we would + // never be able prove are unsatisfiable (e.g. `<T as Foo>` + // 2. We avoid trying to normalize predicates involving generic + // parameters (e.g. `<T as Foo>::MyItem`). This can confuse + // the normalization code (leading to cycle errors), since + // it's usually never invoked in this way. + let predicates = tcx + .predicates_of(body.source.def_id()) + .predicates + .iter() + .filter_map(|(p, _)| if p.is_global() { Some(*p) } else { None }); + if traits::impossible_predicates(tcx, traits::elaborate(tcx, predicates).collect()) { + trace!("found unsatisfiable predicates for {:?}", body.source); + // Clear the body to only contain a single `unreachable` statement. + let bbs = body.basic_blocks.as_mut(); + bbs.raw.truncate(1); + bbs[START_BLOCK].statements.clear(); + bbs[START_BLOCK].terminator_mut().kind = TerminatorKind::Unreachable; + body.var_debug_info.clear(); + body.local_decls.raw.truncate(body.arg_count + 1); + } + + run_analysis_to_runtime_passes(tcx, &mut body); + + // Now that drop elaboration has been performed, we can check for + // unconditional drop recursion. + rustc_mir_build::lints::check_drop_recursion(tcx, &body); + + tcx.alloc_steal_mir(body) +} + +// Made public such that `mir_drops_elaborated_and_const_checked` can be overridden +// by custom rustc drivers, running all the steps by themselves. +pub fn run_analysis_to_runtime_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + assert!(body.phase == MirPhase::Analysis(AnalysisPhase::Initial)); + let did = body.source.def_id(); + + debug!("analysis_mir_cleanup({:?})", did); + run_analysis_cleanup_passes(tcx, body); + assert!(body.phase == MirPhase::Analysis(AnalysisPhase::PostCleanup)); + + // Do a little drop elaboration before const-checking if `const_precise_live_drops` is enabled. + if check_consts::post_drop_elaboration::checking_enabled(&ConstCx::new(tcx, body)) { + pm::run_passes( + tcx, + body, + &[&remove_uninit_drops::RemoveUninitDrops, &simplify::SimplifyCfg::RemoveFalseEdges], + None, + ); + check_consts::post_drop_elaboration::check_live_drops(tcx, body); // FIXME: make this a MIR lint + } + + debug!("runtime_mir_lowering({:?})", did); + run_runtime_lowering_passes(tcx, body); + assert!(body.phase == MirPhase::Runtime(RuntimePhase::Initial)); + + debug!("runtime_mir_cleanup({:?})", did); + run_runtime_cleanup_passes(tcx, body); + assert!(body.phase == MirPhase::Runtime(RuntimePhase::PostCleanup)); +} + +// FIXME(JakobDegen): Can we make these lists of passes consts? + +/// After this series of passes, no lifetime analysis based on borrowing can be done. +fn run_analysis_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let passes: &[&dyn MirPass<'tcx>] = &[ + &cleanup_post_borrowck::CleanupPostBorrowck, + &remove_noop_landing_pads::RemoveNoopLandingPads, + &simplify::SimplifyCfg::EarlyOpt, + &deref_separator::Derefer, + ]; + + pm::run_passes(tcx, body, passes, Some(MirPhase::Analysis(AnalysisPhase::PostCleanup))); +} + +/// Returns the sequence of passes that lowers analysis to runtime MIR. +fn run_runtime_lowering_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let passes: &[&dyn MirPass<'tcx>] = &[ + // These next passes must be executed together + &add_call_guards::CriticalCallEdges, + &reveal_all::RevealAll, // has to be done before drop elaboration, since we need to drop opaque types, too. + &add_subtyping_projections::Subtyper, // calling this after reveal_all ensures that we don't deal with opaque types + &elaborate_drops::ElaborateDrops, + // This will remove extraneous landing pads which are no longer + // necessary as well as well as forcing any call in a non-unwinding + // function calling a possibly-unwinding function to abort the process. + &abort_unwinding_calls::AbortUnwindingCalls, + // AddMovesForPackedDrops needs to run after drop + // elaboration. + &add_moves_for_packed_drops::AddMovesForPackedDrops, + // `AddRetag` needs to run after `ElaborateDrops`. Otherwise it should run fairly late, + // but before optimizations begin. + &elaborate_box_derefs::ElaborateBoxDerefs, + &coroutine::StateTransform, + &add_retag::AddRetag, + &Lint(const_prop_lint::ConstPropLint), + ]; + pm::run_passes_no_validate(tcx, body, passes, Some(MirPhase::Runtime(RuntimePhase::Initial))); +} + +/// Returns the sequence of passes that do the initial cleanup of runtime MIR. +fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let passes: &[&dyn MirPass<'tcx>] = &[ + &lower_intrinsics::LowerIntrinsics, + &remove_place_mention::RemovePlaceMention, + &simplify::SimplifyCfg::ElaborateDrops, + ]; + + pm::run_passes(tcx, body, passes, Some(MirPhase::Runtime(RuntimePhase::PostCleanup))); + + // Clear this by anticipation. Optimizations and runtime MIR have no reason to look + // into this information, which is meant for borrowck diagnostics. + for decl in &mut body.local_decls { + decl.local_info = ClearCrossCrate::Clear; + } +} + +fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + fn o1<T>(x: T) -> WithMinOptLevel<T> { + WithMinOptLevel(1, x) + } + + // The main optimizations that we do on MIR. + pm::run_passes( + tcx, + body, + &[ + &check_alignment::CheckAlignment, + &lower_slice_len::LowerSliceLenCalls, // has to be done before inlining, otherwise actual call will be almost always inlined. Also simple, so can just do first + &inline::Inline, + // Substitutions during inlining may introduce switch on enums with uninhabited branches. + &uninhabited_enum_branching::UninhabitedEnumBranching, + &unreachable_prop::UnreachablePropagation, + &o1(simplify::SimplifyCfg::AfterUninhabitedEnumBranching), + &remove_storage_markers::RemoveStorageMarkers, + &remove_zsts::RemoveZsts, + &normalize_array_len::NormalizeArrayLen, // has to run after `slice::len` lowering + &const_goto::ConstGoto, + &remove_unneeded_drops::RemoveUnneededDrops, + &ref_prop::ReferencePropagation, + &sroa::ScalarReplacementOfAggregates, + &match_branches::MatchBranchSimplification, + // inst combine is after MatchBranchSimplification to clean up Ne(_1, false) + &multiple_return_terminators::MultipleReturnTerminators, + &instsimplify::InstSimplify, + &simplify::SimplifyLocals::BeforeConstProp, + ©_prop::CopyProp, + // Perform `SeparateConstSwitch` after SSA-based analyses, as cloning blocks may + // destroy the SSA property. It should still happen before const-propagation, so the + // latter pass will leverage the created opportunities. + &separate_const_switch::SeparateConstSwitch, + &const_prop::ConstProp, + &gvn::GVN, + &simplify::SimplifyLocals::AfterGVN, + &dataflow_const_prop::DataflowConstProp, + &const_debuginfo::ConstDebugInfo, + &o1(simplify_branches::SimplifyConstCondition::AfterConstProp), + &jump_threading::JumpThreading, + &early_otherwise_branch::EarlyOtherwiseBranch, + &simplify_comparison_integral::SimplifyComparisonIntegral, + &dead_store_elimination::DeadStoreElimination, + &dest_prop::DestinationPropagation, + &o1(simplify_branches::SimplifyConstCondition::Final), + &o1(remove_noop_landing_pads::RemoveNoopLandingPads), + &o1(simplify::SimplifyCfg::Final), + &nrvo::RenameReturnPlace, + &simplify::SimplifyLocals::Final, + &multiple_return_terminators::MultipleReturnTerminators, + &deduplicate_blocks::DeduplicateBlocks, + &large_enums::EnumSizeOpt { discrepancy: 128 }, + // Some cleanup necessary at least for LLVM and potentially other codegen backends. + &add_call_guards::CriticalCallEdges, + // Cleanup for human readability, off by default. + &prettify::ReorderBasicBlocks, + &prettify::ReorderLocals, + // Dump the end result for testing and debugging purposes. + &dump_mir::Marker("PreCodegen"), + ], + Some(MirPhase::Runtime(RuntimePhase::Optimized)), + ); +} + +/// Optimize the MIR and prepare it for codegen. +fn optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> &Body<'_> { + tcx.arena.alloc(inner_optimized_mir(tcx, did)) +} + +fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> { + if tcx.is_constructor(did.to_def_id()) { + // There's no reason to run all of the MIR passes on constructors when + // we can just output the MIR we want directly. This also saves const + // qualification and borrow checking the trouble of special casing + // constructors. + return shim::build_adt_ctor(tcx, did.to_def_id()); + } + + match tcx.hir().body_const_context(did) { + // Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked` + // which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it + // computes and caches its result. + Some(hir::ConstContext::ConstFn) => tcx.ensure_with_value().mir_for_ctfe(did), + None => {} + Some(other) => panic!("do not use `optimized_mir` for constants: {other:?}"), + } + debug!("about to call mir_drops_elaborated..."); + let body = tcx.mir_drops_elaborated_and_const_checked(did).steal(); + let mut body = remap_mir_for_const_eval_select(tcx, body, hir::Constness::NotConst); + debug!("body: {:#?}", body); + + if body.tainted_by_errors.is_some() { + return body; + } + + // If `mir_drops_elaborated_and_const_checked` found that the current body has unsatisfiable + // predicates, it will shrink the MIR to a single `unreachable` terminator. + // More generally, if MIR is a lone `unreachable`, there is nothing to optimize. + if let TerminatorKind::Unreachable = body.basic_blocks[START_BLOCK].terminator().kind + && body.basic_blocks[START_BLOCK].statements.is_empty() + { + return body; + } + + run_optimization_passes(tcx, &mut body); + + body +} + +/// Fetch all the promoteds of an item and prepare their MIR bodies to be ready for +/// constant evaluation once all substitutions become known. +fn promoted_mir(tcx: TyCtxt<'_>, def: LocalDefId) -> &IndexVec<Promoted, Body<'_>> { + if tcx.is_constructor(def.to_def_id()) { + return tcx.arena.alloc(IndexVec::new()); + } + + tcx.ensure_with_value().mir_borrowck(def); + let mut promoted = tcx.mir_promoted(def).1.steal(); + + for body in &mut promoted { + run_analysis_to_runtime_passes(tcx, body); + } + + tcx.arena.alloc(promoted) +} diff --git a/compiler/rustc_mir_transform/src/lower_intrinsics.rs b/compiler/rustc_mir_transform/src/lower_intrinsics.rs new file mode 100644 index 00000000000..18f588dccf6 --- /dev/null +++ b/compiler/rustc_mir_transform/src/lower_intrinsics.rs @@ -0,0 +1,283 @@ +//! Lowers intrinsic calls + +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_span::symbol::sym; + +pub struct LowerIntrinsics; + +impl<'tcx> MirPass<'tcx> for LowerIntrinsics { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let local_decls = &body.local_decls; + for block in body.basic_blocks.as_mut() { + let terminator = block.terminator.as_mut().unwrap(); + if let TerminatorKind::Call { func, args, destination, target, .. } = + &mut terminator.kind + && let ty::FnDef(def_id, generic_args) = *func.ty(local_decls, tcx).kind() + && tcx.is_intrinsic(def_id) + { + let intrinsic_name = tcx.item_name(def_id); + match intrinsic_name { + sym::unreachable => { + terminator.kind = TerminatorKind::Unreachable; + } + sym::forget => { + if let Some(target) = *target { + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::Use(Operand::Constant(Box::new(ConstOperand { + span: terminator.source_info.span, + user_ty: None, + const_: Const::zero_sized(tcx.types.unit), + }))), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + } + sym::copy_nonoverlapping => { + let target = target.unwrap(); + let mut args = args.drain(..); + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Intrinsic(Box::new( + NonDivergingIntrinsic::CopyNonOverlapping( + rustc_middle::mir::CopyNonOverlapping { + src: args.next().unwrap(), + dst: args.next().unwrap(), + count: args.next().unwrap(), + }, + ), + )), + }); + assert_eq!( + args.next(), + None, + "Extra argument for copy_non_overlapping intrinsic" + ); + drop(args); + terminator.kind = TerminatorKind::Goto { target }; + } + sym::assume => { + let target = target.unwrap(); + let mut args = args.drain(..); + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Intrinsic(Box::new( + NonDivergingIntrinsic::Assume(args.next().unwrap()), + )), + }); + assert_eq!( + args.next(), + None, + "Extra argument for copy_non_overlapping intrinsic" + ); + drop(args); + terminator.kind = TerminatorKind::Goto { target }; + } + sym::wrapping_add + | sym::wrapping_sub + | sym::wrapping_mul + | sym::unchecked_add + | sym::unchecked_sub + | sym::unchecked_mul + | sym::unchecked_div + | sym::unchecked_rem + | sym::unchecked_shl + | sym::unchecked_shr => { + let target = target.unwrap(); + let lhs; + let rhs; + { + let mut args = args.drain(..); + lhs = args.next().unwrap(); + rhs = args.next().unwrap(); + } + let bin_op = match intrinsic_name { + sym::wrapping_add => BinOp::Add, + sym::wrapping_sub => BinOp::Sub, + sym::wrapping_mul => BinOp::Mul, + sym::unchecked_add => BinOp::AddUnchecked, + sym::unchecked_sub => BinOp::SubUnchecked, + sym::unchecked_mul => BinOp::MulUnchecked, + sym::unchecked_div => BinOp::Div, + sym::unchecked_rem => BinOp::Rem, + sym::unchecked_shl => BinOp::ShlUnchecked, + sym::unchecked_shr => BinOp::ShrUnchecked, + _ => bug!("unexpected intrinsic"), + }; + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::BinaryOp(bin_op, Box::new((lhs, rhs))), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + sym::add_with_overflow | sym::sub_with_overflow | sym::mul_with_overflow => { + if let Some(target) = *target { + let lhs; + let rhs; + { + let mut args = args.drain(..); + lhs = args.next().unwrap(); + rhs = args.next().unwrap(); + } + let bin_op = match intrinsic_name { + sym::add_with_overflow => BinOp::Add, + sym::sub_with_overflow => BinOp::Sub, + sym::mul_with_overflow => BinOp::Mul, + _ => bug!("unexpected intrinsic"), + }; + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::CheckedBinaryOp(bin_op, Box::new((lhs, rhs))), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + } + sym::size_of | sym::min_align_of => { + if let Some(target) = *target { + let tp_ty = generic_args.type_at(0); + let null_op = match intrinsic_name { + sym::size_of => NullOp::SizeOf, + sym::min_align_of => NullOp::AlignOf, + _ => bug!("unexpected intrinsic"), + }; + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::NullaryOp(null_op, tp_ty), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + } + sym::read_via_copy => { + let [arg] = args.as_slice() else { + span_bug!(terminator.source_info.span, "Wrong number of arguments"); + }; + let derefed_place = if let Some(place) = arg.place() + && let Some(local) = place.as_local() + { + tcx.mk_place_deref(local.into()) + } else { + span_bug!( + terminator.source_info.span, + "Only passing a local is supported" + ); + }; + // Add new statement at the end of the block that does the read, and patch + // up the terminator. + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::Use(Operand::Copy(derefed_place)), + ))), + }); + terminator.kind = match *target { + None => { + // No target means this read something uninhabited, + // so it must be unreachable. + TerminatorKind::Unreachable + } + Some(target) => TerminatorKind::Goto { target }, + } + } + sym::write_via_move => { + let target = target.unwrap(); + let Ok([ptr, val]) = <[_; 2]>::try_from(std::mem::take(args)) else { + span_bug!( + terminator.source_info.span, + "Wrong number of arguments for write_via_move intrinsic", + ); + }; + let derefed_place = if let Some(place) = ptr.place() + && let Some(local) = place.as_local() + { + tcx.mk_place_deref(local.into()) + } else { + span_bug!( + terminator.source_info.span, + "Only passing a local is supported" + ); + }; + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + derefed_place, + Rvalue::Use(val), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + sym::discriminant_value => { + if let (Some(target), Some(arg)) = (*target, args[0].place()) { + let arg = tcx.mk_place_deref(arg); + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::Discriminant(arg), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + } + sym::offset => { + let target = target.unwrap(); + let Ok([ptr, delta]) = <[_; 2]>::try_from(std::mem::take(args)) else { + span_bug!( + terminator.source_info.span, + "Wrong number of arguments for offset intrinsic", + ); + }; + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::BinaryOp(BinOp::Offset, Box::new((ptr, delta))), + ))), + }); + terminator.kind = TerminatorKind::Goto { target }; + } + sym::transmute | sym::transmute_unchecked => { + let dst_ty = destination.ty(local_decls, tcx).ty; + let Ok([arg]) = <[_; 1]>::try_from(std::mem::take(args)) else { + span_bug!( + terminator.source_info.span, + "Wrong number of arguments for transmute intrinsic", + ); + }; + + // Always emit the cast, even if we transmute to an uninhabited type, + // because that lets CTFE and codegen generate better error messages + // when such a transmute actually ends up reachable. + block.statements.push(Statement { + source_info: terminator.source_info, + kind: StatementKind::Assign(Box::new(( + *destination, + Rvalue::Cast(CastKind::Transmute, arg, dst_ty), + ))), + }); + + if let Some(target) = *target { + terminator.kind = TerminatorKind::Goto { target }; + } else { + terminator.kind = TerminatorKind::Unreachable; + } + } + _ => {} + } + } + } + } +} diff --git a/compiler/rustc_mir_transform/src/lower_slice_len.rs b/compiler/rustc_mir_transform/src/lower_slice_len.rs new file mode 100644 index 00000000000..daeb56666f4 --- /dev/null +++ b/compiler/rustc_mir_transform/src/lower_slice_len.rs @@ -0,0 +1,76 @@ +//! This pass lowers calls to core::slice::len to just Len op. +//! It should run before inlining! + +use rustc_hir::def_id::DefId; +use rustc_index::IndexSlice; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; + +pub struct LowerSliceLenCalls; + +impl<'tcx> MirPass<'tcx> for LowerSliceLenCalls { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + lower_slice_len_calls(tcx, body) + } +} + +pub fn lower_slice_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let language_items = tcx.lang_items(); + let Some(slice_len_fn_item_def_id) = language_items.slice_len_fn() else { + // there is no language item to compare to :) + return; + }; + + // The one successor remains unchanged, so no need to invalidate + let basic_blocks = body.basic_blocks.as_mut_preserves_cfg(); + for block in basic_blocks { + // lower `<[_]>::len` calls + lower_slice_len_call(tcx, block, &body.local_decls, slice_len_fn_item_def_id); + } +} + +fn lower_slice_len_call<'tcx>( + tcx: TyCtxt<'tcx>, + block: &mut BasicBlockData<'tcx>, + local_decls: &IndexSlice<Local, LocalDecl<'tcx>>, + slice_len_fn_item_def_id: DefId, +) { + let terminator = block.terminator(); + if let TerminatorKind::Call { + func, + args, + destination, + target: Some(bb), + call_source: CallSource::Normal, + .. + } = &terminator.kind + // some heuristics for fast rejection + && let [arg] = &args[..] + && let Some(arg) = arg.place() + && let ty::FnDef(fn_def_id, _) = func.ty(local_decls, tcx).kind() + && *fn_def_id == slice_len_fn_item_def_id + { + // perform modifications from something like: + // _5 = core::slice::<impl [u8]>::len(move _6) -> bb1 + // into: + // _5 = Len(*_6) + // goto bb1 + + // make new RValue for Len + let deref_arg = tcx.mk_place_deref(arg); + let r_value = Rvalue::Len(deref_arg); + let len_statement_kind = StatementKind::Assign(Box::new((*destination, r_value))); + let add_statement = + Statement { kind: len_statement_kind, source_info: terminator.source_info }; + + // modify terminator into simple Goto + let new_terminator_kind = TerminatorKind::Goto { target: *bb }; + + block.statements.push(add_statement); + block.terminator_mut().kind = new_terminator_kind; + } +} diff --git a/compiler/rustc_mir_transform/src/match_branches.rs b/compiler/rustc_mir_transform/src/match_branches.rs new file mode 100644 index 00000000000..1c4aa37d57f --- /dev/null +++ b/compiler/rustc_mir_transform/src/match_branches.rs @@ -0,0 +1,180 @@ +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use std::iter; + +use super::simplify::simplify_cfg; + +pub struct MatchBranchSimplification; + +/// If a source block is found that switches between two blocks that are exactly +/// the same modulo const bool assignments (e.g., one assigns true another false +/// to the same place), merge a target block statements into the source block, +/// using Eq / Ne comparison with switch value where const bools value differ. +/// +/// For example: +/// +/// ```ignore (MIR) +/// bb0: { +/// switchInt(move _3) -> [42_isize: bb1, otherwise: bb2]; +/// } +/// +/// bb1: { +/// _2 = const true; +/// goto -> bb3; +/// } +/// +/// bb2: { +/// _2 = const false; +/// goto -> bb3; +/// } +/// ``` +/// +/// into: +/// +/// ```ignore (MIR) +/// bb0: { +/// _2 = Eq(move _3, const 42_isize); +/// goto -> bb3; +/// } +/// ``` + +impl<'tcx> MirPass<'tcx> for MatchBranchSimplification { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 1 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let def_id = body.source.def_id(); + let param_env = tcx.param_env_reveal_all_normalized(def_id); + + let bbs = body.basic_blocks.as_mut(); + let mut should_cleanup = false; + 'outer: for bb_idx in bbs.indices() { + if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {def_id:?} ")) { + continue; + } + + let (discr, val, first, second) = match bbs[bb_idx].terminator().kind { + TerminatorKind::SwitchInt { + discr: ref discr @ (Operand::Copy(_) | Operand::Move(_)), + ref targets, + .. + } if targets.iter().len() == 1 => { + let (value, target) = targets.iter().next().unwrap(); + // We require that this block and the two possible target blocks all be + // distinct. + if target == targets.otherwise() + || bb_idx == target + || bb_idx == targets.otherwise() + { + continue; + } + (discr, value, target, targets.otherwise()) + } + // Only optimize switch int statements + _ => continue, + }; + + // Check that destinations are identical, and if not, then don't optimize this block + if bbs[first].terminator().kind != bbs[second].terminator().kind { + continue; + } + + // Check that blocks are assignments of consts to the same place or same statement, + // and match up 1-1, if not don't optimize this block. + let first_stmts = &bbs[first].statements; + let scnd_stmts = &bbs[second].statements; + if first_stmts.len() != scnd_stmts.len() { + continue; + } + for (f, s) in iter::zip(first_stmts, scnd_stmts) { + match (&f.kind, &s.kind) { + // If two statements are exactly the same, we can optimize. + (f_s, s_s) if f_s == s_s => {} + + // If two statements are const bool assignments to the same place, we can optimize. + ( + StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))), + StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))), + ) if lhs_f == lhs_s + && f_c.const_.ty().is_bool() + && s_c.const_.ty().is_bool() + && f_c.const_.try_eval_bool(tcx, param_env).is_some() + && s_c.const_.try_eval_bool(tcx, param_env).is_some() => {} + + // Otherwise we cannot optimize. Try another block. + _ => continue 'outer, + } + } + // Take ownership of items now that we know we can optimize. + let discr = discr.clone(); + let discr_ty = discr.ty(&body.local_decls, tcx); + + // Introduce a temporary for the discriminant value. + let source_info = bbs[bb_idx].terminator().source_info; + let discr_local = body.local_decls.push(LocalDecl::new(discr_ty, source_info.span)); + + // We already checked that first and second are different blocks, + // and bb_idx has a different terminator from both of them. + let (from, first, second) = bbs.pick3_mut(bb_idx, first, second); + + let new_stmts = iter::zip(&first.statements, &second.statements).map(|(f, s)| { + match (&f.kind, &s.kind) { + (f_s, s_s) if f_s == s_s => (*f).clone(), + + ( + StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))), + StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))), + ) => { + // From earlier loop we know that we are dealing with bool constants only: + let f_b = f_c.const_.try_eval_bool(tcx, param_env).unwrap(); + let s_b = s_c.const_.try_eval_bool(tcx, param_env).unwrap(); + if f_b == s_b { + // Same value in both blocks. Use statement as is. + (*f).clone() + } else { + // Different value between blocks. Make value conditional on switch condition. + let size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size; + let const_cmp = Operand::const_from_scalar( + tcx, + discr_ty, + rustc_const_eval::interpret::Scalar::from_uint(val, size), + rustc_span::DUMMY_SP, + ); + let op = if f_b { BinOp::Eq } else { BinOp::Ne }; + let rhs = Rvalue::BinaryOp( + op, + Box::new((Operand::Copy(Place::from(discr_local)), const_cmp)), + ); + Statement { + source_info: f.source_info, + kind: StatementKind::Assign(Box::new((*lhs, rhs))), + } + } + } + + _ => unreachable!(), + } + }); + + from.statements + .push(Statement { source_info, kind: StatementKind::StorageLive(discr_local) }); + from.statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + Place::from(discr_local), + Rvalue::Use(discr), + ))), + }); + from.statements.extend(new_stmts); + from.statements + .push(Statement { source_info, kind: StatementKind::StorageDead(discr_local) }); + from.terminator_mut().kind = first.terminator().kind.clone(); + should_cleanup = true; + } + + if should_cleanup { + simplify_cfg(tcx, body); + } + } +} diff --git a/compiler/rustc_mir_transform/src/multiple_return_terminators.rs b/compiler/rustc_mir_transform/src/multiple_return_terminators.rs new file mode 100644 index 00000000000..64749a4b5b6 --- /dev/null +++ b/compiler/rustc_mir_transform/src/multiple_return_terminators.rs @@ -0,0 +1,43 @@ +//! This pass removes jumps to basic blocks containing only a return, and replaces them with a +//! return instead. + +use crate::simplify; +use rustc_index::bit_set::BitSet; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +pub struct MultipleReturnTerminators; + +impl<'tcx> MirPass<'tcx> for MultipleReturnTerminators { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 4 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // find basic blocks with no statement and a return terminator + let mut bbs_simple_returns = BitSet::new_empty(body.basic_blocks.len()); + let def_id = body.source.def_id(); + let bbs = body.basic_blocks_mut(); + for idx in bbs.indices() { + if bbs[idx].statements.is_empty() + && bbs[idx].terminator().kind == TerminatorKind::Return + { + bbs_simple_returns.insert(idx); + } + } + + for bb in bbs { + if !tcx.consider_optimizing(|| format!("MultipleReturnTerminators {def_id:?} ")) { + break; + } + + if let TerminatorKind::Goto { target } = bb.terminator().kind { + if bbs_simple_returns.contains(target) { + bb.terminator_mut().kind = TerminatorKind::Return; + } + } + } + + simplify::remove_dead_blocks(body) + } +} diff --git a/compiler/rustc_mir_transform/src/normalize_array_len.rs b/compiler/rustc_mir_transform/src/normalize_array_len.rs new file mode 100644 index 00000000000..128634bd7f2 --- /dev/null +++ b/compiler/rustc_mir_transform/src/normalize_array_len.rs @@ -0,0 +1,102 @@ +//! This pass eliminates casting of arrays into slices when their length +//! is taken using `.len()` method. Handy to preserve information in MIR for const prop + +use crate::ssa::SsaLocals; +use rustc_index::IndexVec; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; + +pub struct NormalizeArrayLen; + +impl<'tcx> MirPass<'tcx> for NormalizeArrayLen { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 3 + } + + #[instrument(level = "trace", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + normalize_array_len_calls(tcx, body) + } +} + +fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let ssa = SsaLocals::new(body); + + let slice_lengths = compute_slice_length(tcx, &ssa, body); + debug!(?slice_lengths); + + Replacer { tcx, slice_lengths }.visit_body_preserves_cfg(body); +} + +fn compute_slice_length<'tcx>( + tcx: TyCtxt<'tcx>, + ssa: &SsaLocals, + body: &Body<'tcx>, +) -> IndexVec<Local, Option<ty::Const<'tcx>>> { + let mut slice_lengths = IndexVec::from_elem(None, &body.local_decls); + + for (local, rvalue, _) in ssa.assignments(body) { + match rvalue { + Rvalue::Cast( + CastKind::PointerCoercion(ty::adjustment::PointerCoercion::Unsize), + operand, + cast_ty, + ) => { + let operand_ty = operand.ty(body, tcx); + debug!(?operand_ty); + if let Some(operand_ty) = operand_ty.builtin_deref(true) + && let ty::Array(_, len) = operand_ty.ty.kind() + && let Some(cast_ty) = cast_ty.builtin_deref(true) + && let ty::Slice(..) = cast_ty.ty.kind() + { + slice_lengths[local] = Some(*len); + } + } + // The length information is stored in the fat pointer, so we treat `operand` as a value. + Rvalue::Use(operand) => { + if let Some(rhs) = operand.place() + && let Some(rhs) = rhs.as_local() + { + slice_lengths[local] = slice_lengths[rhs]; + } + } + // The length information is stored in the fat pointer. + // Reborrowing copies length information from one pointer to the other. + Rvalue::Ref(_, _, rhs) | Rvalue::AddressOf(_, rhs) => { + if let [PlaceElem::Deref] = rhs.projection[..] { + slice_lengths[local] = slice_lengths[rhs.local]; + } + } + _ => {} + } + } + + slice_lengths +} + +struct Replacer<'tcx> { + tcx: TyCtxt<'tcx>, + slice_lengths: IndexVec<Local, Option<ty::Const<'tcx>>>, +} + +impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_rvalue(&mut self, rvalue: &mut Rvalue<'tcx>, loc: Location) { + if let Rvalue::Len(place) = rvalue + && let [PlaceElem::Deref] = &place.projection[..] + && let Some(len) = self.slice_lengths[place.local] + { + *rvalue = Rvalue::Use(Operand::Constant(Box::new(ConstOperand { + span: rustc_span::DUMMY_SP, + user_ty: None, + const_: Const::from_ty_const(len, self.tcx), + }))); + } + self.super_rvalue(rvalue, loc); + } +} diff --git a/compiler/rustc_mir_transform/src/nrvo.rs b/compiler/rustc_mir_transform/src/nrvo.rs new file mode 100644 index 00000000000..ff309bd10ec --- /dev/null +++ b/compiler/rustc_mir_transform/src/nrvo.rs @@ -0,0 +1,237 @@ +//! See the docs for [`RenameReturnPlace`]. + +use rustc_hir::Mutability; +use rustc_index::bit_set::HybridBitSet; +use rustc_middle::mir::visit::{MutVisitor, NonUseContext, PlaceContext, Visitor}; +use rustc_middle::mir::{self, BasicBlock, Local, Location}; +use rustc_middle::ty::TyCtxt; + +use crate::MirPass; + +/// This pass looks for MIR that always copies the same local into the return place and eliminates +/// the copy by renaming all uses of that local to `_0`. +/// +/// This allows LLVM to perform an optimization similar to the named return value optimization +/// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the +/// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by +/// value like so: +/// +/// ```rust +/// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] { +/// let mut buf = [0; 1024]; +/// init(&mut buf); +/// buf +/// } +/// ``` +/// +/// For now, this pass is very simple and only capable of eliminating a single copy. A more general +/// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and +/// [#71003], could yield even more benefits. +/// +/// [#47954]: https://github.com/rust-lang/rust/pull/47954 +/// [#71003]: https://github.com/rust-lang/rust/pull/71003 +pub struct RenameReturnPlace; + +impl<'tcx> MirPass<'tcx> for RenameReturnPlace { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // unsound: #111005 + sess.mir_opt_level() > 0 && sess.opts.unstable_opts.unsound_mir_opts + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) { + let def_id = body.source.def_id(); + let Some(returned_local) = local_eligible_for_nrvo(body) else { + debug!("`{:?}` was ineligible for NRVO", def_id); + return; + }; + + if !tcx.consider_optimizing(|| format!("RenameReturnPlace {def_id:?}")) { + return; + } + + debug!( + "`{:?}` was eligible for NRVO, making {:?} the return place", + def_id, returned_local + ); + + RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body_preserves_cfg(body); + + // Clean up the `NOP`s we inserted for statements made useless by our renaming. + for block_data in body.basic_blocks.as_mut_preserves_cfg() { + block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop); + } + + // Overwrite the debuginfo of `_0` with that of the renamed local. + let (renamed_decl, ret_decl) = + body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE); + + // Sometimes, the return place is assigned a local of a different but coercible type, for + // example `&mut T` instead of `&T`. Overwriting the `LocalInfo` for the return place means + // its type may no longer match the return type of its function. This doesn't cause a + // problem in codegen because these two types are layout-compatible, but may be unexpected. + debug!("_0: {:?} = {:?}: {:?}", ret_decl.ty, returned_local, renamed_decl.ty); + ret_decl.clone_from(renamed_decl); + + // The return place is always mutable. + ret_decl.mutability = Mutability::Mut; + } +} + +/// MIR that is eligible for the NRVO must fulfill two conditions: +/// 1. The return place must not be read prior to the `Return` terminator. +/// 2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the +/// only definition of the return place reaching the `Return` terminator. +/// +/// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned +/// to the return place along all possible paths through the control-flow graph. +fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> { + if IsReturnPlaceRead::run(body) { + return None; + } + + let mut copied_to_return_place = None; + for block in body.basic_blocks.indices() { + // Look for blocks with a `Return` terminator. + if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) { + continue; + } + + // Look for an assignment of a single local to the return place prior to the `Return`. + let returned_local = find_local_assigned_to_return_place(block, body)?; + match body.local_kind(returned_local) { + // FIXME: Can we do this for arguments as well? + mir::LocalKind::Arg => return None, + + mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"), + mir::LocalKind::Temp => {} + } + + // If multiple different locals are copied to the return place. We can't pick a + // single one to rename. + if copied_to_return_place.is_some_and(|old| old != returned_local) { + return None; + } + + copied_to_return_place = Some(returned_local); + } + + copied_to_return_place +} + +fn find_local_assigned_to_return_place( + start: BasicBlock, + body: &mut mir::Body<'_>, +) -> Option<Local> { + let mut block = start; + let mut seen = HybridBitSet::new_empty(body.basic_blocks.len()); + + // Iterate as long as `block` has exactly one predecessor that we have not yet visited. + while seen.insert(block) { + trace!("Looking for assignments to `_0` in {:?}", block); + + let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place); + if local.is_some() { + return local; + } + + match body.basic_blocks.predecessors()[block].as_slice() { + &[pred] => block = pred, + _ => return None, + } + } + + None +} + +// If this statement is an assignment of an unprojected local to the return place, +// return that local. +fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> { + if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind { + if lhs.as_local() == Some(mir::RETURN_PLACE) { + if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs { + return rhs.as_local(); + } + } + } + + None +} + +struct RenameToReturnPlace<'tcx> { + to_rename: Local, + tcx: TyCtxt<'tcx>, +} + +/// Replaces all uses of `self.to_rename` with `_0`. +impl<'tcx> MutVisitor<'tcx> for RenameToReturnPlace<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) { + // Remove assignments of the local being replaced to the return place, since it is now the + // return place: + // _0 = _1 + if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) { + stmt.kind = mir::StatementKind::Nop; + return; + } + + // Remove storage annotations for the local being replaced: + // StorageLive(_1) + if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) = + stmt.kind + { + if local == self.to_rename { + stmt.kind = mir::StatementKind::Nop; + return; + } + } + + self.super_statement(stmt, loc) + } + + fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) { + // Ignore the implicit "use" of the return place in a `Return` statement. + if let mir::TerminatorKind::Return = terminator.kind { + return; + } + + self.super_terminator(terminator, loc); + } + + fn visit_local(&mut self, l: &mut Local, ctxt: PlaceContext, _: Location) { + if *l == mir::RETURN_PLACE { + assert_eq!(ctxt, PlaceContext::NonUse(NonUseContext::VarDebugInfo)); + } else if *l == self.to_rename { + *l = mir::RETURN_PLACE; + } + } +} + +struct IsReturnPlaceRead(bool); + +impl IsReturnPlaceRead { + fn run(body: &mir::Body<'_>) -> bool { + let mut vis = IsReturnPlaceRead(false); + vis.visit_body(body); + vis.0 + } +} + +impl<'tcx> Visitor<'tcx> for IsReturnPlaceRead { + fn visit_local(&mut self, l: Local, ctxt: PlaceContext, _: Location) { + if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() { + self.0 = true; + } + } + + fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) { + // Ignore the implicit "use" of the return place in a `Return` statement. + if let mir::TerminatorKind::Return = terminator.kind { + return; + } + + self.super_terminator(terminator, loc); + } +} diff --git a/compiler/rustc_mir_transform/src/pass_manager.rs b/compiler/rustc_mir_transform/src/pass_manager.rs new file mode 100644 index 00000000000..c4eca18ff27 --- /dev/null +++ b/compiler/rustc_mir_transform/src/pass_manager.rs @@ -0,0 +1,195 @@ +use rustc_middle::mir::{self, Body, MirPhase, RuntimePhase}; +use rustc_middle::ty::TyCtxt; +use rustc_session::Session; + +use crate::{validate, MirPass}; + +/// Just like `MirPass`, except it cannot mutate `Body`. +pub trait MirLint<'tcx> { + fn name(&self) -> &'static str { + let name = std::any::type_name::<Self>(); + if let Some((_, tail)) = name.rsplit_once(':') { tail } else { name } + } + + fn is_enabled(&self, _sess: &Session) -> bool { + true + } + + fn run_lint(&self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>); +} + +/// An adapter for `MirLint`s that implements `MirPass`. +#[derive(Debug, Clone)] +pub struct Lint<T>(pub T); + +impl<'tcx, T> MirPass<'tcx> for Lint<T> +where + T: MirLint<'tcx>, +{ + fn name(&self) -> &'static str { + self.0.name() + } + + fn is_enabled(&self, sess: &Session) -> bool { + self.0.is_enabled(sess) + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + self.0.run_lint(tcx, body) + } + + fn is_mir_dump_enabled(&self) -> bool { + false + } +} + +pub struct WithMinOptLevel<T>(pub u32, pub T); + +impl<'tcx, T> MirPass<'tcx> for WithMinOptLevel<T> +where + T: MirPass<'tcx>, +{ + fn name(&self) -> &'static str { + self.1.name() + } + + fn is_enabled(&self, sess: &Session) -> bool { + sess.mir_opt_level() >= self.0 as usize + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + self.1.run_pass(tcx, body) + } +} + +/// Run the sequence of passes without validating the MIR after each pass. The MIR is still +/// validated at the end. +pub fn run_passes_no_validate<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + passes: &[&dyn MirPass<'tcx>], + phase_change: Option<MirPhase>, +) { + run_passes_inner(tcx, body, passes, phase_change, false); +} + +/// The optional `phase_change` is applied after executing all the passes, if present +pub fn run_passes<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + passes: &[&dyn MirPass<'tcx>], + phase_change: Option<MirPhase>, +) { + run_passes_inner(tcx, body, passes, phase_change, true); +} + +pub fn should_run_pass<'tcx, P>(tcx: TyCtxt<'tcx>, pass: &P) -> bool +where + P: MirPass<'tcx> + ?Sized, +{ + let name = pass.name(); + + let overridden_passes = &tcx.sess.opts.unstable_opts.mir_enable_passes; + let overridden = + overridden_passes.iter().rev().find(|(s, _)| s == &*name).map(|(_name, polarity)| { + trace!( + pass = %name, + "{} as requested by flag", + if *polarity { "Running" } else { "Not running" }, + ); + *polarity + }); + overridden.unwrap_or_else(|| pass.is_enabled(tcx.sess)) +} + +fn run_passes_inner<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + passes: &[&dyn MirPass<'tcx>], + phase_change: Option<MirPhase>, + validate_each: bool, +) { + let validate = validate_each & tcx.sess.opts.unstable_opts.validate_mir & !body.should_skip(); + let overridden_passes = &tcx.sess.opts.unstable_opts.mir_enable_passes; + trace!(?overridden_passes); + + let prof_arg = tcx.sess.prof.enabled().then(|| format!("{:?}", body.source.def_id())); + + if !body.should_skip() { + for pass in passes { + let name = pass.name(); + + if !should_run_pass(tcx, *pass) { + continue; + }; + + let dump_enabled = pass.is_mir_dump_enabled(); + + if dump_enabled { + dump_mir_for_pass(tcx, body, name, false); + } + if validate { + validate_body(tcx, body, format!("before pass {name}")); + } + + if let Some(prof_arg) = &prof_arg { + tcx.sess + .prof + .generic_activity_with_arg(pass.profiler_name(), &**prof_arg) + .run(|| pass.run_pass(tcx, body)); + } else { + pass.run_pass(tcx, body); + } + + if dump_enabled { + dump_mir_for_pass(tcx, body, name, true); + } + if validate { + validate_body(tcx, body, format!("after pass {name}")); + } + + body.pass_count += 1; + } + } + + if let Some(new_phase) = phase_change { + if body.phase >= new_phase { + panic!("Invalid MIR phase transition from {:?} to {:?}", body.phase, new_phase); + } + + body.phase = new_phase; + body.pass_count = 0; + + dump_mir_for_phase_change(tcx, body); + if validate || new_phase == MirPhase::Runtime(RuntimePhase::Optimized) { + validate_body(tcx, body, format!("after phase change to {}", new_phase.name())); + } + + body.pass_count = 1; + } +} + +pub fn validate_body<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, when: String) { + validate::Validator { when, mir_phase: body.phase }.run_pass(tcx, body); +} + +pub fn dump_mir_for_pass<'tcx>( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + pass_name: &str, + is_after: bool, +) { + mir::dump_mir( + tcx, + true, + pass_name, + if is_after { &"after" } else { &"before" }, + body, + |_, _| Ok(()), + ); +} + +pub fn dump_mir_for_phase_change<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) { + assert_eq!(body.pass_count, 0); + mir::dump_mir(tcx, true, body.phase.name(), &"after", body, |_, _| Ok(())) +} diff --git a/compiler/rustc_mir_transform/src/prettify.rs b/compiler/rustc_mir_transform/src/prettify.rs new file mode 100644 index 00000000000..7b77d032353 --- /dev/null +++ b/compiler/rustc_mir_transform/src/prettify.rs @@ -0,0 +1,149 @@ +//! These two passes provide no value to the compiler, so are off at every level. +//! +//! However, they can be enabled on the command line +//! (`-Zmir-enable-passes=+ReorderBasicBlocks,+ReorderLocals`) +//! to make the MIR easier to read for humans. + +use rustc_index::{bit_set::BitSet, IndexSlice, IndexVec}; +use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_session::Session; + +/// Rearranges the basic blocks into a *reverse post-order*. +/// +/// Thus after this pass, all the successors of a block are later than it in the +/// `IndexVec`, unless that successor is a back-edge (such as from a loop). +pub struct ReorderBasicBlocks; + +impl<'tcx> MirPass<'tcx> for ReorderBasicBlocks { + fn is_enabled(&self, _session: &Session) -> bool { + false + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let rpo: IndexVec<BasicBlock, BasicBlock> = + body.basic_blocks.reverse_postorder().iter().copied().collect(); + if rpo.iter().is_sorted() { + return; + } + + let mut updater = BasicBlockUpdater { map: rpo.invert_bijective_mapping(), tcx }; + debug_assert_eq!(updater.map[START_BLOCK], START_BLOCK); + updater.visit_body(body); + + permute(body.basic_blocks.as_mut(), &updater.map); + } +} + +/// Rearranges the locals into *use* order. +/// +/// Thus after this pass, a local with a smaller [`Location`] where it was first +/// assigned or referenced will have a smaller number. +/// +/// (Does not reorder arguments nor the [`RETURN_PLACE`].) +pub struct ReorderLocals; + +impl<'tcx> MirPass<'tcx> for ReorderLocals { + fn is_enabled(&self, _session: &Session) -> bool { + false + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let mut finder = + LocalFinder { map: IndexVec::new(), seen: BitSet::new_empty(body.local_decls.len()) }; + + // We can't reorder the return place or the arguments + for local in (0..=body.arg_count).map(Local::from_usize) { + finder.track(local); + } + + for (bb, bbd) in body.basic_blocks.iter_enumerated() { + finder.visit_basic_block_data(bb, bbd); + } + + // track everything in case there are some locals that we never saw, + // such as in non-block things like debug info or in non-uses. + for local in body.local_decls.indices() { + finder.track(local); + } + + if finder.map.iter().is_sorted() { + return; + } + + let mut updater = LocalUpdater { map: finder.map.invert_bijective_mapping(), tcx }; + + for local in (0..=body.arg_count).map(Local::from_usize) { + debug_assert_eq!(updater.map[local], local); + } + + updater.visit_body_preserves_cfg(body); + + permute(&mut body.local_decls, &updater.map); + } +} + +fn permute<I: rustc_index::Idx + Ord, T>(data: &mut IndexVec<I, T>, map: &IndexSlice<I, I>) { + // FIXME: It would be nice to have a less-awkward way to apply permutations, + // but I don't know one that exists. `sort_by_cached_key` has logic for it + // internally, but not in a way that we're allowed to use here. + let mut enumerated: Vec<_> = std::mem::take(data).into_iter_enumerated().collect(); + enumerated.sort_by_key(|p| map[p.0]); + *data = enumerated.into_iter().map(|p| p.1).collect(); +} + +struct BasicBlockUpdater<'tcx> { + map: IndexVec<BasicBlock, BasicBlock>, + tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for BasicBlockUpdater<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, _location: Location) { + for succ in terminator.successors_mut() { + *succ = self.map[*succ]; + } + } +} + +struct LocalFinder { + map: IndexVec<Local, Local>, + seen: BitSet<Local>, +} + +impl LocalFinder { + fn track(&mut self, l: Local) { + if self.seen.insert(l) { + self.map.push(l); + } + } +} + +impl<'tcx> Visitor<'tcx> for LocalFinder { + fn visit_local(&mut self, l: Local, context: PlaceContext, _location: Location) { + // Exclude non-uses to keep `StorageLive` from controlling where we put + // a `Local`, since it might not actually be assigned until much later. + if context.is_use() { + self.track(l); + } + } +} + +struct LocalUpdater<'tcx> { + pub map: IndexVec<Local, Local>, + pub tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) { + *l = self.map[*l]; + } +} diff --git a/compiler/rustc_mir_transform/src/ref_prop.rs b/compiler/rustc_mir_transform/src/ref_prop.rs new file mode 100644 index 00000000000..f13ab5b0f1f --- /dev/null +++ b/compiler/rustc_mir_transform/src/ref_prop.rs @@ -0,0 +1,414 @@ +use rustc_data_structures::fx::FxHashSet; +use rustc_index::bit_set::BitSet; +use rustc_index::IndexVec; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_mir_dataflow::impls::MaybeStorageDead; +use rustc_mir_dataflow::storage::always_storage_live_locals; +use rustc_mir_dataflow::Analysis; + +use crate::ssa::{SsaLocals, StorageLiveLocals}; + +/// Propagate references using SSA analysis. +/// +/// MIR building may produce a lot of borrow-dereference patterns. +/// +/// This pass aims to transform the following pattern: +/// _1 = &raw? mut? PLACE; +/// _3 = *_1; +/// _4 = &raw? mut? *_1; +/// +/// Into +/// _1 = &raw? mut? PLACE; +/// _3 = PLACE; +/// _4 = &raw? mut? PLACE; +/// +/// where `PLACE` is a direct or an indirect place expression. +/// +/// There are 3 properties that need to be upheld for this transformation to be legal: +/// - place stability: `PLACE` must refer to the same memory wherever it appears; +/// - pointer liveness: we must not introduce dereferences of dangling pointers; +/// - `&mut` borrow uniqueness. +/// +/// # Stability +/// +/// If `PLACE` is an indirect projection, if its of the form `(*LOCAL).PROJECTIONS` where: +/// - `LOCAL` is SSA; +/// - all projections in `PROJECTIONS` have a stable offset (no dereference and no indexing). +/// +/// If `PLACE` is a direct projection of a local, we consider it as constant if: +/// - the local is always live, or it has a single `StorageLive`; +/// - all projections have a stable offset. +/// +/// # Liveness +/// +/// When performing a substitution, we must take care not to introduce uses of dangling locals. +/// To ensure this, we walk the body with the `MaybeStorageDead` dataflow analysis: +/// - if we want to replace `*x` by reborrow `*y` and `y` may be dead, we allow replacement and +/// mark storage statements on `y` for removal; +/// - if we want to replace `*x` by non-reborrow `y` and `y` must be live, we allow replacement; +/// - if we want to replace `*x` by non-reborrow `y` and `y` may be dead, we do not replace. +/// +/// # Uniqueness +/// +/// For `&mut` borrows, we also need to preserve the uniqueness property: +/// we must avoid creating a state where we interleave uses of `*_1` and `_2`. +/// To do it, we only perform full substitution of mutable borrows: +/// we replace either all or none of the occurrences of `*_1`. +/// +/// Some care has to be taken when `_1` is copied in other locals. +/// _1 = &raw? mut? _2; +/// _3 = *_1; +/// _4 = _1 +/// _5 = *_4 +/// In such cases, fully substituting `_1` means fully substituting all of the copies. +/// +/// For immutable borrows, we do not need to preserve such uniqueness property, +/// so we perform all the possible substitutions without removing the `_1 = &_2` statement. +pub struct ReferencePropagation; + +impl<'tcx> MirPass<'tcx> for ReferencePropagation { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 2 + } + + #[instrument(level = "trace", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + while propagate_ssa(tcx, body) {} + } +} + +fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool { + let ssa = SsaLocals::new(body); + + let mut replacer = compute_replacement(tcx, body, &ssa); + debug!(?replacer.targets); + debug!(?replacer.allowed_replacements); + debug!(?replacer.storage_to_remove); + + replacer.visit_body_preserves_cfg(body); + + if replacer.any_replacement { + crate::simplify::remove_unused_definitions(body); + } + + replacer.any_replacement +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +enum Value<'tcx> { + /// Not a pointer, or we can't know. + Unknown, + /// We know the value to be a pointer to this place. + /// The boolean indicates whether the reference is mutable, subject the uniqueness rule. + Pointer(Place<'tcx>, bool), +} + +/// For each local, save the place corresponding to `*local`. +#[instrument(level = "trace", skip(tcx, body, ssa))] +fn compute_replacement<'tcx>( + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, + ssa: &SsaLocals, +) -> Replacer<'tcx> { + let always_live_locals = always_storage_live_locals(body); + + // Compute which locals have a single `StorageLive` statement ever. + let storage_live = StorageLiveLocals::new(body, &always_live_locals); + + // Compute `MaybeStorageDead` dataflow to check that we only replace when the pointee is + // definitely live. + let mut maybe_dead = MaybeStorageDead::new(always_live_locals) + .into_engine(tcx, body) + .iterate_to_fixpoint() + .into_results_cursor(body); + + // Map for each local to the pointee. + let mut targets = IndexVec::from_elem(Value::Unknown, &body.local_decls); + // Set of locals for which we will remove their storage statement. This is useful for + // reborrowed references. + let mut storage_to_remove = BitSet::new_empty(body.local_decls.len()); + + let fully_replacable_locals = fully_replacable_locals(ssa); + + // Returns true iff we can use `place` as a pointee. + // + // Note that we only need to verify that there is a single `StorageLive` statement, and we do + // not need to verify that it dominates all uses of that local. + // + // Consider the three statements: + // SL : StorageLive(a) + // DEF: b = &raw? mut? a + // USE: stuff that uses *b + // + // First, we recall that DEF is checked to dominate USE. Now imagine for the sake of + // contradiction there is a DEF -> SL -> USE path. Consider two cases: + // + // - DEF dominates SL. We always have UB the first time control flow reaches DEF, + // because the storage of `a` is dead. Since DEF dominates USE, that means we cannot + // reach USE and so our optimization is ok. + // + // - DEF does not dominate SL. Then there is a `START_BLOCK -> SL` path not including DEF. + // But we can extend this path to USE, meaning there is also a `START_BLOCK -> USE` path not + // including DEF. This violates the DEF dominates USE condition, and so is impossible. + let is_constant_place = |place: Place<'_>| { + // We only allow `Deref` as the first projection, to avoid surprises. + if place.projection.first() == Some(&PlaceElem::Deref) { + // `place == (*some_local).xxx`, it is constant only if `some_local` is constant. + // We approximate constness using SSAness. + ssa.is_ssa(place.local) && place.projection[1..].iter().all(PlaceElem::is_stable_offset) + } else { + storage_live.has_single_storage(place.local) + && place.projection[..].iter().all(PlaceElem::is_stable_offset) + } + }; + + let mut can_perform_opt = |target: Place<'tcx>, loc: Location| { + if target.projection.first() == Some(&PlaceElem::Deref) { + // We are creating a reborrow. As `place.local` is a reference, removing the storage + // statements should not make it much harder for LLVM to optimize. + storage_to_remove.insert(target.local); + true + } else { + // This is a proper dereference. We can only allow it if `target` is live. + maybe_dead.seek_after_primary_effect(loc); + let maybe_dead = maybe_dead.contains(target.local); + !maybe_dead + } + }; + + for (local, rvalue, location) in ssa.assignments(body) { + debug!(?local); + + // Only visit if we have something to do. + let Value::Unknown = targets[local] else { bug!() }; + + let ty = body.local_decls[local].ty; + + // If this is not a reference or pointer, do nothing. + if !ty.is_any_ptr() { + debug!("not a reference or pointer"); + continue; + } + + // Whether the current local is subject to the uniqueness rule. + let needs_unique = ty.is_mutable_ptr(); + + // If this a mutable reference that we cannot fully replace, mark it as unknown. + if needs_unique && !fully_replacable_locals.contains(local) { + debug!("not fully replaceable"); + continue; + } + + debug!(?rvalue); + match rvalue { + // This is a copy, just use the value we have in store for the previous one. + // As we are visiting in `assignment_order`, ie. reverse postorder, `rhs` should + // have been visited before. + Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) + | Rvalue::CopyForDeref(place) => { + if let Some(rhs) = place.as_local() + && ssa.is_ssa(rhs) + { + let target = targets[rhs]; + // Only see through immutable reference and pointers, as we do not know yet if + // mutable references are fully replaced. + if !needs_unique && matches!(target, Value::Pointer(..)) { + targets[local] = target; + } else { + targets[local] = + Value::Pointer(tcx.mk_place_deref(rhs.into()), needs_unique); + } + } + } + Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => { + let mut place = *place; + // Try to see through `place` in order to collapse reborrow chains. + if place.projection.first() == Some(&PlaceElem::Deref) + && let Value::Pointer(target, inner_needs_unique) = targets[place.local] + // Only see through immutable reference and pointers, as we do not know yet if + // mutable references are fully replaced. + && !inner_needs_unique + // Only collapse chain if the pointee is definitely live. + && can_perform_opt(target, location) + { + place = target.project_deeper(&place.projection[1..], tcx); + } + assert_ne!(place.local, local); + if is_constant_place(place) { + targets[local] = Value::Pointer(place, needs_unique); + } + } + // We do not know what to do, so keep as not-a-pointer. + _ => {} + } + } + + debug!(?targets); + + let mut finder = ReplacementFinder { + targets: &mut targets, + can_perform_opt, + allowed_replacements: FxHashSet::default(), + }; + let reachable_blocks = traversal::reachable_as_bitset(body); + for (bb, bbdata) in body.basic_blocks.iter_enumerated() { + // Only visit reachable blocks as we rely on dataflow. + if reachable_blocks.contains(bb) { + finder.visit_basic_block_data(bb, bbdata); + } + } + + let allowed_replacements = finder.allowed_replacements; + return Replacer { + tcx, + targets, + storage_to_remove, + allowed_replacements, + any_replacement: false, + }; + + struct ReplacementFinder<'a, 'tcx, F> { + targets: &'a mut IndexVec<Local, Value<'tcx>>, + can_perform_opt: F, + allowed_replacements: FxHashSet<(Local, Location)>, + } + + impl<'tcx, F> Visitor<'tcx> for ReplacementFinder<'_, 'tcx, F> + where + F: FnMut(Place<'tcx>, Location) -> bool, + { + fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) { + if matches!(ctxt, PlaceContext::NonUse(_)) { + // There is no need to check liveness for non-uses. + return; + } + + if place.projection.first() != Some(&PlaceElem::Deref) { + // This is not a dereference, nothing to do. + return; + } + + let mut place = place.as_ref(); + loop { + if let Value::Pointer(target, needs_unique) = self.targets[place.local] { + let perform_opt = (self.can_perform_opt)(target, loc); + debug!(?place, ?target, ?needs_unique, ?perform_opt); + + // This a reborrow chain, recursively allow the replacement. + // + // This also allows to detect cases where `target.local` is not replacable, + // and mark it as such. + if let &[PlaceElem::Deref] = &target.projection[..] { + assert!(perform_opt); + self.allowed_replacements.insert((target.local, loc)); + place.local = target.local; + continue; + } else if perform_opt { + self.allowed_replacements.insert((target.local, loc)); + } else if needs_unique { + // This mutable reference is not fully replacable, so drop it. + self.targets[place.local] = Value::Unknown; + } + } + + break; + } + } + } +} + +/// Compute the set of locals that can be fully replaced. +/// +/// We consider a local to be replacable iff it's only used in a `Deref` projection `*_local` or +/// non-use position (like storage statements and debuginfo). +fn fully_replacable_locals(ssa: &SsaLocals) -> BitSet<Local> { + let mut replacable = BitSet::new_empty(ssa.num_locals()); + + // First pass: for each local, whether its uses can be fully replaced. + for local in ssa.locals() { + if ssa.num_direct_uses(local) == 0 { + replacable.insert(local); + } + } + + // Second pass: a local can only be fully replaced if all its copies can. + ssa.meet_copy_equivalence(&mut replacable); + + replacable +} + +/// Utility to help performing subtitution of `*pattern` by `target`. +struct Replacer<'tcx> { + tcx: TyCtxt<'tcx>, + targets: IndexVec<Local, Value<'tcx>>, + storage_to_remove: BitSet<Local>, + allowed_replacements: FxHashSet<(Local, Location)>, + any_replacement: bool, +} + +impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_var_debug_info(&mut self, debuginfo: &mut VarDebugInfo<'tcx>) { + // If the debuginfo is a pointer to another place: + // - if it's a reborrow, see through it; + // - if it's a direct borrow, increase `debuginfo.references`. + while let VarDebugInfoContents::Place(ref mut place) = debuginfo.value + && place.projection.is_empty() + && let Value::Pointer(target, _) = self.targets[place.local] + && target.projection.iter().all(|p| p.can_use_in_debuginfo()) + { + if let Some((&PlaceElem::Deref, rest)) = target.projection.split_last() { + *place = Place::from(target.local).project_deeper(rest, self.tcx); + self.any_replacement = true; + } else { + break; + } + } + + // Simplify eventual projections left inside `debuginfo`. + self.super_var_debug_info(debuginfo); + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) { + loop { + if place.projection.first() != Some(&PlaceElem::Deref) { + return; + } + + let Value::Pointer(target, _) = self.targets[place.local] else { return }; + + let perform_opt = match ctxt { + PlaceContext::NonUse(NonUseContext::VarDebugInfo) => { + target.projection.iter().all(|p| p.can_use_in_debuginfo()) + } + PlaceContext::NonUse(_) => true, + _ => self.allowed_replacements.contains(&(target.local, loc)), + }; + + if !perform_opt { + return; + } + + *place = target.project_deeper(&place.projection[1..], self.tcx); + self.any_replacement = true; + } + } + + fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) { + match stmt.kind { + StatementKind::StorageLive(l) | StatementKind::StorageDead(l) + if self.storage_to_remove.contains(l) => + { + stmt.make_nop(); + } + // Do not remove assignments as they may still be useful for debuginfo. + _ => self.super_statement(stmt, loc), + } + } +} diff --git a/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs new file mode 100644 index 00000000000..095119e2e3f --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs @@ -0,0 +1,141 @@ +use rustc_index::bit_set::BitSet; +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use rustc_target::spec::PanicStrategy; + +/// 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. +pub struct RemoveNoopLandingPads; + +impl<'tcx> MirPass<'tcx> for RemoveNoopLandingPads { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.panic_strategy() != PanicStrategy::Abort + } + + fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!("remove_noop_landing_pads({:?})", body); + self.remove_nop_landing_pads(body) + } +} + +impl RemoveNoopLandingPads { + fn is_nop_landing_pad( + &self, + bb: BasicBlock, + body: &Body<'_>, + nop_landing_pads: &BitSet<BasicBlock>, + ) -> bool { + for stmt in &body[bb].statements { + match &stmt.kind { + StatementKind::FakeRead(..) + | StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::PlaceMention(..) + | StatementKind::AscribeUserType(..) + | StatementKind::Coverage(..) + | StatementKind::ConstEvalCounter + | StatementKind::Nop => { + // These are all noops in a landing pad + } + + StatementKind::Assign(box (place, Rvalue::Use(_) | Rvalue::Discriminant(_))) => { + if place.as_local().is_some() { + // Writing to a local (e.g., a drop flag) does not + // turn a landing pad to a non-nop + } else { + return false; + } + } + + StatementKind::Assign { .. } + | StatementKind::SetDiscriminant { .. } + | StatementKind::Deinit(..) + | StatementKind::Intrinsic(..) + | StatementKind::Retag { .. } => { + return false; + } + } + } + + let terminator = body[bb].terminator(); + match terminator.kind { + TerminatorKind::Goto { .. } + | TerminatorKind::UnwindResume + | TerminatorKind::SwitchInt { .. } + | TerminatorKind::FalseEdge { .. } + | TerminatorKind::FalseUnwind { .. } => { + terminator.successors().all(|succ| nop_landing_pads.contains(succ)) + } + TerminatorKind::CoroutineDrop + | TerminatorKind::Yield { .. } + | TerminatorKind::Return + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Unreachable + | TerminatorKind::Call { .. } + | TerminatorKind::Assert { .. } + | TerminatorKind::Drop { .. } + | TerminatorKind::InlineAsm { .. } => false, + } + } + + fn remove_nop_landing_pads(&self, body: &mut Body<'_>) { + debug!("body: {:#?}", body); + + // Skip the pass if there are no blocks with a resume terminator. + let has_resume = body + .basic_blocks + .iter_enumerated() + .any(|(_bb, block)| matches!(block.terminator().kind, TerminatorKind::UnwindResume)); + if !has_resume { + debug!("remove_noop_landing_pads: no resume block in MIR"); + return; + } + + // make sure there's a resume block without any statements + let resume_block = { + let mut patch = MirPatch::new(body); + let resume_block = patch.resume_block(); + patch.apply(body); + resume_block + }; + debug!("remove_noop_landing_pads: resume block is {:?}", resume_block); + + let mut jumps_folded = 0; + let mut landing_pads_removed = 0; + let mut nop_landing_pads = BitSet::new_empty(body.basic_blocks.len()); + + // This is a post-order traversal, so that if A post-dominates B + // then A will be visited before B. + let postorder: Vec<_> = traversal::postorder(body).map(|(bb, _)| bb).collect(); + for bb in postorder { + debug!(" processing {:?}", bb); + if let Some(unwind) = body[bb].terminator_mut().unwind_mut() { + if let UnwindAction::Cleanup(unwind_bb) = *unwind { + if nop_landing_pads.contains(unwind_bb) { + debug!(" removing noop landing pad"); + landing_pads_removed += 1; + *unwind = UnwindAction::Continue; + } + } + } + + for target in body[bb].terminator_mut().successors_mut() { + if *target != resume_block && nop_landing_pads.contains(*target) { + debug!(" folding noop jump to {:?} to resume block", target); + *target = resume_block; + jumps_folded += 1; + } + } + + let is_nop_landing_pad = self.is_nop_landing_pad(bb, body, &nop_landing_pads); + if is_nop_landing_pad { + nop_landing_pads.insert(bb); + } + debug!(" is_nop_landing_pad({:?}) = {}", bb, is_nop_landing_pad); + } + + debug!("removed {:?} jumps and {:?} landing pads", jumps_folded, landing_pads_removed); + } +} diff --git a/compiler/rustc_mir_transform/src/remove_place_mention.rs b/compiler/rustc_mir_transform/src/remove_place_mention.rs new file mode 100644 index 00000000000..78335b3b5e0 --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_place_mention.rs @@ -0,0 +1,22 @@ +//! This pass removes `PlaceMention` statement, which has no effect at codegen. + +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +pub struct RemovePlaceMention; + +impl<'tcx> MirPass<'tcx> for RemovePlaceMention { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + !sess.opts.unstable_opts.mir_keep_place_mention + } + + fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("Running RemovePlaceMention on {:?}", body.source); + for data in body.basic_blocks.as_mut_preserves_cfg() { + data.statements.retain(|statement| match statement.kind { + StatementKind::PlaceMention(..) | StatementKind::Nop => false, + _ => true, + }) + } + } +} diff --git a/compiler/rustc_mir_transform/src/remove_storage_markers.rs b/compiler/rustc_mir_transform/src/remove_storage_markers.rs new file mode 100644 index 00000000000..795f5232ee3 --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_storage_markers.rs @@ -0,0 +1,28 @@ +//! This pass removes storage markers if they won't be emitted during codegen. + +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +pub struct RemoveStorageMarkers; + +impl<'tcx> MirPass<'tcx> for RemoveStorageMarkers { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + if tcx.sess.emit_lifetime_markers() { + return; + } + + trace!("Running RemoveStorageMarkers on {:?}", body.source); + for data in body.basic_blocks.as_mut_preserves_cfg() { + data.statements.retain(|statement| match statement.kind { + StatementKind::StorageLive(..) + | StatementKind::StorageDead(..) + | StatementKind::Nop => false, + _ => true, + }) + } + } +} diff --git a/compiler/rustc_mir_transform/src/remove_uninit_drops.rs b/compiler/rustc_mir_transform/src/remove_uninit_drops.rs new file mode 100644 index 00000000000..7d12bcf2fa1 --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_uninit_drops.rs @@ -0,0 +1,156 @@ +use rustc_index::bit_set::ChunkedBitSet; +use rustc_middle::mir::{Body, TerminatorKind}; +use rustc_middle::ty::GenericArgsRef; +use rustc_middle::ty::{self, ParamEnv, Ty, TyCtxt, VariantDef}; +use rustc_mir_dataflow::impls::MaybeInitializedPlaces; +use rustc_mir_dataflow::move_paths::{LookupResult, MoveData, MovePathIndex}; +use rustc_mir_dataflow::{move_path_children_matching, Analysis, MaybeReachable, MoveDataParamEnv}; +use rustc_target::abi::FieldIdx; + +use crate::MirPass; + +/// Removes `Drop` terminators whose target is known to be uninitialized at +/// that point. +/// +/// This is redundant with drop elaboration, but we need to do it prior to const-checking, and +/// running const-checking after drop elaboration makes it optimization dependent, causing issues +/// like [#90770]. +/// +/// [#90770]: https://github.com/rust-lang/rust/issues/90770 +pub struct RemoveUninitDrops; + +impl<'tcx> MirPass<'tcx> for RemoveUninitDrops { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let param_env = tcx.param_env(body.source.def_id()); + let move_data = + MoveData::gather_moves(body, tcx, param_env, |ty| ty.needs_drop(tcx, param_env)); + + let mdpe = MoveDataParamEnv { move_data, param_env }; + let mut maybe_inits = MaybeInitializedPlaces::new(tcx, body, &mdpe) + .into_engine(tcx, body) + .pass_name("remove_uninit_drops") + .iterate_to_fixpoint() + .into_results_cursor(body); + + let mut to_remove = vec![]; + for (bb, block) in body.basic_blocks.iter_enumerated() { + let terminator = block.terminator(); + let TerminatorKind::Drop { place, .. } = &terminator.kind else { continue }; + + maybe_inits.seek_before_primary_effect(body.terminator_loc(bb)); + let MaybeReachable::Reachable(maybe_inits) = maybe_inits.get() else { continue }; + + // If there's no move path for the dropped place, it's probably a `Deref`. Let it alone. + let LookupResult::Exact(mpi) = mdpe.move_data.rev_lookup.find(place.as_ref()) else { + continue; + }; + + let should_keep = is_needs_drop_and_init( + tcx, + param_env, + maybe_inits, + &mdpe.move_data, + place.ty(body, tcx).ty, + mpi, + ); + if !should_keep { + to_remove.push(bb) + } + } + + for bb in to_remove { + let block = &mut body.basic_blocks_mut()[bb]; + + let TerminatorKind::Drop { target, .. } = &block.terminator().kind else { + unreachable!() + }; + + // Replace block terminator with `Goto`. + block.terminator_mut().kind = TerminatorKind::Goto { target: *target }; + } + } +} + +fn is_needs_drop_and_init<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + maybe_inits: &ChunkedBitSet<MovePathIndex>, + move_data: &MoveData<'tcx>, + ty: Ty<'tcx>, + mpi: MovePathIndex, +) -> bool { + // No need to look deeper if the root is definitely uninit or if it has no `Drop` impl. + if !maybe_inits.contains(mpi) || !ty.needs_drop(tcx, param_env) { + return false; + } + + let field_needs_drop_and_init = |(f, f_ty, mpi)| { + let child = move_path_children_matching(move_data, mpi, |x| x.is_field_to(f)); + let Some(mpi) = child else { + return Ty::needs_drop(f_ty, tcx, param_env); + }; + + is_needs_drop_and_init(tcx, param_env, maybe_inits, move_data, f_ty, mpi) + }; + + // This pass is only needed for const-checking, so it doesn't handle as many cases as + // `DropCtxt::open_drop`, since they aren't relevant in a const-context. + match ty.kind() { + ty::Adt(adt, args) => { + let dont_elaborate = adt.is_union() || adt.is_manually_drop() || adt.has_dtor(tcx); + if dont_elaborate { + return true; + } + + // Look at all our fields, or if we are an enum all our variants and their fields. + // + // If a field's projection *is not* present in `MoveData`, it has the same + // initializedness as its parent (maybe init). + // + // If its projection *is* present in `MoveData`, then the field may have been moved + // from separate from its parent. Recurse. + adt.variants().iter_enumerated().any(|(vid, variant)| { + // Enums have multiple variants, which are discriminated with a `Downcast` projection. + // Structs have a single variant, and don't use a `Downcast` projection. + let mpi = if adt.is_enum() { + let downcast = + move_path_children_matching(move_data, mpi, |x| x.is_downcast_to(vid)); + let Some(dc_mpi) = downcast else { + return variant_needs_drop(tcx, param_env, args, variant); + }; + + dc_mpi + } else { + mpi + }; + + variant + .fields + .iter() + .enumerate() + .map(|(f, field)| (FieldIdx::from_usize(f), field.ty(tcx, args), mpi)) + .any(field_needs_drop_and_init) + }) + } + + ty::Tuple(fields) => fields + .iter() + .enumerate() + .map(|(f, f_ty)| (FieldIdx::from_usize(f), f_ty, mpi)) + .any(field_needs_drop_and_init), + + _ => true, + } +} + +fn variant_needs_drop<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ParamEnv<'tcx>, + args: GenericArgsRef<'tcx>, + variant: &VariantDef, +) -> bool { + variant.fields.iter().any(|field| { + let f_ty = field.ty(tcx, args); + f_ty.needs_drop(tcx, param_env) + }) +} diff --git a/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs b/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs new file mode 100644 index 00000000000..5d528bed356 --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs @@ -0,0 +1,44 @@ +//! This pass replaces a drop of a type that does not need dropping, with a goto. +//! +//! When the MIR is built, we check `needs_drop` before emitting a `Drop` for a place. This pass is +//! useful because (unlike MIR building) it runs after type checking, so it can make use of +//! `Reveal::All` to provide more precise type information. + +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +use super::simplify::simplify_cfg; + +pub struct RemoveUnneededDrops; + +impl<'tcx> MirPass<'tcx> for RemoveUnneededDrops { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("Running RemoveUnneededDrops on {:?}", body.source); + + let did = body.source.def_id(); + let param_env = tcx.param_env_reveal_all_normalized(did); + let mut should_simplify = false; + + for block in body.basic_blocks.as_mut() { + let terminator = block.terminator_mut(); + if let TerminatorKind::Drop { place, target, .. } = terminator.kind { + let ty = place.ty(&body.local_decls, tcx); + if ty.ty.needs_drop(tcx, param_env) { + continue; + } + if !tcx.consider_optimizing(|| format!("RemoveUnneededDrops {did:?} ")) { + continue; + } + debug!("SUCCESS: replacing `drop` with goto({:?})", target); + terminator.kind = TerminatorKind::Goto { target }; + should_simplify = true; + } + } + + // if we applied optimizations, we potentially have some cfg to cleanup to + // make it easier for further passes + if should_simplify { + simplify_cfg(tcx, body); + } + } +} diff --git a/compiler/rustc_mir_transform/src/remove_zsts.rs b/compiler/rustc_mir_transform/src/remove_zsts.rs new file mode 100644 index 00000000000..34d57a45301 --- /dev/null +++ b/compiler/rustc_mir_transform/src/remove_zsts.rs @@ -0,0 +1,138 @@ +//! Removes operations on ZST places, and convert ZST operands to constants. + +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, Ty, TyCtxt}; + +pub struct RemoveZsts; + +impl<'tcx> MirPass<'tcx> for RemoveZsts { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // Avoid query cycles (coroutines require optimized MIR for layout). + if tcx.type_of(body.source.def_id()).instantiate_identity().is_coroutine() { + return; + } + + if !tcx.consider_optimizing(|| format!("RemoveZsts - {:?}", body.source.def_id())) { + return; + } + + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + let local_decls = &body.local_decls; + let mut replacer = Replacer { tcx, param_env, local_decls }; + for var_debug_info in &mut body.var_debug_info { + replacer.visit_var_debug_info(var_debug_info); + } + for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() { + replacer.visit_basic_block_data(bb, data); + } + } +} + +struct Replacer<'a, 'tcx> { + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + local_decls: &'a LocalDecls<'tcx>, +} + +/// A cheap, approximate check to avoid unnecessary `layout_of` calls. +fn maybe_zst(ty: Ty<'_>) -> bool { + match ty.kind() { + // maybe ZST (could be more precise) + ty::Adt(..) + | ty::Array(..) + | ty::Closure(..) + | ty::Tuple(..) + | ty::Alias(ty::Opaque, ..) => true, + // definitely ZST + ty::FnDef(..) | ty::Never => true, + // unreachable or can't be ZST + _ => false, + } +} + +impl<'tcx> Replacer<'_, 'tcx> { + fn known_to_be_zst(&self, ty: Ty<'tcx>) -> bool { + if !maybe_zst(ty) { + return false; + } + let Ok(layout) = self.tcx.layout_of(self.param_env.and(ty)) else { + return false; + }; + layout.is_zst() + } + + fn make_zst(&self, ty: Ty<'tcx>) -> ConstOperand<'tcx> { + debug_assert!(self.known_to_be_zst(ty)); + ConstOperand { + span: rustc_span::DUMMY_SP, + user_ty: None, + const_: Const::Val(ConstValue::ZeroSized, ty), + } + } +} + +impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_var_debug_info(&mut self, var_debug_info: &mut VarDebugInfo<'tcx>) { + match var_debug_info.value { + VarDebugInfoContents::Const(_) => {} + VarDebugInfoContents::Place(place) => { + let place_ty = place.ty(self.local_decls, self.tcx).ty; + if self.known_to_be_zst(place_ty) { + var_debug_info.value = VarDebugInfoContents::Const(self.make_zst(place_ty)) + } + } + } + } + + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) { + if let Operand::Constant(_) = operand { + return; + } + let op_ty = operand.ty(self.local_decls, self.tcx); + if self.known_to_be_zst(op_ty) + && self.tcx.consider_optimizing(|| { + format!("RemoveZsts - Operand: {operand:?} Location: {loc:?}") + }) + { + *operand = Operand::Constant(Box::new(self.make_zst(op_ty))) + } + } + + fn visit_statement(&mut self, statement: &mut Statement<'tcx>, loc: Location) { + let place_for_ty = match statement.kind { + StatementKind::Assign(box (place, ref rvalue)) => { + rvalue.is_safe_to_remove().then_some(place) + } + StatementKind::Deinit(box place) + | StatementKind::SetDiscriminant { box place, variant_index: _ } + | StatementKind::AscribeUserType(box (place, _), _) + | StatementKind::Retag(_, box place) + | StatementKind::PlaceMention(box place) + | StatementKind::FakeRead(box (_, place)) => Some(place), + StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => { + Some(local.into()) + } + StatementKind::Coverage(_) + | StatementKind::Intrinsic(_) + | StatementKind::Nop + | StatementKind::ConstEvalCounter => None, + }; + if let Some(place_for_ty) = place_for_ty + && let ty = place_for_ty.ty(self.local_decls, self.tcx).ty + && self.known_to_be_zst(ty) + { + statement.make_nop(); + } else { + self.super_statement(statement, loc); + } + } +} diff --git a/compiler/rustc_mir_transform/src/required_consts.rs b/compiler/rustc_mir_transform/src/required_consts.rs new file mode 100644 index 00000000000..abde6a47e83 --- /dev/null +++ b/compiler/rustc_mir_transform/src/required_consts.rs @@ -0,0 +1,27 @@ +use rustc_middle::mir::visit::Visitor; +use rustc_middle::mir::{Const, ConstOperand, Location}; +use rustc_middle::ty::ConstKind; + +pub struct RequiredConstsVisitor<'a, 'tcx> { + required_consts: &'a mut Vec<ConstOperand<'tcx>>, +} + +impl<'a, 'tcx> RequiredConstsVisitor<'a, 'tcx> { + pub fn new(required_consts: &'a mut Vec<ConstOperand<'tcx>>) -> Self { + RequiredConstsVisitor { required_consts } + } +} + +impl<'tcx> Visitor<'tcx> for RequiredConstsVisitor<'_, 'tcx> { + fn visit_constant(&mut self, constant: &ConstOperand<'tcx>, _: Location) { + let const_ = constant.const_; + match const_ { + Const::Ty(c) => match c.kind() { + ConstKind::Param(_) | ConstKind::Error(_) | ConstKind::Value(_) => {} + _ => bug!("only ConstKind::Param/Value should be encountered here, got {:#?}", c), + }, + Const::Unevaluated(..) => self.required_consts.push(*constant), + Const::Val(..) => {} + } + } +} diff --git a/compiler/rustc_mir_transform/src/reveal_all.rs b/compiler/rustc_mir_transform/src/reveal_all.rs new file mode 100644 index 00000000000..4d2eca57840 --- /dev/null +++ b/compiler/rustc_mir_transform/src/reveal_all.rs @@ -0,0 +1,71 @@ +//! Normalizes MIR in RevealAll mode. + +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, Ty, TyCtxt}; + +pub struct RevealAll; + +impl<'tcx> MirPass<'tcx> for RevealAll { + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + RevealAllVisitor { tcx, param_env }.visit_body_preserves_cfg(body); + } +} + +struct RevealAllVisitor<'tcx> { + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for RevealAllVisitor<'tcx> { + #[inline] + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + #[inline] + fn visit_place( + &mut self, + place: &mut Place<'tcx>, + _context: PlaceContext, + _location: Location, + ) { + // Performance optimization: don't reintern if there is no `OpaqueCast` to remove. + if place.projection.iter().all(|elem| !matches!(elem, ProjectionElem::OpaqueCast(_))) { + return; + } + // `OpaqueCast` projections are only needed if there are opaque types on which projections are performed. + // After the `RevealAll` pass, all opaque types are replaced with their hidden types, so we don't need these + // projections anymore. + place.projection = self.tcx.mk_place_elems( + &place + .projection + .into_iter() + .filter(|elem| !matches!(elem, ProjectionElem::OpaqueCast(_))) + .collect::<Vec<_>>(), + ); + self.super_place(place, _context, _location); + } + + #[inline] + fn visit_constant(&mut self, constant: &mut ConstOperand<'tcx>, location: Location) { + // We have to use `try_normalize_erasing_regions` here, since it's + // possible that we visit impossible-to-satisfy where clauses here, + // see #91745 + if let Ok(c) = self.tcx.try_normalize_erasing_regions(self.param_env, constant.const_) { + constant.const_ = c; + } + self.super_constant(constant, location); + } + + #[inline] + fn visit_ty(&mut self, ty: &mut Ty<'tcx>, _: TyContext) { + // We have to use `try_normalize_erasing_regions` here, since it's + // possible that we visit impossible-to-satisfy where clauses here, + // see #91745 + if let Ok(t) = self.tcx.try_normalize_erasing_regions(self.param_env, *ty) { + *ty = t; + } + } +} diff --git a/compiler/rustc_mir_transform/src/separate_const_switch.rs b/compiler/rustc_mir_transform/src/separate_const_switch.rs new file mode 100644 index 00000000000..6e22690d8da --- /dev/null +++ b/compiler/rustc_mir_transform/src/separate_const_switch.rs @@ -0,0 +1,343 @@ +//! A pass that duplicates switch-terminated blocks +//! into a new copy for each predecessor, provided +//! the predecessor sets the value being switched +//! over to a constant. +//! +//! The purpose of this pass is to help constant +//! propagation passes to simplify the switch terminator +//! of the copied blocks into gotos when some predecessors +//! statically determine the output of switches. +//! +//! ```text +//! x = 12 --- ---> something +//! \ / 12 +//! --> switch x +//! / \ otherwise +//! x = y --- ---> something else +//! ``` +//! becomes +//! ```text +//! x = 12 ---> switch x ------> something +//! \ / 12 +//! X +//! / \ otherwise +//! x = y ---> switch x ------> something else +//! ``` +//! so it can hopefully later be turned by another pass into +//! ```text +//! x = 12 --------------------> something +//! / 12 +//! / +//! / otherwise +//! x = y ---- switch x ------> something else +//! ``` +//! +//! This optimization is meant to cover simple cases +//! like `?` desugaring. For now, it thus focuses on +//! simplicity rather than completeness (it notably +//! sometimes duplicates abusively). + +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use smallvec::SmallVec; + +pub struct SeparateConstSwitch; + +impl<'tcx> MirPass<'tcx> for SeparateConstSwitch { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // This pass participates in some as-of-yet untested unsoundness found + // in https://github.com/rust-lang/rust/issues/112460 + sess.mir_opt_level() >= 2 && sess.opts.unstable_opts.unsound_mir_opts + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + // If execution did something, applying a simplification layer + // helps later passes optimize the copy away. + if separate_const_switch(body) > 0 { + super::simplify::simplify_cfg(tcx, body); + } + } +} + +/// Returns the amount of blocks that were duplicated +pub fn separate_const_switch(body: &mut Body<'_>) -> usize { + let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new(); + let predecessors = body.basic_blocks.predecessors(); + 'block_iter: for (block_id, block) in body.basic_blocks.iter_enumerated() { + if let TerminatorKind::SwitchInt { + discr: Operand::Copy(switch_place) | Operand::Move(switch_place), + .. + } = block.terminator().kind + { + // If the block is on an unwind path, do not + // apply the optimization as unwind paths + // rely on a unique parent invariant + if block.is_cleanup { + continue 'block_iter; + } + + // If the block has fewer than 2 predecessors, ignore it + // we could maybe chain blocks that have exactly one + // predecessor, but for now we ignore that + if predecessors[block_id].len() < 2 { + continue 'block_iter; + } + + // First, let's find a non-const place + // that determines the result of the switch + if let Some(switch_place) = find_determining_place(switch_place, block) { + // We now have an input place for which it would + // be interesting if predecessors assigned it from a const + + let mut predecessors_left = predecessors[block_id].len(); + 'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() { + let predecessor = &body.basic_blocks[predecessor_id]; + + // First we make sure the predecessor jumps + // in a reasonable way + match &predecessor.terminator().kind { + // The following terminators are + // unconditionally valid + TerminatorKind::Goto { .. } | TerminatorKind::SwitchInt { .. } => {} + + TerminatorKind::FalseEdge { real_target, .. } => { + if *real_target != block_id { + continue 'predec_iter; + } + } + + // The following terminators are not allowed + TerminatorKind::UnwindResume + | TerminatorKind::Drop { .. } + | TerminatorKind::Call { .. } + | TerminatorKind::Assert { .. } + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::Yield { .. } + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::InlineAsm { .. } + | TerminatorKind::CoroutineDrop => { + continue 'predec_iter; + } + } + + if is_likely_const(switch_place, predecessor) { + new_blocks.push((predecessor_id, block_id)); + predecessors_left -= 1; + if predecessors_left < 2 { + // If the original block only has one predecessor left, + // we have nothing left to do + break 'predec_iter; + } + } + } + } + } + } + + // Once the analysis is done, perform the duplication + let body_span = body.span; + let copied_blocks = new_blocks.len(); + let blocks = body.basic_blocks_mut(); + for (pred_id, target_id) in new_blocks { + let new_block = blocks[target_id].clone(); + let new_block_id = blocks.push(new_block); + let terminator = blocks[pred_id].terminator_mut(); + + match terminator.kind { + TerminatorKind::Goto { ref mut target } => { + *target = new_block_id; + } + + TerminatorKind::FalseEdge { ref mut real_target, .. } => { + if *real_target == target_id { + *real_target = new_block_id; + } + } + + TerminatorKind::SwitchInt { ref mut targets, .. } => { + targets.all_targets_mut().iter_mut().for_each(|x| { + if *x == target_id { + *x = new_block_id; + } + }); + } + + TerminatorKind::UnwindResume + | TerminatorKind::UnwindTerminate(_) + | TerminatorKind::Return + | TerminatorKind::Unreachable + | TerminatorKind::CoroutineDrop + | TerminatorKind::Assert { .. } + | TerminatorKind::FalseUnwind { .. } + | TerminatorKind::Drop { .. } + | TerminatorKind::Call { .. } + | TerminatorKind::InlineAsm { .. } + | TerminatorKind::Yield { .. } => { + span_bug!( + body_span, + "basic block terminator had unexpected kind {:?}", + &terminator.kind + ) + } + } + } + + copied_blocks +} + +/// This function describes a rough heuristic guessing +/// whether a place is last set with a const within the block. +/// Notably, it will be overly pessimistic in cases that are already +/// not handled by `separate_const_switch`. +fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool { + for statement in block.statements.iter().rev() { + match &statement.kind { + StatementKind::Assign(assign) => { + if assign.0 == tracked_place { + match assign.1 { + // These rvalues are definitely constant + Rvalue::Use(Operand::Constant(_)) + | Rvalue::Ref(_, _, _) + | Rvalue::AddressOf(_, _) + | Rvalue::Cast(_, Operand::Constant(_), _) + | Rvalue::NullaryOp(_, _) + | Rvalue::ShallowInitBox(_, _) + | Rvalue::UnaryOp(_, Operand::Constant(_)) => return true, + + // These rvalues make things ambiguous + Rvalue::Repeat(_, _) + | Rvalue::ThreadLocalRef(_) + | Rvalue::Len(_) + | Rvalue::BinaryOp(_, _) + | Rvalue::CheckedBinaryOp(_, _) + | Rvalue::Aggregate(_, _) => return false, + + // These rvalues move the place to track + Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _) + | Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) + | Rvalue::CopyForDeref(place) + | Rvalue::UnaryOp(_, Operand::Copy(place) | Operand::Move(place)) + | Rvalue::Discriminant(place) => tracked_place = place, + } + } + } + + // If the discriminant is set, it is always set + // as a constant, so the job is done. + // As we are **ignoring projections**, if the place + // we are tracking sees its discriminant be set, + // that means we had to be tracking the discriminant + // specifically (as it is impossible to switch over + // an enum directly, and if we were switching over + // its content, we would have had to at least cast it to + // some variant first) + StatementKind::SetDiscriminant { place, .. } => { + if **place == tracked_place { + return true; + } + } + + // These statements have no influence on the place + // we are interested in + StatementKind::FakeRead(_) + | StatementKind::Deinit(_) + | StatementKind::StorageLive(_) + | StatementKind::Retag(_, _) + | StatementKind::AscribeUserType(_, _) + | StatementKind::PlaceMention(..) + | StatementKind::Coverage(_) + | StatementKind::StorageDead(_) + | StatementKind::Intrinsic(_) + | StatementKind::ConstEvalCounter + | StatementKind::Nop => {} + } + } + + // If no good reason for the place to be const is found, + // give up. We could maybe go up predecessors, but in + // most cases giving up now should be sufficient. + false +} + +/// Finds a unique place that entirely determines the value +/// of `switch_place`, if it exists. This is only a heuristic. +/// Ideally we would like to track multiple determining places +/// for some edge cases, but one is enough for a lot of situations. +fn find_determining_place<'tcx>( + mut switch_place: Place<'tcx>, + block: &BasicBlockData<'tcx>, +) -> Option<Place<'tcx>> { + for statement in block.statements.iter().rev() { + match &statement.kind { + StatementKind::Assign(op) => { + if op.0 != switch_place { + continue; + } + + match op.1 { + // The following rvalues move the place + // that may be const in the predecessor + Rvalue::Use(Operand::Move(new) | Operand::Copy(new)) + | Rvalue::UnaryOp(_, Operand::Copy(new) | Operand::Move(new)) + | Rvalue::CopyForDeref(new) + | Rvalue::Cast(_, Operand::Move(new) | Operand::Copy(new), _) + | Rvalue::Repeat(Operand::Move(new) | Operand::Copy(new), _) + | Rvalue::Discriminant(new) + => switch_place = new, + + // The following rvalues might still make the block + // be valid but for now we reject them + Rvalue::Len(_) + | Rvalue::Ref(_, _, _) + | Rvalue::BinaryOp(_, _) + | Rvalue::CheckedBinaryOp(_, _) + | Rvalue::Aggregate(_, _) + + // The following rvalues definitely mean we cannot + // or should not apply this optimization + | Rvalue::Use(Operand::Constant(_)) + | Rvalue::Repeat(Operand::Constant(_), _) + | Rvalue::ThreadLocalRef(_) + | Rvalue::AddressOf(_, _) + | Rvalue::NullaryOp(_, _) + | Rvalue::ShallowInitBox(_, _) + | Rvalue::UnaryOp(_, Operand::Constant(_)) + | Rvalue::Cast(_, Operand::Constant(_), _) => return None, + } + } + + // These statements have no influence on the place + // we are interested in + StatementKind::FakeRead(_) + | StatementKind::Deinit(_) + | StatementKind::StorageLive(_) + | StatementKind::StorageDead(_) + | StatementKind::Retag(_, _) + | StatementKind::AscribeUserType(_, _) + | StatementKind::PlaceMention(..) + | StatementKind::Coverage(_) + | StatementKind::Intrinsic(_) + | StatementKind::ConstEvalCounter + | StatementKind::Nop => {} + + // If the discriminant is set, it is always set + // as a constant, so the job is already done. + // As we are **ignoring projections**, if the place + // we are tracking sees its discriminant be set, + // that means we had to be tracking the discriminant + // specifically (as it is impossible to switch over + // an enum directly, and if we were switching over + // its content, we would have had to at least cast it to + // some variant first) + StatementKind::SetDiscriminant { place, .. } => { + if **place == switch_place { + return None; + } + } + } + } + + Some(switch_place) +} diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs new file mode 100644 index 00000000000..fba73d5195b --- /dev/null +++ b/compiler/rustc_mir_transform/src/shim.rs @@ -0,0 +1,981 @@ +use rustc_hir as hir; +use rustc_hir::def_id::DefId; +use rustc_hir::lang_items::LangItem; +use rustc_middle::mir::*; +use rustc_middle::query::Providers; +use rustc_middle::ty::GenericArgs; +use rustc_middle::ty::{self, CoroutineArgs, EarlyBinder, Ty, TyCtxt}; +use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT}; + +use rustc_index::{Idx, IndexVec}; + +use rustc_span::Span; +use rustc_target::spec::abi::Abi; + +use std::fmt; +use std::iter; + +use crate::{ + abort_unwinding_calls, add_call_guards, add_moves_for_packed_drops, deref_separator, + pass_manager as pm, remove_noop_landing_pads, simplify, +}; +use rustc_middle::mir::patch::MirPatch; +use rustc_mir_dataflow::elaborate_drops::{self, DropElaborator, DropFlagMode, DropStyle}; + +pub fn provide(providers: &mut Providers) { + providers.mir_shims = make_shim; +} + +fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'tcx> { + debug!("make_shim({:?})", instance); + + let mut result = match instance { + ty::InstanceDef::Item(..) => bug!("item {:?} passed to make_shim", instance), + ty::InstanceDef::VTableShim(def_id) => { + let adjustment = Adjustment::Deref { source: DerefSource::MutPtr }; + build_call_shim(tcx, instance, Some(adjustment), CallKind::Direct(def_id)) + } + ty::InstanceDef::FnPtrShim(def_id, ty) => { + let trait_ = tcx.trait_of_item(def_id).unwrap(); + let adjustment = match tcx.fn_trait_kind_from_def_id(trait_) { + Some(ty::ClosureKind::FnOnce) => Adjustment::Identity, + Some(ty::ClosureKind::Fn) => Adjustment::Deref { source: DerefSource::ImmRef }, + Some(ty::ClosureKind::FnMut) => Adjustment::Deref { source: DerefSource::MutRef }, + None => bug!("fn pointer {:?} is not an fn", ty), + }; + + build_call_shim(tcx, instance, Some(adjustment), CallKind::Indirect(ty)) + } + // We are generating a call back to our def-id, which the + // codegen backend knows to turn to an actual call, be it + // a virtual call, or a direct call to a function for which + // indirect calls must be codegen'd differently than direct ones + // (such as `#[track_caller]`). + ty::InstanceDef::ReifyShim(def_id) => { + build_call_shim(tcx, instance, None, CallKind::Direct(def_id)) + } + ty::InstanceDef::ClosureOnceShim { call_once: _, track_caller: _ } => { + let fn_mut = tcx.require_lang_item(LangItem::FnMut, None); + let call_mut = tcx + .associated_items(fn_mut) + .in_definition_order() + .find(|it| it.kind == ty::AssocKind::Fn) + .unwrap() + .def_id; + + build_call_shim(tcx, instance, Some(Adjustment::RefMut), CallKind::Direct(call_mut)) + } + + ty::InstanceDef::DropGlue(def_id, ty) => { + // FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end + // of this function. Is this intentional? + if let Some(ty::Coroutine(coroutine_def_id, args, _)) = ty.map(Ty::kind) { + let body = tcx.optimized_mir(*coroutine_def_id).coroutine_drop().unwrap(); + let mut body = EarlyBinder::bind(body.clone()).instantiate(tcx, args); + debug!("make_shim({:?}) = {:?}", instance, body); + + pm::run_passes( + tcx, + &mut body, + &[ + &abort_unwinding_calls::AbortUnwindingCalls, + &add_call_guards::CriticalCallEdges, + ], + Some(MirPhase::Runtime(RuntimePhase::Optimized)), + ); + + return body; + } + + build_drop_shim(tcx, def_id, ty) + } + ty::InstanceDef::ThreadLocalShim(..) => build_thread_local_shim(tcx, instance), + ty::InstanceDef::CloneShim(def_id, ty) => build_clone_shim(tcx, def_id, ty), + ty::InstanceDef::FnPtrAddrShim(def_id, ty) => build_fn_ptr_addr_shim(tcx, def_id, ty), + ty::InstanceDef::Virtual(..) => { + bug!("InstanceDef::Virtual ({:?}) is for direct calls only", instance) + } + ty::InstanceDef::Intrinsic(_) => { + bug!("creating shims from intrinsics ({:?}) is unsupported", instance) + } + }; + debug!("make_shim({:?}) = untransformed {:?}", instance, result); + + // We don't validate MIR here because the shims may generate code that's + // only valid in a reveal-all param-env. However, since we do initial + // validation with the MirBuilt phase, which uses a user-facing param-env. + // This causes validation errors when TAITs are involved. + pm::run_passes_no_validate( + tcx, + &mut result, + &[ + &add_moves_for_packed_drops::AddMovesForPackedDrops, + &deref_separator::Derefer, + &remove_noop_landing_pads::RemoveNoopLandingPads, + &simplify::SimplifyCfg::MakeShim, + &abort_unwinding_calls::AbortUnwindingCalls, + &add_call_guards::CriticalCallEdges, + ], + Some(MirPhase::Runtime(RuntimePhase::Optimized)), + ); + + debug!("make_shim({:?}) = {:?}", instance, result); + + result +} + +#[derive(Copy, Clone, Debug, PartialEq)] +enum DerefSource { + /// `fn shim(&self) { inner(*self )}`. + ImmRef, + /// `fn shim(&mut self) { inner(*self )}`. + MutRef, + /// `fn shim(*mut self) { inner(*self )}`. + MutPtr, +} + +#[derive(Copy, Clone, Debug, PartialEq)] +enum Adjustment { + /// Pass the receiver as-is. + Identity, + + /// We get passed a reference or a raw pointer to `self` and call the target with `*self`. + /// + /// This either copies `self` (if `Self: Copy`, eg. for function items), or moves out of it + /// (for `VTableShim`, which effectively is passed `&own Self`). + Deref { source: DerefSource }, + + /// We get passed `self: Self` and call the target with `&mut self`. + /// + /// In this case we need to ensure that the `Self` is dropped after the call, as the callee + /// won't do it for us. + RefMut, +} + +#[derive(Copy, Clone, Debug, PartialEq)] +enum CallKind<'tcx> { + /// Call the `FnPtr` that was passed as the receiver. + Indirect(Ty<'tcx>), + + /// Call a known `FnDef`. + Direct(DefId), +} + +fn local_decls_for_sig<'tcx>( + sig: &ty::FnSig<'tcx>, + span: Span, +) -> IndexVec<Local, LocalDecl<'tcx>> { + iter::once(LocalDecl::new(sig.output(), span)) + .chain(sig.inputs().iter().map(|ity| LocalDecl::new(*ity, span).immutable())) + .collect() +} + +fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>) -> Body<'tcx> { + debug!("build_drop_shim(def_id={:?}, ty={:?})", def_id, ty); + + assert!(!matches!(ty, Some(ty) if ty.is_coroutine())); + + let args = if let Some(ty) = ty { + tcx.mk_args(&[ty.into()]) + } else { + GenericArgs::identity_for_item(tcx, def_id) + }; + let sig = tcx.fn_sig(def_id).instantiate(tcx, args); + let sig = tcx.instantiate_bound_regions_with_erased(sig); + let span = tcx.def_span(def_id); + + let source_info = SourceInfo::outermost(span); + + let return_block = BasicBlock::new(1); + let mut blocks = IndexVec::with_capacity(2); + let block = |blocks: &mut IndexVec<_, _>, kind| { + blocks.push(BasicBlockData { + statements: vec![], + terminator: Some(Terminator { source_info, kind }), + is_cleanup: false, + }) + }; + block(&mut blocks, TerminatorKind::Goto { target: return_block }); + block(&mut blocks, TerminatorKind::Return); + + let source = MirSource::from_instance(ty::InstanceDef::DropGlue(def_id, ty)); + let mut body = + new_body(source, blocks, local_decls_for_sig(&sig, span), sig.inputs().len(), span); + + // The first argument (index 0), but add 1 for the return value. + let mut dropee_ptr = Place::from(Local::new(1 + 0)); + if tcx.sess.opts.unstable_opts.mir_emit_retag { + // We want to treat the function argument as if it was passed by `&mut`. As such, we + // generate + // ``` + // temp = &mut *arg; + // Retag(temp, FnEntry) + // ``` + // It's important that we do this first, before anything that depends on `dropee_ptr` + // has been put into the body. + let reborrow = Rvalue::Ref( + tcx.lifetimes.re_erased, + BorrowKind::Mut { kind: MutBorrowKind::Default }, + tcx.mk_place_deref(dropee_ptr), + ); + let ref_ty = reborrow.ty(body.local_decls(), tcx); + dropee_ptr = body.local_decls.push(LocalDecl::new(ref_ty, span)).into(); + let new_statements = [ + StatementKind::Assign(Box::new((dropee_ptr, reborrow))), + StatementKind::Retag(RetagKind::FnEntry, Box::new(dropee_ptr)), + ]; + for s in new_statements { + body.basic_blocks_mut()[START_BLOCK] + .statements + .push(Statement { source_info, kind: s }); + } + } + + if ty.is_some() { + let patch = { + let param_env = tcx.param_env_reveal_all_normalized(def_id); + let mut elaborator = + DropShimElaborator { body: &body, patch: MirPatch::new(&body), tcx, param_env }; + let dropee = tcx.mk_place_deref(dropee_ptr); + let resume_block = elaborator.patch.resume_block(); + elaborate_drops::elaborate_drop( + &mut elaborator, + source_info, + dropee, + (), + return_block, + elaborate_drops::Unwind::To(resume_block), + START_BLOCK, + ); + elaborator.patch + }; + patch.apply(&mut body); + } + + body +} + +fn new_body<'tcx>( + source: MirSource<'tcx>, + basic_blocks: IndexVec<BasicBlock, BasicBlockData<'tcx>>, + local_decls: IndexVec<Local, LocalDecl<'tcx>>, + arg_count: usize, + span: Span, +) -> Body<'tcx> { + Body::new( + source, + basic_blocks, + IndexVec::from_elem_n( + SourceScopeData { + span, + parent_scope: None, + inlined: None, + inlined_parent_scope: None, + local_data: ClearCrossCrate::Clear, + }, + 1, + ), + local_decls, + IndexVec::new(), + arg_count, + vec![], + span, + None, + // FIXME(compiler-errors): is this correct? + None, + ) +} + +pub struct DropShimElaborator<'a, 'tcx> { + pub body: &'a Body<'tcx>, + pub patch: MirPatch<'tcx>, + pub tcx: TyCtxt<'tcx>, + pub param_env: ty::ParamEnv<'tcx>, +} + +impl fmt::Debug for DropShimElaborator<'_, '_> { + fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { + Ok(()) + } +} + +impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> { + type Path = (); + + fn patch(&mut self) -> &mut MirPatch<'tcx> { + &mut self.patch + } + fn body(&self) -> &'a Body<'tcx> { + self.body + } + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + fn param_env(&self) -> ty::ParamEnv<'tcx> { + self.param_env + } + + fn drop_style(&self, _path: Self::Path, mode: DropFlagMode) -> DropStyle { + match mode { + DropFlagMode::Shallow => { + // Drops for the contained fields are "shallow" and "static" - they will simply call + // the field's own drop glue. + DropStyle::Static + } + DropFlagMode::Deep => { + // The top-level drop is "deep" and "open" - it will be elaborated to a drop ladder + // dropping each field contained in the value. + DropStyle::Open + } + } + } + + fn get_drop_flag(&mut self, _path: Self::Path) -> Option<Operand<'tcx>> { + None + } + + fn clear_drop_flag(&mut self, _location: Location, _path: Self::Path, _mode: DropFlagMode) {} + + fn field_subpath(&self, _path: Self::Path, _field: FieldIdx) -> Option<Self::Path> { + None + } + fn deref_subpath(&self, _path: Self::Path) -> Option<Self::Path> { + None + } + fn downcast_subpath(&self, _path: Self::Path, _variant: VariantIdx) -> Option<Self::Path> { + Some(()) + } + fn array_subpath(&self, _path: Self::Path, _index: u64, _size: u64) -> Option<Self::Path> { + None + } +} + +fn build_thread_local_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'tcx> { + let def_id = instance.def_id(); + + let span = tcx.def_span(def_id); + let source_info = SourceInfo::outermost(span); + + let mut blocks = IndexVec::with_capacity(1); + blocks.push(BasicBlockData { + statements: vec![Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + Place::return_place(), + Rvalue::ThreadLocalRef(def_id), + ))), + }], + terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }), + is_cleanup: false, + }); + + new_body( + MirSource::from_instance(instance), + blocks, + IndexVec::from_raw(vec![LocalDecl::new(tcx.thread_local_ptr_ty(def_id), span)]), + 0, + span, + ) +} + +/// Builds a `Clone::clone` shim for `self_ty`. Here, `def_id` is `Clone::clone`. +fn build_clone_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'tcx>) -> Body<'tcx> { + debug!("build_clone_shim(def_id={:?})", def_id); + + let param_env = tcx.param_env_reveal_all_normalized(def_id); + + let mut builder = CloneShimBuilder::new(tcx, def_id, self_ty); + let is_copy = self_ty.is_copy_modulo_regions(tcx, param_env); + + let dest = Place::return_place(); + let src = tcx.mk_place_deref(Place::from(Local::new(1 + 0))); + + match self_ty.kind() { + _ if is_copy => builder.copy_shim(), + ty::Closure(_, args) => builder.tuple_like_shim(dest, src, args.as_closure().upvar_tys()), + ty::Tuple(..) => builder.tuple_like_shim(dest, src, self_ty.tuple_fields()), + ty::Coroutine(coroutine_def_id, args, hir::Movability::Movable) => { + builder.coroutine_shim(dest, src, *coroutine_def_id, args.as_coroutine()) + } + _ => bug!("clone shim for `{:?}` which is not `Copy` and is not an aggregate", self_ty), + }; + + builder.into_mir() +} + +struct CloneShimBuilder<'tcx> { + tcx: TyCtxt<'tcx>, + def_id: DefId, + local_decls: IndexVec<Local, LocalDecl<'tcx>>, + blocks: IndexVec<BasicBlock, BasicBlockData<'tcx>>, + span: Span, + sig: ty::FnSig<'tcx>, +} + +impl<'tcx> CloneShimBuilder<'tcx> { + fn new(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'tcx>) -> Self { + // we must subst the self_ty because it's + // otherwise going to be TySelf and we can't index + // or access fields of a Place of type TySelf. + let sig = tcx.fn_sig(def_id).instantiate(tcx, &[self_ty.into()]); + let sig = tcx.instantiate_bound_regions_with_erased(sig); + let span = tcx.def_span(def_id); + + CloneShimBuilder { + tcx, + def_id, + local_decls: local_decls_for_sig(&sig, span), + blocks: IndexVec::new(), + span, + sig, + } + } + + fn into_mir(self) -> Body<'tcx> { + let source = MirSource::from_instance(ty::InstanceDef::CloneShim( + self.def_id, + self.sig.inputs_and_output[0], + )); + new_body(source, self.blocks, self.local_decls, self.sig.inputs().len(), self.span) + } + + fn source_info(&self) -> SourceInfo { + SourceInfo::outermost(self.span) + } + + fn block( + &mut self, + statements: Vec<Statement<'tcx>>, + kind: TerminatorKind<'tcx>, + is_cleanup: bool, + ) -> BasicBlock { + let source_info = self.source_info(); + self.blocks.push(BasicBlockData { + statements, + terminator: Some(Terminator { source_info, kind }), + is_cleanup, + }) + } + + /// Gives the index of an upcoming BasicBlock, with an offset. + /// offset=0 will give you the index of the next BasicBlock, + /// offset=1 will give the index of the next-to-next block, + /// offset=-1 will give you the index of the last-created block + fn block_index_offset(&self, offset: usize) -> BasicBlock { + BasicBlock::new(self.blocks.len() + offset) + } + + fn make_statement(&self, kind: StatementKind<'tcx>) -> Statement<'tcx> { + Statement { source_info: self.source_info(), kind } + } + + fn copy_shim(&mut self) { + let rcvr = self.tcx.mk_place_deref(Place::from(Local::new(1 + 0))); + let ret_statement = self.make_statement(StatementKind::Assign(Box::new(( + Place::return_place(), + Rvalue::Use(Operand::Copy(rcvr)), + )))); + self.block(vec![ret_statement], TerminatorKind::Return, false); + } + + fn make_place(&mut self, mutability: Mutability, ty: Ty<'tcx>) -> Place<'tcx> { + let span = self.span; + let mut local = LocalDecl::new(ty, span); + if mutability.is_not() { + local = local.immutable(); + } + Place::from(self.local_decls.push(local)) + } + + fn make_clone_call( + &mut self, + dest: Place<'tcx>, + src: Place<'tcx>, + ty: Ty<'tcx>, + next: BasicBlock, + cleanup: BasicBlock, + ) { + let tcx = self.tcx; + + // `func == Clone::clone(&ty) -> ty` + let func_ty = Ty::new_fn_def(tcx, self.def_id, [ty]); + let func = Operand::Constant(Box::new(ConstOperand { + span: self.span, + user_ty: None, + const_: Const::zero_sized(func_ty), + })); + + let ref_loc = self.make_place( + Mutability::Not, + Ty::new_ref( + tcx, + tcx.lifetimes.re_erased, + ty::TypeAndMut { ty, mutbl: hir::Mutability::Not }, + ), + ); + + // `let ref_loc: &ty = &src;` + let statement = self.make_statement(StatementKind::Assign(Box::new(( + ref_loc, + Rvalue::Ref(tcx.lifetimes.re_erased, BorrowKind::Shared, src), + )))); + + // `let loc = Clone::clone(ref_loc);` + self.block( + vec![statement], + TerminatorKind::Call { + func, + args: vec![Operand::Move(ref_loc)], + destination: dest, + target: Some(next), + unwind: UnwindAction::Cleanup(cleanup), + call_source: CallSource::Normal, + fn_span: self.span, + }, + false, + ); + } + + fn clone_fields<I>( + &mut self, + dest: Place<'tcx>, + src: Place<'tcx>, + target: BasicBlock, + mut unwind: BasicBlock, + tys: I, + ) -> BasicBlock + where + I: IntoIterator<Item = Ty<'tcx>>, + { + // For an iterator of length n, create 2*n + 1 blocks. + for (i, ity) in tys.into_iter().enumerate() { + // Each iteration creates two blocks, referred to here as block 2*i and block 2*i + 1. + // + // Block 2*i attempts to clone the field. If successful it branches to 2*i + 2 (the + // next clone block). If unsuccessful it branches to the previous unwind block, which + // is initially the `unwind` argument passed to this function. + // + // Block 2*i + 1 is the unwind block for this iteration. It drops the cloned value + // created by block 2*i. We store this block in `unwind` so that the next clone block + // will unwind to it if cloning fails. + + let field = FieldIdx::new(i); + let src_field = self.tcx.mk_place_field(src, field, ity); + + let dest_field = self.tcx.mk_place_field(dest, field, ity); + + let next_unwind = self.block_index_offset(1); + let next_block = self.block_index_offset(2); + self.make_clone_call(dest_field, src_field, ity, next_block, unwind); + self.block( + vec![], + TerminatorKind::Drop { + place: dest_field, + target: unwind, + unwind: UnwindAction::Terminate(UnwindTerminateReason::InCleanup), + replace: false, + }, + /* is_cleanup */ true, + ); + unwind = next_unwind; + } + // If all clones succeed then we end up here. + self.block(vec![], TerminatorKind::Goto { target }, false); + unwind + } + + fn tuple_like_shim<I>(&mut self, dest: Place<'tcx>, src: Place<'tcx>, tys: I) + where + I: IntoIterator<Item = Ty<'tcx>>, + { + self.block(vec![], TerminatorKind::Goto { target: self.block_index_offset(3) }, false); + let unwind = self.block(vec![], TerminatorKind::UnwindResume, true); + let target = self.block(vec![], TerminatorKind::Return, false); + + let _final_cleanup_block = self.clone_fields(dest, src, target, unwind, tys); + } + + fn coroutine_shim( + &mut self, + dest: Place<'tcx>, + src: Place<'tcx>, + coroutine_def_id: DefId, + args: CoroutineArgs<'tcx>, + ) { + self.block(vec![], TerminatorKind::Goto { target: self.block_index_offset(3) }, false); + let unwind = self.block(vec![], TerminatorKind::UnwindResume, true); + // This will get overwritten with a switch once we know the target blocks + let switch = self.block(vec![], TerminatorKind::Unreachable, false); + let unwind = self.clone_fields(dest, src, switch, unwind, args.upvar_tys()); + let target = self.block(vec![], TerminatorKind::Return, false); + let unreachable = self.block(vec![], TerminatorKind::Unreachable, false); + let mut cases = Vec::with_capacity(args.state_tys(coroutine_def_id, self.tcx).count()); + for (index, state_tys) in args.state_tys(coroutine_def_id, self.tcx).enumerate() { + let variant_index = VariantIdx::new(index); + let dest = self.tcx.mk_place_downcast_unnamed(dest, variant_index); + let src = self.tcx.mk_place_downcast_unnamed(src, variant_index); + let clone_block = self.block_index_offset(1); + let start_block = self.block( + vec![self.make_statement(StatementKind::SetDiscriminant { + place: Box::new(Place::return_place()), + variant_index, + })], + TerminatorKind::Goto { target: clone_block }, + false, + ); + cases.push((index as u128, start_block)); + let _final_cleanup_block = self.clone_fields(dest, src, target, unwind, state_tys); + } + let discr_ty = args.discr_ty(self.tcx); + let temp = self.make_place(Mutability::Mut, discr_ty); + let rvalue = Rvalue::Discriminant(src); + let statement = self.make_statement(StatementKind::Assign(Box::new((temp, rvalue)))); + match &mut self.blocks[switch] { + BasicBlockData { statements, terminator: Some(Terminator { kind, .. }), .. } => { + statements.push(statement); + *kind = TerminatorKind::SwitchInt { + discr: Operand::Move(temp), + targets: SwitchTargets::new(cases.into_iter(), unreachable), + }; + } + BasicBlockData { terminator: None, .. } => unreachable!(), + } + } +} + +/// Builds a "call" shim for `instance`. The shim calls the function specified by `call_kind`, +/// first adjusting its first argument according to `rcvr_adjustment`. +#[instrument(level = "debug", skip(tcx), ret)] +fn build_call_shim<'tcx>( + tcx: TyCtxt<'tcx>, + instance: ty::InstanceDef<'tcx>, + rcvr_adjustment: Option<Adjustment>, + call_kind: CallKind<'tcx>, +) -> Body<'tcx> { + // `FnPtrShim` contains the fn pointer type that a call shim is being built for - this is used + // to substitute into the signature of the shim. It is not necessary for users of this + // MIR body to perform further substitutions (see `InstanceDef::has_polymorphic_mir_body`). + let (sig_args, untuple_args) = if let ty::InstanceDef::FnPtrShim(_, ty) = instance { + let sig = tcx.instantiate_bound_regions_with_erased(ty.fn_sig(tcx)); + + let untuple_args = sig.inputs(); + + // Create substitutions for the `Self` and `Args` generic parameters of the shim body. + let arg_tup = Ty::new_tup(tcx, untuple_args); + + (Some([ty.into(), arg_tup.into()]), Some(untuple_args)) + } else { + (None, None) + }; + + let def_id = instance.def_id(); + let sig = tcx.fn_sig(def_id); + let sig = sig.map_bound(|sig| tcx.instantiate_bound_regions_with_erased(sig)); + + assert_eq!(sig_args.is_some(), !instance.has_polymorphic_mir_body()); + let mut sig = if let Some(sig_args) = sig_args { + sig.instantiate(tcx, &sig_args) + } else { + sig.instantiate_identity() + }; + + if let CallKind::Indirect(fnty) = call_kind { + // `sig` determines our local decls, and thus the callee type in the `Call` terminator. This + // can only be an `FnDef` or `FnPtr`, but currently will be `Self` since the types come from + // the implemented `FnX` trait. + + // Apply the opposite adjustment to the MIR input. + let mut inputs_and_output = sig.inputs_and_output.to_vec(); + + // Initial signature is `fn(&? Self, Args) -> Self::Output` where `Args` is a tuple of the + // fn arguments. `Self` may be passed via (im)mutable reference or by-value. + assert_eq!(inputs_and_output.len(), 3); + + // `Self` is always the original fn type `ty`. The MIR call terminator is only defined for + // `FnDef` and `FnPtr` callees, not the `Self` type param. + let self_arg = &mut inputs_and_output[0]; + *self_arg = match rcvr_adjustment.unwrap() { + Adjustment::Identity => fnty, + Adjustment::Deref { source } => match source { + DerefSource::ImmRef => Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, fnty), + DerefSource::MutRef => Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, fnty), + DerefSource::MutPtr => Ty::new_mut_ptr(tcx, fnty), + }, + Adjustment::RefMut => bug!("`RefMut` is never used with indirect calls: {instance:?}"), + }; + sig.inputs_and_output = tcx.mk_type_list(&inputs_and_output); + } + + // FIXME(eddyb) avoid having this snippet both here and in + // `Instance::fn_sig` (introduce `InstanceDef::fn_sig`?). + if let ty::InstanceDef::VTableShim(..) = instance { + // Modify fn(self, ...) to fn(self: *mut Self, ...) + let mut inputs_and_output = sig.inputs_and_output.to_vec(); + let self_arg = &mut inputs_and_output[0]; + debug_assert!(tcx.generics_of(def_id).has_self && *self_arg == tcx.types.self_param); + *self_arg = Ty::new_mut_ptr(tcx, *self_arg); + sig.inputs_and_output = tcx.mk_type_list(&inputs_and_output); + } + + let span = tcx.def_span(def_id); + + debug!(?sig); + + let mut local_decls = local_decls_for_sig(&sig, span); + let source_info = SourceInfo::outermost(span); + + let rcvr_place = || { + assert!(rcvr_adjustment.is_some()); + Place::from(Local::new(1 + 0)) + }; + let mut statements = vec![]; + + let rcvr = rcvr_adjustment.map(|rcvr_adjustment| match rcvr_adjustment { + Adjustment::Identity => Operand::Move(rcvr_place()), + Adjustment::Deref { source: _ } => Operand::Move(tcx.mk_place_deref(rcvr_place())), + Adjustment::RefMut => { + // let rcvr = &mut rcvr; + let ref_rcvr = local_decls.push( + LocalDecl::new( + Ty::new_ref( + tcx, + tcx.lifetimes.re_erased, + ty::TypeAndMut { ty: sig.inputs()[0], mutbl: hir::Mutability::Mut }, + ), + span, + ) + .immutable(), + ); + let borrow_kind = BorrowKind::Mut { kind: MutBorrowKind::Default }; + statements.push(Statement { + source_info, + kind: StatementKind::Assign(Box::new(( + Place::from(ref_rcvr), + Rvalue::Ref(tcx.lifetimes.re_erased, borrow_kind, rcvr_place()), + ))), + }); + Operand::Move(Place::from(ref_rcvr)) + } + }); + + let (callee, mut args) = match call_kind { + // `FnPtr` call has no receiver. Args are untupled below. + CallKind::Indirect(_) => (rcvr.unwrap(), vec![]), + + // `FnDef` call with optional receiver. + CallKind::Direct(def_id) => { + let ty = tcx.type_of(def_id).instantiate_identity(); + ( + Operand::Constant(Box::new(ConstOperand { + span, + user_ty: None, + const_: Const::zero_sized(ty), + })), + rcvr.into_iter().collect::<Vec<_>>(), + ) + } + }; + + let mut arg_range = 0..sig.inputs().len(); + + // Take the `self` ("receiver") argument out of the range (it's adjusted above). + if rcvr_adjustment.is_some() { + arg_range.start += 1; + } + + // Take the last argument, if we need to untuple it (handled below). + if untuple_args.is_some() { + arg_range.end -= 1; + } + + // Pass all of the non-special arguments directly. + args.extend(arg_range.map(|i| Operand::Move(Place::from(Local::new(1 + i))))); + + // Untuple the last argument, if we have to. + if let Some(untuple_args) = untuple_args { + let tuple_arg = Local::new(1 + (sig.inputs().len() - 1)); + args.extend(untuple_args.iter().enumerate().map(|(i, ity)| { + Operand::Move(tcx.mk_place_field(Place::from(tuple_arg), FieldIdx::new(i), *ity)) + })); + } + + let n_blocks = if let Some(Adjustment::RefMut) = rcvr_adjustment { 5 } else { 2 }; + let mut blocks = IndexVec::with_capacity(n_blocks); + let block = |blocks: &mut IndexVec<_, _>, statements, kind, is_cleanup| { + blocks.push(BasicBlockData { + statements, + terminator: Some(Terminator { source_info, kind }), + is_cleanup, + }) + }; + + // BB #0 + block( + &mut blocks, + statements, + TerminatorKind::Call { + func: callee, + args, + destination: Place::return_place(), + target: Some(BasicBlock::new(1)), + unwind: if let Some(Adjustment::RefMut) = rcvr_adjustment { + UnwindAction::Cleanup(BasicBlock::new(3)) + } else { + UnwindAction::Continue + }, + call_source: CallSource::Misc, + fn_span: span, + }, + false, + ); + + if let Some(Adjustment::RefMut) = rcvr_adjustment { + // BB #1 - drop for Self + block( + &mut blocks, + vec![], + TerminatorKind::Drop { + place: rcvr_place(), + target: BasicBlock::new(2), + unwind: UnwindAction::Continue, + replace: false, + }, + false, + ); + } + // BB #1/#2 - return + block(&mut blocks, vec![], TerminatorKind::Return, false); + if let Some(Adjustment::RefMut) = rcvr_adjustment { + // BB #3 - drop if closure panics + block( + &mut blocks, + vec![], + TerminatorKind::Drop { + place: rcvr_place(), + target: BasicBlock::new(4), + unwind: UnwindAction::Terminate(UnwindTerminateReason::InCleanup), + replace: false, + }, + /* is_cleanup */ true, + ); + + // BB #4 - resume + block(&mut blocks, vec![], TerminatorKind::UnwindResume, true); + } + + let mut body = + new_body(MirSource::from_instance(instance), blocks, local_decls, sig.inputs().len(), span); + + if let Abi::RustCall = sig.abi { + body.spread_arg = Some(Local::new(sig.inputs().len())); + } + + body +} + +pub fn build_adt_ctor(tcx: TyCtxt<'_>, ctor_id: DefId) -> Body<'_> { + debug_assert!(tcx.is_constructor(ctor_id)); + + let param_env = tcx.param_env_reveal_all_normalized(ctor_id); + + // Normalize the sig. + let sig = tcx + .fn_sig(ctor_id) + .instantiate_identity() + .no_bound_vars() + .expect("LBR in ADT constructor signature"); + let sig = tcx.normalize_erasing_regions(param_env, sig); + + let ty::Adt(adt_def, args) = sig.output().kind() else { + bug!("unexpected type for ADT ctor {:?}", sig.output()); + }; + + debug!("build_ctor: ctor_id={:?} sig={:?}", ctor_id, sig); + + let span = tcx.def_span(ctor_id); + + let local_decls = local_decls_for_sig(&sig, span); + + let source_info = SourceInfo::outermost(span); + + let variant_index = + if adt_def.is_enum() { adt_def.variant_index_with_ctor_id(ctor_id) } else { FIRST_VARIANT }; + + // Generate the following MIR: + // + // (return as Variant).field0 = arg0; + // (return as Variant).field1 = arg1; + // + // return; + debug!("build_ctor: variant_index={:?}", variant_index); + + let kind = AggregateKind::Adt(adt_def.did(), variant_index, args, None, None); + let variant = adt_def.variant(variant_index); + let statement = Statement { + kind: StatementKind::Assign(Box::new(( + Place::return_place(), + Rvalue::Aggregate( + Box::new(kind), + (0..variant.fields.len()) + .map(|idx| Operand::Move(Place::from(Local::new(idx + 1)))) + .collect(), + ), + ))), + source_info, + }; + + let start_block = BasicBlockData { + statements: vec![statement], + terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }), + is_cleanup: false, + }; + + let source = MirSource::item(ctor_id); + let body = new_body( + source, + IndexVec::from_elem_n(start_block, 1), + local_decls, + sig.inputs().len(), + span, + ); + + crate::pass_manager::dump_mir_for_phase_change(tcx, &body); + + body +} + +/// ```ignore (pseudo-impl) +/// impl FnPtr for fn(u32) { +/// fn addr(self) -> usize { +/// self as usize +/// } +/// } +/// ``` +fn build_fn_ptr_addr_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'tcx>) -> Body<'tcx> { + assert!(matches!(self_ty.kind(), ty::FnPtr(..)), "expected fn ptr, found {self_ty}"); + let span = tcx.def_span(def_id); + let Some(sig) = tcx.fn_sig(def_id).instantiate(tcx, &[self_ty.into()]).no_bound_vars() else { + span_bug!(span, "FnPtr::addr with bound vars for `{self_ty}`"); + }; + let locals = local_decls_for_sig(&sig, span); + + let source_info = SourceInfo::outermost(span); + // FIXME: use `expose_addr` once we figure out whether function pointers have meaningful provenance. + let rvalue = Rvalue::Cast( + CastKind::FnPtrToPtr, + Operand::Move(Place::from(Local::new(1))), + Ty::new_imm_ptr(tcx, tcx.types.unit), + ); + let stmt = Statement { + source_info, + kind: StatementKind::Assign(Box::new((Place::return_place(), rvalue))), + }; + let statements = vec![stmt]; + let start_block = BasicBlockData { + statements, + terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }), + is_cleanup: false, + }; + let source = MirSource::from_instance(ty::InstanceDef::FnPtrAddrShim(def_id, self_ty)); + new_body(source, IndexVec::from_elem_n(start_block, 1), locals, sig.inputs().len(), span) +} diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs new file mode 100644 index 00000000000..856a0f22771 --- /dev/null +++ b/compiler/rustc_mir_transform/src/simplify.rs @@ -0,0 +1,598 @@ +//! A number of passes which remove various redundancies in the CFG. +//! +//! The `SimplifyCfg` pass gets rid of unnecessary blocks in the CFG, whereas the `SimplifyLocals` +//! gets rid of all the unnecessary local variable declarations. +//! +//! The `SimplifyLocals` pass is kinda expensive and therefore not very suitable to be run often. +//! Most of the passes should not care or be impacted in meaningful ways due to extra locals +//! either, so running the pass once, right before codegen, should suffice. +//! +//! On the other side of the spectrum, the `SimplifyCfg` pass is considerably cheap to run, thus +//! one should run it after every pass which may modify CFG in significant ways. This pass must +//! also be run before any analysis passes because it removes dead blocks, and some of these can be +//! ill-typed. +//! +//! The cause of this typing issue is typeck allowing most blocks whose end is not reachable have +//! an arbitrary return type, rather than having the usual () return type (as a note, typeck's +//! notion of reachability is in fact slightly weaker than MIR CFG reachability - see #31617). A +//! standard example of the situation is: +//! +//! ```rust +//! fn example() { +//! let _a: char = { return; }; +//! } +//! ``` +//! +//! Here the block (`{ return; }`) has the return type `char`, rather than `()`, but the MIR we +//! naively generate still contains the `_a = ()` write in the unreachable block "after" the +//! return. + +use rustc_data_structures::fx::FxIndexSet; +use rustc_index::{Idx, IndexSlice, IndexVec}; +use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor}; +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; +use smallvec::SmallVec; + +pub enum SimplifyCfg { + Initial, + PromoteConsts, + RemoveFalseEdges, + EarlyOpt, + ElaborateDrops, + Final, + MakeShim, + AfterUninhabitedEnumBranching, +} + +impl SimplifyCfg { + pub fn name(&self) -> &'static str { + match self { + SimplifyCfg::Initial => "SimplifyCfg-initial", + SimplifyCfg::PromoteConsts => "SimplifyCfg-promote-consts", + SimplifyCfg::RemoveFalseEdges => "SimplifyCfg-remove-false-edges", + SimplifyCfg::EarlyOpt => "SimplifyCfg-early-opt", + SimplifyCfg::ElaborateDrops => "SimplifyCfg-elaborate-drops", + SimplifyCfg::Final => "SimplifyCfg-final", + SimplifyCfg::MakeShim => "SimplifyCfg-make_shim", + SimplifyCfg::AfterUninhabitedEnumBranching => { + "SimplifyCfg-after-uninhabited-enum-branching" + } + } + } +} + +pub fn simplify_cfg<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + CfgSimplifier::new(body).simplify(); + remove_duplicate_unreachable_blocks(tcx, body); + remove_dead_blocks(body); + + // FIXME: Should probably be moved into some kind of pass manager + body.basic_blocks_mut().raw.shrink_to_fit(); +} + +impl<'tcx> MirPass<'tcx> for SimplifyCfg { + fn name(&self) -> &'static str { + self.name() + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!("SimplifyCfg({:?}) - simplifying {:?}", self.name(), body.source); + simplify_cfg(tcx, body); + } +} + +pub struct CfgSimplifier<'a, 'tcx> { + basic_blocks: &'a mut IndexSlice<BasicBlock, BasicBlockData<'tcx>>, + pred_count: IndexVec<BasicBlock, u32>, +} + +impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> { + pub fn new(body: &'a mut Body<'tcx>) -> Self { + let mut pred_count = IndexVec::from_elem(0u32, &body.basic_blocks); + + // we can't use mir.predecessors() here because that counts + // dead blocks, which we don't want to. + pred_count[START_BLOCK] = 1; + + for (_, data) in traversal::preorder(body) { + if let Some(ref term) = data.terminator { + for tgt in term.successors() { + pred_count[tgt] += 1; + } + } + } + + let basic_blocks = body.basic_blocks_mut(); + + CfgSimplifier { basic_blocks, pred_count } + } + + pub fn simplify(mut self) { + self.strip_nops(); + + // Vec of the blocks that should be merged. We store the indices here, instead of the + // statements itself to avoid moving the (relatively) large statements twice. + // We do not push the statements directly into the target block (`bb`) as that is slower + // due to additional reallocations + let mut merged_blocks = Vec::new(); + loop { + let mut changed = false; + + for bb in self.basic_blocks.indices() { + if self.pred_count[bb] == 0 { + continue; + } + + debug!("simplifying {:?}", bb); + + let mut terminator = + self.basic_blocks[bb].terminator.take().expect("invalid terminator state"); + + for successor in terminator.successors_mut() { + self.collapse_goto_chain(successor, &mut changed); + } + + let mut inner_changed = true; + merged_blocks.clear(); + while inner_changed { + inner_changed = false; + inner_changed |= self.simplify_branch(&mut terminator); + inner_changed |= self.merge_successor(&mut merged_blocks, &mut terminator); + changed |= inner_changed; + } + + let statements_to_merge = + merged_blocks.iter().map(|&i| self.basic_blocks[i].statements.len()).sum(); + + if statements_to_merge > 0 { + let mut statements = std::mem::take(&mut self.basic_blocks[bb].statements); + statements.reserve(statements_to_merge); + for &from in &merged_blocks { + statements.append(&mut self.basic_blocks[from].statements); + } + self.basic_blocks[bb].statements = statements; + } + + self.basic_blocks[bb].terminator = Some(terminator); + } + + if !changed { + break; + } + } + } + + /// This function will return `None` if + /// * the block has statements + /// * the block has a terminator other than `goto` + /// * the block has no terminator (meaning some other part of the current optimization stole it) + fn take_terminator_if_simple_goto(&mut self, bb: BasicBlock) -> Option<Terminator<'tcx>> { + match self.basic_blocks[bb] { + BasicBlockData { + ref statements, + terminator: + ref mut terminator @ Some(Terminator { kind: TerminatorKind::Goto { .. }, .. }), + .. + } if statements.is_empty() => terminator.take(), + // if `terminator` is None, this means we are in a loop. In that + // case, let all the loop collapse to its entry. + _ => None, + } + } + + /// Collapse a goto chain starting from `start` + fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) { + // Using `SmallVec` here, because in some logs on libcore oli-obk saw many single-element + // goto chains. We should probably benchmark different sizes. + let mut terminators: SmallVec<[_; 1]> = Default::default(); + let mut current = *start; + while let Some(terminator) = self.take_terminator_if_simple_goto(current) { + let Terminator { kind: TerminatorKind::Goto { target }, .. } = terminator else { + unreachable!(); + }; + terminators.push((current, terminator)); + current = target; + } + let last = current; + *start = last; + while let Some((current, mut terminator)) = terminators.pop() { + let Terminator { kind: TerminatorKind::Goto { ref mut target }, .. } = terminator + else { + unreachable!(); + }; + *changed |= *target != last; + *target = last; + debug!("collapsing goto chain from {:?} to {:?}", current, target); + + if self.pred_count[current] == 1 { + // This is the last reference to current, so the pred-count to + // to target is moved into the current block. + self.pred_count[current] = 0; + } else { + self.pred_count[*target] += 1; + self.pred_count[current] -= 1; + } + self.basic_blocks[current].terminator = Some(terminator); + } + } + + // merge a block with 1 `goto` predecessor to its parent + fn merge_successor( + &mut self, + merged_blocks: &mut Vec<BasicBlock>, + terminator: &mut Terminator<'tcx>, + ) -> bool { + let target = match terminator.kind { + TerminatorKind::Goto { target } if self.pred_count[target] == 1 => target, + _ => return false, + }; + + debug!("merging block {:?} into {:?}", target, terminator); + *terminator = match self.basic_blocks[target].terminator.take() { + Some(terminator) => terminator, + None => { + // unreachable loop - this should not be possible, as we + // don't strand blocks, but handle it correctly. + return false; + } + }; + + merged_blocks.push(target); + self.pred_count[target] = 0; + + true + } + + // turn a branch with all successors identical to a goto + fn simplify_branch(&mut self, terminator: &mut Terminator<'tcx>) -> bool { + match terminator.kind { + TerminatorKind::SwitchInt { .. } => {} + _ => return false, + }; + + let first_succ = { + if let Some(first_succ) = terminator.successors().next() { + if terminator.successors().all(|s| s == first_succ) { + let count = terminator.successors().count(); + self.pred_count[first_succ] -= (count - 1) as u32; + first_succ + } else { + return false; + } + } else { + return false; + } + }; + + debug!("simplifying branch {:?}", terminator); + terminator.kind = TerminatorKind::Goto { target: first_succ }; + true + } + + fn strip_nops(&mut self) { + for blk in self.basic_blocks.iter_mut() { + blk.statements.retain(|stmt| !matches!(stmt.kind, StatementKind::Nop)) + } + } +} + +pub fn simplify_duplicate_switch_targets(terminator: &mut Terminator<'_>) { + if let TerminatorKind::SwitchInt { targets, .. } = &mut terminator.kind { + let otherwise = targets.otherwise(); + if targets.iter().any(|t| t.1 == otherwise) { + *targets = SwitchTargets::new( + targets.iter().filter(|t| t.1 != otherwise), + targets.otherwise(), + ); + } + } +} + +pub fn remove_duplicate_unreachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + struct OptApplier<'tcx> { + tcx: TyCtxt<'tcx>, + duplicates: FxIndexSet<BasicBlock>, + } + + impl<'tcx> MutVisitor<'tcx> for OptApplier<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) { + for target in terminator.successors_mut() { + // We don't have to check whether `target` is a cleanup block, because have + // entirely excluded cleanup blocks in building the set of duplicates. + if self.duplicates.contains(target) { + *target = self.duplicates[0]; + } + } + + simplify_duplicate_switch_targets(terminator); + + self.super_terminator(terminator, location); + } + } + + let unreachable_blocks = body + .basic_blocks + .iter_enumerated() + .filter(|(_, bb)| { + // CfgSimplifier::simplify leaves behind some unreachable basic blocks without a + // terminator. Those blocks will be deleted by remove_dead_blocks, but we run just + // before then so we need to handle missing terminators. + // We also need to prevent confusing cleanup and non-cleanup blocks. In practice we + // don't emit empty unreachable cleanup blocks, so this simple check suffices. + bb.terminator.is_some() && bb.is_empty_unreachable() && !bb.is_cleanup + }) + .map(|(block, _)| block) + .collect::<FxIndexSet<_>>(); + + if unreachable_blocks.len() > 1 { + OptApplier { tcx, duplicates: unreachable_blocks }.visit_body(body); + } +} + +pub fn remove_dead_blocks(body: &mut Body<'_>) { + let reachable = traversal::reachable_as_bitset(body); + let num_blocks = body.basic_blocks.len(); + if num_blocks == reachable.count() { + return; + } + + let basic_blocks = body.basic_blocks.as_mut(); + + let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect(); + let mut orig_index = 0; + let mut used_index = 0; + basic_blocks.raw.retain(|_| { + let keep = reachable.contains(BasicBlock::new(orig_index)); + if keep { + replacements[orig_index] = BasicBlock::new(used_index); + used_index += 1; + } + orig_index += 1; + keep + }); + + for block in basic_blocks { + for target in block.terminator_mut().successors_mut() { + *target = replacements[target.index()]; + } + } +} + +pub enum SimplifyLocals { + BeforeConstProp, + AfterGVN, + Final, +} + +impl<'tcx> MirPass<'tcx> for SimplifyLocals { + fn name(&self) -> &'static str { + match &self { + SimplifyLocals::BeforeConstProp => "SimplifyLocals-before-const-prop", + SimplifyLocals::AfterGVN => "SimplifyLocals-after-value-numbering", + SimplifyLocals::Final => "SimplifyLocals-final", + } + } + + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("running SimplifyLocals on {:?}", body.source); + simplify_locals(body, tcx); + } +} + +pub fn remove_unused_definitions<'tcx>(body: &mut Body<'tcx>) { + // First, we're going to get a count of *actual* uses for every `Local`. + let mut used_locals = UsedLocals::new(body); + + // Next, we're going to remove any `Local` with zero actual uses. When we remove those + // `Locals`, we're also going to subtract any uses of other `Locals` from the `used_locals` + // count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from + // `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a + // fixedpoint where there are no more unused locals. + remove_unused_definitions_helper(&mut used_locals, body); +} + +pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) { + // First, we're going to get a count of *actual* uses for every `Local`. + let mut used_locals = UsedLocals::new(body); + + // Next, we're going to remove any `Local` with zero actual uses. When we remove those + // `Locals`, we're also going to subtract any uses of other `Locals` from the `used_locals` + // count. For example, if we removed `_2 = discriminant(_1)`, then we'll subtract one from + // `use_counts[_1]`. That in turn might make `_1` unused, so we loop until we hit a + // fixedpoint where there are no more unused locals. + remove_unused_definitions_helper(&mut used_locals, body); + + // Finally, we'll actually do the work of shrinking `body.local_decls` and remapping the `Local`s. + let map = make_local_map(&mut body.local_decls, &used_locals); + + // Only bother running the `LocalUpdater` if we actually found locals to remove. + if map.iter().any(Option::is_none) { + // Update references to all vars and tmps now + let mut updater = LocalUpdater { map, tcx }; + updater.visit_body_preserves_cfg(body); + + body.local_decls.shrink_to_fit(); + } +} + +/// Construct the mapping while swapping out unused stuff out from the `vec`. +fn make_local_map<V>( + local_decls: &mut IndexVec<Local, V>, + used_locals: &UsedLocals, +) -> IndexVec<Local, Option<Local>> { + let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, local_decls); + let mut used = Local::new(0); + + for alive_index in local_decls.indices() { + // `is_used` treats the `RETURN_PLACE` and arguments as used. + if !used_locals.is_used(alive_index) { + continue; + } + + map[alive_index] = Some(used); + if alive_index != used { + local_decls.swap(alive_index, used); + } + used.increment_by(1); + } + local_decls.truncate(used.index()); + map +} + +/// Keeps track of used & unused locals. +struct UsedLocals { + increment: bool, + arg_count: u32, + use_count: IndexVec<Local, u32>, +} + +impl UsedLocals { + /// Determines which locals are used & unused in the given body. + fn new(body: &Body<'_>) -> Self { + let mut this = Self { + increment: true, + arg_count: body.arg_count.try_into().unwrap(), + use_count: IndexVec::from_elem(0, &body.local_decls), + }; + this.visit_body(body); + this + } + + /// Checks if local is used. + /// + /// Return place and arguments are always considered used. + fn is_used(&self, local: Local) -> bool { + trace!("is_used({:?}): use_count: {:?}", local, self.use_count[local]); + local.as_u32() <= self.arg_count || self.use_count[local] != 0 + } + + /// Updates the use counts to reflect the removal of given statement. + fn statement_removed(&mut self, statement: &Statement<'_>) { + self.increment = false; + + // The location of the statement is irrelevant. + let location = Location::START; + self.visit_statement(statement, location); + } + + /// Visits a left-hand side of an assignment. + fn visit_lhs(&mut self, place: &Place<'_>, location: Location) { + if place.is_indirect() { + // A use, not a definition. + self.visit_place(place, PlaceContext::MutatingUse(MutatingUseContext::Store), location); + } else { + // A definition. The base local itself is not visited, so this occurrence is not counted + // toward its use count. There might be other locals still, used in an indexing + // projection. + self.super_projection( + place.as_ref(), + PlaceContext::MutatingUse(MutatingUseContext::Projection), + location, + ); + } + } +} + +impl<'tcx> Visitor<'tcx> for UsedLocals { + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + match statement.kind { + StatementKind::Intrinsic(..) + | StatementKind::Retag(..) + | StatementKind::Coverage(..) + | StatementKind::FakeRead(..) + | StatementKind::PlaceMention(..) + | StatementKind::AscribeUserType(..) => { + self.super_statement(statement, location); + } + + StatementKind::ConstEvalCounter | StatementKind::Nop => {} + + StatementKind::StorageLive(_local) | StatementKind::StorageDead(_local) => {} + + StatementKind::Assign(box (ref place, ref rvalue)) => { + if rvalue.is_safe_to_remove() { + self.visit_lhs(place, location); + self.visit_rvalue(rvalue, location); + } else { + self.super_statement(statement, location); + } + } + + StatementKind::SetDiscriminant { ref place, variant_index: _ } + | StatementKind::Deinit(ref place) => { + self.visit_lhs(place, location); + } + } + } + + fn visit_local(&mut self, local: Local, _ctx: PlaceContext, _location: Location) { + if self.increment { + self.use_count[local] += 1; + } else { + assert_ne!(self.use_count[local], 0); + self.use_count[local] -= 1; + } + } +} + +/// Removes unused definitions. Updates the used locals to reflect the changes made. +fn remove_unused_definitions_helper(used_locals: &mut UsedLocals, body: &mut Body<'_>) { + // The use counts are updated as we remove the statements. A local might become unused + // during the retain operation, leading to a temporary inconsistency (storage statements or + // definitions referencing the local might remain). For correctness it is crucial that this + // computation reaches a fixed point. + + let mut modified = true; + while modified { + modified = false; + + for data in body.basic_blocks.as_mut_preserves_cfg() { + // Remove unnecessary StorageLive and StorageDead annotations. + data.statements.retain(|statement| { + let keep = match &statement.kind { + StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => { + used_locals.is_used(*local) + } + StatementKind::Assign(box (place, _)) => used_locals.is_used(place.local), + + StatementKind::SetDiscriminant { ref place, .. } + | StatementKind::Deinit(ref place) => used_locals.is_used(place.local), + StatementKind::Nop => false, + _ => true, + }; + + if !keep { + trace!("removing statement {:?}", statement); + modified = true; + used_locals.statement_removed(statement); + } + + keep + }); + } + } +} + +struct LocalUpdater<'tcx> { + map: IndexVec<Local, Option<Local>>, + tcx: TyCtxt<'tcx>, +} + +impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) { + *l = self.map[*l].unwrap(); + } +} diff --git a/compiler/rustc_mir_transform/src/simplify_branches.rs b/compiler/rustc_mir_transform/src/simplify_branches.rs new file mode 100644 index 00000000000..35a052166bd --- /dev/null +++ b/compiler/rustc_mir_transform/src/simplify_branches.rs @@ -0,0 +1,60 @@ +use rustc_middle::mir::*; +use rustc_middle::ty::TyCtxt; + +pub enum SimplifyConstCondition { + AfterConstProp, + Final, +} +/// A pass that replaces a branch with a goto when its condition is known. +impl<'tcx> MirPass<'tcx> for SimplifyConstCondition { + fn name(&self) -> &'static str { + match self { + SimplifyConstCondition::AfterConstProp => "SimplifyConstCondition-after-const-prop", + SimplifyConstCondition::Final => "SimplifyConstCondition-final", + } + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("Running SimplifyConstCondition on {:?}", body.source); + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + 'blocks: for block in body.basic_blocks_mut() { + for stmt in block.statements.iter_mut() { + if let StatementKind::Intrinsic(box ref intrinsic) = stmt.kind + && let NonDivergingIntrinsic::Assume(discr) = intrinsic + && let Operand::Constant(ref c) = discr + && let Some(constant) = c.const_.try_eval_bool(tcx, param_env) + { + if constant { + stmt.make_nop(); + } else { + block.statements.clear(); + block.terminator_mut().kind = TerminatorKind::Unreachable; + continue 'blocks; + } + } + } + + let terminator = block.terminator_mut(); + terminator.kind = match terminator.kind { + TerminatorKind::SwitchInt { + discr: Operand::Constant(ref c), ref targets, .. + } => { + let constant = c.const_.try_eval_bits(tcx, param_env); + if let Some(constant) = constant { + let target = targets.target_for_value(constant); + TerminatorKind::Goto { target } + } else { + continue; + } + } + TerminatorKind::Assert { + target, cond: Operand::Constant(ref c), expected, .. + } => match c.const_.try_eval_bool(tcx, param_env) { + Some(v) if v == expected => TerminatorKind::Goto { target }, + _ => continue, + }, + _ => continue, + }; + } + } +} diff --git a/compiler/rustc_mir_transform/src/simplify_comparison_integral.rs b/compiler/rustc_mir_transform/src/simplify_comparison_integral.rs new file mode 100644 index 00000000000..1a8cfc41178 --- /dev/null +++ b/compiler/rustc_mir_transform/src/simplify_comparison_integral.rs @@ -0,0 +1,239 @@ +use std::iter; + +use super::MirPass; +use rustc_middle::{ + mir::{ + interpret::Scalar, BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement, + StatementKind, SwitchTargets, TerminatorKind, + }, + ty::{Ty, TyCtxt}, +}; + +/// Pass to convert `if` conditions on integrals into switches on the integral. +/// For an example, it turns something like +/// +/// ```ignore (MIR) +/// _3 = Eq(move _4, const 43i32); +/// StorageDead(_4); +/// switchInt(_3) -> [false: bb2, otherwise: bb3]; +/// ``` +/// +/// into: +/// +/// ```ignore (MIR) +/// switchInt(_4) -> [43i32: bb3, otherwise: bb2]; +/// ``` +pub struct SimplifyComparisonIntegral; + +impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("Running SimplifyComparisonIntegral on {:?}", body.source); + + let helper = OptimizationFinder { body }; + let opts = helper.find_optimizations(); + let mut storage_deads_to_insert = vec![]; + let mut storage_deads_to_remove: Vec<(usize, BasicBlock)> = vec![]; + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + for opt in opts { + trace!("SUCCESS: Applying {:?}", opt); + // replace terminator with a switchInt that switches on the integer directly + let bbs = &mut body.basic_blocks_mut(); + let bb = &mut bbs[opt.bb_idx]; + let new_value = match opt.branch_value_scalar { + Scalar::Int(int) => { + let layout = tcx + .layout_of(param_env.and(opt.branch_value_ty)) + .expect("if we have an evaluated constant we must know the layout"); + int.assert_bits(layout.size) + } + Scalar::Ptr(..) => continue, + }; + const FALSE: u128 = 0; + + let mut new_targets = opt.targets; + let first_value = new_targets.iter().next().unwrap().0; + let first_is_false_target = first_value == FALSE; + match opt.op { + BinOp::Eq => { + // if the assignment was Eq we want the true case to be first + if first_is_false_target { + new_targets.all_targets_mut().swap(0, 1); + } + } + BinOp::Ne => { + // if the assignment was Ne we want the false case to be first + if !first_is_false_target { + new_targets.all_targets_mut().swap(0, 1); + } + } + _ => unreachable!(), + } + + // delete comparison statement if it the value being switched on was moved, which means it can not be user later on + if opt.can_remove_bin_op_stmt { + bb.statements[opt.bin_op_stmt_idx].make_nop(); + } else { + // if the integer being compared to a const integral is being moved into the comparison, + // e.g `_2 = Eq(move _3, const 'x');` + // we want to avoid making a double move later on in the switchInt on _3. + // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`, + // we convert the move in the comparison statement to a copy. + + // unwrap is safe as we know this statement is an assign + let (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap(); + + use Operand::*; + match rhs { + Rvalue::BinaryOp(_, box (ref mut left @ Move(_), Constant(_))) => { + *left = Copy(opt.to_switch_on); + } + Rvalue::BinaryOp(_, box (Constant(_), ref mut right @ Move(_))) => { + *right = Copy(opt.to_switch_on); + } + _ => (), + } + } + + let terminator = bb.terminator(); + + // remove StorageDead (if it exists) being used in the assign of the comparison + for (stmt_idx, stmt) in bb.statements.iter().enumerate() { + if !matches!(stmt.kind, StatementKind::StorageDead(local) if local == opt.to_switch_on.local) + { + continue; + } + storage_deads_to_remove.push((stmt_idx, opt.bb_idx)); + // if we have StorageDeads to remove then make sure to insert them at the top of each target + for bb_idx in new_targets.all_targets() { + storage_deads_to_insert.push(( + *bb_idx, + Statement { + source_info: terminator.source_info, + kind: StatementKind::StorageDead(opt.to_switch_on.local), + }, + )); + } + } + + let [bb_cond, bb_otherwise] = match new_targets.all_targets() { + [a, b] => [*a, *b], + e => bug!("expected 2 switch targets, got: {:?}", e), + }; + + let targets = SwitchTargets::new(iter::once((new_value, bb_cond)), bb_otherwise); + + let terminator = bb.terminator_mut(); + terminator.kind = + TerminatorKind::SwitchInt { discr: Operand::Move(opt.to_switch_on), targets }; + } + + for (idx, bb_idx) in storage_deads_to_remove { + body.basic_blocks_mut()[bb_idx].statements[idx].make_nop(); + } + + for (idx, stmt) in storage_deads_to_insert { + body.basic_blocks_mut()[idx].statements.insert(0, stmt); + } + } +} + +struct OptimizationFinder<'a, 'tcx> { + body: &'a Body<'tcx>, +} + +impl<'tcx> OptimizationFinder<'_, 'tcx> { + fn find_optimizations(&self) -> Vec<OptimizationInfo<'tcx>> { + self.body + .basic_blocks + .iter_enumerated() + .filter_map(|(bb_idx, bb)| { + // find switch + let (place_switched_on, targets, place_switched_on_moved) = + match &bb.terminator().kind { + rustc_middle::mir::TerminatorKind::SwitchInt { discr, targets, .. } => { + Some((discr.place()?, targets, discr.is_move())) + } + _ => None, + }?; + + // find the statement that assigns the place being switched on + bb.statements.iter().enumerate().rev().find_map(|(stmt_idx, stmt)| { + match &stmt.kind { + rustc_middle::mir::StatementKind::Assign(box (lhs, rhs)) + if *lhs == place_switched_on => + { + match rhs { + Rvalue::BinaryOp( + op @ (BinOp::Eq | BinOp::Ne), + box (left, right), + ) => { + let (branch_value_scalar, branch_value_ty, to_switch_on) = + find_branch_value_info(left, right)?; + + Some(OptimizationInfo { + bin_op_stmt_idx: stmt_idx, + bb_idx, + can_remove_bin_op_stmt: place_switched_on_moved, + to_switch_on, + branch_value_scalar, + branch_value_ty, + op: *op, + targets: targets.clone(), + }) + } + _ => None, + } + } + _ => None, + } + }) + }) + .collect() + } +} + +fn find_branch_value_info<'tcx>( + left: &Operand<'tcx>, + right: &Operand<'tcx>, +) -> Option<(Scalar, Ty<'tcx>, Place<'tcx>)> { + // check that either left or right is a constant. + // if any are, we can use the other to switch on, and the constant as a value in a switch + use Operand::*; + match (left, right) { + (Constant(branch_value), Copy(to_switch_on) | Move(to_switch_on)) + | (Copy(to_switch_on) | Move(to_switch_on), Constant(branch_value)) => { + let branch_value_ty = branch_value.const_.ty(); + // we only want to apply this optimization if we are matching on integrals (and chars), as it is not possible to switch on floats + if !branch_value_ty.is_integral() && !branch_value_ty.is_char() { + return None; + }; + let branch_value_scalar = branch_value.const_.try_to_scalar()?; + Some((branch_value_scalar, branch_value_ty, *to_switch_on)) + } + _ => None, + } +} + +#[derive(Debug)] +struct OptimizationInfo<'tcx> { + /// Basic block to apply the optimization + bb_idx: BasicBlock, + /// Statement index of Eq/Ne assignment that can be removed. None if the assignment can not be removed - i.e the statement is used later on + bin_op_stmt_idx: usize, + /// Can remove Eq/Ne assignment + can_remove_bin_op_stmt: bool, + /// Place that needs to be switched on. This place is of type integral + to_switch_on: Place<'tcx>, + /// Constant to use in switch target value + branch_value_scalar: Scalar, + /// Type of the constant value + branch_value_ty: Ty<'tcx>, + /// Either Eq or Ne + op: BinOp, + /// Current targets used in the switch + targets: SwitchTargets, +} diff --git a/compiler/rustc_mir_transform/src/sroa.rs b/compiler/rustc_mir_transform/src/sroa.rs new file mode 100644 index 00000000000..06d5e17fdd6 --- /dev/null +++ b/compiler/rustc_mir_transform/src/sroa.rs @@ -0,0 +1,448 @@ +use rustc_data_structures::flat_map_in_place::FlatMapInPlace; +use rustc_index::bit_set::{BitSet, GrowableBitSet}; +use rustc_index::IndexVec; +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 rustc_target::abi::{FieldIdx, FIRST_VARIANT}; + +pub struct ScalarReplacementOfAggregates; + +impl<'tcx> MirPass<'tcx> for ScalarReplacementOfAggregates { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() >= 2 + } + + #[instrument(level = "debug", skip(self, tcx, body))] + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + debug!(def_id = ?body.source.def_id()); + + // Avoid query cycles (coroutines require optimized MIR for layout). + if tcx.type_of(body.source.def_id()).instantiate_identity().is_coroutine() { + return; + } + + let mut excluded = excluded_locals(body); + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); + loop { + debug!(?excluded); + let escaping = escaping_locals(tcx, param_env, &excluded, body); + debug!(?escaping); + let replacements = compute_flattening(tcx, param_env, body, escaping); + debug!(?replacements); + let all_dead_locals = replace_flattened_locals(tcx, body, replacements); + if !all_dead_locals.is_empty() { + excluded.union(&all_dead_locals); + excluded = { + let mut growable = GrowableBitSet::from(excluded); + growable.ensure(body.local_decls.len()); + growable.into() + }; + } else { + break; + } + } + } +} + +/// Identify all locals that are not eligible for SROA. +/// +/// There are 3 cases: +/// - the aggregated local is used or passed to other code (function parameters and arguments); +/// - the locals is a union or an enum; +/// - the local's address is taken, and thus the relative addresses of the fields are observable to +/// client code. +fn escaping_locals<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + excluded: &BitSet<Local>, + body: &Body<'tcx>, +) -> BitSet<Local> { + let is_excluded_ty = |ty: Ty<'tcx>| { + if ty.is_union() || ty.is_enum() { + return true; + } + if let ty::Adt(def, _args) = ty.kind() { + if def.repr().simd() { + // Exclude #[repr(simd)] types so that they are not de-optimized into an array + return true; + } + // We already excluded unions and enums, so this ADT must have one variant + let variant = def.variant(FIRST_VARIANT); + if variant.fields.len() > 1 { + // If this has more than one field, it cannot be a wrapper that only provides a + // niche, so we do not want to automatically exclude it. + return false; + } + let Ok(layout) = tcx.layout_of(param_env.and(ty)) else { + // We can't get the layout + return true; + }; + if layout.layout.largest_niche().is_some() { + // This type has a niche + return true; + } + } + // Default for non-ADTs + false + }; + + let mut set = BitSet::new_empty(body.local_decls.len()); + set.insert_range(RETURN_PLACE..=Local::from_usize(body.arg_count)); + for (local, decl) in body.local_decls().iter_enumerated() { + if excluded.contains(local) || is_excluded_ty(decl.ty) { + set.insert(local); + } + } + let mut visitor = EscapeVisitor { set }; + visitor.visit_body(body); + return visitor.set; + + struct EscapeVisitor { + set: BitSet<Local>, + } + + impl<'tcx> Visitor<'tcx> for EscapeVisitor { + fn visit_local(&mut self, local: Local, _: PlaceContext, _: Location) { + self.set.insert(local); + } + + fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) { + // Mirror the implementation in PreFlattenVisitor. + if let &[PlaceElem::Field(..), ..] = &place.projection[..] { + return; + } + self.super_place(place, context, location); + } + + fn visit_assign( + &mut self, + lvalue: &Place<'tcx>, + rvalue: &Rvalue<'tcx>, + location: Location, + ) { + if lvalue.as_local().is_some() { + match rvalue { + // Aggregate assignments are expanded in run_pass. + Rvalue::Aggregate(..) | Rvalue::Use(..) => { + self.visit_rvalue(rvalue, location); + return; + } + _ => {} + } + } + self.super_assign(lvalue, rvalue, location) + } + + fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { + match statement.kind { + // Storage statements are expanded in run_pass. + StatementKind::StorageLive(..) + | StatementKind::StorageDead(..) + | StatementKind::Deinit(..) => return, + _ => self.super_statement(statement, location), + } + } + + // We ignore anything that happens in debuginfo, since we expand it using + // `VarDebugInfoFragment`. + fn visit_var_debug_info(&mut self, _: &VarDebugInfo<'tcx>) {} + } +} + +#[derive(Default, Debug)] +struct ReplacementMap<'tcx> { + /// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage + /// and deinit statement and debuginfo. + fragments: IndexVec<Local, Option<IndexVec<FieldIdx, Option<(Ty<'tcx>, Local)>>>>, +} + +impl<'tcx> ReplacementMap<'tcx> { + fn replace_place(&self, tcx: TyCtxt<'tcx>, place: PlaceRef<'tcx>) -> Option<Place<'tcx>> { + let &[PlaceElem::Field(f, _), ref rest @ ..] = place.projection else { + return None; + }; + let fields = self.fragments[place.local].as_ref()?; + let (_, new_local) = fields[f]?; + Some(Place { local: new_local, projection: tcx.mk_place_elems(rest) }) + } + + fn place_fragments( + &self, + place: Place<'tcx>, + ) -> Option<impl Iterator<Item = (FieldIdx, Ty<'tcx>, Local)> + '_> { + let local = place.as_local()?; + let fields = self.fragments[local].as_ref()?; + Some(fields.iter_enumerated().filter_map(|(field, &opt_ty_local)| { + let (ty, local) = opt_ty_local?; + Some((field, ty, local)) + })) + } +} + +/// Compute the replacement of flattened places into locals. +/// +/// For each eligible place, we assign a new local to each accessed field. +/// The replacement will be done later in `ReplacementVisitor`. +fn compute_flattening<'tcx>( + tcx: TyCtxt<'tcx>, + param_env: ty::ParamEnv<'tcx>, + body: &mut Body<'tcx>, + escaping: BitSet<Local>, +) -> ReplacementMap<'tcx> { + let mut fragments = IndexVec::from_elem(None, &body.local_decls); + + for local in body.local_decls.indices() { + if escaping.contains(local) { + continue; + } + let decl = body.local_decls[local].clone(); + let ty = decl.ty; + iter_fields(ty, tcx, param_env, |variant, field, field_ty| { + if variant.is_some() { + // Downcasts are currently not supported. + return; + }; + let new_local = + body.local_decls.push(LocalDecl { ty: field_ty, user_ty: None, ..decl.clone() }); + fragments.get_or_insert_with(local, IndexVec::new).insert(field, (field_ty, new_local)); + }); + } + ReplacementMap { fragments } +} + +/// Perform the replacement computed by `compute_flattening`. +fn replace_flattened_locals<'tcx>( + tcx: TyCtxt<'tcx>, + body: &mut Body<'tcx>, + replacements: ReplacementMap<'tcx>, +) -> BitSet<Local> { + let mut all_dead_locals = BitSet::new_empty(replacements.fragments.len()); + for (local, replacements) in replacements.fragments.iter_enumerated() { + if replacements.is_some() { + all_dead_locals.insert(local); + } + } + debug!(?all_dead_locals); + if all_dead_locals.is_empty() { + return all_dead_locals; + } + + let mut visitor = ReplacementVisitor { + tcx, + local_decls: &body.local_decls, + replacements: &replacements, + all_dead_locals, + patch: MirPatch::new(body), + }; + for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() { + visitor.visit_basic_block_data(bb, data); + } + for scope in &mut body.source_scopes { + visitor.visit_source_scope_data(scope); + } + for (index, annotation) in body.user_type_annotations.iter_enumerated_mut() { + visitor.visit_user_type_annotation(index, annotation); + } + visitor.expand_var_debug_info(&mut body.var_debug_info); + let ReplacementVisitor { patch, all_dead_locals, .. } = visitor; + patch.apply(body); + all_dead_locals +} + +struct ReplacementVisitor<'tcx, 'll> { + tcx: TyCtxt<'tcx>, + /// This is only used to compute the type for `VarDebugInfoFragment`. + local_decls: &'ll LocalDecls<'tcx>, + /// Work to do. + replacements: &'ll ReplacementMap<'tcx>, + /// This is used to check that we are not leaving references to replaced locals behind. + all_dead_locals: BitSet<Local>, + patch: MirPatch<'tcx>, +} + +impl<'tcx> ReplacementVisitor<'tcx, '_> { + #[instrument(level = "trace", skip(self))] + fn expand_var_debug_info(&mut self, var_debug_info: &mut Vec<VarDebugInfo<'tcx>>) { + var_debug_info.flat_map_in_place(|mut var_debug_info| { + let place = match var_debug_info.value { + VarDebugInfoContents::Const(_) => return vec![var_debug_info], + VarDebugInfoContents::Place(ref mut place) => place, + }; + + if let Some(repl) = self.replacements.replace_place(self.tcx, place.as_ref()) { + *place = repl; + return vec![var_debug_info]; + } + + let Some(parts) = self.replacements.place_fragments(*place) else { + return vec![var_debug_info]; + }; + + let ty = place.ty(self.local_decls, self.tcx).ty; + + parts + .map(|(field, field_ty, replacement_local)| { + let mut var_debug_info = var_debug_info.clone(); + let composite = var_debug_info.composite.get_or_insert_with(|| { + Box::new(VarDebugInfoFragment { ty, projection: Vec::new() }) + }); + composite.projection.push(PlaceElem::Field(field, field_ty)); + + var_debug_info.value = VarDebugInfoContents::Place(replacement_local.into()); + var_debug_info + }) + .collect() + }); + } +} + +impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> { + fn tcx(&self) -> TyCtxt<'tcx> { + self.tcx + } + + fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) { + if let Some(repl) = self.replacements.replace_place(self.tcx, place.as_ref()) { + *place = repl + } else { + self.super_place(place, context, location) + } + } + + #[instrument(level = "trace", skip(self))] + fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) { + match statement.kind { + // Duplicate storage and deinit statements, as they pretty much apply to all fields. + StatementKind::StorageLive(l) => { + if let Some(final_locals) = self.replacements.place_fragments(l.into()) { + for (_, _, fl) in final_locals { + self.patch.add_statement(location, StatementKind::StorageLive(fl)); + } + statement.make_nop(); + } + return; + } + StatementKind::StorageDead(l) => { + if let Some(final_locals) = self.replacements.place_fragments(l.into()) { + for (_, _, fl) in final_locals { + self.patch.add_statement(location, StatementKind::StorageDead(fl)); + } + statement.make_nop(); + } + return; + } + StatementKind::Deinit(box place) => { + if let Some(final_locals) = self.replacements.place_fragments(place) { + for (_, _, fl) in final_locals { + self.patch + .add_statement(location, StatementKind::Deinit(Box::new(fl.into()))); + } + statement.make_nop(); + return; + } + } + + // We have `a = Struct { 0: x, 1: y, .. }`. + // We replace it by + // ``` + // a_0 = x + // a_1 = y + // ... + // ``` + StatementKind::Assign(box (place, Rvalue::Aggregate(_, ref mut operands))) => { + if let Some(local) = place.as_local() + && let Some(final_locals) = &self.replacements.fragments[local] + { + // This is ok as we delete the statement later. + let operands = std::mem::take(operands); + for (&opt_ty_local, mut operand) in final_locals.iter().zip(operands) { + if let Some((_, new_local)) = opt_ty_local { + // Replace mentions of SROA'd locals that appear in the operand. + self.visit_operand(&mut operand, location); + + let rvalue = Rvalue::Use(operand); + self.patch.add_statement( + location, + StatementKind::Assign(Box::new((new_local.into(), rvalue))), + ); + } + } + statement.make_nop(); + return; + } + } + + // We have `a = some constant` + // We add the projections. + // ``` + // a_0 = a.0 + // a_1 = a.1 + // ... + // ``` + // ConstProp will pick up the pieces and replace them by actual constants. + StatementKind::Assign(box (place, Rvalue::Use(Operand::Constant(_)))) => { + if let Some(final_locals) = self.replacements.place_fragments(place) { + // Put the deaggregated statements *after* the original one. + let location = location.successor_within_block(); + for (field, ty, new_local) in final_locals { + let rplace = self.tcx.mk_place_field(place, field, ty); + let rvalue = Rvalue::Use(Operand::Move(rplace)); + self.patch.add_statement( + location, + StatementKind::Assign(Box::new((new_local.into(), rvalue))), + ); + } + // We still need `place.local` to exist, so don't make it nop. + return; + } + } + + // We have `a = move? place` + // We replace it by + // ``` + // a_0 = move? place.0 + // a_1 = move? place.1 + // ... + // ``` + StatementKind::Assign(box (lhs, Rvalue::Use(ref op))) => { + let (rplace, copy) = match *op { + Operand::Copy(rplace) => (rplace, true), + Operand::Move(rplace) => (rplace, false), + Operand::Constant(_) => bug!(), + }; + if let Some(final_locals) = self.replacements.place_fragments(lhs) { + for (field, ty, new_local) in final_locals { + let rplace = self.tcx.mk_place_field(rplace, field, ty); + debug!(?rplace); + let rplace = self + .replacements + .replace_place(self.tcx, rplace.as_ref()) + .unwrap_or(rplace); + debug!(?rplace); + let rvalue = if copy { + Rvalue::Use(Operand::Copy(rplace)) + } else { + Rvalue::Use(Operand::Move(rplace)) + }; + self.patch.add_statement( + location, + StatementKind::Assign(Box::new((new_local.into(), rvalue))), + ); + } + statement.make_nop(); + return; + } + } + + _ => {} + } + self.super_statement(statement, location) + } + + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { + assert!(!self.all_dead_locals.contains(*local)); + } +} diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs new file mode 100644 index 00000000000..3a6e1ef3488 --- /dev/null +++ b/compiler/rustc_mir_transform/src/ssa.rs @@ -0,0 +1,385 @@ +//! We denote as "SSA" the set of locals that verify the following properties: +//! 1/ They are only assigned-to once, either as a function parameter, or in an assign statement; +//! 2/ This single assignment dominates all uses; +//! +//! As a consequence of rule 2, we consider that borrowed locals are not SSA, even if they are +//! `Freeze`, as we do not track that the assignment dominates all uses of the borrow. + +use rustc_data_structures::graph::dominators::Dominators; +use rustc_index::bit_set::BitSet; +use rustc_index::{IndexSlice, IndexVec}; +use rustc_middle::middle::resolve_bound_vars::Set1; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; + +pub struct SsaLocals { + /// Assignments to each local. This defines whether the local is SSA. + assignments: IndexVec<Local, Set1<DefLocation>>, + /// We visit the body in reverse postorder, to ensure each local is assigned before it is used. + /// We remember the order in which we saw the assignments to compute the SSA values in a single + /// pass. + assignment_order: Vec<Local>, + /// Copy equivalence classes between locals. See `copy_classes` for documentation. + copy_classes: IndexVec<Local, Local>, + /// Number of "direct" uses of each local, ie. uses that are not dereferences. + /// We ignore non-uses (Storage statements, debuginfo). + direct_uses: IndexVec<Local, u32>, +} + +pub enum AssignedValue<'a, 'tcx> { + Arg, + Rvalue(&'a mut Rvalue<'tcx>), + Terminator(&'a mut TerminatorKind<'tcx>), +} + +impl SsaLocals { + pub fn new<'tcx>(body: &Body<'tcx>) -> SsaLocals { + let assignment_order = Vec::with_capacity(body.local_decls.len()); + + let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls); + let dominators = body.basic_blocks.dominators(); + + let direct_uses = IndexVec::from_elem(0, &body.local_decls); + let mut visitor = + SsaVisitor { body, assignments, assignment_order, dominators, direct_uses }; + + for local in body.args_iter() { + visitor.assignments[local] = Set1::One(DefLocation::Argument); + visitor.assignment_order.push(local); + } + + // For SSA assignments, a RPO visit will see the assignment before it sees any use. + // We only visit reachable nodes: computing `dominates` on an unreachable node ICEs. + for (bb, data) in traversal::reverse_postorder(body) { + visitor.visit_basic_block_data(bb, data); + } + + for var_debug_info in &body.var_debug_info { + visitor.visit_var_debug_info(var_debug_info); + } + + debug!(?visitor.assignments); + debug!(?visitor.direct_uses); + + visitor + .assignment_order + .retain(|&local| matches!(visitor.assignments[local], Set1::One(_))); + debug!(?visitor.assignment_order); + + let mut ssa = SsaLocals { + assignments: visitor.assignments, + assignment_order: visitor.assignment_order, + direct_uses: visitor.direct_uses, + // This is filled by `compute_copy_classes`. + copy_classes: IndexVec::default(), + }; + compute_copy_classes(&mut ssa, body); + ssa + } + + pub fn num_locals(&self) -> usize { + self.assignments.len() + } + + pub fn locals(&self) -> impl Iterator<Item = Local> { + self.assignments.indices() + } + + pub fn is_ssa(&self, local: Local) -> bool { + matches!(self.assignments[local], Set1::One(_)) + } + + /// Return the number of uses if a local that are not "Deref". + pub fn num_direct_uses(&self, local: Local) -> u32 { + self.direct_uses[local] + } + + pub fn assignment_dominates( + &self, + dominators: &Dominators<BasicBlock>, + local: Local, + location: Location, + ) -> bool { + match self.assignments[local] { + Set1::One(def) => def.dominates(location, dominators), + _ => false, + } + } + + pub fn assignments<'a, 'tcx>( + &'a self, + body: &'a Body<'tcx>, + ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a { + self.assignment_order.iter().filter_map(|&local| { + if let Set1::One(DefLocation::Assignment(loc)) = self.assignments[local] { + let stmt = body.stmt_at(loc).left()?; + // `loc` must point to a direct assignment to `local`. + let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() }; + assert_eq!(target.as_local(), Some(local)); + Some((local, rvalue, loc)) + } else { + None + } + }) + } + + pub fn for_each_assignment_mut<'tcx>( + &self, + basic_blocks: &mut IndexSlice<BasicBlock, BasicBlockData<'tcx>>, + mut f: impl FnMut(Local, AssignedValue<'_, 'tcx>, Location), + ) { + for &local in &self.assignment_order { + match self.assignments[local] { + Set1::One(DefLocation::Argument) => f( + local, + AssignedValue::Arg, + Location { block: START_BLOCK, statement_index: 0 }, + ), + Set1::One(DefLocation::Assignment(loc)) => { + let bb = &mut basic_blocks[loc.block]; + // `loc` must point to a direct assignment to `local`. + let stmt = &mut bb.statements[loc.statement_index]; + let StatementKind::Assign(box (target, ref mut rvalue)) = stmt.kind else { + bug!() + }; + assert_eq!(target.as_local(), Some(local)); + f(local, AssignedValue::Rvalue(rvalue), loc) + } + Set1::One(DefLocation::CallReturn { call, .. }) => { + let bb = &mut basic_blocks[call]; + let loc = Location { block: call, statement_index: bb.statements.len() }; + let term = bb.terminator_mut(); + f(local, AssignedValue::Terminator(&mut term.kind), loc) + } + _ => {} + } + } + } + + /// Compute the equivalence classes for locals, based on copy statements. + /// + /// The returned vector maps each local to the one it copies. In the following case: + /// _a = &mut _0 + /// _b = move? _a + /// _c = move? _a + /// _d = move? _c + /// We return the mapping + /// _a => _a // not a copy so, represented by itself + /// _b => _a + /// _c => _a + /// _d => _a // transitively through _c + /// + /// Exception: we do not see through the return place, as it cannot be substituted. + pub fn copy_classes(&self) -> &IndexSlice<Local, Local> { + &self.copy_classes + } + + /// Make a property uniform on a copy equivalence class by removing elements. + pub fn meet_copy_equivalence(&self, property: &mut BitSet<Local>) { + // Consolidate to have a local iff all its copies are. + // + // `copy_classes` defines equivalence classes between locals. The `local`s that recursively + // move/copy the same local all have the same `head`. + for (local, &head) in self.copy_classes.iter_enumerated() { + // If any copy does not have `property`, then the head is not. + if !property.contains(local) { + property.remove(head); + } + } + for (local, &head) in self.copy_classes.iter_enumerated() { + // If any copy does not have `property`, then the head doesn't either, + // then no copy has `property`. + if !property.contains(head) { + property.remove(local); + } + } + + // Verify that we correctly computed equivalence classes. + #[cfg(debug_assertions)] + for (local, &head) in self.copy_classes.iter_enumerated() { + assert_eq!(property.contains(local), property.contains(head)); + } + } +} + +struct SsaVisitor<'tcx, 'a> { + body: &'a Body<'tcx>, + dominators: &'a Dominators<BasicBlock>, + assignments: IndexVec<Local, Set1<DefLocation>>, + assignment_order: Vec<Local>, + direct_uses: IndexVec<Local, u32>, +} + +impl SsaVisitor<'_, '_> { + fn check_dominates(&mut self, local: Local, loc: Location) { + let set = &mut self.assignments[local]; + let assign_dominates = match *set { + Set1::Empty | Set1::Many => false, + Set1::One(def) => def.dominates(loc, self.dominators), + }; + // We are visiting a use that is not dominated by an assignment. + // Either there is a cycle involved, or we are reading for uninitialized local. + // Bail out. + if !assign_dominates { + *set = Set1::Many; + } + } +} + +impl<'tcx> Visitor<'tcx> for SsaVisitor<'tcx, '_> { + fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) { + match ctxt { + PlaceContext::MutatingUse(MutatingUseContext::Projection) + | PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection) => bug!(), + // Anything can happen with raw pointers, so remove them. + // We do not verify that all uses of the borrow dominate the assignment to `local`, + // so we have to remove them too. + PlaceContext::NonMutatingUse( + NonMutatingUseContext::SharedBorrow + | NonMutatingUseContext::FakeBorrow + | NonMutatingUseContext::AddressOf, + ) + | PlaceContext::MutatingUse(_) => { + self.assignments[local] = Set1::Many; + } + PlaceContext::NonMutatingUse(_) => { + self.check_dominates(local, loc); + self.direct_uses[local] += 1; + } + PlaceContext::NonUse(_) => {} + } + } + + fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) { + let location = match ctxt { + PlaceContext::MutatingUse(MutatingUseContext::Store) => { + Some(DefLocation::Assignment(loc)) + } + PlaceContext::MutatingUse(MutatingUseContext::Call) => { + let call = loc.block; + let TerminatorKind::Call { target, .. } = + self.body.basic_blocks[call].terminator().kind + else { + bug!() + }; + Some(DefLocation::CallReturn { call, target }) + } + _ => None, + }; + if let Some(location) = location + && let Some(local) = place.as_local() + { + self.assignments[local].insert(location); + if let Set1::One(_) = self.assignments[local] { + // Only record if SSA-like, to avoid growing the vector needlessly. + self.assignment_order.push(local); + } + } else if place.projection.first() == Some(&PlaceElem::Deref) { + // Do not do anything for debuginfo. + if ctxt.is_use() { + // Only change the context if it is a real use, not a "use" in debuginfo. + let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy); + + self.visit_projection(place.as_ref(), new_ctxt, loc); + self.check_dominates(place.local, loc); + } + } else { + self.visit_projection(place.as_ref(), ctxt, loc); + self.visit_local(place.local, ctxt, loc); + } + } +} + +#[instrument(level = "trace", skip(ssa, body))] +fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) { + let mut direct_uses = std::mem::take(&mut ssa.direct_uses); + let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len()); + + for (local, rvalue, _) in ssa.assignments(body) { + let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) + | Rvalue::CopyForDeref(place)) = rvalue + else { + continue; + }; + + let Some(rhs) = place.as_local() else { continue }; + let local_ty = body.local_decls()[local].ty; + let rhs_ty = body.local_decls()[rhs].ty; + if local_ty != rhs_ty { + // FIXME(#112651): This can be removed afterwards. + trace!("skipped `{local:?} = {rhs:?}` due to subtyping: {local_ty} != {rhs_ty}"); + continue; + } + + if !ssa.is_ssa(rhs) { + continue; + } + + // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been + // visited before `local`, and we just have to copy the representing local. + let head = copies[rhs]; + + if local == RETURN_PLACE { + // `_0` is special, we cannot rename it. Instead, rename the class of `rhs` to + // `RETURN_PLACE`. This is only possible if the class head is a temporary, not an + // argument. + if body.local_kind(head) != LocalKind::Temp { + continue; + } + for h in copies.iter_mut() { + if *h == head { + *h = RETURN_PLACE; + } + } + } else { + copies[local] = head; + } + direct_uses[rhs] -= 1; + } + + debug!(?copies); + debug!(?direct_uses); + + // Invariant: `copies` must point to the head of an equivalence class. + #[cfg(debug_assertions)] + for &head in copies.iter() { + assert_eq!(copies[head], head); + } + debug_assert_eq!(copies[RETURN_PLACE], RETURN_PLACE); + + ssa.direct_uses = direct_uses; + ssa.copy_classes = copies; +} + +#[derive(Debug)] +pub(crate) struct StorageLiveLocals { + /// Set of "StorageLive" statements for each local. + storage_live: IndexVec<Local, Set1<DefLocation>>, +} + +impl StorageLiveLocals { + pub(crate) fn new( + body: &Body<'_>, + always_storage_live_locals: &BitSet<Local>, + ) -> StorageLiveLocals { + let mut storage_live = IndexVec::from_elem(Set1::Empty, &body.local_decls); + for local in always_storage_live_locals.iter() { + storage_live[local] = Set1::One(DefLocation::Argument); + } + for (block, bbdata) in body.basic_blocks.iter_enumerated() { + for (statement_index, statement) in bbdata.statements.iter().enumerate() { + if let StatementKind::StorageLive(local) = statement.kind { + storage_live[local] + .insert(DefLocation::Assignment(Location { block, statement_index })); + } + } + } + debug!(?storage_live); + StorageLiveLocals { storage_live } + } + + #[inline] + pub(crate) fn has_single_storage(&self, local: Local) -> bool { + matches!(self.storage_live[local], Set1::One(_)) + } +} diff --git a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs new file mode 100644 index 00000000000..e68d37f4c70 --- /dev/null +++ b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs @@ -0,0 +1,137 @@ +//! A pass that eliminates branches on uninhabited enum variants. + +use crate::MirPass; +use rustc_data_structures::fx::FxHashSet; +use rustc_middle::mir::{ + BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, Terminator, TerminatorKind, +}; +use rustc_middle::ty::layout::TyAndLayout; +use rustc_middle::ty::{Ty, TyCtxt}; +use rustc_target::abi::{Abi, Variants}; + +pub struct UninhabitedEnumBranching; + +fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> { + if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator { + p.as_local() + } else { + None + } +} + +/// If the basic block terminates by switching on a discriminant, this returns the `Ty` the +/// discriminant is read from. Otherwise, returns None. +fn get_switched_on_type<'tcx>( + block_data: &BasicBlockData<'tcx>, + tcx: TyCtxt<'tcx>, + body: &Body<'tcx>, +) -> Option<Ty<'tcx>> { + let terminator = block_data.terminator(); + + // Only bother checking blocks which terminate by switching on a local. + let local = get_discriminant_local(&terminator.kind)?; + + let stmt_before_term = block_data.statements.last()?; + + if let StatementKind::Assign(box (l, Rvalue::Discriminant(place))) = stmt_before_term.kind + && l.as_local() == Some(local) + { + let ty = place.ty(body, tcx).ty; + if ty.is_enum() { + return Some(ty); + } + } + + None +} + +fn variant_discriminants<'tcx>( + layout: &TyAndLayout<'tcx>, + ty: Ty<'tcx>, + tcx: TyCtxt<'tcx>, +) -> FxHashSet<u128> { + match &layout.variants { + Variants::Single { index } => { + let mut res = FxHashSet::default(); + res.insert( + ty.discriminant_for_variant(tcx, *index) + .map_or(index.as_u32() as u128, |discr| discr.val), + ); + res + } + Variants::Multiple { variants, .. } => variants + .iter_enumerated() + .filter_map(|(idx, layout)| { + (layout.abi != Abi::Uninhabited) + .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val) + }) + .collect(), + } +} + +impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + sess.mir_opt_level() > 0 + } + + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + trace!("UninhabitedEnumBranching starting for {:?}", body.source); + + let mut removable_switchs = Vec::new(); + + for (bb, bb_data) in body.basic_blocks.iter_enumerated() { + trace!("processing block {:?}", bb); + + if bb_data.is_cleanup { + continue; + } + + let Some(discriminant_ty) = get_switched_on_type(bb_data, tcx, body) else { continue }; + + let layout = tcx.layout_of( + tcx.param_env_reveal_all_normalized(body.source.def_id()).and(discriminant_ty), + ); + + let allowed_variants = if let Ok(layout) = layout { + variant_discriminants(&layout, discriminant_ty, tcx) + } else { + continue; + }; + + trace!("allowed_variants = {:?}", allowed_variants); + + let terminator = bb_data.terminator(); + let TerminatorKind::SwitchInt { targets, .. } = &terminator.kind else { bug!() }; + + let mut reachable_count = 0; + for (index, (val, _)) in targets.iter().enumerate() { + if allowed_variants.contains(&val) { + reachable_count += 1; + } else { + removable_switchs.push((bb, index)); + } + } + + if reachable_count == allowed_variants.len() { + removable_switchs.push((bb, targets.iter().count())); + } + } + + if removable_switchs.is_empty() { + return; + } + + let new_block = BasicBlockData::new(Some(Terminator { + source_info: body.basic_blocks[removable_switchs[0].0].terminator().source_info, + kind: TerminatorKind::Unreachable, + })); + let unreachable_block = body.basic_blocks.as_mut().push(new_block); + + for (bb, index) in removable_switchs { + let bb = &mut body.basic_blocks.as_mut()[bb]; + let terminator = bb.terminator_mut(); + let TerminatorKind::SwitchInt { targets, .. } = &mut terminator.kind else { bug!() }; + targets.all_targets_mut()[index] = unreachable_block; + } + } +} diff --git a/compiler/rustc_mir_transform/src/unreachable_prop.rs b/compiler/rustc_mir_transform/src/unreachable_prop.rs new file mode 100644 index 00000000000..f12a6aa2429 --- /dev/null +++ b/compiler/rustc_mir_transform/src/unreachable_prop.rs @@ -0,0 +1,154 @@ +//! A pass that propagates the unreachable terminator of a block to its predecessors +//! when all of their successors are unreachable. This is achieved through a +//! post-order traversal of the blocks. + +use rustc_data_structures::fx::FxHashSet; +use rustc_middle::mir::interpret::Scalar; +use rustc_middle::mir::patch::MirPatch; +use rustc_middle::mir::*; +use rustc_middle::ty::{self, TyCtxt}; +use rustc_target::abi::Size; + +pub struct UnreachablePropagation; + +impl MirPass<'_> for UnreachablePropagation { + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { + // Enable only under -Zmir-opt-level=2 as this can make programs less debuggable. + + // FIXME(#116171) Coverage gets confused by MIR passes that can remove all + // coverage statements from an instrumented function. This pass can be + // re-enabled when coverage codegen is robust against that happening. + sess.mir_opt_level() >= 2 && !sess.instrument_coverage() + } + + fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { + let mut patch = MirPatch::new(body); + let mut unreachable_blocks = FxHashSet::default(); + + for (bb, bb_data) in traversal::postorder(body) { + let terminator = bb_data.terminator(); + let is_unreachable = match &terminator.kind { + TerminatorKind::Unreachable => true, + // This will unconditionally run into an unreachable and is therefore unreachable as well. + TerminatorKind::Goto { target } if unreachable_blocks.contains(target) => { + patch.patch_terminator(bb, TerminatorKind::Unreachable); + true + } + // Try to remove unreachable targets from the switch. + TerminatorKind::SwitchInt { .. } => { + remove_successors_from_switch(tcx, bb, &unreachable_blocks, body, &mut patch) + } + _ => false, + }; + if is_unreachable { + unreachable_blocks.insert(bb); + } + } + + if !tcx + .consider_optimizing(|| format!("UnreachablePropagation {:?} ", body.source.def_id())) + { + return; + } + + patch.apply(body); + + // We do want do keep some unreachable blocks, but make them empty. + for bb in unreachable_blocks { + body.basic_blocks_mut()[bb].statements.clear(); + } + } +} + +/// Return whether the current terminator is fully unreachable. +fn remove_successors_from_switch<'tcx>( + tcx: TyCtxt<'tcx>, + bb: BasicBlock, + unreachable_blocks: &FxHashSet<BasicBlock>, + body: &Body<'tcx>, + patch: &mut MirPatch<'tcx>, +) -> bool { + let terminator = body.basic_blocks[bb].terminator(); + let TerminatorKind::SwitchInt { discr, targets } = &terminator.kind else { bug!() }; + let source_info = terminator.source_info; + let location = body.terminator_loc(bb); + + let is_unreachable = |bb| unreachable_blocks.contains(&bb); + + // If there are multiple targets, we want to keep information about reachability for codegen. + // For example (see tests/codegen/match-optimizes-away.rs) + // + // pub enum Two { A, B } + // pub fn identity(x: Two) -> Two { + // match x { + // Two::A => Two::A, + // Two::B => Two::B, + // } + // } + // + // This generates a `switchInt() -> [0: 0, 1: 1, otherwise: unreachable]`, which allows us or LLVM to + // turn it into just `x` later. Without the unreachable, such a transformation would be illegal. + // + // In order to preserve this information, we record reachable and unreachable targets as + // `Assume` statements in MIR. + + let discr_ty = discr.ty(body, tcx); + let discr_size = Size::from_bits(match discr_ty.kind() { + ty::Uint(uint) => uint.normalize(tcx.sess.target.pointer_width).bit_width().unwrap(), + ty::Int(int) => int.normalize(tcx.sess.target.pointer_width).bit_width().unwrap(), + ty::Char => 32, + ty::Bool => 1, + other => bug!("unhandled type: {:?}", other), + }); + + let mut add_assumption = |binop, value| { + let local = patch.new_temp(tcx.types.bool, source_info.span); + let value = Operand::Constant(Box::new(ConstOperand { + span: source_info.span, + user_ty: None, + const_: Const::from_scalar(tcx, Scalar::from_uint(value, discr_size), discr_ty), + })); + let cmp = Rvalue::BinaryOp(binop, Box::new((discr.to_copy(), value))); + patch.add_assign(location, local.into(), cmp); + + let assume = NonDivergingIntrinsic::Assume(Operand::Move(local.into())); + patch.add_statement(location, StatementKind::Intrinsic(Box::new(assume))); + }; + + let otherwise = targets.otherwise(); + let otherwise_unreachable = is_unreachable(otherwise); + + let reachable_iter = targets.iter().filter(|&(value, bb)| { + let is_unreachable = is_unreachable(bb); + // We remove this target from the switch, so record the inequality using `Assume`. + if is_unreachable && !otherwise_unreachable { + add_assumption(BinOp::Ne, value); + } + !is_unreachable + }); + + let new_targets = SwitchTargets::new(reachable_iter, otherwise); + + let num_targets = new_targets.all_targets().len(); + let fully_unreachable = num_targets == 1 && otherwise_unreachable; + + let terminator = match (num_targets, otherwise_unreachable) { + // If all targets are unreachable, we can be unreachable as well. + (1, true) => TerminatorKind::Unreachable, + (1, false) => TerminatorKind::Goto { target: otherwise }, + (2, true) => { + // All targets are unreachable except one. Record the equality, and make it a goto. + let (value, target) = new_targets.iter().next().unwrap(); + add_assumption(BinOp::Eq, value); + TerminatorKind::Goto { target } + } + _ if num_targets == targets.all_targets().len() => { + // Nothing has changed. + return false; + } + _ => TerminatorKind::SwitchInt { discr: discr.clone(), targets: new_targets }, + }; + + patch.patch_terminator(bb, terminator); + fully_unreachable +} |