[mir-opt] Allow SRoA when an object is transmuted to a field type

Author: scottmcm

compiler/rustc_mir_transform/src/sroa.rs

@@ -32,11 +32,11 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
         let typing_env = body.typing_env(tcx);
         loop {
             debug!(?excluded);
-            let escaping = escaping_locals(tcx, &excluded, body);
+            let escaping = escaping_locals(tcx, &excluded, typing_env, body);
             debug!(?escaping);
             let replacements = compute_flattening(tcx, typing_env, body, escaping);
             debug!(?replacements);
-            let all_dead_locals = replace_flattened_locals(tcx, body, replacements);
+            let all_dead_locals = replace_flattened_locals(tcx, typing_env, body, replacements);
             if !all_dead_locals.is_empty() {
                 excluded.union(&all_dead_locals);
                 excluded = {
@@ -65,6 +65,7 @@ impl<'tcx> crate::MirPass<'tcx> for ScalarReplacementOfAggregates {
 fn escaping_locals<'tcx>(
     tcx: TyCtxt<'tcx>,
     excluded: &DenseBitSet<Local>,
+    typing_env: ty::TypingEnv<'tcx>,
     body: &Body<'tcx>,
 ) -> DenseBitSet<Local> {
     let is_excluded_ty = |ty: Ty<'tcx>| {
@@ -88,20 +89,24 @@ fn escaping_locals<'tcx>(
 
     let mut set = DenseBitSet::new_empty(body.local_decls.len());
     set.insert_range(RETURN_PLACE..Local::arg(body.arg_count));
-    for (local, decl) in body.local_decls().iter_enumerated() {
+    for (local, decl) in body.local_decls.iter_enumerated() {
         if excluded.contains(local) || is_excluded_ty(decl.ty) {
             set.insert(local);
         }
     }
-    let mut visitor = EscapeVisitor { set };
+    let mut visitor = EscapeVisitor { tcx, typing_env, set, decls: &body.local_decls };
     visitor.visit_body(body);
     return visitor.set;
 
-    struct EscapeVisitor {
+    struct EscapeVisitor<'tcx, 'a> {
+        tcx: TyCtxt<'tcx>,
+        typing_env: ty::TypingEnv<'tcx>,
         set: DenseBitSet<Local>,
+        /// This is used to look at the field types of a transmuted local.
+        decls: &'a LocalDecls<'tcx>,
     }
 
-    impl<'tcx> Visitor<'tcx> for EscapeVisitor {
+    impl<'tcx> Visitor<'tcx> for EscapeVisitor<'tcx, '_> {
         fn visit_local(&mut self, local: Local, _: PlaceContext, _: Location) {
             self.set.insert(local);
         }
@@ -114,6 +119,28 @@ fn escaping_locals<'tcx>(
             self.super_place(place, context, location);
         }
 
+        fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
+            // A transmute to a field type is either the same as a read of that
+            // field or it's UB for a size mismatch, so we can allow SRoA the same
+            // as if it had been written `Use(op)` with a field projection.
+            if let Rvalue::Cast(CastKind::Transmute, op, to_ty) = rvalue
+                && let Some(place) = op.place()
+                && let Some(local) = place.as_local()
+                && !self.set.contains(local)
+                && find_matching_struct_field(
+                    self.tcx,
+                    self.typing_env,
+                    *to_ty,
+                    self.decls[local].ty,
+                )
+                .is_some()
+            {
+                return;
+            }
+
+            self.super_rvalue(rvalue, location)
+        }
+
         fn visit_assign(
             &mut self,
             lvalue: &Place<'tcx>,
@@ -147,6 +174,27 @@ fn escaping_locals<'tcx>(
     }
 }
 
+fn find_matching_struct_field<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    typing_env: ty::TypingEnv<'tcx>,
+    needle_ty: Ty<'tcx>,
+    haystack_ty: Ty<'tcx>,
+) -> Option<FieldIdx> {
+    if let ty::Adt(adt_def, adt_args) = haystack_ty.kind()
+        && adt_def.is_struct()
+    {
+        for (idx, data) in adt_def.non_enum_variant().fields.iter_enumerated() {
+            let field_ty = data.ty(tcx, adt_args);
+            let field_ty = tcx.normalize_erasing_regions(typing_env, field_ty);
+            if field_ty == needle_ty {
+                return Some(idx);
+            }
+        }
+    }
+
+    None
+}
+
 #[derive(Default, Debug)]
 struct ReplacementMap<'tcx> {
     /// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage
@@ -211,6 +259,7 @@ fn compute_flattening<'tcx>(
 /// Perform the replacement computed by `compute_flattening`.
 fn replace_flattened_locals<'tcx>(
     tcx: TyCtxt<'tcx>,
+    typing_env: ty::TypingEnv<'tcx>,
     body: &mut Body<'tcx>,
     replacements: ReplacementMap<'tcx>,
 ) -> DenseBitSet<Local> {
@@ -227,6 +276,7 @@ fn replace_flattened_locals<'tcx>(
 
     let mut visitor = ReplacementVisitor {
         tcx,
+        typing_env,
         local_decls: &body.local_decls,
         replacements: &replacements,
         all_dead_locals,
@@ -249,7 +299,9 @@ fn replace_flattened_locals<'tcx>(
 
 struct ReplacementVisitor<'tcx, 'll> {
     tcx: TyCtxt<'tcx>,
-    /// This is only used to compute the type for `VarDebugInfoFragment`.
+    typing_env: ty::TypingEnv<'tcx>,
+    /// This is used to compute the type for `VarDebugInfoFragment`
+    /// and to look at the field types of a transmuted local.
     local_decls: &'ll LocalDecls<'tcx>,
     /// Work to do.
     replacements: &'ll ReplacementMap<'tcx>,
@@ -430,6 +482,37 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
         self.super_statement(statement, location)
     }
 
+    fn visit_rvalue(&mut self, rvalue: &mut Rvalue<'tcx>, location: Location) {
+        // We have `other = transmute(move? a)`
+        // We replace it with
+        // ```
+        // other = move? a_i
+        // ```
+        // for the one relevant field.
+        if let Rvalue::Cast(CastKind::Transmute, ref op, to_ty) = *rvalue
+            && let Some(op_place) = op.place()
+            && let Some(op_local) = op_place.as_local()
+            && let is_move = matches!(op, Operand::Move(..))
+            && let Some(op_final_locals) = &self.replacements.fragments[op_local]
+        {
+            let field_idx = find_matching_struct_field(
+                self.tcx,
+                self.typing_env,
+                to_ty,
+                self.local_decls[op_local].ty,
+            )
+            .unwrap();
+            let (new_local_ty, new_local) = op_final_locals[field_idx].unwrap();
+            assert_eq!(new_local_ty, to_ty);
+            let new_place = Place::from(new_local);
+            let new_op = if is_move { Operand::Move(new_place) } else { Operand::Copy(new_place) };
+            *rvalue = Rvalue::Use(new_op, WithRetag::Yes);
+            return;
+        }
+
+        self.super_rvalue(rvalue, location);
+    }
+
     fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
         assert!(!self.all_dead_locals.contains(*local));
     }

tests/codegen-llvm/read_write_unaligned.rs

@@ -22,8 +22,8 @@ unsafe fn read_unaligned_ptr(ptr: *const NonNull<i16>) -> NonNull<i16> {
 unsafe fn read_unaligned_i16(ptr: *const NonZero<i16>) -> NonZero<i16> {
     // CHECK: start:
     // CHECK-NEXT: [[TEMP:%.+]] = load i16, ptr %ptr, align 1
-    // CHECK-NOT: !noundef
-    // CHECK-NOT: !range
+    // CHECK-SAME: !range [[R16:![0-9]+]]
+    // CHECK-SAME: !noundef
     // CHECK-NEXT: ret i16 [[TEMP]]
     ptr.read_unaligned()
 }
@@ -33,8 +33,8 @@ unsafe fn read_unaligned_i16(ptr: *const NonZero<i16>) -> NonZero<i16> {
 unsafe fn typed_copy_unaligned_i32(src: *const NonZero<i32>, dst: *mut NonZero<i32>) {
     // CHECK: start:
     // CHECK-NEXT: [[TEMP:%.+]] = load i32, ptr %src, align 1
-    // CHECK-NOT: !noundef
-    // CHECK-NOT: !range
+    // CHECK-SAME: !range [[R32:![0-9]+]]
+    // CHECK-SAME: !noundef
     // CHECK-NEXT: store i32 [[TEMP]], ptr %dst, align 1
     // CHECK-NEXT: ret void
     dst.write_unaligned(src.read_unaligned())
@@ -69,3 +69,6 @@ unsafe fn write_unaligned_huge(ptr: *mut HugeBuffer, val: HugeBuffer) {
     // CHECK-NEXT: ret void
     ptr.write_unaligned(val)
 }
+
+// CHECK: [[R16]] = !{i16 1, i16 0}
+// CHECK: [[R32]] = !{i32 1, i32 0}

tests/mir-opt/pre-codegen/unaligned.rs

@@ -25,11 +25,10 @@ pub unsafe fn unaligned_copy_generic<T>(src: *const T, dst: *mut T) {
     // CHECK: debug dst => _2;
     // CHECK: debug val => [[VAL:_.+]];
     // CHECK: [[SRC_P:_.+]] = copy _1 as *const {{.+}}::Packed<T> (PtrToPtr);
-    // CHECK: [[PACKED1:_.+]] = copy (*[[SRC_P]]);
-    // CHECK: [[VAL]] = copy [[PACKED1]] as T (Transmute);
+    // CHECK: [[VAL]] = copy ((*[[SRC_P]]).0: T);
     // CHECK: [[DST_P:_.+]] = copy _2 as *mut {{.+}}::Packed<T> (PtrToPtr);
-    // CHECK: [[PACKED2:_.+]] = {{.+}}::Packed::<T>(copy [[VAL]]);
-    // CHECK: (*[[DST_P]]) = copy [[PACKED2]];
+    // CHECK: [[PACKED:_.+]] = {{.+}}::Packed::<T>(copy [[VAL]]);
+    // CHECK: (*[[DST_P]]) = copy [[PACKED]];
     // CHECK-NOT: copy_nonoverlapping
     // CHECK-NOT: drop
     unsafe {

tests/mir-opt/pre-codegen/unaligned.unaligned_copy_generic.runtime-optimized.after.mir

@@ -4,13 +4,13 @@ fn unaligned_copy_generic(_1: *const T, _2: *mut T) -> () {
     debug src => _1;
     debug dst => _2;
     let mut _0: ();
-    let _5: T;
+    let _4: T;
     scope 1 {
-        debug val => _5;
+        debug val => _4;
         scope 8 (inlined #[track_caller] write_unaligned::<T>) {
-            let _6: *mut std::ptr::Packed<T>;
+            let _5: *mut std::ptr::Packed<T>;
             scope 9 {
-                let _7: std::ptr::Packed<T>;
+                let _6: std::ptr::Packed<T>;
                 scope 10 {
                     scope 12 (inlined #[track_caller] std::ptr::write::<std::ptr::Packed<T>>) {
                     }
@@ -23,7 +23,6 @@ fn unaligned_copy_generic(_1: *const T, _2: *mut T) -> () {
     scope 2 (inlined #[track_caller] read_unaligned::<T>) {
         let _3: *const std::ptr::Packed<T>;
         scope 3 {
-            let _4: std::ptr::Packed<T>;
             scope 4 {
                 scope 7 (inlined transmute_neo::<std::ptr::Packed<T>, T>) {
                 }
@@ -36,22 +35,19 @@ fn unaligned_copy_generic(_1: *const T, _2: *mut T) -> () {
     }
 
     bb0: {
-        StorageLive(_5);
+        StorageLive(_4);
         StorageLive(_3);
         _3 = copy _1 as *const std::ptr::Packed<T> (PtrToPtr);
-        StorageLive(_4);
-        _4 = copy (*_3);
-        _5 = copy _4 as T (Transmute);
-        StorageDead(_4);
+        _4 = copy ((*_3).0: T);
         StorageDead(_3);
+        StorageLive(_5);
+        _5 = copy _2 as *mut std::ptr::Packed<T> (PtrToPtr);
         StorageLive(_6);
-        _6 = copy _2 as *mut std::ptr::Packed<T> (PtrToPtr);
-        StorageLive(_7);
-        _7 = std::ptr::Packed::<T>(copy _5);
-        (*_6) = copy _7;
-        StorageDead(_7);
+        _6 = std::ptr::Packed::<T>(copy _4);
+        (*_5) = copy _6;
         StorageDead(_6);
         StorageDead(_5);
+        StorageDead(_4);
         return;
     }
 }

tests/mir-opt/sroa/read_packed.read_unaligned.ScalarReplacementOfAggregates.diff

@@ -0,0 +1,59 @@
+- // MIR for `read_unaligned` before ScalarReplacementOfAggregates
++ // MIR for `read_unaligned` after ScalarReplacementOfAggregates
+  
+  fn read_unaligned(_1: *const T) -> T {
+      debug ptr => _1;
+      let mut _0: T;
+      let _2: *const Packed<T>;
+      let mut _3: *const Packed<T>;
+      let mut _4: *const T;
+      let mut _6: *const Packed<T>;
+      let mut _7: Packed<T>;
++     let mut _8: T;
+      scope 1 {
+          debug packed_ptr => _2;
+          let _5: Packed<T>;
++         let _9: T;
+          scope 2 {
+-             debug packed_val => _5;
++             debug ((packed_val: Packed<T>).0: T) => _9;
+          }
+      }
+  
+      bb0: {
+          StorageLive(_2);
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          _3 = move _4 as *const Packed<T> (PtrToPtr);
+          _2 = copy _3;
+          StorageDead(_4);
+          StorageDead(_3);
+-         StorageLive(_5);
++         StorageLive(_9);
++         nop;
+          StorageLive(_6);
+          _6 = copy _2;
+-         _5 = copy (*_6);
++         _9 = copy ((*_6).0: T);
++         nop;
+          StorageDead(_6);
+-         StorageLive(_7);
+-         _7 = move _5;
+-         _0 = move _7 as T (Transmute);
+-         StorageDead(_7);
+-         StorageDead(_5);
++         StorageLive(_8);
++         nop;
++         _8 = move _9;
++         nop;
++         _0 = move _8;
++         StorageDead(_8);
++         nop;
++         StorageDead(_9);
++         nop;
+          StorageDead(_2);
+          return;
+      }
+  }
+  

tests/mir-opt/sroa/read_packed.rs

@@ -0,0 +1,26 @@
+//@ test-mir-pass: ScalarReplacementOfAggregates
+//@ compile-flags: -Cpanic=abort
+//@ no-prefer-dynamic
+
+#![crate_type = "lib"]
+#![feature(core_intrinsics)]
+
+use std::intrinsics::{read_via_copy, transmute_unchecked};
+
+#[repr(packed)]
+struct Packed<T>(T);
+
+// EMIT_MIR read_packed.read_unaligned.ScalarReplacementOfAggregates.diff
+pub const unsafe fn read_unaligned<T>(ptr: *const T) -> T {
+    // CHECK-LABEL: fn read_unaligned(_1: *const T) -> T
+    // CHECK: debug packed_ptr => [[PPTR:_.+]];
+    // CHECK: debug ((packed_val: Packed<T>).0: T) => [[VAL:_.+]];
+    // CHECK: [[TEMP:_.+]] = copy [[PPTR]];
+    // CHECK: [[VAL]] = copy ((*{{_.+}}).0: T);
+    unsafe {
+        let packed_ptr = ptr as *const Packed<T>;
+        let packed_val = read_via_copy(packed_ptr);
+        // transmute because you can't destructure it in a `const fn`
+        transmute_unchecked(packed_val)
+    }
+}