Skip to content

Commit 2dcf095

Browse files
saethlinsaethlin
committed
Avoid lowering code under dead SwitchInt targets
1 parent d7bd9cd commit 2dcf095

File tree

6 files changed

+210
-2
lines changed

6 files changed

+210
-2
lines changed

‎compiler/rustc_codegen_ssa/src/mir/block.rs

+10
Original file line numberDiff line numberDiff line change
@@ -1179,6 +1179,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
11791179
}
11801180
}
11811181

1182+
pub fn codegen_block_as_unreachable(&mut self, bb: mir::BasicBlock) {
1183+
let llbb = match self.try_llbb(bb) {
1184+
Some(llbb) => llbb,
1185+
None => return,
1186+
};
1187+
let bx = &mut Bx::build(self.cx, llbb);
1188+
debug!("codegen_block_as_unreachable({:?})", bb);
1189+
bx.unreachable();
1190+
}
1191+
11821192
fn codegen_terminator(
11831193
&mut self,
11841194
bx: &mut Bx,

‎compiler/rustc_codegen_ssa/src/mir/mod.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -256,13 +256,22 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
256256
// Apply debuginfo to the newly allocated locals.
257257
fx.debug_introduce_locals(&mut start_bx);
258258

259+
let reachable_blocks = mir.reachable_blocks_in_mono(cx.tcx(), instance);
260+
259261
// The builders will be created separately for each basic block at `codegen_block`.
260262
// So drop the builder of `start_llbb` to avoid having two at the same time.
261263
drop(start_bx);
262264

263265
// Codegen the body of each block using reverse postorder
264266
for (bb, _) in traversal::reverse_postorder(mir) {
265-
fx.codegen_block(bb);
267+
if reachable_blocks.contains(bb) {
268+
fx.codegen_block(bb);
269+
} else {
270+
// This may have references to things we didn't monomorphize, so we
271+
// don't actually codegen the body. We still create the block so
272+
// terminators in other blocks can reference it without worry.
273+
fx.codegen_block_as_unreachable(bb);
274+
}
266275
}
267276
}
268277

‎compiler/rustc_middle/src/mir/mod.rs

+69-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use crate::ty::print::{pretty_print_const, with_no_trimmed_paths};
1010
use crate::ty::print::{FmtPrinter, Printer};
1111
use crate::ty::visit::TypeVisitableExt;
1212
use crate::ty::{self, List, Ty, TyCtxt};
13-
use crate::ty::{AdtDef, InstanceDef, UserTypeAnnotationIndex};
13+
use crate::ty::{AdtDef, Instance, InstanceDef, UserTypeAnnotationIndex};
1414
use crate::ty::{GenericArg, GenericArgsRef};
1515

1616
use rustc_data_structures::captures::Captures;
@@ -29,6 +29,7 @@ pub use rustc_ast::Mutability;
2929
use rustc_data_structures::fx::FxHashMap;
3030
use rustc_data_structures::fx::FxHashSet;
3131
use rustc_data_structures::graph::dominators::Dominators;
32+
use rustc_index::bit_set::BitSet;
3233
use rustc_index::{Idx, IndexSlice, IndexVec};
3334
use rustc_serialize::{Decodable, Encodable};
3435
use rustc_span::symbol::Symbol;
@@ -642,6 +643,73 @@ impl<'tcx> Body<'tcx> {
642643
self.injection_phase.is_some()
643644
}
644645

646+
/// Finds which basic blocks are actually reachable for a specific
647+
/// monomorphization of this body.
648+
///
649+
/// This is allowed to have false positives; just because this says a block
650+
/// is reachable doesn't mean that's necessarily true. It's thus always
651+
/// legal for this to return a filled set.
652+
///
653+
/// Regardless, the [`BitSet::domain_size`] of the returned set will always
654+
/// exactly match the number of blocks in the body so that `contains`
655+
/// checks can be done without worrying about panicking.
656+
///
657+
/// The main case this supports is filtering out `if <T as Trait>::CONST`
658+
/// bodies that can't be removed in generic MIR, but *can* be removed once
659+
/// the specific `T` is known.
660+
///
661+
/// This is used in the monomorphization collector as well as in codegen.
662+
pub fn reachable_blocks_in_mono(
663+
&self,
664+
tcx: TyCtxt<'tcx>,
665+
instance: Instance<'tcx>,
666+
) -> BitSet<BasicBlock> {
667+
if instance.args.non_erasable_generics(tcx, instance.def_id()).next().is_none() {
668+
// If it's non-generic, then mir-opt const prop has already run, meaning it's
669+
// probably not worth doing any further filtering. So call everything reachable.
670+
return BitSet::new_filled(self.basic_blocks.len());
671+
}
672+
673+
let mut set = BitSet::new_empty(self.basic_blocks.len());
674+
self.reachable_blocks_in_mono_from(tcx, instance, &mut set, START_BLOCK);
675+
set
676+
}
677+
678+
fn reachable_blocks_in_mono_from(
679+
&self,
680+
tcx: TyCtxt<'tcx>,
681+
instance: Instance<'tcx>,
682+
set: &mut BitSet<BasicBlock>,
683+
bb: BasicBlock,
684+
) {
685+
if !set.insert(bb) {
686+
return;
687+
}
688+
689+
let data = &self.basic_blocks[bb];
690+
691+
if let TerminatorKind::SwitchInt { discr: Operand::Constant(constant), targets } =
692+
&data.terminator().kind
693+
{
694+
let env = ty::ParamEnv::reveal_all();
695+
let mono_literal = instance.instantiate_mir_and_normalize_erasing_regions(
696+
tcx,
697+
env,
698+
crate::ty::EarlyBinder::bind(constant.const_),
699+
);
700+
if let Some(bits) = mono_literal.try_eval_bits(tcx, env) {
701+
let target = targets.target_for_value(bits);
702+
return self.reachable_blocks_in_mono_from(tcx, instance, set, target);
703+
} else {
704+
bug!("Couldn't evaluate constant {:?} in mono {:?}", constant, instance);
705+
}
706+
}
707+
708+
for target in data.terminator().successors() {
709+
self.reachable_blocks_in_mono_from(tcx, instance, set, target);
710+
}
711+
}
712+
645713
/// For a `Location` in this scope, determine what the "caller location" at that point is. This
646714
/// is interesting because of inlining: the `#[track_caller]` attribute of inlined functions
647715
/// must be honored. Falls back to the `tracked_caller` value for `#[track_caller]` functions,

‎compiler/rustc_mir_transform/src/lib.rs

+2
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ mod check_alignment;
108108
pub mod simplify;
109109
mod simplify_branches;
110110
mod simplify_comparison_integral;
111+
mod simplify_if_const;
111112
mod sroa;
112113
mod uninhabited_enum_branching;
113114
mod unreachable_prop;
@@ -616,6 +617,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
616617
&large_enums::EnumSizeOpt { discrepancy: 128 },
617618
// Some cleanup necessary at least for LLVM and potentially other codegen backends.
618619
&add_call_guards::CriticalCallEdges,
620+
&simplify_if_const::SimplifyIfConst,
619621
// Cleanup for human readability, off by default.
620622
&prettify::ReorderBasicBlocks,
621623
&prettify::ReorderLocals,

‎compiler/rustc_mir_transform/src/simplify_if_const.rs

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
//! A pass that simplifies branches when their condition is known.
2+
3+
use crate::MirPass;
4+
use rustc_middle::mir::*;
5+
use rustc_middle::ty::TyCtxt;
6+
7+
/// The lowering for `if CONST` produces
8+
/// ```
9+
/// _1 = Const(...);
10+
/// switchInt (move _1)
11+
/// ```
12+
/// so this pass replaces that with
13+
/// ```
14+
/// switchInt (Const(...))
15+
/// ```
16+
/// so that further MIR consumers can special-case it more easily.
17+
///
18+
/// Unlike ConstProp, this supports generic constants too, not just concrete ones.
19+
pub struct SimplifyIfConst;
20+
21+
impl<'tcx> MirPass<'tcx> for SimplifyIfConst {
22+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
23+
for block in body.basic_blocks_mut() {
24+
simplify_assign_move_switch(tcx, block);
25+
}
26+
}
27+
}
28+
29+
fn simplify_assign_move_switch(tcx: TyCtxt<'_>, block: &mut BasicBlockData<'_>) {
30+
let Some(Terminator { kind: TerminatorKind::SwitchInt { discr: switch_desc, .. }, .. }) =
31+
&mut block.terminator
32+
else {
33+
return;
34+
};
35+
36+
let &mut Operand::Move(switch_place) = &mut *switch_desc else { return };
37+
38+
let Some(switch_local) = switch_place.as_local() else { return };
39+
40+
let Some(last_statement) = block.statements.last_mut() else { return };
41+
42+
let StatementKind::Assign(boxed_place_rvalue) = &last_statement.kind else { return };
43+
44+
let Some(assigned_local) = boxed_place_rvalue.0.as_local() else { return };
45+
46+
if switch_local != assigned_local {
47+
return;
48+
}
49+
50+
if !matches!(boxed_place_rvalue.1, Rvalue::Use(Operand::Constant(_))) {
51+
return;
52+
}
53+
54+
let should_optimize = tcx.consider_optimizing(|| {
55+
format!(
56+
"SimplifyBranches - Assignment: {:?} SourceInfo: {:?}",
57+
boxed_place_rvalue, last_statement.source_info
58+
)
59+
});
60+
61+
if should_optimize {
62+
let Some(last_statement) = block.statements.pop() else {
63+
bug!("Somehow the statement disappeared?");
64+
};
65+
66+
let StatementKind::Assign(boxed_place_rvalue) = last_statement.kind else {
67+
bug!("Somehow it's not an assignment any more?");
68+
};
69+
70+
let Rvalue::Use(assigned_constant @ Operand::Constant(_)) = boxed_place_rvalue.1 else {
71+
bug!("Somehow it's not a use of a constant any more?");
72+
};
73+
74+
*switch_desc = assigned_constant;
75+
}
76+
}

‎tests/codegen/skip-mono-inside-if-false.rs

+43
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
// compile-flags: -O -C no-prepopulate-passes
2+
3+
#![crate_type = "lib"]
4+
5+
#[no_mangle]
6+
pub fn demo_for_i32() {
7+
generic_impl::<i32>();
8+
}
9+
10+
// Two important things here:
11+
// - We replace the "then" block with `unreachable` to avoid linking problems
12+
// - We neither declare nor define the `big_impl` that said block "calls".
13+
14+
// CHECK-LABEL: ; skip_mono_inside_if_false::generic_impl
15+
// CHECK: start:
16+
// CHECK-NEXT: br i1 false, label %[[THEN_BRANCH:bb[0-9]+]], label %[[ELSE_BRANCH:bb[0-9]+]]
17+
// CHECK: [[ELSE_BRANCH]]:
18+
// CHECK-NEXT: call skip_mono_inside_if_false::small_impl
19+
// CHECK: [[THEN_BRANCH]]:
20+
// CHECK-NEXT: unreachable
21+
22+
// CHECK-NOT: @_ZN25skip_mono_inside_if_false8big_impl
23+
// CHECK: define internal void @_ZN25skip_mono_inside_if_false10small_impl
24+
// CHECK-NOT: @_ZN25skip_mono_inside_if_false8big_impl
25+
26+
fn generic_impl<T>() {
27+
trait MagicTrait {
28+
const IS_BIG: bool;
29+
}
30+
impl<T> MagicTrait for T {
31+
const IS_BIG: bool = std::mem::size_of::<T>() > 10;
32+
}
33+
if T::IS_BIG {
34+
big_impl::<T>();
35+
} else {
36+
small_impl::<T>();
37+
}
38+
}
39+
40+
#[inline(never)]
41+
fn small_impl<T>() {}
42+
#[inline(never)]
43+
fn big_impl<T>() {}

0 commit comments

Comments
 (0)