GVN: Support unions.

compiler/rustc_mir_transform/src/gvn.rs

@@ -213,6 +213,8 @@ enum Value<'a, 'tcx> {
     /// An aggregate value, either tuple/closure/struct/enum.
     /// This does not contain unions, as we cannot reason with the value.
     Aggregate(VariantIdx, &'a [VnIndex]),
+    /// A union aggregate value.
+    Union(FieldIdx, VnIndex),
     /// A raw pointer aggregate built from a thin pointer and metadata.
     RawPtr {
         /// Thin pointer component. This is field 0 in MIR.
@@ -600,6 +602,21 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                     return None;
                 }
             }
+            Union(active_field, field) => {
+                let field = self.evaluated[field].as_ref()?;
+                if matches!(ty.backend_repr, BackendRepr::Scalar(..) | BackendRepr::ScalarPair(..))
+                {
+                    let dest = self.ecx.allocate(ty, MemoryKind::Stack).discard_err()?;
+                    let field_dest = self.ecx.project_field(&dest, active_field).discard_err()?;
+                    self.ecx.copy_op(field, &field_dest).discard_err()?;
+                    self.ecx
+                        .alloc_mark_immutable(dest.ptr().provenance.unwrap().alloc_id())
+                        .discard_err()?;
+                    dest.into()
+                } else {
+                    return None;
+                }
+            }
             RawPtr { pointer, metadata } => {
                 let pointer = self.evaluated[pointer].as_ref()?;
                 let metadata = self.evaluated[metadata].as_ref()?;
@@ -802,11 +819,11 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                 }
             }
             ProjectionElem::Downcast(name, index) => ProjectionElem::Downcast(name, index),
-            ProjectionElem::Field(f, _) => {
-                if let Value::Aggregate(_, fields) = self.get(value) {
-                    return Some((projection_ty, fields[f.as_usize()]));
-                } else if let Value::Projection(outer_value, ProjectionElem::Downcast(_, read_variant)) = self.get(value)
-                    && let Value::Aggregate(written_variant, fields) = self.get(outer_value)
+            ProjectionElem::Field(f, _) => match self.get(value) {
+                Value::Aggregate(_, fields) => return Some((projection_ty, fields[f.as_usize()])),
+                Value::Union(active, field) if active == f => return Some((projection_ty, field)),
+                Value::Projection(outer_value, ProjectionElem::Downcast(_, read_variant))
+                    if let Value::Aggregate(written_variant, fields) = self.get(outer_value)
                     // This pass is not aware of control-flow, so we do not know whether the
                     // replacement we are doing is actually reachable. We could be in any arm of
                     // ```
@@ -822,12 +839,12 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
                     // accessing the wrong variant is not UB if the enum has repr.
                     // So it's not impossible for a series of MIR opts to generate
                     // a downcast to an inactive variant.
-                    && written_variant == read_variant
+                    && written_variant == read_variant =>
                 {
                     return Some((projection_ty, fields[f.as_usize()]));
                 }
-                ProjectionElem::Field(f, ())
-            }
+                _ => ProjectionElem::Field(f, ()),
+            },
             ProjectionElem::Index(idx) => {
                 if let Value::Repeat(inner, _) = self.get(value) {
                     return Some((projection_ty, inner));
@@ -1167,7 +1184,10 @@ impl<'body, 'a, 'tcx> VnState<'body, 'a, 'tcx> {
             | AggregateKind::Coroutine(..) => FIRST_VARIANT,
             AggregateKind::Adt(_, variant_index, _, _, None) => variant_index,
             // Do not track unions.
-            AggregateKind::Adt(_, _, _, _, Some(_)) => return None,
+            AggregateKind::Adt(_, _, _, _, Some(active_field)) => {
+                let field = *fields.first()?;
+                return Some(self.insert(ty, Value::Union(active_field, field)));
+            }
             AggregateKind::RawPtr(..) => {
                 assert_eq!(field_ops.len(), 2);
                 let [mut pointer, metadata] = fields.try_into().unwrap();

tests/mir-opt/const_prop/invalid_constant.main.GVN.diff

@@ -28,21 +28,27 @@
       bb0: {
           StorageLive(_1);
           StorageLive(_2);
-          _2 = InvalidChar { int: const 1114113_u32 };
-          _1 = copy (_2.1: char);
+-         _2 = InvalidChar { int: const 1114113_u32 };
+-         _1 = copy (_2.1: char);
++         _2 = const InvalidChar {{ int: 1114113_u32, chr: {transmute(0x00110001): char} }};
++         _1 = const {transmute(0x00110001): char};
           StorageDead(_2);
           StorageLive(_3);
           StorageLive(_4);
           StorageLive(_5);
-          _5 = InvalidTag { int: const 4_u32 };
-          _4 = copy (_5.1: E);
-          _3 = [move _4];
+-         _5 = InvalidTag { int: const 4_u32 };
+-         _4 = copy (_5.1: E);
+-         _3 = [move _4];
++         _5 = const InvalidTag {{ int: 4_u32, e: Scalar(0x00000004): E }};
++         _4 = const Scalar(0x00000004): E;
++         _3 = [const Scalar(0x00000004): E];
           StorageDead(_4);
           StorageDead(_5);
           nop;
           nop;
           StorageLive(_8);
-          _8 = NoVariants { int: const 0_u32 };
+-         _8 = NoVariants { int: const 0_u32 };
++         _8 = const NoVariants {{ int: 0_u32, empty: ZeroSized: Empty }};
           nop;
           nop;
           nop;
@@ -55,5 +61,17 @@
           StorageDead(_1);
           return;
       }
++ }
++ 
++ ALLOC0 (size: 4, align: 4) {
++     00 00 00 00                                     │ ....
++ }
++ 
++ ALLOC1 (size: 4, align: 4) {
++     04 00 00 00                                     │ ....
++ }
++ 
++ ALLOC2 (size: 4, align: 4) {
++     01 00 11 00                                     │ ....
   }
 

tests/mir-opt/const_prop/transmute.rs

@@ -43,8 +43,8 @@ pub unsafe fn invalid_bool() -> bool {
 // EMIT_MIR transmute.undef_union_as_integer.GVN.diff
 pub unsafe fn undef_union_as_integer() -> u32 {
     // CHECK-LABEL: fn undef_union_as_integer(
-    // CHECK: _1 = Union32 {
-    // CHECK: _0 = move _1 as u32 (Transmute);
+    // CHECK: _1 = const Union32
+    // CHECK: _0 = const {{.*}}: u32;
     union Union32 {
         value: u32,
         unit: (),

tests/mir-opt/const_prop/transmute.undef_union_as_integer.GVN.32bit.diff

@@ -12,11 +12,16 @@
 -         _2 = ();
 -         _1 = Union32 { value: move _2 };
 +         _2 = const ();
-+         _1 = Union32 { value: const () };
++         _1 = const Union32 {{ value: Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: u32, unit: () }};
           StorageDead(_2);
-          _0 = move _1 as u32 (Transmute);
+-         _0 = move _1 as u32 (Transmute);
++         _0 = const Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: u32;
           StorageDead(_1);
           return;
       }
++ }
++ 
++ ALLOC0 (size: 4, align: 4) {
++     __ __ __ __                                     │ ░░░░
   }
 

tests/mir-opt/const_prop/transmute.undef_union_as_integer.GVN.64bit.diff

@@ -12,11 +12,16 @@
 -         _2 = ();
 -         _1 = Union32 { value: move _2 };
 +         _2 = const ();
-+         _1 = Union32 { value: const () };
++         _1 = const Union32 {{ value: Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: u32, unit: () }};
           StorageDead(_2);
-          _0 = move _1 as u32 (Transmute);
+-         _0 = move _1 as u32 (Transmute);
++         _0 = const Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: u32;
           StorageDead(_1);
           return;
       }
++ }
++ 
++ ALLOC0 (size: 4, align: 4) {
++     __ __ __ __                                     │ ░░░░
   }
 

tests/mir-opt/const_prop/union.main.GVN.diff

@@ -0,0 +1,42 @@
+- // MIR for `main` before GVN
++ // MIR for `main` after GVN
+
+  fn main() -> () {
+      let mut _0: ();
+      let _1: main::Un;
+      let mut _2: u32;
+      let _3: ();
+      let mut _4: u32;
+      scope 1 {
+          debug un => _1;
+          scope 3 (inlined std::mem::drop::<u32>) {
+          }
+      }
+      scope 2 (inlined val) {
+      }
+
+      bb0: {
+          StorageLive(_1);
+-         StorageLive(_2);
++         nop;
+          _2 = const 1_u32;
+-         _1 = Un { us: move _2 };
+-         StorageDead(_2);
++         _1 = const Un {{ us: 1_u32 }};
++         nop;
+          StorageLive(_3);
+          StorageLive(_4);
+-         _4 = copy (_1.0: u32);
++         _4 = const 1_u32;
+          StorageDead(_4);
+          StorageDead(_3);
+          _0 = const ();
+          StorageDead(_1);
+          return;
+      }
++ }
++ 
++ ALLOC0 (size: 4, align: 4) {
++     01 00 00 00                                     │ ....
+  }
+

tests/mir-opt/dest-prop/union.rs → tests/mir-opt/const_prop/union.rs

@@ -1,14 +1,17 @@
-// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
 //! Tests that we can propagate into places that are projections into unions
-//@ compile-flags: -Zunsound-mir-opts -C debuginfo=full
+//@ test-mir-pass: GVN
+//@ compile-flags: -Zinline-mir
+
 fn val() -> u32 {
     1
 }
 
-// EMIT_MIR union.main.DestinationPropagation.diff
+// EMIT_MIR union.main.GVN.diff
 fn main() {
     // CHECK-LABEL: fn main(
-    // CHECK: {{_.*}} = Un { us: const 1_u32 };
+    // CHECK: debug un => [[un:_.*]];
+    // CHECK: bb0: {
+    // CHECK: [[un]] = const Un {{{{ us: 1_u32 }}}};
     union Un {
         us: u32,
     }