Auto merge of #102570 - cjgillot:deagg-debuginfo, r=oli-obk

Perform simple scalar replacement of aggregates (SROA) MIR opt

This is a re-open of https://github.com/rust-lang/rust/pull/85796

I copied the debuginfo implementation (first commit) from `@eddyb's` own SROA PR.

This pass replaces plain field accesses by simple locals when possible.
To be eligible, the replaced locals:
- must not be enums or unions;
- must not be used whole;
- must not have their address taken.

The storage and deinit statements are duplicated on each created local.

cc `@tmiasko` who reviewed the former version of this PR.

10 files changed, 549 insertions, 16 deletions

diff --git a/compiler/rustc_codegen_gcc/src/debuginfo.rs b/compiler/rustc_codegen_gcc/src/debuginfo.rs
index 266759ed6cf..a81585d4128 100644
--- a/compiler/rustc_codegen_gcc/src/debuginfo.rs
+++ b/compiler/rustc_codegen_gcc/src/debuginfo.rs
@@ -4,8 +4,9 @@ use rustc_codegen_ssa::traits::{DebugInfoBuilderMethods, DebugInfoMethods};
 use rustc_middle::mir;
 use rustc_middle::ty::{Instance, PolyExistentialTraitRef, Ty};
 use rustc_span::{SourceFile, Span, Symbol};
-use rustc_target::abi::Size;
 use rustc_target::abi::call::FnAbi;
+use rustc_target::abi::Size;
+use std::ops::Range;
 
 use crate::builder::Builder;
 use crate::context::CodegenCx;
@@ -13,7 +14,15 @@ use crate::context::CodegenCx;
 impl<'a, 'gcc, 'tcx> DebugInfoBuilderMethods for Builder<'a, 'gcc, 'tcx> {
     // FIXME(eddyb) find a common convention for all of the debuginfo-related
     // names (choose between `dbg`, `debug`, `debuginfo`, `debug_info` etc.).
-    fn dbg_var_addr(&mut self, _dbg_var: Self::DIVariable, _scope_metadata: Self::DIScope, _variable_alloca: Self::Value, _direct_offset: Size, _indirect_offsets: &[Size]) {
+    fn dbg_var_addr(
+        &mut self,
+        _dbg_var: Self::DIVariable,
+        _scope_metadata: Self::DIScope,
+        _variable_alloca: Self::Value,
+        _direct_offset: Size,
+        _indirect_offsets: &[Size],
+        _fragment: Option<Range<Size>>,
+    ) {
         unimplemented!();
     }
 
@@ -31,16 +40,31 @@ impl<'a, 'gcc, 'tcx> DebugInfoBuilderMethods for Builder<'a, 'gcc, 'tcx> {
 }
 
 impl<'gcc, 'tcx> DebugInfoMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
-    fn create_vtable_debuginfo(&self, _ty: Ty<'tcx>, _trait_ref: Option<PolyExistentialTraitRef<'tcx>>, _vtable: Self::Value) {
+    fn create_vtable_debuginfo(
+        &self,
+        _ty: Ty<'tcx>,
+        _trait_ref: Option<PolyExistentialTraitRef<'tcx>>,
+        _vtable: Self::Value,
+    ) {
         // TODO(antoyo)
     }
 
-    fn create_function_debug_context(&self, _instance: Instance<'tcx>, _fn_abi: &FnAbi<'tcx, Ty<'tcx>>, _llfn: RValue<'gcc>, _mir: &mir::Body<'tcx>) -> Option<FunctionDebugContext<Self::DIScope, Self::DILocation>> {
+    fn create_function_debug_context(
+        &self,
+        _instance: Instance<'tcx>,
+        _fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        _llfn: RValue<'gcc>,
+        _mir: &mir::Body<'tcx>,
+    ) -> Option<FunctionDebugContext<Self::DIScope, Self::DILocation>> {
         // TODO(antoyo)
         None
     }
 
-    fn extend_scope_to_file(&self, _scope_metadata: Self::DIScope, _file: &SourceFile) -> Self::DIScope {
+    fn extend_scope_to_file(
+        &self,
+        _scope_metadata: Self::DIScope,
+        _file: &SourceFile,
+    ) -> Self::DIScope {
         unimplemented!();
     }
 
@@ -48,15 +72,32 @@ impl<'gcc, 'tcx> DebugInfoMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         // TODO(antoyo)
     }
 
-    fn create_dbg_var(&self, _variable_name: Symbol, _variable_type: Ty<'tcx>, _scope_metadata: Self::DIScope, _variable_kind: VariableKind, _span: Span) -> Self::DIVariable {
+    fn create_dbg_var(
+        &self,
+        _variable_name: Symbol,
+        _variable_type: Ty<'tcx>,
+        _scope_metadata: Self::DIScope,
+        _variable_kind: VariableKind,
+        _span: Span,
+    ) -> Self::DIVariable {
         unimplemented!();
     }
 
-    fn dbg_scope_fn(&self, _instance: Instance<'tcx>, _fn_abi: &FnAbi<'tcx, Ty<'tcx>>, _maybe_definition_llfn: Option<RValue<'gcc>>) -> Self::DIScope {
+    fn dbg_scope_fn(
+        &self,
+        _instance: Instance<'tcx>,
+        _fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        _maybe_definition_llfn: Option<RValue<'gcc>>,
+    ) -> Self::DIScope {
         unimplemented!();
     }
 
-    fn dbg_loc(&self, _scope: Self::DIScope, _inlined_at: Option<Self::DILocation>, _span: Span) -> Self::DILocation {
+    fn dbg_loc(
+        &self,
+        _scope: Self::DIScope,
+        _inlined_at: Option<Self::DILocation>,
+        _span: Span,
+    ) -> Self::DILocation {
         unimplemented!();
     }
 }
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
index b23fe3fc9d5..ca7a07d8391 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
@@ -39,6 +39,7 @@ use smallvec::SmallVec;
 use std::cell::OnceCell;
 use std::cell::RefCell;
 use std::iter;
+use std::ops::Range;
 
 mod create_scope_map;
 pub mod gdb;
@@ -163,12 +164,14 @@ impl<'ll> DebugInfoBuilderMethods for Builder<'_, 'll, '_> {
         variable_alloca: Self::Value,
         direct_offset: Size,
         indirect_offsets: &[Size],
+        fragment: Option<Range<Size>>,
     ) {
-        // Convert the direct and indirect offsets to address ops.
+        // Convert the direct and indirect offsets and fragment byte range to address ops.
         // FIXME(eddyb) use `const`s instead of getting the values via FFI,
         // the values should match the ones in the DWARF standard anyway.
         let op_deref = || unsafe { llvm::LLVMRustDIBuilderCreateOpDeref() };
         let op_plus_uconst = || unsafe { llvm::LLVMRustDIBuilderCreateOpPlusUconst() };
+        let op_llvm_fragment = || unsafe { llvm::LLVMRustDIBuilderCreateOpLLVMFragment() };
         let mut addr_ops = SmallVec::<[u64; 8]>::new();
 
         if direct_offset.bytes() > 0 {
@@ -182,6 +185,13 @@ impl<'ll> DebugInfoBuilderMethods for Builder<'_, 'll, '_> {
                 addr_ops.push(offset.bytes() as u64);
             }
         }
+        if let Some(fragment) = fragment {
+            // `DW_OP_LLVM_fragment` takes as arguments the fragment's
+            // offset and size, both of them in bits.
+            addr_ops.push(op_llvm_fragment());
+            addr_ops.push(fragment.start.bits() as u64);
+            addr_ops.push((fragment.end - fragment.start).bits() as u64);
+        }
 
         unsafe {
             // FIXME(eddyb) replace `llvm.dbg.declare` with `llvm.dbg.addr`.
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index e2d0390821d..8f7728da9dd 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -2210,6 +2210,7 @@ extern "C" {
     ) -> &'a DILocation;
     pub fn LLVMRustDIBuilderCreateOpDeref() -> u64;
     pub fn LLVMRustDIBuilderCreateOpPlusUconst() -> u64;
+    pub fn LLVMRustDIBuilderCreateOpLLVMFragment() -> u64;
 
     #[allow(improper_ctypes)]
     pub fn LLVMRustWriteTypeToString(Type: &Type, s: &RustString);
diff --git a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
index 157c1c82311..99283d3bb29 100644
--- a/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/debuginfo.rs
@@ -14,6 +14,8 @@ use super::operand::{OperandRef, OperandValue};
 use super::place::PlaceRef;
 use super::{FunctionCx, LocalRef};
 
+use std::ops::Range;
+
 pub struct FunctionDebugContext<S, L> {
     pub scopes: IndexVec<mir::SourceScope, DebugScope<S, L>>,
 }
@@ -25,7 +27,7 @@ pub enum VariableKind {
 }
 
 /// Like `mir::VarDebugInfo`, but within a `mir::Local`.
-#[derive(Copy, Clone)]
+#[derive(Clone)]
 pub struct PerLocalVarDebugInfo<'tcx, D> {
     pub name: Symbol,
     pub source_info: mir::SourceInfo,
@@ -33,6 +35,10 @@ pub struct PerLocalVarDebugInfo<'tcx, D> {
     /// `DIVariable` returned by `create_dbg_var`.
     pub dbg_var: Option<D>,
 
+    /// Byte range in the `dbg_var` covered by this fragment,
+    /// if this is a fragment of a composite `VarDebugInfo`.
+    pub fragment: Option<Range<Size>>,
+
     /// `.place.projection` from `mir::VarDebugInfo`.
     pub projection: &'tcx ty::List<mir::PlaceElem<'tcx>>,
 }
@@ -145,7 +151,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             Some(per_local) => &per_local[local],
             None => return,
         };
-        let whole_local_var = vars.iter().find(|var| var.projection.is_empty()).copied();
+        let whole_local_var = vars.iter().find(|var| var.projection.is_empty()).cloned();
         let has_proj = || vars.iter().any(|var| !var.projection.is_empty());
 
         let fallback_var = if self.mir.local_kind(local) == mir::LocalKind::Arg {
@@ -187,6 +193,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     name,
                     source_info: decl.source_info,
                     dbg_var,
+                    fragment: None,
                     projection: ty::List::empty(),
                 })
             }
@@ -199,7 +206,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let name = if bx.sess().fewer_names() {
             None
         } else {
-            Some(match whole_local_var.or(fallback_var) {
+            Some(match whole_local_var.or(fallback_var.clone()) {
                 Some(var) if var.name != kw::Empty => var.name.to_string(),
                 _ => format!("{:?}", local),
             })
@@ -249,7 +256,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             LocalRef::UnsizedPlace(_) => return,
         };
 
-        let vars = vars.iter().copied().chain(fallback_var);
+        let vars = vars.iter().cloned().chain(fallback_var);
 
         for var in vars {
             let Some(dbg_var) = var.dbg_var else { continue };
@@ -312,9 +319,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 bx.store(place.llval, alloca.llval, alloca.align);
 
                 // Point the debug info to `*alloca` for the current variable
-                bx.dbg_var_addr(dbg_var, dbg_loc, alloca.llval, Size::ZERO, &[Size::ZERO]);
+                bx.dbg_var_addr(dbg_var, dbg_loc, alloca.llval, Size::ZERO, &[Size::ZERO], None);
             } else {
-                bx.dbg_var_addr(dbg_var, dbg_loc, base.llval, direct_offset, &indirect_offsets);
+                bx.dbg_var_addr(
+                    dbg_var,
+                    dbg_loc,
+                    base.llval,
+                    direct_offset,
+                    &indirect_offsets,
+                    None,
+                );
             }
         }
     }
@@ -382,6 +396,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         let ty = self.monomorphize(c.ty());
                         (ty, VariableKind::LocalVariable)
                     }
+                    mir::VarDebugInfoContents::Composite { ty, fragments: _ } => {
+                        let ty = self.monomorphize(ty);
+                        (ty, VariableKind::LocalVariable)
+                    }
                 };
 
                 self.cx.create_dbg_var(var.name, var_ty, dbg_scope, var_kind, span)
@@ -393,6 +411,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         name: var.name,
                         source_info: var.source_info,
                         dbg_var,
+                        fragment: None,
                         projection: place.projection,
                     });
                 }
@@ -407,10 +426,48 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                                 bx,
                             );
 
-                            bx.dbg_var_addr(dbg_var, dbg_loc, base.llval, Size::ZERO, &[]);
+                            bx.dbg_var_addr(dbg_var, dbg_loc, base.llval, Size::ZERO, &[], None);
                         }
                     }
                 }
+                mir::VarDebugInfoContents::Composite { ty, ref fragments } => {
+                    let var_ty = self.monomorphize(ty);
+                    let var_layout = self.cx.layout_of(var_ty);
+                    for fragment in fragments {
+                        let mut fragment_start = Size::ZERO;
+                        let mut fragment_layout = var_layout;
+
+                        for elem in &fragment.projection {
+                            match *elem {
+                                mir::ProjectionElem::Field(field, _) => {
+                                    let i = field.index();
+                                    fragment_start += fragment_layout.fields.offset(i);
+                                    fragment_layout = fragment_layout.field(self.cx, i);
+                                }
+                                _ => span_bug!(
+                                    var.source_info.span,
+                                    "unsupported fragment projection `{:?}`",
+                                    elem,
+                                ),
+                            }
+                        }
+
+                        let place = fragment.contents;
+                        per_local[place.local].push(PerLocalVarDebugInfo {
+                            name: var.name,
+                            source_info: var.source_info,
+                            dbg_var,
+                            fragment: if fragment_layout.size == var_layout.size {
+                                // Fragment covers entire variable, so as far as
+                                // DWARF is concerned, it's not really a fragment.
+                                None
+                            } else {
+                                Some(fragment_start..fragment_start + fragment_layout.size)
+                            },
+                            projection: place.projection,
+                        });
+                    }
+                }
             }
         }
         Some(per_local)
diff --git a/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs b/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
index f310789d144..63fecaf34fd 100644
--- a/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
@@ -6,6 +6,8 @@ use rustc_span::{SourceFile, Span, Symbol};
 use rustc_target::abi::call::FnAbi;
 use rustc_target::abi::Size;
 
+use std::ops::Range;
+
 pub trait DebugInfoMethods<'tcx>: BackendTypes {
     fn create_vtable_debuginfo(
         &self,
@@ -72,6 +74,9 @@ pub trait DebugInfoBuilderMethods: BackendTypes {
         direct_offset: Size,
         // NB: each offset implies a deref (i.e. they're steps in a pointer chain).
         indirect_offsets: &[Size],
+        // Byte range in the `dbg_var` covered by this fragment,
+        // if this is a fragment of a composite `DIVariable`.
+        fragment: Option<Range<Size>>,
     );
     fn set_dbg_loc(&mut self, dbg_loc: Self::DILocation);
     fn insert_reference_to_gdb_debug_scripts_section_global(&mut self);
diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
index 85c520a7911..6f8bb676104 100644
--- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
+++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp
@@ -1111,6 +1111,10 @@ extern "C" uint64_t LLVMRustDIBuilderCreateOpPlusUconst() {
   return dwarf::DW_OP_plus_uconst;
 }
 
+extern "C" int64_t LLVMRustDIBuilderCreateOpLLVMFragment() {
+  return dwarf::DW_OP_LLVM_fragment;
+}
+
 extern "C" void LLVMRustWriteTypeToString(LLVMTypeRef Ty, RustStringRef Str) {
   RawRustStringOstream OS(Str);
   unwrap<llvm::Type>(Ty)->print(OS);
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index 4781651071d..12e9ebfeaec 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -1071,6 +1071,18 @@ pub enum VarDebugInfoContents<'tcx> {
     /// based on a `Local`, not a `Static`, and contains no indexing.
     Place(Place<'tcx>),
     Const(Constant<'tcx>),
+    /// The user variable's data is split across several fragments,
+    /// each described by a `VarDebugInfoFragment`.
+    /// See DWARF 5's "2.6.1.2 Composite Location Descriptions"
+    /// and LLVM's `DW_OP_LLVM_fragment` for more details on
+    /// the underlying debuginfo feature this relies on.
+    Composite {
+        /// Type of the original user variable.
+        ty: Ty<'tcx>,
+        /// All the parts of the original user variable, which ended
+        /// up in disjoint places, due to optimizations.
+        fragments: Vec<VarDebugInfoFragment<'tcx>>,
+    },
 }
 
 impl<'tcx> Debug for VarDebugInfoContents<'tcx> {
@@ -1078,7 +1090,48 @@ impl<'tcx> Debug for VarDebugInfoContents<'tcx> {
         match self {
             VarDebugInfoContents::Const(c) => write!(fmt, "{}", c),
             VarDebugInfoContents::Place(p) => write!(fmt, "{:?}", p),
+            VarDebugInfoContents::Composite { ty, fragments } => {
+                write!(fmt, "{:?}{{ ", ty)?;
+                for f in fragments.iter() {
+                    write!(fmt, "{:?}, ", f)?;
+                }
+                write!(fmt, "}}")
+            }
+        }
+    }
+}
+
+#[derive(Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub struct VarDebugInfoFragment<'tcx> {
+    /// Where in the composite user variable this fragment is,
+    /// represented as a "projection" into the composite variable.
+    /// At lower levels, this corresponds to a byte/bit range.
+    // NOTE(eddyb) there's an unenforced invariant that this contains
+    // only `Field`s, and not into `enum` variants or `union`s.
+    // FIXME(eddyb) support this for `enum`s by either using DWARF's
+    // more advanced control-flow features (unsupported by LLVM?)
+    // to match on the discriminant, or by using custom type debuginfo
+    // with non-overlapping variants for the composite variable.
+    pub projection: Vec<PlaceElem<'tcx>>,
+
+    /// Where the data for this fragment can be found.
+    // NOTE(eddyb) There's an unenforced invariant that this `Place` is
+    // contains no indexing (with a non-constant index).
+    pub contents: Place<'tcx>,
+}
+
+impl Debug for VarDebugInfoFragment<'_> {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
+        for elem in self.projection.iter() {
+            match elem {
+                ProjectionElem::Field(field, _) => {
+                    write!(fmt, ".{:?}", field.index())?;
+                }
+                _ => bug!("unsupported fragment projection `{:?}`", elem),
+            }
         }
+
+        write!(fmt, " => {:?}", self.contents)
     }
 }
 
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
index d87eb28970e..b21f50ae5ea 100644
--- a/compiler/rustc_middle/src/mir/visit.rs
+++ b/compiler/rustc_middle/src/mir/visit.rs
@@ -847,6 +847,17 @@ macro_rules! make_mir_visitor {
                             PlaceContext::NonUse(NonUseContext::VarDebugInfo),
                             location
                         ),
+                    VarDebugInfoContents::Composite { ty, fragments } => {
+                        // FIXME(eddyb) use a better `TyContext` here.
+                        self.visit_ty($(& $mutability)? *ty, TyContext::Location(location));
+                        for VarDebugInfoFragment { projection: _, contents } in fragments {
+                            self.visit_place(
+                                contents,
+                                PlaceContext::NonUse(NonUseContext::VarDebugInfo),
+                                location,
+                            );
+                        }
+                    }
                 }
             }
 
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 692eeddfb98..5e85d1f0db4 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -1,5 +1,6 @@
 #![allow(rustc::potential_query_instability)]
 #![feature(box_patterns)]
+#![feature(drain_filter)]
 #![feature(let_chains)]
 #![feature(map_try_insert)]
 #![feature(min_specialization)]
@@ -92,6 +93,7 @@ pub mod simplify;
 mod simplify_branches;
 mod simplify_comparison_integral;
 mod simplify_try;
+mod sroa;
 mod uninhabited_enum_branching;
 mod unreachable_prop;
 
@@ -562,6 +564,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &remove_zsts::RemoveZsts,
             &const_goto::ConstGoto,
             &remove_unneeded_drops::RemoveUnneededDrops,
+            &sroa::ScalarReplacementOfAggregates,
             &match_branches::MatchBranchSimplification,
             // inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
             &multiple_return_terminators::MultipleReturnTerminators,
diff --git a/compiler/rustc_mir_transform/src/sroa.rs b/compiler/rustc_mir_transform/src/sroa.rs
new file mode 100644
index 00000000000..558a372fb1e
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/sroa.rs
@@ -0,0 +1,348 @@
+use crate::MirPass;
+use rustc_data_structures::fx::{FxIndexMap, IndexEntry};
+use rustc_index::bit_set::BitSet;
+use rustc_index::vec::IndexVec;
+use rustc_middle::mir::visit::*;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+
+pub struct ScalarReplacementOfAggregates;
+
+impl<'tcx> MirPass<'tcx> for ScalarReplacementOfAggregates {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() >= 3
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let escaping = escaping_locals(&*body);
+        debug!(?escaping);
+        let replacements = compute_flattening(tcx, body, escaping);
+        debug!(?replacements);
+        replace_flattened_locals(tcx, body, replacements);
+    }
+}
+
+/// Identify all locals that are not eligible for SROA.
+///
+/// There are 3 cases:
+/// - the aggegated local is used or passed to other code (function parameters and arguments);
+/// - the locals is a union or an enum;
+/// - the local's address is taken, and thus the relative addresses of the fields are observable to
+///   client code.
+fn escaping_locals(body: &Body<'_>) -> BitSet<Local> {
+    let mut set = BitSet::new_empty(body.local_decls.len());
+    set.insert_range(RETURN_PLACE..=Local::from_usize(body.arg_count));
+    for (local, decl) in body.local_decls().iter_enumerated() {
+        if decl.ty.is_union() || decl.ty.is_enum() {
+            set.insert(local);
+        }
+    }
+    let mut visitor = EscapeVisitor { set };
+    visitor.visit_body(body);
+    return visitor.set;
+
+    struct EscapeVisitor {
+        set: BitSet<Local>,
+    }
+
+    impl<'tcx> Visitor<'tcx> for EscapeVisitor {
+        fn visit_local(&mut self, local: Local, _: PlaceContext, _: Location) {
+            self.set.insert(local);
+        }
+
+        fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) {
+            // Mirror the implementation in PreFlattenVisitor.
+            if let &[PlaceElem::Field(..), ..] = &place.projection[..] {
+                return;
+            }
+            self.super_place(place, context, location);
+        }
+
+        fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, location: Location) {
+            if let Rvalue::AddressOf(.., place) | Rvalue::Ref(.., place) = rvalue {
+                if !place.is_indirect() {
+                    // Raw pointers may be used to access anything inside the enclosing place.
+                    self.set.insert(place.local);
+                    return;
+                }
+            }
+            self.super_rvalue(rvalue, location)
+        }
+
+        fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
+            if let StatementKind::StorageLive(..)
+            | StatementKind::StorageDead(..)
+            | StatementKind::Deinit(..) = statement.kind
+            {
+                // Storage statements are expanded in run_pass.
+                return;
+            }
+            self.super_statement(statement, location)
+        }
+
+        fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
+            // Drop implicitly calls `drop_in_place`, which takes a `&mut`.
+            // This implies that `Drop` implicitly takes the address of the place.
+            if let TerminatorKind::Drop { place, .. }
+            | TerminatorKind::DropAndReplace { place, .. } = terminator.kind
+            {
+                if !place.is_indirect() {
+                    // Raw pointers may be used to access anything inside the enclosing place.
+                    self.set.insert(place.local);
+                    return;
+                }
+            }
+            self.super_terminator(terminator, location);
+        }
+
+        // We ignore anything that happens in debuginfo, since we expand it using
+        // `VarDebugInfoContents::Composite`.
+        fn visit_var_debug_info(&mut self, _: &VarDebugInfo<'tcx>) {}
+    }
+}
+
+#[derive(Default, Debug)]
+struct ReplacementMap<'tcx> {
+    fields: FxIndexMap<PlaceRef<'tcx>, Local>,
+}
+
+/// Compute the replacement of flattened places into locals.
+///
+/// For each eligible place, we assign a new local to each accessed field.
+/// The replacement will be done later in `ReplacementVisitor`.
+fn compute_flattening<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &mut Body<'tcx>,
+    escaping: BitSet<Local>,
+) -> ReplacementMap<'tcx> {
+    let mut visitor = PreFlattenVisitor {
+        tcx,
+        escaping,
+        local_decls: &mut body.local_decls,
+        map: Default::default(),
+    };
+    for (block, bbdata) in body.basic_blocks.iter_enumerated() {
+        visitor.visit_basic_block_data(block, bbdata);
+    }
+    return visitor.map;
+
+    struct PreFlattenVisitor<'tcx, 'll> {
+        tcx: TyCtxt<'tcx>,
+        local_decls: &'ll mut LocalDecls<'tcx>,
+        escaping: BitSet<Local>,
+        map: ReplacementMap<'tcx>,
+    }
+
+    impl<'tcx, 'll> PreFlattenVisitor<'tcx, 'll> {
+        fn create_place(&mut self, place: PlaceRef<'tcx>) {
+            if self.escaping.contains(place.local) {
+                return;
+            }
+
+            match self.map.fields.entry(place) {
+                IndexEntry::Occupied(_) => {}
+                IndexEntry::Vacant(v) => {
+                    let ty = place.ty(&*self.local_decls, self.tcx).ty;
+                    let local = self.local_decls.push(LocalDecl {
+                        ty,
+                        user_ty: None,
+                        ..self.local_decls[place.local].clone()
+                    });
+                    v.insert(local);
+                }
+            }
+        }
+    }
+
+    impl<'tcx, 'll> Visitor<'tcx> for PreFlattenVisitor<'tcx, 'll> {
+        fn visit_place(&mut self, place: &Place<'tcx>, _: PlaceContext, _: Location) {
+            if let &[PlaceElem::Field(..), ..] = &place.projection[..] {
+                let pr = PlaceRef { local: place.local, projection: &place.projection[..1] };
+                self.create_place(pr)
+            }
+        }
+    }
+}
+
+/// Perform the replacement computed by `compute_flattening`.
+fn replace_flattened_locals<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &mut Body<'tcx>,
+    replacements: ReplacementMap<'tcx>,
+) {
+    let mut all_dead_locals = BitSet::new_empty(body.local_decls.len());
+    for p in replacements.fields.keys() {
+        all_dead_locals.insert(p.local);
+    }
+    debug!(?all_dead_locals);
+    if all_dead_locals.is_empty() {
+        return;
+    }
+
+    let mut fragments = IndexVec::new();
+    for (k, v) in &replacements.fields {
+        fragments.ensure_contains_elem(k.local, || Vec::new());
+        fragments[k.local].push((&k.projection[..], *v));
+    }
+    debug!(?fragments);
+
+    let mut visitor = ReplacementVisitor {
+        tcx,
+        local_decls: &body.local_decls,
+        replacements,
+        all_dead_locals,
+        fragments,
+    };
+    for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
+        visitor.visit_basic_block_data(bb, data);
+    }
+    for scope in &mut body.source_scopes {
+        visitor.visit_source_scope_data(scope);
+    }
+    for (index, annotation) in body.user_type_annotations.iter_enumerated_mut() {
+        visitor.visit_user_type_annotation(index, annotation);
+    }
+    for var_debug_info in &mut body.var_debug_info {
+        visitor.visit_var_debug_info(var_debug_info);
+    }
+}
+
+struct ReplacementVisitor<'tcx, 'll> {
+    tcx: TyCtxt<'tcx>,
+    /// This is only used to compute the type for `VarDebugInfoContents::Composite`.
+    local_decls: &'ll LocalDecls<'tcx>,
+    /// Work to do.
+    replacements: ReplacementMap<'tcx>,
+    /// This is used to check that we are not leaving references to replaced locals behind.
+    all_dead_locals: BitSet<Local>,
+    /// Pre-computed list of all "new" locals for each "old" local.  This is used to expand storage
+    /// and deinit statement and debuginfo.
+    fragments: IndexVec<Local, Vec<(&'tcx [PlaceElem<'tcx>], Local)>>,
+}
+
+impl<'tcx, 'll> ReplacementVisitor<'tcx, 'll> {
+    fn gather_debug_info_fragments(
+        &self,
+        place: PlaceRef<'tcx>,
+    ) -> Vec<VarDebugInfoFragment<'tcx>> {
+        let mut fragments = Vec::new();
+        let parts = &self.fragments[place.local];
+        for (proj, replacement_local) in parts {
+            if proj.starts_with(place.projection) {
+                fragments.push(VarDebugInfoFragment {
+                    projection: proj[place.projection.len()..].to_vec(),
+                    contents: Place::from(*replacement_local),
+                });
+            }
+        }
+        fragments
+    }
+
+    fn replace_place(&self, place: PlaceRef<'tcx>) -> Option<Place<'tcx>> {
+        if let &[PlaceElem::Field(..), ref rest @ ..] = place.projection {
+            let pr = PlaceRef { local: place.local, projection: &place.projection[..1] };
+            let local = self.replacements.fields.get(&pr)?;
+            Some(Place { local: *local, projection: self.tcx.intern_place_elems(&rest) })
+        } else {
+            None
+        }
+    }
+}
+
+impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
+        if let StatementKind::StorageLive(..)
+        | StatementKind::StorageDead(..)
+        | StatementKind::Deinit(..) = statement.kind
+        {
+            // Storage statements are expanded in run_pass.
+            return;
+        }
+        self.super_statement(statement, location)
+    }
+
+    fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
+        if let Some(repl) = self.replace_place(place.as_ref()) {
+            *place = repl
+        } else {
+            self.super_place(place, context, location)
+        }
+    }
+
+    fn visit_var_debug_info(&mut self, var_debug_info: &mut VarDebugInfo<'tcx>) {
+        match &mut var_debug_info.value {
+            VarDebugInfoContents::Place(ref mut place) => {
+                if let Some(repl) = self.replace_place(place.as_ref()) {
+                    *place = repl;
+                } else if self.all_dead_locals.contains(place.local) {
+                    let ty = place.ty(self.local_decls, self.tcx).ty;
+                    let fragments = self.gather_debug_info_fragments(place.as_ref());
+                    var_debug_info.value = VarDebugInfoContents::Composite { ty, fragments };
+                }
+            }
+            VarDebugInfoContents::Composite { ty: _, ref mut fragments } => {
+                let mut new_fragments = Vec::new();
+                fragments
+                    .drain_filter(|fragment| {
+                        if let Some(repl) = self.replace_place(fragment.contents.as_ref()) {
+                            fragment.contents = repl;
+                            true
+                        } else if self.all_dead_locals.contains(fragment.contents.local) {
+                            let frg = self.gather_debug_info_fragments(fragment.contents.as_ref());
+                            new_fragments.extend(frg.into_iter().map(|mut f| {
+                                f.projection.splice(0..0, fragment.projection.iter().copied());
+                                f
+                            }));
+                            false
+                        } else {
+                            true
+                        }
+                    })
+                    .for_each(drop);
+                fragments.extend(new_fragments);
+            }
+            VarDebugInfoContents::Const(_) => {}
+        }
+    }
+
+    fn visit_basic_block_data(&mut self, bb: BasicBlock, bbdata: &mut BasicBlockData<'tcx>) {
+        self.super_basic_block_data(bb, bbdata);
+
+        #[derive(Debug)]
+        enum Stmt {
+            StorageLive,
+            StorageDead,
+            Deinit,
+        }
+
+        bbdata.expand_statements(|stmt| {
+            let source_info = stmt.source_info;
+            let (stmt, origin_local) = match &stmt.kind {
+                StatementKind::StorageLive(l) => (Stmt::StorageLive, *l),
+                StatementKind::StorageDead(l) => (Stmt::StorageDead, *l),
+                StatementKind::Deinit(p) if let Some(l) = p.as_local() => (Stmt::Deinit, l),
+                _ => return None,
+            };
+            if !self.all_dead_locals.contains(origin_local) {
+                return None;
+            }
+            let final_locals = self.fragments.get(origin_local)?;
+            Some(final_locals.iter().map(move |&(_, l)| {
+                let kind = match stmt {
+                    Stmt::StorageLive => StatementKind::StorageLive(l),
+                    Stmt::StorageDead => StatementKind::StorageDead(l),
+                    Stmt::Deinit => StatementKind::Deinit(Box::new(l.into())),
+                };
+                Statement { source_info, kind }
+            }))
+        });
+    }
+
+    fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
+        assert!(!self.all_dead_locals.contains(*local));
+    }
+}