@@ -51,7 +51,107 @@ pub fn provide(providers: &mut Providers) {
5151 } ;
5252}
5353
54+ use rustc_middle:: mir:: {
55+ BasicBlockData , ConstOperand , NullOp , Operand , Rvalue , StatementKind , SwitchTargets ,
56+ TerminatorKind ,
57+ } ;
58+ use rustc_middle:: ty:: { Instance , TyCtxt } ;
59+
5460/// `rustc_driver::main` installs a handler that will set this to `true` if
5561/// the compiler has been sent a request to shut down, such as by a Ctrl-C.
5662/// This static lives here because it is only read by the interpreter.
5763pub static CTRL_C_RECEIVED : AtomicBool = AtomicBool :: new ( false ) ;
64+
65+ /// If this basic block ends with a [`TerminatorKind::SwitchInt`] for which we can evaluate the
66+ /// dimscriminant in monomorphization, we return the discriminant bits and the
67+ /// [`SwitchTargets`], just so the caller doesn't also have to match on the terminator.
68+ pub fn try_const_mono_switchint < ' a , ' tcx > (
69+ tcx : TyCtxt < ' tcx > ,
70+ instance : Instance < ' tcx > ,
71+ block : & ' a BasicBlockData < ' tcx > ,
72+ ) -> Option < ( u128 , & ' a SwitchTargets ) > {
73+ // There are two places here we need to evaluate a constant.
74+ let eval_mono_const = |constant : & ConstOperand < ' tcx > | {
75+ let env = ty:: ParamEnv :: reveal_all ( ) ;
76+ let mono_literal = instance. instantiate_mir_and_normalize_erasing_regions (
77+ tcx,
78+ env,
79+ crate :: ty:: EarlyBinder :: bind ( constant. const_ ) ,
80+ ) ;
81+ mono_literal. try_eval_bits ( tcx, env)
82+ } ;
83+
84+ let TerminatorKind :: SwitchInt { discr, targets } = & block. terminator ( ) . kind else {
85+ return None ;
86+ } ;
87+
88+ // If this is a SwitchInt(const _), then we can just evaluate the constant and return.
89+ let discr = match discr {
90+ Operand :: Constant ( constant) => {
91+ let bits = eval_mono_const ( constant) ?;
92+ return Some ( ( bits, targets) ) ;
93+ }
94+ Operand :: Move ( place) | Operand :: Copy ( place) => place,
95+ } ;
96+
97+ // MIR for `if false` actually looks like this:
98+ // _1 = const _
99+ // SwitchInt(_1)
100+ //
101+ // And MIR for if intrinsics::ub_checks() looks like this:
102+ // _1 = UbChecks()
103+ // SwitchInt(_1)
104+ //
105+ // So we're going to try to recognize this pattern.
106+ //
107+ // If we have a SwitchInt on a non-const place, we find the most recent statement that
108+ // isn't a storage marker. If that statement is an assignment of a const to our
109+ // discriminant place, we evaluate and return the const, as if we've const-propagated it
110+ // into the SwitchInt.
111+
112+ let last_stmt = block. statements . iter ( ) . rev ( ) . find ( |stmt| {
113+ !matches ! ( stmt. kind, StatementKind :: StorageDead ( _) | StatementKind :: StorageLive ( _) )
114+ } ) ?;
115+
116+ let ( place, rvalue) = last_stmt. kind . as_assign ( ) ?;
117+
118+ if discr != place {
119+ return None ;
120+ }
121+
122+ use rustc_span:: DUMMY_SP ;
123+
124+ use crate :: const_eval:: DummyMachine ;
125+ use crate :: interpret:: InterpCx ;
126+ match rvalue {
127+ Rvalue :: NullaryOp ( NullOp :: UbChecks , _) => Some ( ( tcx. sess . ub_checks ( ) as u128 , targets) ) ,
128+ Rvalue :: Use ( Operand :: Constant ( constant) ) => {
129+ let bits = eval_mono_const ( constant) ?;
130+ Some ( ( bits, targets) )
131+ }
132+ Rvalue :: BinaryOp ( binop, box ( Operand :: Constant ( lhs) , Operand :: Constant ( rhs) ) ) => {
133+ let env = ty:: ParamEnv :: reveal_all ( ) ;
134+ let lhs = instance. instantiate_mir_and_normalize_erasing_regions (
135+ tcx,
136+ env,
137+ crate :: ty:: EarlyBinder :: bind ( lhs. const_ ) ,
138+ ) ;
139+ let ecx = InterpCx :: new ( tcx, DUMMY_SP , env, DummyMachine ) ;
140+ let lhs = ecx. eval_mir_constant ( & lhs, DUMMY_SP , None ) . ok ( ) ?;
141+ let lhs = ecx. read_immediate ( & lhs) . ok ( ) ?;
142+
143+ let rhs = instance. instantiate_mir_and_normalize_erasing_regions (
144+ tcx,
145+ env,
146+ crate :: ty:: EarlyBinder :: bind ( rhs. const_ ) ,
147+ ) ;
148+ let rhs = ecx. eval_mir_constant ( & rhs, DUMMY_SP , None ) . ok ( ) ?;
149+ let rhs = ecx. read_immediate ( & rhs) . ok ( ) ?;
150+
151+ let res = ecx. binary_op ( * binop, & lhs, & rhs) . ok ( ) ?;
152+ let res = res. to_scalar_int ( ) . unwrap ( ) . to_bits_unchecked ( ) ;
153+ Some ( ( res, targets) )
154+ }
155+ _ => None ,
156+ }
157+ }
0 commit comments