|
| 1 | +//! A pass that checks for and simplifies calls to `pow` where the receiver is a power of |
| 2 | +//! two. This can be done with `<<` instead. |
| 3 | +
|
| 4 | +use crate::MirPass; |
| 5 | +use rustc_const_eval::interpret::{ConstValue, Scalar}; |
| 6 | +use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData}; |
| 7 | +use rustc_middle::mir::patch::MirPatch; |
| 8 | +use rustc_middle::mir::*; |
| 9 | +use rustc_middle::ty::{self, Ty, TyCtxt, UintTy}; |
| 10 | +use rustc_span::sym; |
| 11 | +use rustc_target::abi::FieldIdx; |
| 12 | + |
| 13 | +pub struct SimplifyPowOfTwo; |
| 14 | + |
| 15 | +impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo { |
| 16 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 17 | + let mut patch = MirPatch::new(body); |
| 18 | + |
| 19 | + for (i, bb) in body.basic_blocks.iter_enumerated() { |
| 20 | + let term = bb.terminator(); |
| 21 | + let source_info = term.source_info; |
| 22 | + let span = source_info.span; |
| 23 | + |
| 24 | + if let TerminatorKind::Call { |
| 25 | + func, |
| 26 | + args, |
| 27 | + destination, |
| 28 | + target: Some(target), |
| 29 | + unwind, |
| 30 | + call_source: CallSource::Normal, |
| 31 | + .. |
| 32 | + } = &term.kind |
| 33 | + && let Some(def_id) = func.const_fn_def().map(|def| def.0) |
| 34 | + && let def_path = tcx.def_path(def_id) |
| 35 | + && tcx.crate_name(def_path.krate) == sym::core |
| 36 | + // FIXME(Centri3): I feel like we should do this differently... |
| 37 | + && let [ |
| 38 | + DisambiguatedDefPathData { data: DefPathData::TypeNs(sym::num), disambiguator: 0 }, |
| 39 | + DisambiguatedDefPathData { data: DefPathData::Impl, .. }, |
| 40 | + DisambiguatedDefPathData { data: DefPathData::ValueNs(sym::pow), .. }, |
| 41 | + ] = &*def_path.data |
| 42 | + && let [recv, exp] = args.as_slice() |
| 43 | + && let Some(recv_const) = recv.constant() |
| 44 | + && let ConstantKind::Val( |
| 45 | + ConstValue::Scalar(Scalar::Int(recv_int)), |
| 46 | + recv_ty, |
| 47 | + ) = recv_const.literal |
| 48 | + && let Ok(recv_val) = match recv_ty.kind() { |
| 49 | + ty::Int(_) => { |
| 50 | + let result = recv_int.try_to_int(recv_int.size()).unwrap_or(-1).max(0); |
| 51 | + if result > 0 { |
| 52 | + Ok(result as u128) |
| 53 | + } else { |
| 54 | + continue; |
| 55 | + } |
| 56 | + }, |
| 57 | + ty::Uint(_) => recv_int.try_to_uint(recv_int.size()), |
| 58 | + _ => continue, |
| 59 | + } |
| 60 | + && let power_used = f32::log2(recv_val as f32) |
| 61 | + // Precision loss means it's not a power of two |
| 62 | + && power_used == (power_used as u32) as f32 |
| 63 | + // `0` would be `1.pow()`, which we shouldn't try to optimize as it's |
| 64 | + // already entirely optimized away |
| 65 | + && power_used != 0.0 |
| 66 | + // Same here |
| 67 | + && recv_val != 0 |
| 68 | + { |
| 69 | + let power_used = power_used as u32; |
| 70 | + let loc = Location { block: i, statement_index: bb.statements.len() }; |
| 71 | + let exp_ty = Ty::new(tcx, ty::Uint(UintTy::U32)); |
| 72 | + let checked_mul = |
| 73 | + patch.new_temp(Ty::new_tup(tcx, &[exp_ty, Ty::new_bool(tcx)]), span); |
| 74 | + |
| 75 | + // If this is not `2.pow(...)`, we need to multiply the number of times we |
| 76 | + // shift the bits left by the receiver's power of two used, e.g.: |
| 77 | + // |
| 78 | + // > 2 -> 1 |
| 79 | + // > 4 -> 2 |
| 80 | + // > 16 -> 4 |
| 81 | + // > 256 -> 8 |
| 82 | + // |
| 83 | + // If this is `1`, then we *could* remove this entirely but it'll be |
| 84 | + // optimized out anyway by later passes (or perhaps LLVM) so it's entirely |
| 85 | + // unnecessary to do so. |
| 86 | + patch.add_assign( |
| 87 | + loc, |
| 88 | + checked_mul.into(), |
| 89 | + Rvalue::CheckedBinaryOp( |
| 90 | + BinOp::Mul, |
| 91 | + Box::new(( |
| 92 | + exp.clone(), |
| 93 | + Operand::Constant(Box::new(Constant { |
| 94 | + span, |
| 95 | + user_ty: None, |
| 96 | + literal: ConstantKind::Val( |
| 97 | + ConstValue::from_u32(power_used), |
| 98 | + exp_ty, |
| 99 | + ), |
| 100 | + })), |
| 101 | + )), |
| 102 | + ), |
| 103 | + ); |
| 104 | + |
| 105 | + let num_shl = tcx.mk_place_field(checked_mul.into(), FieldIdx::from_u32(0), exp_ty); |
| 106 | + let mul_result = |
| 107 | + tcx.mk_place_field(checked_mul.into(), FieldIdx::from_u32(1), Ty::new_bool(tcx)); |
| 108 | + let shl_result = patch.new_temp(Ty::new_bool(tcx), span); |
| 109 | + |
| 110 | + // Whether the shl will overflow, if so we return 0 |
| 111 | + patch.add_assign( |
| 112 | + loc, |
| 113 | + shl_result.into(), |
| 114 | + Rvalue::BinaryOp( |
| 115 | + BinOp::Lt, |
| 116 | + Box::new(( |
| 117 | + Operand::Copy(num_shl), |
| 118 | + Operand::Constant(Box::new(Constant { |
| 119 | + span, |
| 120 | + user_ty: None, |
| 121 | + literal: ConstantKind::Val(ConstValue::from_u32(32), exp_ty), |
| 122 | + })), |
| 123 | + )), |
| 124 | + ), |
| 125 | + ); |
| 126 | + |
| 127 | + let should_be_zero_bool = patch.new_temp(Ty::new_bool(tcx), span); |
| 128 | + let should_be_zero = patch.new_temp(recv_ty, span); |
| 129 | + |
| 130 | + patch.add_assign( |
| 131 | + loc, |
| 132 | + should_be_zero_bool.into(), |
| 133 | + Rvalue::BinaryOp( |
| 134 | + BinOp::BitOr, |
| 135 | + Box::new(( |
| 136 | + Operand::Copy(mul_result.into()), |
| 137 | + Operand::Copy(shl_result.into()), |
| 138 | + )), |
| 139 | + ), |
| 140 | + ); |
| 141 | + |
| 142 | + patch.add_assign( |
| 143 | + loc, |
| 144 | + should_be_zero.into(), |
| 145 | + Rvalue::Cast( |
| 146 | + CastKind::IntToInt, |
| 147 | + Operand::Copy(should_be_zero_bool.into()), |
| 148 | + recv_ty, |
| 149 | + ), |
| 150 | + ); |
| 151 | + |
| 152 | + let shl_exp_ty = patch.new_temp(exp_ty, span); |
| 153 | + let shl = patch.new_temp(recv_ty, span); |
| 154 | + |
| 155 | + patch.add_assign( |
| 156 | + loc, |
| 157 | + shl_exp_ty.into(), |
| 158 | + Rvalue::BinaryOp( |
| 159 | + BinOp::Shl, |
| 160 | + Box::new(( |
| 161 | + Operand::Constant(Box::new(Constant { |
| 162 | + span, |
| 163 | + user_ty: None, |
| 164 | + literal: ConstantKind::Val(ConstValue::from_u32(1), exp_ty), |
| 165 | + })), |
| 166 | + Operand::Copy(num_shl.into()), |
| 167 | + )), |
| 168 | + ), |
| 169 | + ); |
| 170 | + |
| 171 | + patch.add_assign( |
| 172 | + loc, |
| 173 | + shl.into(), |
| 174 | + Rvalue::Cast( |
| 175 | + CastKind::IntToInt, |
| 176 | + Operand::Copy(shl_exp_ty.into()), |
| 177 | + recv_ty, |
| 178 | + ), |
| 179 | + ); |
| 180 | + |
| 181 | + patch.add_assign( |
| 182 | + loc, |
| 183 | + *destination, |
| 184 | + Rvalue::BinaryOp( |
| 185 | + BinOp::MulUnchecked, |
| 186 | + Box::new((Operand::Copy(shl.into()), Operand::Copy(should_be_zero.into()))), |
| 187 | + ), |
| 188 | + ); |
| 189 | + |
| 190 | + // shl doesn't set the overflow flag on x86_64 or even in Rust, so shr to |
| 191 | + // see if it overflowed. If it equals 1, it did not, but we also need to |
| 192 | + // check `shl_result` to ensure that if this is a multiple of the type's |
| 193 | + // size it won't wrap back over to 1 |
| 194 | + // |
| 195 | + // FIXME(Centri3): Do we use `debug_assertions` or `overflow_checks` here? |
| 196 | + if tcx.sess.opts.debug_assertions { |
| 197 | + let shr = patch.new_temp(recv_ty, span); |
| 198 | + let shl_eq_shr = patch.new_temp(Ty::new_bool(tcx), span); |
| 199 | + let overflowed = patch.new_temp(Ty::new_bool(tcx), span); |
| 200 | + |
| 201 | + patch.add_assign( |
| 202 | + loc, |
| 203 | + shr.into(), |
| 204 | + Rvalue::BinaryOp( |
| 205 | + BinOp::Shr, |
| 206 | + Box::new((Operand::Copy(shl.into()), Operand::Copy(num_shl.into()))), |
| 207 | + ), |
| 208 | + ); |
| 209 | + |
| 210 | + patch.add_assign( |
| 211 | + loc, |
| 212 | + shl_eq_shr.into(), |
| 213 | + Rvalue::BinaryOp( |
| 214 | + BinOp::Eq, |
| 215 | + Box::new((Operand::Copy(shl.into()), Operand::Copy(shr.into()))), |
| 216 | + ), |
| 217 | + ); |
| 218 | + |
| 219 | + patch.add_assign( |
| 220 | + loc, |
| 221 | + overflowed.into(), |
| 222 | + Rvalue::BinaryOp( |
| 223 | + BinOp::BitAnd, |
| 224 | + Box::new((Operand::Copy(shl_eq_shr.into()), Operand::Copy(shl_result.into()))), |
| 225 | + ), |
| 226 | + ); |
| 227 | + |
| 228 | + patch.patch_terminator( |
| 229 | + i, |
| 230 | + TerminatorKind::Assert { |
| 231 | + cond: Operand::Copy(overflowed.into()), |
| 232 | + expected: true, |
| 233 | + msg: Box::new(AssertMessage::Overflow( |
| 234 | + // For consistency with the previous error message, though |
| 235 | + // it's technically incorrect |
| 236 | + BinOp::Mul, |
| 237 | + Operand::Constant(Box::new(Constant { |
| 238 | + span, |
| 239 | + user_ty: None, |
| 240 | + literal: ConstantKind::Val(ConstValue::Scalar(Scalar::from_u32(1)), exp_ty), |
| 241 | + })), |
| 242 | + Operand::Copy(num_shl.into()), |
| 243 | + )), |
| 244 | + target: *target, |
| 245 | + unwind: *unwind, |
| 246 | + }, |
| 247 | + ); |
| 248 | + } else { |
| 249 | + patch.patch_terminator(i, TerminatorKind::Goto { target: *target }); |
| 250 | + } |
| 251 | + } |
| 252 | + } |
| 253 | + |
| 254 | + patch.apply(body); |
| 255 | + } |
| 256 | +} |
0 commit comments