Introduce StatementKind::as_assign

Author: simonvandel

compiler/rustc_middle/src/mir/mod.rs

@@ -1504,6 +1504,13 @@ impl<'tcx> StatementKind<'tcx> {
             _ => None,
         }
     }
+
+    pub fn as_assign(&self) -> Option<&Box<(Place<'tcx>, Rvalue<'tcx>)>> {
+        match self {
+            StatementKind::Assign(x) => Some(x),
+            _ => None,
+        }
+    }
 }
 
 /// Describes what kind of retag is to be performed.

compiler/rustc_mir/src/transform/const_goto.rs

@@ -59,55 +59,52 @@ impl<'a, 'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'a, 'tcx> {
                     let last_statement =
                         self.body.basic_blocks()[location.block].statements.last()?;
 
-                    match &last_statement.kind {
-                        StatementKind::Assign(box (place, Rvalue::Use(op))) => {
-                            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();
-                            match &target_bb_terminator.kind {
-                                TerminatorKind::SwitchInt { discr, switch_ty, targets }
-                                    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();
+                        match &target_bb_terminator.kind {
+                            TerminatorKind::SwitchInt { discr, switch_ty, targets }
+                                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,
+                                });
                             }
+                            _ => {}
                         }
-                        _ => {}
                     }
                 }
                 _ => {}