Auto merge of #60745 - wesleywiser:const_prop_into_terminators, r=<try>

[WIP] Const prop into terminators

Rebased on top of #60597. New commits start at eea75dc.

Last commit is just for testing purposes and I'll remove it before we merge this.

r? @ghost

Author: bors

src/librustc_mir/transform/const_prop.rs

@@ -3,12 +3,14 @@
 
 use rustc::hir::def::DefKind;
 use rustc::mir::{
-    Constant, Location, Place, PlaceBase, Mir, Operand, Rvalue, Local,
+    AggregateKind, Constant, Location, Place, PlaceBase, Mir, Operand, Rvalue, Local,
     NullOp, UnOp, StatementKind, Statement, LocalKind, Static, StaticKind,
     TerminatorKind, Terminator,  ClearCrossCrate, SourceInfo, BinOp, ProjectionElem,
     SourceScope, SourceScopeLocalData, LocalDecl, Promoted,
 };
-use rustc::mir::visit::{Visitor, PlaceContext, MutatingUseContext, NonMutatingUseContext};
+use rustc::mir::visit::{
+    Visitor, PlaceContext, MutatingUseContext, MutVisitor, NonMutatingUseContext,
+};
 use rustc::mir::interpret::{InterpError, Scalar, GlobalId, EvalResult};
 use rustc::ty::{self, Instance, ParamEnv, Ty, TyCtxt};
 use syntax::source_map::DUMMY_SP;
@@ -19,7 +21,7 @@ use rustc::ty::layout::{
     HasTyCtxt, TargetDataLayout, HasDataLayout,
 };
 
-use crate::interpret::{InterpretCx, ScalarMaybeUndef, Immediate, OpTy, ImmTy, MemoryKind};
+use crate::interpret::{self, InterpretCx, ScalarMaybeUndef, Immediate, OpTy, ImmTy, MemoryKind};
 use crate::const_eval::{
     CompileTimeInterpreter, error_to_const_error, eval_promoted, mk_eval_cx,
 };
@@ -497,6 +499,50 @@ impl<'a, 'mir, 'tcx> ConstPropagator<'a, 'mir, 'tcx> {
             },
         }
     }
+
+    fn operand_from_scalar(&self, scalar: Scalar, ty: Ty<'tcx>) -> Operand<'tcx> {
+        Operand::Constant(Box::new(
+            Constant {
+                span: DUMMY_SP,
+                ty,
+                user_ty: None,
+                literal: self.tcx.mk_const(ty::Const::from_scalar(
+                    scalar,
+                    ty,
+                ))
+            }
+        ))
+    }
+
+    fn replace_with_const(&self, rval: &mut Rvalue<'tcx>, value: Const<'tcx>) {
+        if let interpret::Operand::Immediate(im) = *value {
+            match im {
+                interpret::Immediate::Scalar(ScalarMaybeUndef::Scalar(scalar)) => {
+                    *rval = Rvalue::Use(self.operand_from_scalar(scalar, value.layout.ty));
+                },
+                Immediate::ScalarPair(
+                    ScalarMaybeUndef::Scalar(one),
+                    ScalarMaybeUndef::Scalar(two)
+                ) => {
+                    let ty = &value.layout.ty.sty;
+                    if let ty::Tuple(substs) = ty {
+                        *rval = Rvalue::Aggregate(
+                            Box::new(AggregateKind::Tuple),
+                            vec![
+                                self.operand_from_scalar(one, substs[0].expect_ty()),
+                                self.operand_from_scalar(two, substs[1].expect_ty()),
+                            ],
+                        );
+                    }
+                },
+                _ => { }
+            }
+        }
+    }
+
+    fn should_const_prop(&self) -> bool {
+        true
+    }
 }
 
 fn type_size_of<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
@@ -560,10 +606,10 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
     }
 }
 
-impl<'b, 'a, 'tcx> Visitor<'tcx> for ConstPropagator<'b, 'a, 'tcx> {
+impl<'b, 'a, 'tcx> MutVisitor<'tcx> for ConstPropagator<'b, 'a, 'tcx> {
     fn visit_constant(
         &mut self,
-        constant: &Constant<'tcx>,
+        constant: &mut Constant<'tcx>,
         location: Location,
     ) {
         trace!("visit_constant: {:?}", constant);
@@ -573,11 +619,11 @@ impl<'b, 'a, 'tcx> Visitor<'tcx> for ConstPropagator<'b, 'a, 'tcx> {
 
     fn visit_statement(
         &mut self,
-        statement: &Statement<'tcx>,
+        statement: &mut Statement<'tcx>,
         location: Location,
     ) {
         trace!("visit_statement: {:?}", statement);
-        if let StatementKind::Assign(ref place, ref rval) = statement.kind {
+        if let StatementKind::Assign(ref place, ref mut rval) = statement.kind {
             let place_ty: Ty<'tcx> = place
                 .ty(&self.local_decls, self.tcx)
                 .ty;
@@ -589,6 +635,10 @@ impl<'b, 'a, 'tcx> Visitor<'tcx> for ConstPropagator<'b, 'a, 'tcx> {
                             trace!("storing {:?} to {:?}", value, local);
                             assert!(self.places[local].is_none());
                             self.places[local] = Some(value);
+
+                            if self.should_const_prop() {
+                                self.replace_with_const(rval, value);
+                            }
                         }
                     }
                 }
@@ -599,79 +649,111 @@ impl<'b, 'a, 'tcx> Visitor<'tcx> for ConstPropagator<'b, 'a, 'tcx> {
 
     fn visit_terminator(
         &mut self,
-        terminator: &Terminator<'tcx>,
+        terminator: &mut Terminator<'tcx>,
         location: Location,
     ) {
         self.super_terminator(terminator, location);
-        let source_info = terminator.source_info;;
-        if let TerminatorKind::Assert { expected, msg, cond, .. } = &terminator.kind {
-            if let Some(value) = self.eval_operand(&cond, source_info) {
-                trace!("assertion on {:?} should be {:?}", value, expected);
-                let expected = ScalarMaybeUndef::from(Scalar::from_bool(*expected));
-                if expected != self.ecx.read_scalar(value).unwrap() {
-                    // poison all places this operand references so that further code
-                    // doesn't use the invalid value
-                    match cond {
-                        Operand::Move(ref place) | Operand::Copy(ref place) => {
-                            let mut place = place;
-                            while let Place::Projection(ref proj) = *place {
-                                place = &proj.base;
-                            }
-                            if let Place::Base(PlaceBase::Local(local)) = *place {
-                                self.places[local] = None;
+        let source_info = terminator.source_info;
+        match &mut terminator.kind {
+            TerminatorKind::Assert { expected, msg, ref mut cond, .. } => {
+                if let Some(value) = self.eval_operand(&cond, source_info) {
+                    trace!("assertion on {:?} should be {:?}", value, expected);
+                    let expected = ScalarMaybeUndef::from(Scalar::from_bool(*expected));
+                    let value_const = self.ecx.read_scalar(value).unwrap();
+                    if expected != value_const {
+                        // poison all places this operand references so that further code
+                        // doesn't use the invalid value
+                        match cond {
+                            Operand::Move(ref place) | Operand::Copy(ref place) => {
+                                let mut place = place;
+                                while let Place::Projection(ref proj) = *place {
+                                    place = &proj.base;
+                                }
+                                if let Place::Base(PlaceBase::Local(local)) = *place {
+                                    self.places[local] = None;
+                                }
+                            },
+                            Operand::Constant(_) => {}
+                        }
+                        let span = terminator.source_info.span;
+                        let hir_id = self
+                            .tcx
+                            .hir()
+                            .as_local_hir_id(self.source.def_id())
+                            .expect("some part of a failing const eval must be local");
+                        use rustc::mir::interpret::InterpError::*;
+                        let msg = match msg {
+                            Overflow(_) |
+                            OverflowNeg |
+                            DivisionByZero |
+                            RemainderByZero => msg.description().to_owned(),
+                            BoundsCheck { ref len, ref index } => {
+                                let len = self
+                                    .eval_operand(len, source_info)
+                                    .expect("len must be const");
+                                let len = match self.ecx.read_scalar(len) {
+                                    Ok(ScalarMaybeUndef::Scalar(Scalar::Bits {
+                                        bits, ..
+                                    })) => bits,
+                                    other => bug!("const len not primitive: {:?}", other),
+                                };
+                                let index = self
+                                    .eval_operand(index, source_info)
+                                    .expect("index must be const");
+                                let index = match self.ecx.read_scalar(index) {
+                                    Ok(ScalarMaybeUndef::Scalar(Scalar::Bits {
+                                        bits, ..
+                                    })) => bits,
+                                    other => bug!("const index not primitive: {:?}", other),
+                                };
+                                format!(
+                                    "index out of bounds: \
+                                    the len is {} but the index is {}",
+                                    len,
+                                    index,
+                                )
+                            },
+                            // Need proper const propagator for these
+                            _ => return,
+                        };
+                        self.tcx.lint_hir(
+                            ::rustc::lint::builtin::CONST_ERR,
+                            hir_id,
+                            span,
+                            &msg,
+                        );
+                    } else {
+                        if self.should_const_prop() {
+                            if let ScalarMaybeUndef::Scalar(scalar) = value_const {
+                                *cond = self.operand_from_scalar(scalar, self.tcx.types.bool);
                             }
-                        },
-                        Operand::Constant(_) => {}
+                        }
                     }
-                    let span = terminator.source_info.span;
-                    let hir_id = self
-                        .tcx
-                        .hir()
-                        .as_local_hir_id(self.source.def_id())
-                        .expect("some part of a failing const eval must be local");
-                    use rustc::mir::interpret::InterpError::*;
-                    let msg = match msg {
-                        Overflow(_) |
-                        OverflowNeg |
-                        DivisionByZero |
-                        RemainderByZero => msg.description().to_owned(),
-                        BoundsCheck { ref len, ref index } => {
-                            let len = self
-                                .eval_operand(len, source_info)
-                                .expect("len must be const");
-                            let len = match self.ecx.read_scalar(len) {
-                                Ok(ScalarMaybeUndef::Scalar(Scalar::Bits {
-                                    bits, ..
-                                })) => bits,
-                                other => bug!("const len not primitive: {:?}", other),
-                            };
-                            let index = self
-                                .eval_operand(index, source_info)
-                                .expect("index must be const");
-                            let index = match self.ecx.read_scalar(index) {
-                                Ok(ScalarMaybeUndef::Scalar(Scalar::Bits {
-                                    bits, ..
-                                })) => bits,
-                                other => bug!("const index not primitive: {:?}", other),
-                            };
-                            format!(
-                                "index out of bounds: \
-                                the len is {} but the index is {}",
-                                len,
-                                index,
-                            )
-                        },
-                        // Need proper const propagator for these
-                        _ => return,
-                    };
-                    self.tcx.lint_hir(
-                        ::rustc::lint::builtin::CONST_ERR,
-                        hir_id,
-                        span,
-                        &msg,
-                    );
                 }
-            }
+            },
+            TerminatorKind::SwitchInt { ref mut discr, switch_ty, .. } => {
+                if self.should_const_prop() {
+                    if let Some(value) = self.eval_operand(&discr, source_info) {
+                        if let Ok(ScalarMaybeUndef::Scalar(scalar)) = self.ecx.read_scalar(value) {
+                            *discr = self.operand_from_scalar(scalar, switch_ty);
+                        }
+                    }
+                }
+            },
+            //none of these have Operands to const-propagate
+            TerminatorKind::Goto { .. } |
+            TerminatorKind::Resume |
+            TerminatorKind::Abort |
+            TerminatorKind::Return |
+            TerminatorKind::Unreachable |
+            TerminatorKind::Drop { .. } |
+            TerminatorKind::DropAndReplace { .. } |
+            TerminatorKind::Yield { .. } |
+            TerminatorKind::GeneratorDrop |
+            TerminatorKind::FalseEdges { .. } |
+            TerminatorKind::FalseUnwind { .. } => { }
+            //FIXME(wesleywiser) Call does have Operands that could be const-propagated
+            TerminatorKind::Call { .. } => { }
         }
     }
 }

src/test/codegen/inline-always-works-always.rs

@@ -14,7 +14,7 @@ pub extern "C" fn callee() -> u32 {
 // CHECK-LABEL: caller
 // SIZE-OPT: ret i32 8
 // SPEED-OPT: ret i32 8
-// NO-OPT: ret i32 8
+// NO-OPT: ret i32 %3
 #[no_mangle]
 pub extern "C" fn caller() -> u32 {
     callee()

src/test/codegen/optimize-attr-1.rs

@@ -8,7 +8,7 @@
 
 // CHECK-LABEL: define i32 @nothing
 // CHECK-SAME: [[NOTHING_ATTRS:#[0-9]+]]
-// NO-OPT: ret i32 %1
+// NO-OPT: ret i32 %3
 // SIZE-OPT: ret i32 4
 // SPEEC-OPT: ret i32 4
 #[no_mangle]
@@ -18,7 +18,7 @@ pub fn nothing() -> i32 {
 
 // CHECK-LABEL: define i32 @size
 // CHECK-SAME: [[SIZE_ATTRS:#[0-9]+]]
-// NO-OPT: ret i32 %1
+// NO-OPT: ret i32 %3
 // SIZE-OPT: ret i32 6
 // SPEED-OPT: ret i32 6
 #[optimize(size)]
@@ -31,7 +31,7 @@ pub fn size() -> i32 {
 // NO-OPT-SAME: [[NOTHING_ATTRS]]
 // SPEED-OPT-SAME: [[NOTHING_ATTRS]]
 // SIZE-OPT-SAME: [[SPEED_ATTRS:#[0-9]+]]
-// NO-OPT: ret i32 %1
+// NO-OPT: ret i32 %3
 // SIZE-OPT: ret i32 8
 // SPEED-OPT: ret i32 8
 #[optimize(speed)]

src/test/mir-opt/const_prop/array_index.rs

@@ -0,0 +1,33 @@
+fn main() {
+    let x: u32 = [0, 1, 2, 3][2];
+}
+
+// END RUST SOURCE
+// START rustc.main.ConstProp.before.mir
+//  bb0: {
+//      ...
+//      _2 = [const 0u32, const 1u32, const 2u32, const 3u32];
+//      ...
+//      _3 = const 2usize;
+//      _4 = const 4usize;
+//      _5 = Lt(_3, _4);
+//      assert(move _5, "index out of bounds: the len is move _4 but the index is _3") -> bb1;
+//  }
+//  bb1: {
+//      _1 = _2[_3];
+//      ...
+//      return;
+//  }
+// END rustc.main.ConstProp.before.mir
+// START rustc.main.ConstProp.after.mir
+//  bb0: {
+//      ...
+//      _5 = const true;
+//      assert(const true, "index out of bounds: the len is move _4 but the index is _3") -> bb1;
+//  }
+//  bb1: {
+//      _1 = _2[_3];
+//      ...
+//      return;
+//  }
+// END rustc.main.ConstProp.after.mir

src/test/mir-opt/const_prop/checked_add.rs

@@ -0,0 +1,21 @@
+// compile-flags: -C overflow-checks=on
+
+fn main() {
+    let x: u32 = 1 + 1;
+}
+
+// END RUST SOURCE
+// START rustc.main.ConstProp.before.mir
+//  bb0: {
+//      ...
+//      _2 = CheckedAdd(const 1u32, const 1u32);
+//      assert(!move (_2.1: bool), "attempt to add with overflow") -> bb1;
+//  }
+// END rustc.main.ConstProp.before.mir
+// START rustc.main.ConstProp.after.mir
+//  bb0: {
+//      ...
+//      _2 = (const 2u32, const false);
+//      assert(!const false, "attempt to add with overflow") -> bb1;
+//  }
+// END rustc.main.ConstProp.after.mir