compiler/rustc_mir/src/transform/early_otherwise_branch.rs

   1 use crate::{transform::MirPass, util::patch::MirPatch};
   2 use rustc_middle::mir::*;
   3 use rustc_middle::ty::{Ty, TyCtxt};
   4 use std::fmt::Debug;
   5
   6 use super::simplify::simplify_cfg;
   7
   8 /// This pass optimizes something like
   9 /// ```text
  10 /// let x: Option<()>;
  11 /// let y: Option<()>;
  12 /// match (x,y) {
  13 ///     (Some(_), Some(_)) => {0},
  14 ///     _ => {1}
  15 /// }
  16 /// ```
  17 /// into something like
  18 /// ```text
  19 /// let x: Option<()>;
  20 /// let y: Option<()>;
  21 /// let discriminant_x = // get discriminant of x
  22 /// let discriminant_y = // get discriminant of y
  23 /// if discriminant_x != discriminant_y || discriminant_x == None {1} else {0}
  24 /// ```
  25 pub struct EarlyOtherwiseBranch;
  26
  27 impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
  28     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  29         if tcx.sess.opts.debugging_opts.mir_opt_level < 2 {
  30             return;
  31         }
  32         trace!("running EarlyOtherwiseBranch on {:?}", body.source);
  33         // we are only interested in this bb if the terminator is a switchInt
  34         let bbs_with_switch =
  35             body.basic_blocks().iter_enumerated().filter(|(_, bb)| is_switch(bb.terminator()));
  36
  37         let opts_to_apply: Vec<OptimizationToApply<'tcx>> = bbs_with_switch
  38             .flat_map(|(bb_idx, bb)| {
  39                 let switch = bb.terminator();
  40                 let helper = Helper { body, tcx };
  41                 let infos = helper.go(bb, switch)?;
  42                 Some(OptimizationToApply { infos, basic_block_first_switch: bb_idx })
  43             })
  44             .collect();
  45
  46         let should_cleanup = !opts_to_apply.is_empty();
  47
  48         for opt_to_apply in opts_to_apply {
  49             trace!("SUCCESS: found optimization possibility to apply: {:?}", &opt_to_apply);
  50
  51             let statements_before =
  52                 body.basic_blocks()[opt_to_apply.basic_block_first_switch].statements.len();
  53             let end_of_block_location = Location {
  54                 block: opt_to_apply.basic_block_first_switch,
  55                 statement_index: statements_before,
  56             };
  57
  58             let mut patch = MirPatch::new(body);
  59
  60             // create temp to store second discriminant in
  61             let discr_type = opt_to_apply.infos[0].second_switch_info.discr_ty;
  62             let discr_span = opt_to_apply.infos[0].second_switch_info.discr_source_info.span;
  63             let second_discriminant_temp = patch.new_temp(discr_type, discr_span);
  64
  65             patch.add_statement(
  66                 end_of_block_location,
  67                 StatementKind::StorageLive(second_discriminant_temp),
  68             );
  69
  70             // create assignment of discriminant
  71             let place_of_adt_to_get_discriminant_of =
  72                 opt_to_apply.infos[0].second_switch_info.place_of_adt_discr_read;
  73             patch.add_assign(
  74                 end_of_block_location,
  75                 Place::from(second_discriminant_temp),
  76                 Rvalue::Discriminant(place_of_adt_to_get_discriminant_of),
  77             );
  78
  79             // create temp to store NotEqual comparison between the two discriminants
  80             let not_equal = BinOp::Ne;
  81             let not_equal_res_type = not_equal.ty(tcx, discr_type, discr_type);
  82             let not_equal_temp = patch.new_temp(not_equal_res_type, discr_span);
  83             patch.add_statement(end_of_block_location, StatementKind::StorageLive(not_equal_temp));
  84
  85             // create NotEqual comparison between the two discriminants
  86             let first_descriminant_place =
  87                 opt_to_apply.infos[0].first_switch_info.discr_used_in_switch;
  88             let not_equal_rvalue = Rvalue::BinaryOp(
  89                 not_equal,
  90                 Operand::Copy(Place::from(second_discriminant_temp)),
  91                 Operand::Copy(first_descriminant_place),
  92             );
  93             patch.add_statement(
  94                 end_of_block_location,
  95                 StatementKind::Assign(box (Place::from(not_equal_temp), not_equal_rvalue)),
  96             );
  97
  98             let new_targets = opt_to_apply
  99                 .infos
 100                 .iter()
 101                 .flat_map(|x| x.second_switch_info.targets_with_values.iter())
 102                 .cloned();
 103
 104             let targets = SwitchTargets::new(
 105                 new_targets,
 106                 opt_to_apply.infos[0].first_switch_info.otherwise_bb,
 107             );
 108
 109             // new block that jumps to the correct discriminant case. This block is switched to if the discriminants are equal
 110             let new_switch_data = BasicBlockData::new(Some(Terminator {
 111                 source_info: opt_to_apply.infos[0].second_switch_info.discr_source_info,
 112                 kind: TerminatorKind::SwitchInt {
 113                     // the first and second discriminants are equal, so just pick one
 114                     discr: Operand::Copy(first_descriminant_place),
 115                     switch_ty: discr_type,
 116                     targets,
 117                 },
 118             }));
 119
 120             let new_switch_bb = patch.new_block(new_switch_data);
 121
 122             // switch on the NotEqual. If true, then jump to the `otherwise` case.
 123             // If false, then jump to a basic block that then jumps to the correct disciminant case
 124             let true_case = opt_to_apply.infos[0].first_switch_info.otherwise_bb;
 125             let false_case = new_switch_bb;
 126             patch.patch_terminator(
 127                 opt_to_apply.basic_block_first_switch,
 128                 TerminatorKind::if_(
 129                     tcx,
 130                     Operand::Move(Place::from(not_equal_temp)),
 131                     true_case,
 132                     false_case,
 133                 ),
 134             );
 135
 136             // generate StorageDead for the second_discriminant_temp not in use anymore
 137             patch.add_statement(
 138                 end_of_block_location,
 139                 StatementKind::StorageDead(second_discriminant_temp),
 140             );
 141
 142             // Generate a StorageDead for not_equal_temp in each of the targets, since we moved it into the switch
 143             for bb in [false_case, true_case].iter() {
 144                 patch.add_statement(
 145                     Location { block: *bb, statement_index: 0 },
 146                     StatementKind::StorageDead(not_equal_temp),
 147                 );
 148             }
 149
 150             patch.apply(body);
 151         }
 152
 153         // Since this optimization adds new basic blocks and invalidates others,
 154         // clean up the cfg to make it nicer for other passes
 155         if should_cleanup {
 156             simplify_cfg(body);
 157         }
 158     }
 159 }
 160
 161 fn is_switch<'tcx>(terminator: &Terminator<'tcx>) -> bool {
 162     match terminator.kind {
 163         TerminatorKind::SwitchInt { .. } => true,
 164         _ => false,
 165     }
 166 }
 167
 168 struct Helper<'a, 'tcx> {
 169     body: &'a Body<'tcx>,
 170     tcx: TyCtxt<'tcx>,
 171 }
 172
 173 #[derive(Debug, Clone)]
 174 struct SwitchDiscriminantInfo<'tcx> {
 175     /// Type of the discriminant being switched on
 176     discr_ty: Ty<'tcx>,
 177     /// The basic block that the otherwise branch points to
 178     otherwise_bb: BasicBlock,
 179     /// Target along with the value being branched from. Otherwise is not included
 180     targets_with_values: Vec<(u128, BasicBlock)>,
 181     discr_source_info: SourceInfo,
 182     /// The place of the discriminant used in the switch
 183     discr_used_in_switch: Place<'tcx>,
 184     /// The place of the adt that has its discriminant read
 185     place_of_adt_discr_read: Place<'tcx>,
 186     /// The type of the adt that has its discriminant read
 187     type_adt_matched_on: Ty<'tcx>,
 188 }
 189
 190 #[derive(Debug)]
 191 struct OptimizationToApply<'tcx> {
 192     infos: Vec<OptimizationInfo<'tcx>>,
 193     /// Basic block of the original first switch
 194     basic_block_first_switch: BasicBlock,
 195 }
 196
 197 #[derive(Debug)]
 198 struct OptimizationInfo<'tcx> {
 199     /// Info about the first switch and discriminant
 200     first_switch_info: SwitchDiscriminantInfo<'tcx>,
 201     /// Info about the second switch and discriminant
 202     second_switch_info: SwitchDiscriminantInfo<'tcx>,
 203 }
 204
 205 impl<'a, 'tcx> Helper<'a, 'tcx> {
 206     pub fn go(
 207         &self,
 208         bb: &BasicBlockData<'tcx>,
 209         switch: &Terminator<'tcx>,
 210     ) -> Option<Vec<OptimizationInfo<'tcx>>> {
 211         // try to find the statement that defines the discriminant that is used for the switch
 212         let discr = self.find_switch_discriminant_info(bb, switch)?;
 213
 214         // go through each target, finding a discriminant read, and a switch
 215         let results = discr.targets_with_values.iter().map(|(value, target)| {
 216             self.find_discriminant_switch_pairing(&discr, target.clone(), value.clone())
 217         });
 218
 219         // if the optimization did not apply for one of the targets, then abort
 220         if results.clone().any(|x| x.is_none()) || results.len() == 0 {
 221             trace!("NO: not all of the targets matched the pattern for optimization");
 222             return None;
 223         }
 224
 225         Some(results.flatten().collect())
 226     }
 227
 228     fn find_discriminant_switch_pairing(
 229         &self,
 230         discr_info: &SwitchDiscriminantInfo<'tcx>,
 231         target: BasicBlock,
 232         value: u128,
 233     ) -> Option<OptimizationInfo<'tcx>> {
 234         let bb = &self.body.basic_blocks()[target];
 235         // find switch
 236         let terminator = bb.terminator();
 237         if is_switch(terminator) {
 238             let this_bb_discr_info = self.find_switch_discriminant_info(bb, terminator)?;
 239
 240             // the types of the two adts matched on have to be equalfor this optimization to apply
 241             if discr_info.type_adt_matched_on != this_bb_discr_info.type_adt_matched_on {
 242                 trace!(
 243                     "NO: types do not match. LHS: {:?}, RHS: {:?}",
 244                     discr_info.type_adt_matched_on,
 245                     this_bb_discr_info.type_adt_matched_on
 246                 );
 247                 return None;
 248             }
 249
 250             // the otherwise branch of the two switches have to point to the same bb
 251             if discr_info.otherwise_bb != this_bb_discr_info.otherwise_bb {
 252                 trace!("NO: otherwise target is not the same");
 253                 return None;
 254             }
 255
 256             // check that the value being matched on is the same. The
 257             if this_bb_discr_info.targets_with_values.iter().find(|x| x.0 == value).is_none() {
 258                 trace!("NO: values being matched on are not the same");
 259                 return None;
 260             }
 261
 262             // only allow optimization if the left and right of the tuple being matched are the same variants.
 263             // so the following should not optimize
 264             //  ```rust
 265             // let x: Option<()>;
 266             // let y: Option<()>;
 267             // match (x,y) {
 268             //     (Some(_), None) => {},
 269             //     _ => {}
 270             // }
 271             //  ```
 272             // We check this by seeing that the value of the first discriminant is the only other discriminant value being used as a target in the second switch
 273             if !(this_bb_discr_info.targets_with_values.len() == 1
 274                 && this_bb_discr_info.targets_with_values[0].0 == value)
 275             {
 276                 trace!(
 277                     "NO: The second switch did not have only 1 target (besides otherwise) that had the same value as the value from the first switch that got us here"
 278                 );
 279                 return None;
 280             }
 281
 282             // if we reach this point, the optimization applies, and we should be able to optimize this case
 283             // store the info that is needed to apply the optimization
 284
 285             Some(OptimizationInfo {
 286                 first_switch_info: discr_info.clone(),
 287                 second_switch_info: this_bb_discr_info,
 288             })
 289         } else {
 290             None
 291         }
 292     }
 293
 294     fn find_switch_discriminant_info(
 295         &self,
 296         bb: &BasicBlockData<'tcx>,
 297         switch: &Terminator<'tcx>,
 298     ) -> Option<SwitchDiscriminantInfo<'tcx>> {
 299         match &switch.kind {
 300             TerminatorKind::SwitchInt { discr, targets, .. } => {
 301                 let discr_local = discr.place()?.as_local()?;
 302                 // the declaration of the discriminant read. Place of this read is being used in the switch
 303                 let discr_decl = &self.body.local_decls()[discr_local];
 304                 let discr_ty = discr_decl.ty;
 305                 // the otherwise target lies as the last element
 306                 let otherwise_bb = targets.otherwise();
 307                 let targets_with_values = targets.iter().collect();
 308
 309                 // find the place of the adt where the discriminant is being read from
 310                 // assume this is the last statement of the block
 311                 let place_of_adt_discr_read = match bb.statements.last()?.kind {
 312                     StatementKind::Assign(box (_, Rvalue::Discriminant(adt_place))) => {
 313                         Some(adt_place)
 314                     }
 315                     _ => None,
 316                 }?;
 317
 318                 let type_adt_matched_on = place_of_adt_discr_read.ty(self.body, self.tcx).ty;
 319
 320                 Some(SwitchDiscriminantInfo {
 321                     discr_used_in_switch: discr.place()?,
 322                     discr_ty,
 323                     otherwise_bb,
 324                     targets_with_values,
 325                     discr_source_info: discr_decl.source_info,
 326                     place_of_adt_discr_read,
 327                     type_adt_matched_on,
 328                 })
 329             }
 330             _ => unreachable!("must only be passed terminator that is a switch"),
 331         }
 332     }
 333 }