GVN: Support unions.

cjgillot · cjgillot · commit 37cd90df6e7a · 2025-09-25T23:52:48.000Z
diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
@@ -202,6 +202,8 @@ enum Value<'tcx> {
     /// An aggregate value, either tuple/closure/struct/enum.
     /// This does not contain unions, as we cannot reason with the value.
     Aggregate(VariantIdx, Vec<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.
@@ -563,6 +565,21 @@ impl<'body, 'tcx> VnState<'body, '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()?;
@@ -735,6 +752,10 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             ProjectionElem::Field(f, _) => {
                 if let Value::Aggregate(_, fields) = self.get(value) {
                     return Some((projection_ty, fields[f.as_usize()]));
+                } else if let Value::Union(active, field) = *self.get(value)
+                    && active == f
+                {
+                    return Some((projection_ty, field));
                 } else if let Value::Projection(outer_value, ProjectionElem::Downcast(_, read_variant)) = self.get(value)
                     && 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
@@ -1106,7 +1127,10 @@ impl<'body, 'tcx> VnState<'body, '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();
diff --git a/tests/mir-opt/const_prop/invalid_constant.main.GVN.diff b/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                                     │ ....
   }
   
diff --git a/tests/mir-opt/const_prop/transmute.rs b/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: (),
diff --git a/tests/mir-opt/const_prop/transmute.undef_union_as_integer.GVN.32bit.diff b/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) {
++     __ __ __ __                                     │ ░░░░
   }
   
diff --git a/tests/mir-opt/const_prop/transmute.undef_union_as_integer.GVN.64bit.diff b/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) {
++     __ __ __ __                                     │ ░░░░
   }
   
diff --git a/tests/mir-opt/dest-prop/union.main.DestinationPropagation.panic-abort.diff b/tests/mir-opt/dest-prop/union.main.DestinationPropagation.panic-abort.diff
@@ -6,22 +6,26 @@
       let _1: main::Un;
       let mut _2: u32;
       scope 1 {
-          debug un => _1;
+          debug un => const Un {{ us: 1_u32 }};
           scope 3 (inlined std::mem::drop::<u32>) {
-              debug _x => _2;
+              debug _x => const 1_u32;
           }
       }
       scope 2 (inlined val) {
       }
   
       bb0: {
           StorageLive(_1);
-          _1 = Un { us: const 1_u32 };
+          nop;
           StorageLive(_2);
-          _2 = copy (_1.0: u32);
+          nop;
           StorageDead(_2);
           StorageDead(_1);
           return;
       }
   }
   
+  ALLOC0 (size: 4, align: 4) {
+      01 00 00 00                                     │ ....
+  }
+  
diff --git a/tests/mir-opt/dest-prop/union.main.DestinationPropagation.panic-unwind.diff b/tests/mir-opt/dest-prop/union.main.DestinationPropagation.panic-unwind.diff
@@ -6,22 +6,26 @@
       let _1: main::Un;
       let mut _2: u32;
       scope 1 {
-          debug un => _1;
+          debug un => const Un {{ us: 1_u32 }};
           scope 3 (inlined std::mem::drop::<u32>) {
-              debug _x => _2;
+              debug _x => const 1_u32;
           }
       }
       scope 2 (inlined val) {
       }
   
       bb0: {
           StorageLive(_1);
-          _1 = Un { us: const 1_u32 };
+          nop;
           StorageLive(_2);
-          _2 = copy (_1.0: u32);
+          nop;
           StorageDead(_2);
           StorageDead(_1);
           return;
       }
   }
   
+  ALLOC0 (size: 4, align: 4) {
+      01 00 00 00                                     │ ....
+  }
+  
diff --git a/tests/mir-opt/dest-prop/union.rs b/tests/mir-opt/dest-prop/union.rs
@@ -8,7 +8,10 @@ fn val() -> u32 {
 // EMIT_MIR union.main.DestinationPropagation.diff
 fn main() {
     // CHECK-LABEL: fn main(
-    // CHECK: {{_.*}} = Un { us: const 1_u32 };
+    // CHECK: debug un => const Un
+    // CHECK: debug _x => const 1_u32;
+    // CHECK: bb0: {
+    // CHECK-NEXT: return;
     union Un {
         us: u32,
     }

Original file line number	Diff line number	Diff line change
`@@ -6,22 +6,26 @@`
`6`	`6`	`let _1: main::Un;`
`7`	`7`	`let mut _2: u32;`
`8`	`8`	`scope 1 {`
`9`		`- debug un => _1;`
	`9`	`+ debug un => const Un {{ us: 1_u32 }};`
`10`	`10`	`scope 3 (inlined std::mem::drop::<u32>) {`
`11`		`- debug _x => _2;`
	`11`	`+ debug _x => const 1_u32;`
`12`	`12`	`}`
`13`	`13`	`}`
`14`	`14`	`scope 2 (inlined val) {`
`15`	`15`	`}`
`16`	`16`
`17`	`17`	`bb0: {`
`18`	`18`	`StorageLive(_1);`
`19`		`- _1 = Un { us: const 1_u32 };`
	`19`	`+ nop;`
`20`	`20`	`StorageLive(_2);`
`21`		`- _2 = copy (_1.0: u32);`
	`21`	`+ nop;`
`22`	`22`	`StorageDead(_2);`
`23`	`23`	`StorageDead(_1);`
`24`	`24`	`return;`
`25`	`25`	`}`
`26`	`26`	`}`
`27`	`27`
	`28`	`+ ALLOC0 (size: 4, align: 4) {`
	`29`	`+ 01 00 00 00 │ ....`
	`30`	`+ }`
	`31`	`+`