Introduce TerminatorKind::as_goto

Author: simonvandel

compiler/rustc_middle/src/mir/terminator.rs

@@ -416,6 +416,13 @@ impl<'tcx> TerminatorKind<'tcx> {
             _ => None,
         }
     }
+
+    pub fn as_goto(&self) -> Option<BasicBlock> {
+        match self {
+            TerminatorKind::Goto { target } => Some(*target),
+            _ => None,
+        }
+    }
 }
 
 impl<'tcx> Debug for TerminatorKind<'tcx> {

compiler/rustc_mir/src/transform/const_goto.rs

@@ -53,57 +53,54 @@ impl<'tcx> MirPass<'tcx> for ConstGoto {
 impl<'a, 'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'a, 'tcx> {
     fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
         let _:Option<_> = try {
-            match terminator.kind {
-                TerminatorKind::Goto { target } => {
-                    // We only apply this optimization if the last statement is a const assignment
-                    let last_statement =
-                        self.body.basic_blocks()[location.block].statements.last()?;
+            let target = terminator.kind.as_goto()?;
+            // We only apply this optimization if the last statement is a const assignment
+            let last_statement =
+                self.body.basic_blocks()[location.block].statements.last()?;
 
-                    if let Some(box (place, Rvalue::Use(op))) = last_statement.kind.as_assign(){
-                        let _const = op.constant()?;
-                        // We found a constant being assigned to `place`.
-                        // Now check that the target of this Goto switches on this place.
-                        let target_bb = &self.body.basic_blocks()[target];
-                        // FIXME(simonvandel): We are conservative here when we don't allow 
-                        // any statements in the target basic block.
-                        // This could probably be relaxed to allow `StorageDead`s which could be 
-                        // copied to the predecessor of this block.
-                        if !target_bb.statements.is_empty() {
-                            None?
-                        }
+            if let Some(box (place, Rvalue::Use(op))) = last_statement.kind.as_assign(){
+                let _const = op.constant()?;
+                // We found a constant being assigned to `place`.
+                // Now check that the target of this Goto switches on this place.
+                let target_bb = &self.body.basic_blocks()[target];
+                
+                // FIXME(simonvandel): We are conservative here when we don't allow 
+                // any statements in the target basic block.
+                // This could probably be relaxed to allow `StorageDead`s which could be 
+                // copied to the predecessor of this block.
+                if !target_bb.statements.is_empty() {
+                    None?
+                }
 
-                        let target_bb_terminator = target_bb.terminator();
-                        let (discr, switch_ty, targets) = target_bb_terminator.kind.as_switch()?;
-                        if discr.place() == Some(*place) {
-                            // We now know that the Switch matches on the const place, and it is statementless
-                            // Now find which value in the Switch matches the const value.
-                            let const_value = _const.literal.try_eval_bits(
-                                self.tcx,
-                                self.param_env,
-                                switch_ty,
-                            )?;
-                            let found_value_idx_option = targets
-                                .iter()
-                                .enumerate()
-                                .find(|(_, x)| const_value == x.0)
-                                .map(|(idx, _)| idx);
+                let target_bb_terminator = target_bb.terminator();
+                let (discr, switch_ty, targets) = target_bb_terminator.kind.as_switch()?;
+                if discr.place() == Some(*place) {
+                    // We now know that the Switch matches on the const place, and it is statementless
+                    // Now find which value in the Switch matches the const value.
+                    let const_value = _const.literal.try_eval_bits(
+                        self.tcx,
+                        self.param_env,
+                        switch_ty,
+                    )?;
+                    let found_value_idx_option = targets
+                        .iter()
+                        .enumerate()
+                        .find(|(_, x)| const_value == x.0)
+                        .map(|(idx, _)| idx);
 
-                            let target_to_use_in_goto =
-                                if let Some(found_value_idx) = found_value_idx_option {
-                                    targets.iter().nth(found_value_idx).unwrap().1
-                                } else {
-                                    // If we did not find the const value in values, it must be the otherwise case
-                                    targets.otherwise()
-                                };
+                    let target_to_use_in_goto =
+                        if let Some(found_value_idx) = found_value_idx_option {
+                            targets.iter().nth(found_value_idx).unwrap().1
+                        } else {
+                            // If we did not find the const value in values, it must be the otherwise case
+                            targets.otherwise()
+                        };
 
-                            self.optimizations.push(OptimizationToApply {
-                                bb_with_goto: location.block,
-                                target_to_use_in_goto,
-                            });
-                        }
-                    }
+                    self.optimizations.push(OptimizationToApply {
+                        bb_with_goto: location.block,
+                        target_to_use_in_goto,
+                    });
                 }
-                _ => {}
             }
             Some(())
         };