1 files changed, 36 insertions, 57 deletions

diff --git a/compiler/rustc_mir_transform/src/sroa.rs b/compiler/rustc_mir_transform/src/sroa.rs
index e66ae8ff8..c21b1724c 100644
--- a/compiler/rustc_mir_transform/src/sroa.rs
+++ b/compiler/rustc_mir_transform/src/sroa.rs
@@ -1,4 +1,5 @@
 use crate::MirPass;
+use rustc_data_structures::flat_map_in_place::FlatMapInPlace;
 use rustc_index::bit_set::{BitSet, GrowableBitSet};
 use rustc_index::IndexVec;
 use rustc_middle::mir::patch::MirPatch;
@@ -147,7 +148,7 @@ fn escaping_locals<'tcx>(
         }
 
         // We ignore anything that happens in debuginfo, since we expand it using
-        // `VarDebugInfoContents::Composite`.
+        // `VarDebugInfoFragment`.
         fn visit_var_debug_info(&mut self, _: &VarDebugInfo<'tcx>) {}
     }
 }
@@ -246,9 +247,7 @@ fn replace_flattened_locals<'tcx>(
     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);
-    }
+    visitor.expand_var_debug_info(&mut body.var_debug_info);
     let ReplacementVisitor { patch, all_dead_locals, .. } = visitor;
     patch.apply(body);
     all_dead_locals
@@ -256,7 +255,7 @@ fn replace_flattened_locals<'tcx>(
 
 struct ReplacementVisitor<'tcx, 'll> {
     tcx: TyCtxt<'tcx>,
-    /// This is only used to compute the type for `VarDebugInfoContents::Composite`.
+    /// This is only used to compute the type for `VarDebugInfoFragment`.
     local_decls: &'ll LocalDecls<'tcx>,
     /// Work to do.
     replacements: &'ll ReplacementMap<'tcx>,
@@ -266,16 +265,38 @@ struct ReplacementVisitor<'tcx, 'll> {
 }
 
 impl<'tcx> ReplacementVisitor<'tcx, '_> {
-    fn gather_debug_info_fragments(&self, local: Local) -> Option<Vec<VarDebugInfoFragment<'tcx>>> {
-        let mut fragments = Vec::new();
-        let parts = self.replacements.place_fragments(local.into())?;
-        for (field, ty, replacement_local) in parts {
-            fragments.push(VarDebugInfoFragment {
-                projection: vec![PlaceElem::Field(field, ty)],
-                contents: Place::from(replacement_local),
-            });
-        }
-        Some(fragments)
+    #[instrument(level = "trace", skip(self))]
+    fn expand_var_debug_info(&mut self, var_debug_info: &mut Vec<VarDebugInfo<'tcx>>) {
+        var_debug_info.flat_map_in_place(|mut var_debug_info| {
+            let place = match var_debug_info.value {
+                VarDebugInfoContents::Const(_) => return vec![var_debug_info],
+                VarDebugInfoContents::Place(ref mut place) => place,
+            };
+
+            if let Some(repl) = self.replacements.replace_place(self.tcx, place.as_ref()) {
+                *place = repl;
+                return vec![var_debug_info];
+            }
+
+            let Some(parts) = self.replacements.place_fragments(*place) else {
+                return vec![var_debug_info];
+            };
+
+            let ty = place.ty(self.local_decls, self.tcx).ty;
+
+            parts
+                .map(|(field, field_ty, replacement_local)| {
+                    let mut var_debug_info = var_debug_info.clone();
+                    let composite = var_debug_info.composite.get_or_insert_with(|| {
+                        Box::new(VarDebugInfoFragment { ty, projection: Vec::new() })
+                    });
+                    composite.projection.push(PlaceElem::Field(field, field_ty));
+
+                    var_debug_info.value = VarDebugInfoContents::Place(replacement_local.into());
+                    var_debug_info
+                })
+                .collect()
+        });
     }
 }
 
@@ -422,48 +443,6 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
         self.super_statement(statement, location)
     }
 
-    #[instrument(level = "trace", skip(self))]
-    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.replacements.replace_place(self.tcx, place.as_ref()) {
-                    *place = repl;
-                } else if let Some(local) = place.as_local()
-                    && let Some(fragments) = self.gather_debug_info_fragments(local)
-                {
-                    let ty = place.ty(self.local_decls, self.tcx).ty;
-                    var_debug_info.value = VarDebugInfoContents::Composite { ty, fragments };
-                }
-            }
-            VarDebugInfoContents::Composite { ty: _, ref mut fragments } => {
-                let mut new_fragments = Vec::new();
-                debug!(?fragments);
-                fragments.retain_mut(|fragment| {
-                    if let Some(repl) =
-                            self.replacements.replace_place(self.tcx, fragment.contents.as_ref())
-                        {
-                            fragment.contents = repl;
-                            true
-                        } else if let Some(local) = fragment.contents.as_local()
-                            && let Some(frg) = self.gather_debug_info_fragments(local)
-                        {
-                            new_fragments.extend(frg.into_iter().map(|mut f| {
-                                f.projection.splice(0..0, fragment.projection.iter().copied());
-                                f
-                            }));
-                            false
-                        } else {
-                            true
-                        }
-                });
-                debug!(?fragments);
-                debug!(?new_fragments);
-                fragments.extend(new_fragments);
-            }
-            VarDebugInfoContents::Const(_) => {}
-        }
-    }
-
     fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
         assert!(!self.all_dead_locals.contains(*local));
     }