compiler/rustc_mir_transform/src/separate_const_switch.rs

   1 //! A pass that duplicates switch-terminated blocks
   2 //! into a new copy for each predecessor, provided
   3 //! the predecessor sets the value being switched
   4 //! over to a constant.
   5 //!
   6 //! The purpose of this pass is to help constant
   7 //! propagation passes to simplify the switch terminator
   8 //! of the copied blocks into gotos when some predecessors
   9 //! statically determine the output of switches.
  10 //!
  11 //! ```text
  12 //!     x = 12 ---              ---> something
  13 //!               \            / 12
  14 //!                --> switch x
  15 //!               /            \ otherwise
  16 //!     x = y  ---              ---> something else
  17 //! ```
  18 //! becomes
  19 //! ```text
  20 //!     x = 12 ---> switch x ------> something
  21 //!                          \ / 12
  22 //!                           X
  23 //!                          / \ otherwise
  24 //!     x = y  ---> switch x ------> something else
  25 //! ```
  26 //! so it can hopefully later be turned by another pass into
  27 //! ```text
  28 //!     x = 12 --------------------> something
  29 //!                            / 12
  30 //!                           /
  31 //!                          /   otherwise
  32 //!     x = y  ---- switch x ------> something else
  33 //! ```
  34 //!
  35 //! This optimization is meant to cover simple cases
  36 //! like `?` desugaring. For now, it thus focuses on
  37 //! simplicity rather than completeness (it notably
  38 //! sometimes duplicates abusively).
  39
  40 use crate::MirPass;
  41 use rustc_middle::mir::*;
  42 use rustc_middle::ty::TyCtxt;
  43 use smallvec::SmallVec;
  44
  45 pub struct SeparateConstSwitch;
  46
  47 impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
  48     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  49         sess.mir_opt_level() >= 4
  50     }
  51
  52     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  53         // If execution did something, applying a simplification layer
  54         // helps later passes optimize the copy away.
  55         if separate_const_switch(body) > 0 {
  56             super::simplify::simplify_cfg(tcx, body);
  57         }
  58     }
  59 }
  60
  61 /// Returns the amount of blocks that were duplicated
  62 pub fn separate_const_switch(body: &mut Body<'_>) -> usize {
  63     let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new();
  64     let predecessors = body.basic_blocks.predecessors();
  65     'block_iter: for (block_id, block) in body.basic_blocks.iter_enumerated() {
  66         if let TerminatorKind::SwitchInt {
  67             discr: Operand::Copy(switch_place) | Operand::Move(switch_place),
  68             ..
  69         } = block.terminator().kind
  70         {
  71             // If the block is on an unwind path, do not
  72             // apply the optimization as unwind paths
  73             // rely on a unique parent invariant
  74             if block.is_cleanup {
  75                 continue 'block_iter;
  76             }
  77
  78             // If the block has fewer than 2 predecessors, ignore it
  79             // we could maybe chain blocks that have exactly one
  80             // predecessor, but for now we ignore that
  81             if predecessors[block_id].len() < 2 {
  82                 continue 'block_iter;
  83             }
  84
  85             // First, let's find a non-const place
  86             // that determines the result of the switch
  87             if let Some(switch_place) = find_determining_place(switch_place, block) {
  88                 // We now have an input place for which it would
  89                 // be interesting if predecessors assigned it from a const
  90
  91                 let mut predecessors_left = predecessors[block_id].len();
  92                 'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() {
  93                     let predecessor = &body.basic_blocks[predecessor_id];
  94
  95                     // First we make sure the predecessor jumps
  96                     // in a reasonable way
  97                     match &predecessor.terminator().kind {
  98                         // The following terminators are
  99                         // unconditionally valid
 100                         TerminatorKind::Goto { .. } | TerminatorKind::SwitchInt { .. } => {}
 101
 102                         TerminatorKind::FalseEdge { real_target, .. } => {
 103                             if *real_target != block_id {
 104                                 continue 'predec_iter;
 105                             }
 106                         }
 107
 108                         // The following terminators are not allowed
 109                         TerminatorKind::Resume
 110                         | TerminatorKind::Drop { .. }
 111                         | TerminatorKind::DropAndReplace { .. }
 112                         | TerminatorKind::Call { .. }
 113                         | TerminatorKind::Assert { .. }
 114                         | TerminatorKind::FalseUnwind { .. }
 115                         | TerminatorKind::Yield { .. }
 116                         | TerminatorKind::Abort
 117                         | TerminatorKind::Return
 118                         | TerminatorKind::Unreachable
 119                         | TerminatorKind::InlineAsm { .. }
 120                         | TerminatorKind::GeneratorDrop => {
 121                             continue 'predec_iter;
 122                         }
 123                     }
 124
 125                     if is_likely_const(switch_place, predecessor) {
 126                         new_blocks.push((predecessor_id, block_id));
 127                         predecessors_left -= 1;
 128                         if predecessors_left < 2 {
 129                             // If the original block only has one predecessor left,
 130                             // we have nothing left to do
 131                             break 'predec_iter;
 132                         }
 133                     }
 134                 }
 135             }
 136         }
 137     }
 138
 139     // Once the analysis is done, perform the duplication
 140     let body_span = body.span;
 141     let copied_blocks = new_blocks.len();
 142     let blocks = body.basic_blocks_mut();
 143     for (pred_id, target_id) in new_blocks {
 144         let new_block = blocks[target_id].clone();
 145         let new_block_id = blocks.push(new_block);
 146         let terminator = blocks[pred_id].terminator_mut();
 147
 148         match terminator.kind {
 149             TerminatorKind::Goto { ref mut target } => {
 150                 *target = new_block_id;
 151             }
 152
 153             TerminatorKind::FalseEdge { ref mut real_target, .. } => {
 154                 if *real_target == target_id {
 155                     *real_target = new_block_id;
 156                 }
 157             }
 158
 159             TerminatorKind::SwitchInt { ref mut targets, .. } => {
 160                 targets.all_targets_mut().iter_mut().for_each(|x| {
 161                     if *x == target_id {
 162                         *x = new_block_id;
 163                     }
 164                 });
 165             }
 166
 167             TerminatorKind::Resume
 168             | TerminatorKind::Abort
 169             | TerminatorKind::Return
 170             | TerminatorKind::Unreachable
 171             | TerminatorKind::GeneratorDrop
 172             | TerminatorKind::Assert { .. }
 173             | TerminatorKind::DropAndReplace { .. }
 174             | TerminatorKind::FalseUnwind { .. }
 175             | TerminatorKind::Drop { .. }
 176             | TerminatorKind::Call { .. }
 177             | TerminatorKind::InlineAsm { .. }
 178             | TerminatorKind::Yield { .. } => {
 179                 span_bug!(
 180                     body_span,
 181                     "basic block terminator had unexpected kind {:?}",
 182                     &terminator.kind
 183                 )
 184             }
 185         }
 186     }
 187
 188     copied_blocks
 189 }
 190
 191 /// This function describes a rough heuristic guessing
 192 /// whether a place is last set with a const within the block.
 193 /// Notably, it will be overly pessimistic in cases that are already
 194 /// not handled by `separate_const_switch`.
 195 fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool {
 196     for statement in block.statements.iter().rev() {
 197         match &statement.kind {
 198             StatementKind::Assign(assign) => {
 199                 if assign.0 == tracked_place {
 200                     match assign.1 {
 201                         // These rvalues are definitely constant
 202                         Rvalue::Use(Operand::Constant(_))
 203                         | Rvalue::Ref(_, _, _)
 204                         | Rvalue::AddressOf(_, _)
 205                         | Rvalue::Cast(_, Operand::Constant(_), _)
 206                         | Rvalue::NullaryOp(_, _)
 207                         | Rvalue::ShallowInitBox(_, _)
 208                         | Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,
 209
 210                         // These rvalues make things ambiguous
 211                         Rvalue::Repeat(_, _)
 212                         | Rvalue::ThreadLocalRef(_)
 213                         | Rvalue::Len(_)
 214                         | Rvalue::BinaryOp(_, _)
 215                         | Rvalue::CheckedBinaryOp(_, _)
 216                         | Rvalue::Aggregate(_, _) => return false,
 217
 218                         // These rvalues move the place to track
 219                         Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _)
 220                         | Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
 221                         | Rvalue::CopyForDeref(place)
 222                         | Rvalue::UnaryOp(_, Operand::Copy(place) | Operand::Move(place))
 223                         | Rvalue::Discriminant(place) => tracked_place = place,
 224                     }
 225                 }
 226             }
 227
 228             // If the discriminant is set, it is always set
 229             // as a constant, so the job is done.
 230             // As we are **ignoring projections**, if the place
 231             // we are tracking sees its discriminant be set,
 232             // that means we had to be tracking the discriminant
 233             // specifically (as it is impossible to switch over
 234             // an enum directly, and if we were switching over
 235             // its content, we would have had to at least cast it to
 236             // some variant first)
 237             StatementKind::SetDiscriminant { place, .. } => {
 238                 if **place == tracked_place {
 239                     return true;
 240                 }
 241             }
 242
 243             // These statements have no influence on the place
 244             // we are interested in
 245             StatementKind::FakeRead(_)
 246             | StatementKind::Deinit(_)
 247             | StatementKind::StorageLive(_)
 248             | StatementKind::Retag(_, _)
 249             | StatementKind::AscribeUserType(_, _)
 250             | StatementKind::Coverage(_)
 251             | StatementKind::StorageDead(_)
 252             | StatementKind::Intrinsic(_)
 253             | StatementKind::Nop => {}
 254         }
 255     }
 256
 257     // If no good reason for the place to be const is found,
 258     // give up. We could maybe go up predecessors, but in
 259     // most cases giving up now should be sufficient.
 260     false
 261 }
 262
 263 /// Finds a unique place that entirely determines the value
 264 /// of `switch_place`, if it exists. This is only a heuristic.
 265 /// Ideally we would like to track multiple determining places
 266 /// for some edge cases, but one is enough for a lot of situations.
 267 fn find_determining_place<'tcx>(
 268     mut switch_place: Place<'tcx>,
 269     block: &BasicBlockData<'tcx>,
 270 ) -> Option<Place<'tcx>> {
 271     for statement in block.statements.iter().rev() {
 272         match &statement.kind {
 273             StatementKind::Assign(op) => {
 274                 if op.0 != switch_place {
 275                     continue;
 276                 }
 277
 278                 match op.1 {
 279                     // The following rvalues move the place
 280                     // that may be const in the predecessor
 281                     Rvalue::Use(Operand::Move(new) | Operand::Copy(new))
 282                     | Rvalue::UnaryOp(_, Operand::Copy(new) | Operand::Move(new))
 283                     | Rvalue::CopyForDeref(new)
 284                     | Rvalue::Cast(_, Operand::Move(new) | Operand::Copy(new), _)
 285                     | Rvalue::Repeat(Operand::Move(new) | Operand::Copy(new), _)
 286                     | Rvalue::Discriminant(new)
 287                     => switch_place = new,
 288
 289                     // The following rvalues might still make the block
 290                     // be valid but for now we reject them
 291                     Rvalue::Len(_)
 292                     | Rvalue::Ref(_, _, _)
 293                     | Rvalue::BinaryOp(_, _)
 294                     | Rvalue::CheckedBinaryOp(_, _)
 295                     | Rvalue::Aggregate(_, _)
 296
 297                     // The following rvalues definitely mean we cannot
 298                     // or should not apply this optimization
 299                     | Rvalue::Use(Operand::Constant(_))
 300                     | Rvalue::Repeat(Operand::Constant(_), _)
 301                     | Rvalue::ThreadLocalRef(_)
 302                     | Rvalue::AddressOf(_, _)
 303                     | Rvalue::NullaryOp(_, _)
 304                     | Rvalue::ShallowInitBox(_, _)
 305                     | Rvalue::UnaryOp(_, Operand::Constant(_))
 306                     | Rvalue::Cast(_, Operand::Constant(_), _)
 307                     => return None,
 308                 }
 309             }
 310
 311             // These statements have no influence on the place
 312             // we are interested in
 313             StatementKind::FakeRead(_)
 314             | StatementKind::Deinit(_)
 315             | StatementKind::StorageLive(_)
 316             | StatementKind::StorageDead(_)
 317             | StatementKind::Retag(_, _)
 318             | StatementKind::AscribeUserType(_, _)
 319             | StatementKind::Coverage(_)
 320             | StatementKind::Intrinsic(_)
 321             | StatementKind::Nop => {}
 322
 323             // If the discriminant is set, it is always set
 324             // as a constant, so the job is already done.
 325             // As we are **ignoring projections**, if the place
 326             // we are tracking sees its discriminant be set,
 327             // that means we had to be tracking the discriminant
 328             // specifically (as it is impossible to switch over
 329             // an enum directly, and if we were switching over
 330             // its content, we would have had to at least cast it to
 331             // some variant first)
 332             StatementKind::SetDiscriminant { place, .. } => {
 333                 if **place == switch_place {
 334                     return None;
 335                 }
 336             }
 337         }
 338     }
 339
 340     Some(switch_place)
 341 }