Let MonoReachable traversal evaluate BinOps

Author: saethlin

Cargo.lock

@@ -3479,6 +3479,7 @@ dependencies = [
  "rustc_macros",
  "rustc_metadata",
  "rustc_middle",
+ "rustc_mir_transform",
  "rustc_monomorphize",
  "rustc_query_system",
  "rustc_serialize",
@@ -4147,6 +4148,7 @@ dependencies = [
  "rustc_hir",
  "rustc_macros",
  "rustc_middle",
+ "rustc_mir_transform",
  "rustc_session",
  "rustc_span",
  "rustc_target",

compiler/rustc_codegen_cranelift/src/base.rs

@@ -299,7 +299,8 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
         .generic_activity("codegen prelude")
         .run(|| crate::abi::codegen_fn_prelude(fx, start_block));
 
-    let reachable_blocks = traversal::mono_reachable_as_bitset(fx.mir, fx.tcx, fx.instance);
+    let reachable_blocks =
+        rustc_mir_transform::mono_reachable_as_bitset(fx.mir, fx.tcx, fx.instance);
 
     for (bb, bb_data) in fx.mir.basic_blocks.iter_enumerated() {
         let block = fx.get_block(bb);

compiler/rustc_codegen_cranelift/src/lib.rs

@@ -24,6 +24,7 @@ extern crate rustc_hir;
 extern crate rustc_incremental;
 extern crate rustc_index;
 extern crate rustc_metadata;
+extern crate rustc_mir_transform;
 extern crate rustc_session;
 extern crate rustc_span;
 extern crate rustc_target;

compiler/rustc_codegen_ssa/Cargo.toml

@@ -27,6 +27,7 @@ rustc_index = { path = "../rustc_index" }
 rustc_macros = { path = "../rustc_macros" }
 rustc_metadata = { path = "../rustc_metadata" }
 rustc_middle = { path = "../rustc_middle" }
+rustc_mir_transform = { path = "../rustc_mir_transform" }
 rustc_monomorphize = { path = "../rustc_monomorphize" }
 rustc_query_system = { path = "../rustc_query_system" }
 rustc_serialize = { path = "../rustc_serialize" }

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -277,7 +277,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     // So drop the builder of `start_llbb` to avoid having two at the same time.
     drop(start_bx);
 
-    let reachable_blocks = traversal::mono_reachable_as_bitset(mir, cx.tcx(), instance);
+    let reachable_blocks = rustc_mir_transform::mono_reachable_as_bitset(mir, cx.tcx(), instance);
 
     // Codegen the body of each block using reverse postorder
     for (bb, _) in traversal::reverse_postorder(mir) {

compiler/rustc_const_eval/src/lib.rs

@@ -51,7 +51,107 @@ pub fn provide(providers: &mut Providers) {
     };
 }
 
+use rustc_middle::mir::{
+    BasicBlockData, ConstOperand, NullOp, Operand, Rvalue, StatementKind, SwitchTargets,
+    TerminatorKind,
+};
+use rustc_middle::ty::{Instance, TyCtxt};
+
 /// `rustc_driver::main` installs a handler that will set this to `true` if
 /// the compiler has been sent a request to shut down, such as by a Ctrl-C.
 /// This static lives here because it is only read by the interpreter.
 pub static CTRL_C_RECEIVED: AtomicBool = AtomicBool::new(false);
+
+/// If this basic block ends with a [`TerminatorKind::SwitchInt`] for which we can evaluate the
+/// dimscriminant in monomorphization, we return the discriminant bits and the
+/// [`SwitchTargets`], just so the caller doesn't also have to match on the terminator.
+pub fn try_const_mono_switchint<'a, 'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    block: &'a BasicBlockData<'tcx>,
+) -> Option<(u128, &'a SwitchTargets)> {
+    // There are two places here we need to evaluate a constant.
+    let eval_mono_const = |constant: &ConstOperand<'tcx>| {
+        let env = ty::ParamEnv::reveal_all();
+        let mono_literal = instance.instantiate_mir_and_normalize_erasing_regions(
+            tcx,
+            env,
+            crate::ty::EarlyBinder::bind(constant.const_),
+        );
+        mono_literal.try_eval_bits(tcx, env)
+    };
+
+    let TerminatorKind::SwitchInt { discr, targets } = &block.terminator().kind else {
+        return None;
+    };
+
+    // If this is a SwitchInt(const _), then we can just evaluate the constant and return.
+    let discr = match discr {
+        Operand::Constant(constant) => {
+            let bits = eval_mono_const(constant)?;
+            return Some((bits, targets));
+        }
+        Operand::Move(place) | Operand::Copy(place) => place,
+    };
+
+    // MIR for `if false` actually looks like this:
+    // _1 = const _
+    // SwitchInt(_1)
+    //
+    // And MIR for if intrinsics::ub_checks() looks like this:
+    // _1 = UbChecks()
+    // SwitchInt(_1)
+    //
+    // So we're going to try to recognize this pattern.
+    //
+    // If we have a SwitchInt on a non-const place, we find the most recent statement that
+    // isn't a storage marker. If that statement is an assignment of a const to our
+    // discriminant place, we evaluate and return the const, as if we've const-propagated it
+    // into the SwitchInt.
+
+    let last_stmt = block.statements.iter().rev().find(|stmt| {
+        !matches!(stmt.kind, StatementKind::StorageDead(_) | StatementKind::StorageLive(_))
+    })?;
+
+    let (place, rvalue) = last_stmt.kind.as_assign()?;
+
+    if discr != place {
+        return None;
+    }
+
+    use rustc_span::DUMMY_SP;
+
+    use crate::const_eval::DummyMachine;
+    use crate::interpret::InterpCx;
+    match rvalue {
+        Rvalue::NullaryOp(NullOp::UbChecks, _) => Some((tcx.sess.ub_checks() as u128, targets)),
+        Rvalue::Use(Operand::Constant(constant)) => {
+            let bits = eval_mono_const(constant)?;
+            Some((bits, targets))
+        }
+        Rvalue::BinaryOp(binop, box (Operand::Constant(lhs), Operand::Constant(rhs))) => {
+            let env = ty::ParamEnv::reveal_all();
+            let lhs = instance.instantiate_mir_and_normalize_erasing_regions(
+                tcx,
+                env,
+                crate::ty::EarlyBinder::bind(lhs.const_),
+            );
+            let ecx = InterpCx::new(tcx, DUMMY_SP, env, DummyMachine);
+            let lhs = ecx.eval_mir_constant(&lhs, DUMMY_SP, None).ok()?;
+            let lhs = ecx.read_immediate(&lhs).ok()?;
+
+            let rhs = instance.instantiate_mir_and_normalize_erasing_regions(
+                tcx,
+                env,
+                crate::ty::EarlyBinder::bind(rhs.const_),
+            );
+            let rhs = ecx.eval_mir_constant(&rhs, DUMMY_SP, None).ok()?;
+            let rhs = ecx.read_immediate(&rhs).ok()?;
+
+            let res = ecx.binary_op(*binop, &lhs, &rhs).ok()?;
+            let res = res.to_scalar_int().unwrap().to_bits_unchecked();
+            Some((res, targets))
+        }
+        _ => None,
+    }
+}

compiler/rustc_middle/src/mir/mod.rs

@@ -42,7 +42,7 @@ use crate::ty::fold::{FallibleTypeFolder, TypeFoldable};
 use crate::ty::print::{pretty_print_const, with_no_trimmed_paths, FmtPrinter, Printer};
 use crate::ty::visit::TypeVisitableExt;
 use crate::ty::{
-    self, AdtDef, GenericArg, GenericArgsRef, Instance, InstanceKind, List, Ty, TyCtxt,
+    self, AdtDef, GenericArg, GenericArgsRef, InstanceKind, List, Ty, TyCtxt,
     UserTypeAnnotationIndex,
 };
 
@@ -687,73 +687,6 @@ impl<'tcx> Body<'tcx> {
         self.injection_phase.is_some()
     }
 
-    /// If this basic block ends with a [`TerminatorKind::SwitchInt`] for which we can evaluate the
-    /// dimscriminant in monomorphization, we return the discriminant bits and the
-    /// [`SwitchTargets`], just so the caller doesn't also have to match on the terminator.
-    fn try_const_mono_switchint<'a>(
-        tcx: TyCtxt<'tcx>,
-        instance: Instance<'tcx>,
-        block: &'a BasicBlockData<'tcx>,
-    ) -> Option<(u128, &'a SwitchTargets)> {
-        // There are two places here we need to evaluate a constant.
-        let eval_mono_const = |constant: &ConstOperand<'tcx>| {
-            let env = ty::ParamEnv::reveal_all();
-            let mono_literal = instance.instantiate_mir_and_normalize_erasing_regions(
-                tcx,
-                env,
-                crate::ty::EarlyBinder::bind(constant.const_),
-            );
-            mono_literal.try_eval_bits(tcx, env)
-        };
-
-        let TerminatorKind::SwitchInt { discr, targets } = &block.terminator().kind else {
-            return None;
-        };
-
-        // If this is a SwitchInt(const _), then we can just evaluate the constant and return.
-        let discr = match discr {
-            Operand::Constant(constant) => {
-                let bits = eval_mono_const(constant)?;
-                return Some((bits, targets));
-            }
-            Operand::Move(place) | Operand::Copy(place) => place,
-        };
-
-        // MIR for `if false` actually looks like this:
-        // _1 = const _
-        // SwitchInt(_1)
-        //
-        // And MIR for if intrinsics::ub_checks() looks like this:
-        // _1 = UbChecks()
-        // SwitchInt(_1)
-        //
-        // So we're going to try to recognize this pattern.
-        //
-        // If we have a SwitchInt on a non-const place, we find the most recent statement that
-        // isn't a storage marker. If that statement is an assignment of a const to our
-        // discriminant place, we evaluate and return the const, as if we've const-propagated it
-        // into the SwitchInt.
-
-        let last_stmt = block.statements.iter().rev().find(|stmt| {
-            !matches!(stmt.kind, StatementKind::StorageDead(_) | StatementKind::StorageLive(_))
-        })?;
-
-        let (place, rvalue) = last_stmt.kind.as_assign()?;
-
-        if discr != place {
-            return None;
-        }
-
-        match rvalue {
-            Rvalue::NullaryOp(NullOp::UbChecks, _) => Some((tcx.sess.ub_checks() as u128, targets)),
-            Rvalue::Use(Operand::Constant(constant)) => {
-                let bits = eval_mono_const(constant)?;
-                Some((bits, targets))
-            }
-            _ => None,
-        }
-    }
-
     /// For a `Location` in this scope, determine what the "caller location" at that point is. This
     /// is interesting because of inlining: the `#[track_caller]` attribute of inlined functions
     /// must be honored. Falls back to the `tracked_caller` value for `#[track_caller]` functions,

compiler/rustc_middle/src/mir/traversal.rs

@@ -279,97 +279,3 @@ pub fn reverse_postorder<'a, 'tcx>(
 {
     body.basic_blocks.reverse_postorder().iter().map(|&bb| (bb, &body.basic_blocks[bb]))
 }
-
-/// Traversal of a [`Body`] that tries to avoid unreachable blocks in a monomorphized [`Instance`].
-///
-/// This is allowed to have false positives; blocks may be visited even if they are not actually
-/// reachable.
-///
-/// Such a traversal is mostly useful because it lets us skip lowering the `false` side
-/// of `if <T as Trait>::CONST`, as well as [`NullOp::UbChecks`].
-///
-/// [`NullOp::UbChecks`]: rustc_middle::mir::NullOp::UbChecks
-pub fn mono_reachable<'a, 'tcx>(
-    body: &'a Body<'tcx>,
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
-) -> MonoReachable<'a, 'tcx> {
-    MonoReachable::new(body, tcx, instance)
-}
-
-/// [`MonoReachable`] internally accumulates a [`BitSet`] of visited blocks. This is just a
-/// convenience function to run that traversal then extract its set of reached blocks.
-pub fn mono_reachable_as_bitset<'a, 'tcx>(
-    body: &'a Body<'tcx>,
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
-) -> BitSet<BasicBlock> {
-    let mut iter = mono_reachable(body, tcx, instance);
-    while let Some(_) = iter.next() {}
-    iter.visited
-}
-
-pub struct MonoReachable<'a, 'tcx> {
-    body: &'a Body<'tcx>,
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
-    visited: BitSet<BasicBlock>,
-    // Other traversers track their worklist in a Vec. But we don't care about order, so we can
-    // store ours in a BitSet and thus save allocations because BitSet has a small size
-    // optimization.
-    worklist: BitSet<BasicBlock>,
-}
-
-impl<'a, 'tcx> MonoReachable<'a, 'tcx> {
-    pub fn new(
-        body: &'a Body<'tcx>,
-        tcx: TyCtxt<'tcx>,
-        instance: Instance<'tcx>,
-    ) -> MonoReachable<'a, 'tcx> {
-        let mut worklist = BitSet::new_empty(body.basic_blocks.len());
-        worklist.insert(START_BLOCK);
-        MonoReachable {
-            body,
-            tcx,
-            instance,
-            visited: BitSet::new_empty(body.basic_blocks.len()),
-            worklist,
-        }
-    }
-
-    fn add_work(&mut self, blocks: impl IntoIterator<Item = BasicBlock>) {
-        for block in blocks.into_iter() {
-            if !self.visited.contains(block) {
-                self.worklist.insert(block);
-            }
-        }
-    }
-}
-
-impl<'a, 'tcx> Iterator for MonoReachable<'a, 'tcx> {
-    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
-
-    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
-        while let Some(idx) = self.worklist.iter().next() {
-            self.worklist.remove(idx);
-            if !self.visited.insert(idx) {
-                continue;
-            }
-
-            let data = &self.body[idx];
-
-            if let Some((bits, targets)) =
-                Body::try_const_mono_switchint(self.tcx, self.instance, data)
-            {
-                let target = targets.target_for_value(bits);
-                self.add_work([target]);
-            } else {
-                self.add_work(data.terminator().successors());
-            }
-
-            return Some((idx, data));
-        }
-
-        None
-    }
-}