GVN: Track known value ranges through assert and switchInt.

Author: dianqk

compiler/rustc_mir_transform/src/gvn.rs

@@ -90,7 +90,9 @@ use std::hash::{Hash, Hasher};
 use either::Either;
 use hashbrown::hash_table::{Entry, HashTable};
 use itertools::Itertools as _;
-use rustc_abi::{self as abi, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, VariantIdx};
+use rustc_abi::{
+    self as abi, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, VariantIdx, WrappingRange,
+};
 use rustc_arena::DroplessArena;
 use rustc_const_eval::const_eval::DummyMachine;
 use rustc_const_eval::interpret::{
@@ -380,6 +382,8 @@ struct VnState<'body, 'a, 'tcx> {
     dominators: Dominators<BasicBlock>,
     reused_locals: DenseBitSet<Local>,
     arena: &'a DroplessArena,
+    /// Known ranges at each locations.
+    ranges: IndexVec<VnIndex, Vec<(Location, WrappingRange)>>,
 }
 
 impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
@@ -413,6 +417,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             dominators,
             reused_locals: DenseBitSet::new_empty(local_decls.len()),
             arena,
+            ranges: IndexVec::with_capacity(num_values),
         }
     }
 
@@ -429,6 +434,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             debug_assert_eq!(index, _index);
             let _index = self.rev_locals.push(SmallVec::new());
             debug_assert_eq!(index, _index);
+            let _index = self.ranges.push(Vec::new());
+            debug_assert_eq!(index, _index);
         }
         index
     }
@@ -442,6 +449,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         debug_assert_eq!(index, _index);
         let _index = self.rev_locals.push(SmallVec::new());
         debug_assert_eq!(index, _index);
+        let _index = self.ranges.push(Vec::new());
+        debug_assert_eq!(index, _index);
         index
     }
 
@@ -480,6 +489,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         debug_assert_eq!(index, _index);
         let _index = self.rev_locals.push(SmallVec::new());
         debug_assert_eq!(index, _index);
+        let _index = self.ranges.push(Vec::new());
+        debug_assert_eq!(index, _index);
 
         Some(index)
     }
@@ -504,6 +515,8 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             debug_assert_eq!(index, _index);
             let _index = self.rev_locals.push(SmallVec::new());
             debug_assert_eq!(index, _index);
+            let _index = self.ranges.push(Vec::new());
+            debug_assert_eq!(index, _index);
         }
         index
     }
@@ -526,6 +539,18 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
         self.rev_locals[value].push(local);
     }
 
+    fn insert_range(&mut self, value: VnIndex, location: Location, range: WrappingRange) {
+        self.ranges[value].push((location, range));
+    }
+
+    fn get_range(&self, value: VnIndex, location: Location) -> Option<WrappingRange> {
+        // FIXME: This should use the intersection of all valid ranges.
+        let (_, range) = self.ranges[value]
+            .iter()
+            .find(|(range_loc, _)| range_loc.dominates(location, &self.dominators))?;
+        Some(*range)
+    }
+
     fn insert_bool(&mut self, flag: bool) -> VnIndex {
         // Booleans are deterministic.
         let value = Const::from_bool(self.tcx, flag);
@@ -1010,7 +1035,7 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             Operand::Constant(ref constant) => Some(self.insert_constant(constant.const_)),
             Operand::Copy(ref mut place) | Operand::Move(ref mut place) => {
                 let value = self.simplify_place_value(place, location)?;
-                if let Some(const_) = self.try_as_constant(value) {
+                if let Some(const_) = self.try_as_constant(value, location) {
                     *operand = Operand::Constant(Box::new(const_));
                 }
                 Some(value)
@@ -1780,7 +1805,7 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
     /// If either [`Self::try_as_constant`] as [`Self::try_as_place`] succeeds,
     /// returns that result as an [`Operand`].
     fn try_as_operand(&mut self, index: VnIndex, location: Location) -> Option<Operand<'tcx>> {
-        if let Some(const_) = self.try_as_constant(index) {
+        if let Some(const_) = self.try_as_constant(index, location) {
             Some(Operand::Constant(Box::new(const_)))
         } else if let Some(place) = self.try_as_place(index, location, false) {
             self.reused_locals.insert(place.local);
@@ -1791,7 +1816,11 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
     }
 
     /// If `index` is a `Value::Constant`, return the `Constant` to be put in the MIR.
-    fn try_as_constant(&mut self, index: VnIndex) -> Option<ConstOperand<'tcx>> {
+    fn try_as_constant(
+        &mut self,
+        index: VnIndex,
+        location: Location,
+    ) -> Option<ConstOperand<'tcx>> {
         // This was already constant in MIR, do not change it. If the constant is not
         // deterministic, adding an additional mention of it in MIR will not give the same value as
         // the former mention.
@@ -1800,21 +1829,34 @@ impl<'tcx> VnState<'_, '_, 'tcx> {
             return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_: value });
         }
 
-        let op = self.eval_to_const(index)?;
-        if op.layout.is_unsized() {
-            // Do not attempt to propagate unsized locals.
-            return None;
-        }
+        if let Some(op) = self.eval_to_const(index) {
+            if op.layout.is_unsized() {
+                // Do not attempt to propagate unsized locals.
+                return None;
+            }
 
-        let value = op_to_prop_const(&mut self.ecx, op)?;
+            let value = op_to_prop_const(&mut self.ecx, op)?;
 
-        // Check that we do not leak a pointer.
-        // Those pointers may lose part of their identity in codegen.
-        // FIXME: remove this hack once https://github.com/rust-lang/rust/issues/79738 is fixed.
-        assert!(!value.may_have_provenance(self.tcx, op.layout.size));
+            // Check that we do not leak a pointer.
+            // Those pointers may lose part of their identity in codegen.
+            // FIXME: remove this hack once https://github.com/rust-lang/rust/issues/79738 is fixed.
+            assert!(!value.may_have_provenance(self.tcx, op.layout.size));
+
+            let const_ = Const::Val(value, op.layout.ty);
+            return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ });
+        }
 
-        let const_ = Const::Val(value, op.layout.ty);
-        Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ })
+        if let Some(range) = self.get_range(index, location)
+            && range.start == range.end
+        {
+            let ty = self.ty(index);
+            let layout = self.ecx.layout_of(ty).ok()?;
+            let value = ConstValue::Scalar(Scalar::from_uint(range.start, layout.size));
+            let const_ = Const::Val(value, self.ty(index));
+            return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ });
+        }
+
+        None
     }
 
     /// Construct a place which holds the same value as `index` and for which all locals strictly
@@ -1891,7 +1933,7 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
 
         let value = self.simplify_rvalue(lhs, rvalue, location);
         if let Some(value) = value {
-            if let Some(const_) = self.try_as_constant(value) {
+            if let Some(const_) = self.try_as_constant(value, location) {
                 *rvalue = Rvalue::Use(Operand::Constant(Box::new(const_)));
             } else if let Some(place) = self.try_as_place(value, location, false)
                 && *rvalue != Rvalue::Use(Operand::Move(place))
@@ -1920,14 +1962,56 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
     }
 
     fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
-        if let Terminator { kind: TerminatorKind::Call { destination, .. }, .. } = terminator {
-            if let Some(local) = destination.as_local()
-                && self.ssa.is_ssa(local)
-            {
-                let ty = self.local_decls[local].ty;
-                let opaque = self.new_opaque(ty);
-                self.assign(local, opaque);
+        match &mut terminator.kind {
+            TerminatorKind::Call { destination, .. } => {
+                if let Some(local) = destination.as_local()
+                    && self.ssa.is_ssa(local)
+                {
+                    let ty = self.local_decls[local].ty;
+                    let opaque = self.new_opaque(ty);
+                    self.assign(local, opaque);
+                }
+            }
+            TerminatorKind::Assert { cond, expected, target, .. } => {
+                if let Some(value) = self.simplify_operand(cond, location) {
+                    let successor = Location { block: *target, statement_index: 0 };
+                    if location.block != successor.block
+                        && location.dominates(successor, &self.dominators)
+                    {
+                        let val = *expected as u128;
+                        let range = WrappingRange { start: val, end: val };
+                        self.insert_range(value, successor, range);
+                    }
+                }
+            }
+            TerminatorKind::SwitchInt { discr, targets } => {
+                if let Some(value) = self.simplify_operand(discr, location) {
+                    for (val, target) in targets.iter() {
+                        let successor = Location { block: target, statement_index: 0 };
+                        if location.block != successor.block
+                            && location.dominates(successor, &self.dominators)
+                        {
+                            let range = WrappingRange { start: val, end: val };
+                            self.insert_range(value, successor, range);
+                        }
+                    }
+
+                    let otherwise = Location { block: targets.otherwise(), statement_index: 0 };
+                    if self.ty(value).is_bool()
+                        && let [val] = targets.all_values()
+                        && location.block != otherwise.block
+                        && location.dominates(otherwise, &self.dominators)
+                    {
+                        let range = if val.get() == 0 {
+                            WrappingRange { start: 1, end: 1 }
+                        } else {
+                            WrappingRange { start: 0, end: 0 }
+                        };
+                        self.insert_range(value, otherwise, range);
+                    }
+                }
             }
+            _ => {}
         }
         // Terminators that can write to memory may invalidate (nested) derefs.
         if terminator.kind.can_write_to_memory() {

tests/mir-opt/early_otherwise_branch_unwind.poll.EarlyOtherwiseBranch.diff

@@ -39,7 +39,7 @@
 
       bb3: {
           _2 = discriminant(((((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>) as Ok).0: std::option::Option<std::vec::Vec<u8>>));
-          switchInt(copy _2) -> [0: bb5, 1: bb7, otherwise: bb1];
+          switchInt(move _2) -> [0: bb5, 1: bb7, otherwise: bb1];
       }
 
       bb4: {
@@ -112,15 +112,15 @@
       }
 
       bb18 (cleanup): {
-          switchInt(copy _3) -> [0: bb19, otherwise: bb9];
+          switchInt(const 0_isize) -> [0: bb19, otherwise: bb9];
       }
 
       bb19 (cleanup): {
           goto -> bb9;
       }
 
       bb20 (cleanup): {
-          switchInt(copy _4) -> [0: bb18, otherwise: bb9];
+          switchInt(const 0_isize) -> [0: bb18, otherwise: bb9];
       }
   }
 

tests/mir-opt/early_otherwise_branch_unwind.unwind.EarlyOtherwiseBranch.diff

@@ -35,7 +35,7 @@
 
       bb3: {
           _2 = discriminant(((((_1 as Some).0: std::option::Option<std::option::Option<T>>) as Some).0: std::option::Option<T>));
-          switchInt(copy _2) -> [0: bb6, 1: bb7, otherwise: bb1];
+          switchInt(move _2) -> [0: bb6, 1: bb7, otherwise: bb1];
       }
 
       bb4: {
@@ -105,15 +105,15 @@
       }
 
       bb18 (cleanup): {
-          switchInt(copy _3) -> [1: bb19, otherwise: bb9];
+          switchInt(const 1_isize) -> [1: bb19, otherwise: bb9];
       }
 
       bb19 (cleanup): {
           goto -> bb9;
       }
 
       bb20 (cleanup): {
-          switchInt(copy _4) -> [1: bb18, otherwise: bb9];
+          switchInt(const 1_isize) -> [1: bb18, otherwise: bb9];
       }
   }
 

tests/mir-opt/gvn.arithmetic.GVN.panic-abort.diff

@@ -222,8 +222,8 @@
           _32 = copy _1;
 -         _33 = Eq(copy _32, const 0_u64);
 -         assert(!move _33, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind unreachable];
-+         _33 = copy _29;
-+         assert(!copy _29, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind unreachable];
++         _33 = const false;
++         assert(!const false, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind unreachable];
       }
 
       bb12: {
@@ -283,8 +283,8 @@
           _44 = copy _1;
 -         _45 = Eq(copy _44, const 0_u64);
 -         assert(!move _45, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind unreachable];
-+         _45 = copy _29;
-+         assert(!copy _29, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind unreachable];
++         _45 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind unreachable];
       }
 
       bb18: {
@@ -304,8 +304,8 @@
           _48 = copy _1;
 -         _49 = Eq(copy _48, const 0_u64);
 -         assert(!move _49, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind unreachable];
-+         _49 = copy _29;
-+         assert(!copy _29, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind unreachable];
++         _49 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind unreachable];
       }
 
       bb20: {

tests/mir-opt/gvn.arithmetic.GVN.panic-unwind.diff

@@ -222,8 +222,8 @@
           _32 = copy _1;
 -         _33 = Eq(copy _32, const 0_u64);
 -         assert(!move _33, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind continue];
-+         _33 = copy _29;
-+         assert(!copy _29, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind continue];
++         _33 = const false;
++         assert(!const false, "attempt to divide `{}` by zero", const 1_u64) -> [success: bb12, unwind continue];
       }
 
       bb12: {
@@ -283,8 +283,8 @@
           _44 = copy _1;
 -         _45 = Eq(copy _44, const 0_u64);
 -         assert(!move _45, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind continue];
-+         _45 = copy _29;
-+         assert(!copy _29, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind continue];
++         _45 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", const 0_u64) -> [success: bb18, unwind continue];
       }
 
       bb18: {
@@ -304,8 +304,8 @@
           _48 = copy _1;
 -         _49 = Eq(copy _48, const 0_u64);
 -         assert(!move _49, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind continue];
-+         _49 = copy _29;
-+         assert(!copy _29, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind continue];
++         _49 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", const 1_u64) -> [success: bb20, unwind continue];
       }
 
       bb20: {

tests/mir-opt/gvn.subexpression_elimination.GVN.panic-abort.diff

@@ -283,8 +283,8 @@
           _24 = copy _2;
 -         _25 = Eq(copy _24, const 0_u64);
 -         assert(!move _25, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _23) -> [success: bb6, unwind unreachable];
-+         _25 = copy _20;
-+         assert(!copy _20, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _1) -> [success: bb6, unwind unreachable];
++         _25 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _1) -> [success: bb6, unwind unreachable];
       }
 
       bb6: {
@@ -520,8 +520,8 @@
           _74 = copy _2;
 -         _75 = Eq(copy _74, const 0_u64);
 -         assert(!move _75, "attempt to divide `{}` by zero", copy _73) -> [success: bb20, unwind unreachable];
-+         _75 = copy _20;
-+         assert(!copy _20, "attempt to divide `{}` by zero", copy _1) -> [success: bb20, unwind unreachable];
++         _75 = const false;
++         assert(!const false, "attempt to divide `{}` by zero", copy _1) -> [success: bb20, unwind unreachable];
       }
 
       bb20: {
@@ -544,8 +544,8 @@
           _79 = copy _2;
 -         _80 = Eq(copy _79, const 0_u64);
 -         assert(!move _80, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _78) -> [success: bb22, unwind unreachable];
-+         _80 = copy _20;
-+         assert(!copy _20, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _1) -> [success: bb22, unwind unreachable];
++         _80 = const false;
++         assert(!const false, "attempt to calculate the remainder of `{}` with a divisor of zero", copy _1) -> [success: bb22, unwind unreachable];
       }
 
       bb22: {