Generate ValTrees in DataflowConstProp.

cjgillot · cjgillot · commit 4d0a7c13b54d · 2023-09-18T17:17:15.000Z
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -550,11 +550,105 @@ impl<'tcx, 'locals> Collector<'tcx, 'locals> {
         state: &State<FlatSet<Scalar>>,
         map: &Map,
     ) -> Option<ConstantKind<'tcx>> {
-        let FlatSet::Elem(Scalar::Int(value)) = state.get(place.as_ref(), &map) else {
-            return None;
-        };
         let ty = place.ty(self.local_decls, self.patch.tcx).ty;
-        Some(ConstantKind::Val(ConstValue::Scalar(value.into()), ty))
+        let place = map.find(place.as_ref())?;
+        if let FlatSet::Elem(Scalar::Int(value)) = state.get_idx(place, map) {
+            Some(ConstantKind::Val(ConstValue::Scalar(value.into()), ty))
+        } else {
+            let valtree = self.try_make_valtree(place, ty, state, map)?;
+            let constant = ty::Const::new_value(self.patch.tcx, valtree, ty);
+            Some(ConstantKind::Ty(constant))
+        }
+    }
+
+    fn try_make_valtree(
+        &self,
+        place: PlaceIndex,
+        ty: Ty<'tcx>,
+        state: &State<FlatSet<Scalar>>,
+        map: &Map,
+    ) -> Option<ty::ValTree<'tcx>> {
+        let tcx = self.patch.tcx;
+        match ty.kind() {
+            // ZSTs.
+            ty::FnDef(..) => Some(ty::ValTree::zst()),
+
+            // Scalars.
+            ty::Bool | ty::Int(_) | ty::Uint(_) | ty::Float(_) | ty::Char => {
+                if let FlatSet::Elem(Scalar::Int(value)) = state.get_idx(place, map) {
+                    Some(ty::ValTree::Leaf(value))
+                } else {
+                    None
+                }
+            }
+
+            // Unsupported for now.
+            ty::Array(_, _) => None,
+
+            ty::Tuple(elem_tys) => {
+                let branches = elem_tys
+                    .iter()
+                    .enumerate()
+                    .map(|(i, ty)| {
+                        let field = map.apply(place, TrackElem::Field(FieldIdx::from_usize(i)))?;
+                        self.try_make_valtree(field, ty, state, map)
+                    })
+                    .collect::<Option<Vec<_>>>()?;
+                Some(ty::ValTree::Branch(tcx.arena.alloc_from_iter(branches.into_iter())))
+            }
+
+            ty::Adt(def, args) => {
+                if def.is_union() {
+                    return None;
+                }
+
+                let (variant_idx, variant_def, variant_place) = if def.is_enum() {
+                    let discr = map.apply(place, TrackElem::Discriminant)?;
+                    let FlatSet::Elem(Scalar::Int(discr)) = state.get_idx(discr, map) else {
+                        return None;
+                    };
+                    let discr_bits = discr.assert_bits(discr.size());
+                    let (variant, _) =
+                        def.discriminants(tcx).find(|(_, var)| discr_bits == var.val)?;
+                    let variant_place = map.apply(place, TrackElem::Variant(variant))?;
+                    let variant_int = ty::ValTree::Leaf(variant.as_u32().into());
+                    (Some(variant_int), def.variant(variant), variant_place)
+                } else {
+                    (None, def.non_enum_variant(), place)
+                };
+
+                let branches = variant_def
+                    .fields
+                    .iter_enumerated()
+                    .map(|(i, field)| {
+                        let ty = field.ty(tcx, args);
+                        let field = map.apply(variant_place, TrackElem::Field(i))?;
+                        self.try_make_valtree(field, ty, state, map)
+                    })
+                    .collect::<Option<Vec<_>>>()?;
+                Some(ty::ValTree::Branch(
+                    tcx.arena.alloc_from_iter(variant_idx.into_iter().chain(branches)),
+                ))
+            }
+
+            // Do not attempt to support indirection in constants.
+            ty::Ref(..) | ty::RawPtr(..) | ty::FnPtr(..) | ty::Str | ty::Slice(_) => None,
+
+            ty::Never
+            | ty::Foreign(..)
+            | ty::Alias(..)
+            | ty::Param(_)
+            | ty::Bound(..)
+            | ty::Placeholder(..)
+            | ty::Closure(..)
+            | ty::Generator(..)
+            | ty::Dynamic(..) => None,
+
+            ty::Error(_)
+            | ty::Infer(..)
+            | ty::GeneratorWitness(..)
+            | ty::GeneratorWitnessMIR(..) => bug!(),
+        }
     }
 }
 
diff --git a/tests/mir-opt/const_debuginfo.main.ConstDebugInfo.diff b/tests/mir-opt/const_debuginfo.main.ConstDebugInfo.diff
@@ -41,7 +41,8 @@
 +                             debug ((f: (bool, bool, u32)).2: u32) => const 123_u32;
                               let _10: std::option::Option<u16>;
                               scope 7 {
-                                  debug o => _10;
+-                                 debug o => _10;
++                                 debug o => const Option::<u16>::Some(99);
                                   let _17: u32;
                                   let _18: u32;
                                   scope 8 {
@@ -81,7 +82,7 @@
           _15 = const false;
           _16 = const 123_u32;
           StorageLive(_10);
-          _10 = Option::<u16>::Some(const 99_u16);
+          _10 = const Option::<u16>::Some(99);
           _17 = const 32_u32;
           _18 = const 32_u32;
           StorageLive(_11);
diff --git a/tests/mir-opt/dataflow-const-prop/checked.main.DataflowConstProp.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/checked.main.DataflowConstProp.panic-abort.diff
@@ -43,7 +43,7 @@
 -         _6 = CheckedAdd(_4, _5);
 -         assert(!move (_6.1: bool), "attempt to compute `{} + {}`, which would overflow", move _4, move _5) -> [success: bb1, unwind unreachable];
 +         _5 = const 2_i32;
-+         _6 = CheckedAdd(const 1_i32, const 2_i32);
++         _6 = const (3, false);
 +         assert(!const false, "attempt to compute `{} + {}`, which would overflow", const 1_i32, const 2_i32) -> [success: bb1, unwind unreachable];
       }
   
@@ -60,7 +60,7 @@
 -         _10 = CheckedAdd(_9, const 1_i32);
 -         assert(!move (_10.1: bool), "attempt to compute `{} + {}`, which would overflow", move _9, const 1_i32) -> [success: bb2, unwind unreachable];
 +         _9 = const i32::MAX;
-+         _10 = CheckedAdd(const i32::MAX, const 1_i32);
++         _10 = const (i32::MIN, true);
 +         assert(!const true, "attempt to compute `{} + {}`, which would overflow", const i32::MAX, const 1_i32) -> [success: bb2, unwind unreachable];
       }
   
diff --git a/tests/mir-opt/dataflow-const-prop/checked.main.DataflowConstProp.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/checked.main.DataflowConstProp.panic-unwind.diff
@@ -43,7 +43,7 @@
 -         _6 = CheckedAdd(_4, _5);
 -         assert(!move (_6.1: bool), "attempt to compute `{} + {}`, which would overflow", move _4, move _5) -> [success: bb1, unwind continue];
 +         _5 = const 2_i32;
-+         _6 = CheckedAdd(const 1_i32, const 2_i32);
++         _6 = const (3, false);
 +         assert(!const false, "attempt to compute `{} + {}`, which would overflow", const 1_i32, const 2_i32) -> [success: bb1, unwind continue];
       }
   
@@ -60,7 +60,7 @@
 -         _10 = CheckedAdd(_9, const 1_i32);
 -         assert(!move (_10.1: bool), "attempt to compute `{} + {}`, which would overflow", move _9, const 1_i32) -> [success: bb2, unwind continue];
 +         _9 = const i32::MAX;
-+         _10 = CheckedAdd(const i32::MAX, const 1_i32);
++         _10 = const (i32::MIN, true);
 +         assert(!const true, "attempt to compute `{} + {}`, which would overflow", const i32::MAX, const 1_i32) -> [success: bb2, unwind continue];
       }
   
diff --git a/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.32bit.diff b/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.32bit.diff
@@ -23,7 +23,8 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = E::V1(const 0_i32);
+-         _1 = E::V1(const 0_i32);
++         _1 = const E::V1(0);
           StorageLive(_2);
 -         _3 = discriminant(_1);
 -         switchInt(move _3) -> [0: bb3, 1: bb1, otherwise: bb2];
diff --git a/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.64bit.diff b/tests/mir-opt/dataflow-const-prop/enum.simple.DataflowConstProp.64bit.diff
@@ -23,7 +23,8 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = E::V1(const 0_i32);
+-         _1 = E::V1(const 0_i32);
++         _1 = const E::V1(0);
           StorageLive(_2);
 -         _3 = discriminant(_1);
 -         switchInt(move _3) -> [0: bb3, 1: bb1, otherwise: bb2];
diff --git a/tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.32bit.diff b/tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.32bit.diff
@@ -44,7 +44,8 @@
           StorageLive(_1);
           StorageLive(_2);
           _2 = const {alloc1: &E};
-          _1 = (*_2);
+-         _1 = (*_2);
++         _1 = const E::V1(0);
           StorageDead(_2);
           StorageLive(_3);
 -         _4 = discriminant(_1);
diff --git a/tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.64bit.diff b/tests/mir-opt/dataflow-const-prop/enum.statics.DataflowConstProp.64bit.diff
@@ -44,7 +44,8 @@
           StorageLive(_1);
           StorageLive(_2);
           _2 = const {alloc1: &E};
-          _1 = (*_2);
+-         _1 = (*_2);
++         _1 = const E::V1(0);
           StorageDead(_2);
           StorageLive(_3);
 -         _4 = discriminant(_1);
diff --git a/tests/mir-opt/dataflow-const-prop/inherit_overflow.main.DataflowConstProp.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/inherit_overflow.main.DataflowConstProp.panic-abort.diff
@@ -22,7 +22,7 @@
           _3 = const 1_u8;
 -         _4 = CheckedAdd(_2, _3);
 -         assert(!move (_4.1: bool), "attempt to compute `{} + {}`, which would overflow", _2, _3) -> [success: bb1, unwind unreachable];
-+         _4 = CheckedAdd(const u8::MAX, const 1_u8);
++         _4 = const (0, true);
 +         assert(!const true, "attempt to compute `{} + {}`, which would overflow", const u8::MAX, const 1_u8) -> [success: bb1, unwind unreachable];
       }
   
diff --git a/tests/mir-opt/dataflow-const-prop/inherit_overflow.main.DataflowConstProp.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/inherit_overflow.main.DataflowConstProp.panic-unwind.diff
@@ -22,7 +22,7 @@
           _3 = const 1_u8;
 -         _4 = CheckedAdd(_2, _3);
 -         assert(!move (_4.1: bool), "attempt to compute `{} + {}`, which would overflow", _2, _3) -> [success: bb1, unwind continue];
-+         _4 = CheckedAdd(const u8::MAX, const 1_u8);
++         _4 = const (0, true);
 +         assert(!const true, "attempt to compute `{} + {}`, which would overflow", const u8::MAX, const 1_u8) -> [success: bb1, unwind continue];
       }
   
diff --git a/tests/mir-opt/dataflow-const-prop/repr_transparent.main.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/repr_transparent.main.DataflowConstProp.diff
@@ -17,7 +17,8 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = I32(const 0_i32);
+-         _1 = I32(const 0_i32);
++         _1 = const I32(0);
           StorageLive(_2);
           StorageLive(_3);
           StorageLive(_4);
@@ -31,7 +32,7 @@
           StorageDead(_5);
           StorageDead(_4);
 -         _2 = I32(move _3);
-+         _2 = I32(const 0_i32);
++         _2 = const I32(0);
           StorageDead(_3);
           _0 = const ();
           StorageDead(_2);
diff --git a/tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.32bit.diff b/tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.32bit.diff
@@ -9,11 +9,15 @@
       let mut _6: i32;
       let mut _11: BigStruct;
       let mut _16: &&BigStruct;
-      let mut _17: &BigStruct;
-      let mut _18: &BigStruct;
-      let mut _19: &BigStruct;
-      let mut _20: &BigStruct;
-      let mut _21: &BigStruct;
+      let mut _18: S;
+      let mut _19: u8;
+      let mut _20: f32;
+      let mut _21: S;
+      let mut _22: &BigStruct;
+      let mut _23: &BigStruct;
+      let mut _24: &BigStruct;
+      let mut _25: &BigStruct;
+      let mut _26: &BigStruct;
       scope 1 {
           debug s => _1;
           let _2: i32;
@@ -40,6 +44,10 @@
                           debug b => _13;
                           debug c => _14;
                           debug d => _15;
+                          let _17: BigStruct;
+                          scope 6 {
+                              debug bs => _17;
+                          }
                       }
                   }
               }
@@ -48,7 +56,8 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = S(const 1_i32);
+-         _1 = S(const 1_i32);
++         _1 = const S(1);
           StorageLive(_2);
           StorageLive(_3);
 -         _3 = (_1.0: i32);
@@ -85,25 +94,45 @@
           StorageDead(_11);
           StorageLive(_16);
           _16 = const {alloc1: &&BigStruct};
-          _17 = deref_copy (*_16);
+          _22 = deref_copy (*_16);
           StorageLive(_12);
-          _18 = deref_copy (*_16);
--         _12 = ((*_18).0: S);
+          _23 = deref_copy (*_16);
+-         _12 = ((*_23).0: S);
 +         _12 = const S(1_i32);
           StorageLive(_13);
-          _19 = deref_copy (*_16);
--         _13 = ((*_19).1: u8);
+          _24 = deref_copy (*_16);
+-         _13 = ((*_24).1: u8);
 +         _13 = const 5_u8;
           StorageLive(_14);
-          _20 = deref_copy (*_16);
--         _14 = ((*_20).2: f32);
+          _25 = deref_copy (*_16);
+-         _14 = ((*_25).2: f32);
 +         _14 = const 7f32;
           StorageLive(_15);
-          _21 = deref_copy (*_16);
--         _15 = ((*_21).3: S);
+          _26 = deref_copy (*_16);
+-         _15 = ((*_26).3: S);
 +         _15 = const S(13_i32);
           StorageDead(_16);
+          StorageLive(_17);
+          StorageLive(_18);
+-         _18 = _12;
++         _18 = const S(1_i32);
+          StorageLive(_19);
+-         _19 = _13;
++         _19 = const 5_u8;
+          StorageLive(_20);
+-         _20 = _14;
++         _20 = const 7f32;
+          StorageLive(_21);
+-         _21 = _15;
+-         _17 = BigStruct(move _18, move _19, move _20, move _21);
++         _21 = const S(13_i32);
++         _17 = const BigStruct(S(1), 5, 7f32, S(13));
+          StorageDead(_21);
+          StorageDead(_20);
+          StorageDead(_19);
+          StorageDead(_18);
           _0 = const ();
+          StorageDead(_17);
           StorageDead(_15);
           StorageDead(_14);
           StorageDead(_13);
diff --git a/tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.64bit.diff b/tests/mir-opt/dataflow-const-prop/struct.main.DataflowConstProp.64bit.diff
diff --git a/tests/mir-opt/dataflow-const-prop/struct.rs b/tests/mir-opt/dataflow-const-prop/struct.rs
diff --git a/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.diff b/tests/mir-opt/dataflow-const-prop/tuple.main.DataflowConstProp.diff
diff --git a/tests/mir-opt/dataflow-const-prop/tuple.rs b/tests/mir-opt/dataflow-const-prop/tuple.rs

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@`
`22`	`22`	`_3 = const 1_u8;`
`23`	`23`	`- _4 = CheckedAdd(_2, _3);`
`24`	`24`	- assert(!move (_4.1: bool), "attempt to compute `{} + {}`, which would overflow", _2, _3) -> [success: bb1, unwind unreachable];
`25`		`-+ _4 = CheckedAdd(const u8::MAX, const 1_u8);`
	`25`	`++ _4 = const (0, true);`
`26`	`26`	+ assert(!const true, "attempt to compute `{} + {}`, which would overflow", const u8::MAX, const 1_u8) -> [success: bb1, unwind unreachable];
`27`	`27`	`}`
`28`	`28`