Refactor the cast-then-cast cases together, and support transmute-then-transmute

Author: scottmcm

compiler/rustc_mir_transform/src/gvn.rs

@@ -1366,110 +1366,112 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             return self.new_opaque();
         }
 
-        let mut was_updated = false;
-
-        // Transmuting `*const T` <=> `*mut T` is just a pointer cast,
-        // which we might be able to merge with other ones later.
-        if let Transmute = kind
-            && let ty::RawPtr(from_pointee, _) = from.kind()
-            && let ty::RawPtr(to_pointee, _) = to.kind()
-            && from_pointee == to_pointee
-        {
-            *kind = PtrToPtr;
-            was_updated = true;
-        }
-
-        // If a cast just casts away the metadata again, then we can get it by
-        // casting the original thin pointer passed to `from_raw_parts`
-        if let PtrToPtr = kind
-            && let Value::Aggregate(AggregateTy::RawPtr { data_pointer_ty, .. }, _, fields) =
-                self.get(value)
-            && let ty::RawPtr(to_pointee, _) = to.kind()
-            && to_pointee.is_sized(self.tcx, self.typing_env())
-        {
-            from = *data_pointer_ty;
-            value = fields[0];
-            was_updated = true;
-            if *data_pointer_ty == to {
-                return Some(fields[0]);
+        let mut was_ever_updated = false;
+        loop {
+            let mut was_updated_this_iteration = false;
+
+            // Transmuting `*const T` <=> `*mut T` is just a pointer cast,
+            // which we might be able to merge with other ones later.
+            if let Transmute = kind
+                && let ty::RawPtr(from_pointee, _) = from.kind()
+                && let ty::RawPtr(to_pointee, _) = to.kind()
+                && from_pointee == to_pointee
+            {
+                *kind = PtrToPtr;
+                was_updated_this_iteration = true;
             }
-        }
 
-        // PtrToPtr-then-PtrToPtr can skip the intermediate step
-        if let PtrToPtr = kind
-            && let Value::Cast { kind: inner_kind, value: inner_value, from: inner_from, to: _ } =
-                *self.get(value)
-            && let PtrToPtr = inner_kind
-        {
-            from = inner_from;
-            value = inner_value;
-            was_updated = true;
-            if inner_from == to {
-                return Some(inner_value);
+            // If a cast just casts away the metadata again, then we can get it by
+            // casting the original thin pointer passed to `from_raw_parts`
+            if let PtrToPtr = kind
+                && let Value::Aggregate(AggregateTy::RawPtr { data_pointer_ty, .. }, _, fields) =
+                    self.get(value)
+                && let ty::RawPtr(to_pointee, _) = to.kind()
+                && to_pointee.is_sized(self.tcx, self.typing_env())
+            {
+                from = *data_pointer_ty;
+                value = fields[0];
+                was_updated_this_iteration = true;
+                if *data_pointer_ty == to {
+                    return Some(fields[0]);
+                }
             }
-        }
 
-        // Aggregate-then-Transmute can just transmute the original field value,
-        // 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),
-                variant_idx,
-                field_values,
-            ) = self.get(value)
-            && 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)
-        {
-            from = field_ty;
-            value = field_values[field_idx.as_usize()];
-            was_updated = true;
-            if field_ty == to {
-                return Some(value);
+            // Aggregate-then-Transmute can just transmute the original field value,
+            // 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),
+                    variant_idx,
+                    field_values,
+                ) = self.get(value)
+                && 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)
+            {
+                from = field_ty;
+                value = field_values[field_idx.as_usize()];
+                was_updated_this_iteration = true;
+                if field_ty == to {
+                    return Some(value);
+                }
             }
-        }
 
-        // PtrToPtr-then-Transmute can just transmute the original, so long as the
-        // PtrToPtr didn't change metadata (and thus the size of the pointer)
-        if let Transmute = kind
-            && let Value::Cast {
-                kind: PtrToPtr,
+            // Various cast-then-cast cases can be simplified.
+            if let Value::Cast {
+                kind: inner_kind,
                 value: inner_value,
                 from: inner_from,
                 to: inner_to,
             } = *self.get(value)
-            && self.pointers_have_same_metadata(inner_from, inner_to)
-        {
-            from = inner_from;
-            value = inner_value;
-            was_updated = true;
-            if inner_from == to {
-                return Some(inner_value);
+            {
+                let new_kind = match (inner_kind, *kind) {
+                    // Even if there's a narrowing cast in here that's fine, because
+                    // things like `*mut [i32] -> *mut i32 -> *const i32` and
+                    // `*mut [i32] -> *const [i32] -> *const i32` can skip the middle in MIR.
+                    (PtrToPtr, PtrToPtr) => Some(PtrToPtr),
+                    // PtrToPtr-then-Transmute is fine so long as the pointer cast is identity:
+                    // `*const T -> *mut T -> NonNull<T>` is fine, but we need to check for narrowing
+                    // to skip things like `*const [i32] -> *const i32 -> NonNull<T>`.
+                    (PtrToPtr, Transmute)
+                        if self.pointers_have_same_metadata(inner_from, inner_to) =>
+                    {
+                        Some(Transmute)
+                    }
+                    // Similarly, for Transmute-then-PtrToPtr. Note that we need to check different
+                    // variables for their metadata, and thus this can't merge with the previous arm.
+                    (Transmute, PtrToPtr) if self.pointers_have_same_metadata(from, to) => {
+                        Some(Transmute)
+                    }
+                    // If would be legal to always do this, but we don't want to hide information
+                    // from the backend that it'd otherwise be able to use for optimizations.
+                    (Transmute, Transmute)
+                        if !self.type_may_have_niche_of_interest_to_backend(inner_to) =>
+                    {
+                        Some(Transmute)
+                    }
+                    _ => None,
+                };
+                if let Some(new_kind) = new_kind {
+                    *kind = new_kind;
+                    from = inner_from;
+                    value = inner_value;
+                    was_updated_this_iteration = true;
+                    if inner_from == to {
+                        return Some(inner_value);
+                    }
+                }
             }
-        }
 
-        // Transmute-then-PtrToPtr can just transmute the original, so long as the
-        // PtrToPtr won't change metadata (and thus the size of the pointer)
-        if let PtrToPtr = kind
-            && let Value::Cast {
-                kind: Transmute,
-                value: inner_value,
-                from: inner_from,
-                to: _inner_to,
-            } = *self.get(value)
-            && self.pointers_have_same_metadata(from, to)
-        {
-            *kind = Transmute;
-            from = inner_from;
-            value = inner_value;
-            was_updated = true;
-            if inner_from == to {
-                return Some(inner_value);
+            if was_updated_this_iteration {
+                was_ever_updated = true;
+            } else {
+                break;
             }
         }
 
-        if was_updated && let Some(op) = self.try_as_operand(value, location) {
+        if was_ever_updated && let Some(op) = self.try_as_operand(value, location) {
             *operand = op;
         }
 
@@ -1492,6 +1494,28 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         }
     }
 
+    /// Returns `false` if we know for sure that this type has no interesting niche,
+    /// and thus we can skip transmuting through it without worrying.
+    ///
+    /// The backend will emit `assume`s when transmuting between types with niches,
+    /// so we want to preserve `i32 -> char -> u32` so that that data is around,
+    /// but it's fine to skip whole-range-is-value steps like `A -> u32 -> B`.
+    fn type_may_have_niche_of_interest_to_backend(&self, ty: Ty<'tcx>) -> bool {
+        let Ok(layout) = self.ecx.layout_of(ty) else {
+            // If it's too generic or something, then assume it might be interesting later.
+            return true;
+        };
+
+        match layout.backend_repr {
+            BackendRepr::Uninhabited => true,
+            BackendRepr::Scalar(a) => !a.is_always_valid(&self.ecx),
+            BackendRepr::ScalarPair(a, b) => {
+                !a.is_always_valid(&self.ecx) || !b.is_always_valid(&self.ecx)
+            }
+            BackendRepr::Vector { .. } | BackendRepr::Memory { .. } => false,
+        }
+    }
+
     fn value_is_all_in_one_field(
         &self,
         ty: Ty<'tcx>,

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

@@ -17,15 +17,49 @@
       let mut _14: u16;
       let _15: ();
       let mut _16: u16;
-      let mut _17: std::result::Result<aggregate_struct_then_transmute::Never, u16>;
+      let mut _17: std::result::Result<Never, u16>;
+      let mut _19: u16;
+      let _20: ();
+      let mut _21: u32;
+      let mut _22: std::option::Option<u16>;
+      let mut _24: u16;
+      let _25: ();
+      let mut _26: i16;
+      let mut _27: MyId;
+      let mut _29: u16;
+      let mut _30: u16;
+      let _31: ();
+      let mut _32: u32;
+      let mut _33: aggregate_struct_then_transmute::Pair;
+      let mut _35: u16;
+      let mut _36: u16;
+      let _37: ();
+      let mut _38: u16;
+      let mut _39: aggregate_struct_then_transmute::Pair;
       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>;
+              let _13: std::result::Result<Never, u16>;
               scope 3 {
                   debug c => _13;
+                  let _18: std::option::Option<u16>;
+                  scope 4 {
+                      debug d => _18;
+                      let _23: MyId;
+                      scope 5 {
+                          debug e => _23;
+                          let _28: aggregate_struct_then_transmute::Pair;
+                          scope 6 {
+                              debug f => _28;
+                              let _34: aggregate_struct_then_transmute::Pair;
+                              scope 7 {
+                                  debug g => _34;
+                              }
+                          }
+                      }
+                  }
               }
           }
       }
@@ -101,12 +135,105 @@
       bb3: {
           StorageDead(_16);
           StorageDead(_15);
+-         StorageLive(_18);
++         nop;
+          StorageLive(_19);
+          _19 = copy _1;
+-         _18 = Option::<u16>::Some(move _19);
++         _18 = Option::<u16>::Some(copy _1);
+          StorageDead(_19);
+          StorageLive(_20);
+          StorageLive(_21);
+          StorageLive(_22);
+          _22 = copy _18;
+-         _21 = move _22 as u32 (Transmute);
++         _21 = copy _18 as u32 (Transmute);
+          StorageDead(_22);
+          _20 = opaque::<u32>(move _21) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+          StorageDead(_21);
+          StorageDead(_20);
+          StorageLive(_23);
+          StorageLive(_24);
+          _24 = copy _1;
+-         _23 = MyId(move _24);
++         _23 = copy _2;
+          StorageDead(_24);
+          StorageLive(_25);
+          StorageLive(_26);
+          StorageLive(_27);
+-         _27 = move _23;
+-         _26 = move _27 as i16 (Transmute);
++         _27 = copy _2;
++         _26 = copy _1 as i16 (Transmute);
+          StorageDead(_27);
+          _25 = opaque::<i16>(move _26) -> [return: bb5, unwind unreachable];
+      }
+  
+      bb5: {
+          StorageDead(_26);
+          StorageDead(_25);
+-         StorageLive(_28);
++         nop;
+          StorageLive(_29);
+          _29 = copy _1;
+          StorageLive(_30);
+          _30 = copy _1;
+-         _28 = Pair(move _29, move _30);
++         _28 = Pair(copy _1, copy _1);
+          StorageDead(_30);
+          StorageDead(_29);
+          StorageLive(_31);
+          StorageLive(_32);
+          StorageLive(_33);
+-         _33 = move _28;
+-         _32 = move _33 as u32 (Transmute);
++         _33 = copy _28;
++         _32 = copy _28 as u32 (Transmute);
+          StorageDead(_33);
+          _31 = opaque::<u32>(move _32) -> [return: bb6, unwind unreachable];
+      }
+  
+      bb6: {
+          StorageDead(_32);
+          StorageDead(_31);
+          StorageLive(_34);
+          StorageLive(_35);
+          _35 = copy _1;
+          StorageLive(_36);
+          _36 = copy _1;
+-         _34 = Pair(move _35, move _36);
++         _34 = copy _28;
+          StorageDead(_36);
+          StorageDead(_35);
+          StorageLive(_37);
+          StorageLive(_38);
+          StorageLive(_39);
+-         _39 = move _34;
+-         _38 = move _39 as u16 (Transmute);
++         _39 = copy _28;
++         _38 = copy _28 as u16 (Transmute);
+          StorageDead(_39);
+          _37 = opaque::<u16>(move _38) -> [return: bb7, unwind unreachable];
+      }
+  
+      bb7: {
+          StorageDead(_38);
+          StorageDead(_37);
           _0 = const ();
+          StorageDead(_34);
+-         StorageDead(_28);
++         nop;
+          StorageDead(_23);
+-         StorageDead(_18);
 -         StorageDead(_13);
 -         StorageDead(_7);
 -         StorageDead(_2);
 +         nop;
 +         nop;
++         nop;
 +         nop;
           return;
       }