From a576ef93af24b6dc0f1f8853e41cd9a2ee19a29a Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Tue, 23 Jun 2026 00:26:45 -0700 Subject: [PATCH] [mir-opt] Allow SRoA when an object is transmuted to a field type --- compiler/rustc_mir_transform/src/sroa.rs | 97 +++++++++++++++++-- tests/codegen-llvm/read_write_unaligned.rs | 11 ++- tests/mir-opt/pre-codegen/unaligned.rs | 7 +- ...d_copy_generic.runtime-optimized.after.mir | 26 +++-- ...aligned.ScalarReplacementOfAggregates.diff | 59 +++++++++++ tests/mir-opt/sroa/read_packed.rs | 26 +++++ 6 files changed, 196 insertions(+), 30 deletions(-) create mode 100644 tests/mir-opt/sroa/read_packed.read_unaligned.ScalarReplacementOfAggregates.diff create mode 100644 tests/mir-opt/sroa/read_packed.rs diff --git a/compiler/rustc_mir_transform/src/sroa.rs b/compiler/rustc_mir_transform/src/sroa.rs index b18aaa829afd3..de7ca74504bf2 100644 --- a/compiler/rustc_mir_transform/src/sroa.rs +++ b/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, + typing_env: ty::TypingEnv<'tcx>, body: &Body<'tcx>, ) -> DenseBitSet { 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, + /// 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 { + 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 { @@ -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)); } diff --git a/tests/codegen-llvm/read_write_unaligned.rs b/tests/codegen-llvm/read_write_unaligned.rs index beba4400dc54a..90bce95d67ee1 100644 --- a/tests/codegen-llvm/read_write_unaligned.rs +++ b/tests/codegen-llvm/read_write_unaligned.rs @@ -22,8 +22,8 @@ unsafe fn read_unaligned_ptr(ptr: *const NonNull) -> NonNull { unsafe fn read_unaligned_i16(ptr: *const NonZero) -> NonZero { // 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) -> NonZero { unsafe fn typed_copy_unaligned_i32(src: *const NonZero, dst: *mut NonZero) { // 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} diff --git a/tests/mir-opt/pre-codegen/unaligned.rs b/tests/mir-opt/pre-codegen/unaligned.rs index 4aeb10c261e21..5423c7faf44fd 100644 --- a/tests/mir-opt/pre-codegen/unaligned.rs +++ b/tests/mir-opt/pre-codegen/unaligned.rs @@ -25,11 +25,10 @@ pub unsafe fn unaligned_copy_generic(src: *const T, dst: *mut T) { // CHECK: debug dst => _2; // CHECK: debug val => [[VAL:_.+]]; // CHECK: [[SRC_P:_.+]] = copy _1 as *const {{.+}}::Packed (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 (PtrToPtr); - // CHECK: [[PACKED2:_.+]] = {{.+}}::Packed::(copy [[VAL]]); - // CHECK: (*[[DST_P]]) = copy [[PACKED2]]; + // CHECK: [[PACKED:_.+]] = {{.+}}::Packed::(copy [[VAL]]); + // CHECK: (*[[DST_P]]) = copy [[PACKED]]; // CHECK-NOT: copy_nonoverlapping // CHECK-NOT: drop unsafe { diff --git a/tests/mir-opt/pre-codegen/unaligned.unaligned_copy_generic.runtime-optimized.after.mir b/tests/mir-opt/pre-codegen/unaligned.unaligned_copy_generic.runtime-optimized.after.mir index 4d0cfa6a4dfa3..33d3ce39cc53e 100644 --- a/tests/mir-opt/pre-codegen/unaligned.unaligned_copy_generic.runtime-optimized.after.mir +++ b/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::) { - let _6: *mut std::ptr::Packed; + let _5: *mut std::ptr::Packed; scope 9 { - let _7: std::ptr::Packed; + let _6: std::ptr::Packed; scope 10 { scope 12 (inlined #[track_caller] std::ptr::write::>) { } @@ -23,7 +23,6 @@ fn unaligned_copy_generic(_1: *const T, _2: *mut T) -> () { scope 2 (inlined #[track_caller] read_unaligned::) { let _3: *const std::ptr::Packed; scope 3 { - let _4: std::ptr::Packed; scope 4 { scope 7 (inlined transmute_neo::, 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 (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 (PtrToPtr); StorageLive(_6); - _6 = copy _2 as *mut std::ptr::Packed (PtrToPtr); - StorageLive(_7); - _7 = std::ptr::Packed::(copy _5); - (*_6) = copy _7; - StorageDead(_7); + _6 = std::ptr::Packed::(copy _4); + (*_5) = copy _6; StorageDead(_6); StorageDead(_5); + StorageDead(_4); return; } } diff --git a/tests/mir-opt/sroa/read_packed.read_unaligned.ScalarReplacementOfAggregates.diff b/tests/mir-opt/sroa/read_packed.read_unaligned.ScalarReplacementOfAggregates.diff new file mode 100644 index 0000000000000..4e0c1c653b555 --- /dev/null +++ b/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; + let mut _3: *const Packed; + let mut _4: *const T; + let mut _6: *const Packed; + let mut _7: Packed; ++ let mut _8: T; + scope 1 { + debug packed_ptr => _2; + let _5: Packed; ++ let _9: T; + scope 2 { +- debug packed_val => _5; ++ debug ((packed_val: Packed).0: T) => _9; + } + } + + bb0: { + StorageLive(_2); + StorageLive(_3); + StorageLive(_4); + _4 = copy _1; + _3 = move _4 as *const Packed (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; + } + } + diff --git a/tests/mir-opt/sroa/read_packed.rs b/tests/mir-opt/sroa/read_packed.rs new file mode 100644 index 0000000000000..7138dd6b0654b --- /dev/null +++ b/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); + +// EMIT_MIR read_packed.read_unaligned.ScalarReplacementOfAggregates.diff +pub const unsafe fn read_unaligned(ptr: *const T) -> T { + // CHECK-LABEL: fn read_unaligned(_1: *const T) -> T + // CHECK: debug packed_ptr => [[PPTR:_.+]]; + // CHECK: debug ((packed_val: Packed).0: T) => [[VAL:_.+]]; + // CHECK: [[TEMP:_.+]] = copy [[PPTR]]; + // CHECK: [[VAL]] = copy ((*{{_.+}}).0: T); + unsafe { + let packed_ptr = ptr as *const Packed; + let packed_val = read_via_copy(packed_ptr); + // transmute because you can't destructure it in a `const fn` + transmute_unchecked(packed_val) + } +}