Use layout information to detect transparent transmutes

scottmcm · scottmcm · commit 9574787559fe · 2024-12-08T11:48:44.000-08:00
diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
@@ -1420,24 +1420,23 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         }
 
         // Aggregate-then-Transmute can just transmute the original field value,
-        // so long as the type is transparent over only that one single field.
+        // so long as the bytes of a value from only from a single field.
         if let Transmute = kind
             && let Value::Aggregate(
                 AggregateTy::Def(aggregate_did, aggregate_args),
-                FIRST_VARIANT,
+                variant_idx,
                 field_values,
             ) = self.get(value)
-            && let [single_field_value] = **field_values
-            && let adt = self.tcx.adt_def(aggregate_did)
-            && adt.is_struct()
-            && adt.repr().transparent()
+            && let aggregate_ty =
+                self.tcx.type_of(aggregate_did).instantiate(self.tcx, aggregate_args)
+            && let Some((field_idx, field_ty)) =
+                self.value_is_all_in_one_field(aggregate_ty, *variant_idx)
         {
-            let field_ty = adt.non_enum_variant().single_field().ty(self.tcx, aggregate_args);
             from = field_ty;
-            value = single_field_value;
+            value = field_values[field_idx.as_usize()];
             was_updated = true;
             if field_ty == to {
-                return Some(single_field_value);
+                return Some(value);
             }
         }
 
@@ -1543,6 +1542,32 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             false
         }
     }
+
+    fn value_is_all_in_one_field(
+        &self,
+        ty: Ty<'tcx>,
+        variant: VariantIdx,
+    ) -> Option<(FieldIdx, Ty<'tcx>)> {
+        if let Ok(layout) = self.ecx.layout_of(ty)
+            && let abi::Variants::Single { index } = layout.variants
+            && index == variant
+            && let Some((field_idx, field_layout)) = layout.non_1zst_field(&self.ecx)
+            && layout.size == field_layout.size
+        {
+            // We needed to check the variant to avoid trying to read the tag
+            // field from an enum where no fields have variants, since that tag
+            // field isn't in the `Aggregate` from which we're getting values.
+            Some((FieldIdx::from_usize(field_idx), field_layout.ty))
+        } else if let ty::Adt(adt, args) = ty.kind()
+            && adt.is_struct()
+            && adt.repr().transparent()
+            && let [single_field] = adt.non_enum_variant().fields.raw.as_slice()
+        {
+            Some((FieldIdx::ZERO, single_field.ty(self.tcx, args)))
+        } else {
+            None
+        }
+    }
 }
 
 fn op_to_prop_const<'tcx>(
diff --git a/tests/mir-opt/gvn.aggregate_struct_then_transmute.GVN.panic-abort.diff b/tests/mir-opt/gvn.aggregate_struct_then_transmute.GVN.panic-abort.diff
@@ -14,11 +14,19 @@
       let _10: ();
       let mut _11: u16;
       let mut _12: TypedId<std::string::String>;
+      let mut _14: u16;
+      let _15: ();
+      let mut _16: u16;
+      let mut _17: std::result::Result<aggregate_struct_then_transmute::Never, u16>;
       scope 1 {
           debug a => _2;
           let _7: TypedId<std::string::String>;
           scope 2 {
               debug b => _7;
+              let _13: std::result::Result<aggregate_struct_then_transmute::Never, u16>;
+              scope 3 {
+                  debug c => _13;
+              }
           }
       }
   
@@ -62,18 +70,43 @@
 -         _12 = move _7;
 -         _11 = move _12 as u16 (Transmute);
 +         _12 = copy _7;
-+         _11 = copy _7 as u16 (Transmute);
++         _11 = copy _1;
           StorageDead(_12);
-          _10 = opaque::<u16>(move _11) -> [return: bb2, unwind unreachable];
+-         _10 = opaque::<u16>(move _11) -> [return: bb2, unwind unreachable];
++         _10 = opaque::<u16>(copy _1) -> [return: bb2, unwind unreachable];
       }
   
       bb2: {
           StorageDead(_11);
           StorageDead(_10);
+-         StorageLive(_13);
++         nop;
+          StorageLive(_14);
+          _14 = copy _1;
+-         _13 = Result::<Never, u16>::Err(move _14);
++         _13 = Result::<Never, u16>::Err(copy _1);
+          StorageDead(_14);
+          StorageLive(_15);
+          StorageLive(_16);
+          StorageLive(_17);
+-         _17 = move _13;
+-         _16 = move _17 as u16 (Transmute);
++         _17 = copy _13;
++         _16 = copy _1;
+          StorageDead(_17);
+-         _15 = opaque::<u16>(move _16) -> [return: bb3, unwind unreachable];
++         _15 = opaque::<u16>(copy _1) -> [return: bb3, unwind unreachable];
+      }
+  
+      bb3: {
+          StorageDead(_16);
+          StorageDead(_15);
           _0 = const ();
+-         StorageDead(_13);
 -         StorageDead(_7);
 -         StorageDead(_2);
 +         nop;
++         nop;
 +         nop;
           return;
       }
diff --git a/tests/mir-opt/gvn.aggregate_struct_then_transmute.GVN.panic-unwind.diff b/tests/mir-opt/gvn.aggregate_struct_then_transmute.GVN.panic-unwind.diff
@@ -14,11 +14,19 @@
       let _10: ();
       let mut _11: u16;
       let mut _12: TypedId<std::string::String>;
+      let mut _14: u16;
+      let _15: ();
+      let mut _16: u16;
+      let mut _17: std::result::Result<aggregate_struct_then_transmute::Never, u16>;
       scope 1 {
           debug a => _2;
           let _7: TypedId<std::string::String>;
           scope 2 {
               debug b => _7;
+              let _13: std::result::Result<aggregate_struct_then_transmute::Never, u16>;
+              scope 3 {
+                  debug c => _13;
+              }
           }
       }
   
@@ -62,18 +70,43 @@
 -         _12 = move _7;
 -         _11 = move _12 as u16 (Transmute);
 +         _12 = copy _7;
-+         _11 = copy _7 as u16 (Transmute);
++         _11 = copy _1;
           StorageDead(_12);
-          _10 = opaque::<u16>(move _11) -> [return: bb2, unwind continue];
+-         _10 = opaque::<u16>(move _11) -> [return: bb2, unwind continue];
++         _10 = opaque::<u16>(copy _1) -> [return: bb2, unwind continue];
       }
   
       bb2: {
           StorageDead(_11);
           StorageDead(_10);
+-         StorageLive(_13);
++         nop;
+          StorageLive(_14);
+          _14 = copy _1;
+-         _13 = Result::<Never, u16>::Err(move _14);
++         _13 = Result::<Never, u16>::Err(copy _1);
+          StorageDead(_14);
+          StorageLive(_15);
+          StorageLive(_16);
+          StorageLive(_17);
+-         _17 = move _13;
+-         _16 = move _17 as u16 (Transmute);
++         _17 = copy _13;
++         _16 = copy _1;
+          StorageDead(_17);
+-         _15 = opaque::<u16>(move _16) -> [return: bb3, unwind continue];
++         _15 = opaque::<u16>(copy _1) -> [return: bb3, unwind continue];
+      }
+  
+      bb3: {
+          StorageDead(_16);
+          StorageDead(_15);
           _0 = const ();
+-         StorageDead(_13);
 -         StorageDead(_7);
 -         StorageDead(_2);
 +         nop;
++         nop;
 +         nop;
           return;
       }
diff --git a/tests/mir-opt/gvn.rs b/tests/mir-opt/gvn.rs
@@ -919,13 +919,15 @@ unsafe fn aggregate_struct_then_transmute(id: u16) {
     let a = MyId(id);
     opaque(std::intrinsics::transmute::<_, u16>(a));
 
-    // GVN can't do this yet because it doesn't know which field is the ZST,
-    // but future changes might enable it.
-    // CHECK: [[AGG:_.+]] = TypedId::<String>(copy _1, const PhantomData::<String>);
-    // CHECK: [[INT:_.+]] = copy [[AGG]] as u16 (Transmute);
-    // CHECK: opaque::<u16>(move [[INT]])
+    // CHECK: opaque::<u16>(copy _1)
     let b = TypedId::<String>(id, PhantomData);
     opaque(std::intrinsics::transmute::<_, u16>(b));
+
+    // CHECK: opaque::<u16>(copy _1)
+    let c = Err::<Never, u16>(id);
+    opaque(std::intrinsics::transmute::<_, u16>(c));
+
+    enum Never {}
 }
 
 // Transmuting can skip a pointer cast so long as it wasn't a fat-to-thin cast.
