fixup! simplifycfg: Preserve debuginfos when merging bbs

Move all debuginfos to the last basic block if each bb has
only one predecessor.

Author: dianqk

compiler/rustc_mir_transform/src/simplify.rs

@@ -158,17 +158,9 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 let mut terminator =
                     self.basic_blocks[bb].terminator.take().expect("invalid terminator state");
 
-                let identical_succ = terminator.identical_successor();
                 terminator.successors_mut(|successor| {
                     self.collapse_goto_chain(successor, &mut changed);
                 });
-                if changed && let Some(identical_succ) = identical_succ {
-                    // Add debugging information from the goto chain only when all successors are identical,
-                    // otherwise, we may provide misleading debugging information within a branch.
-                    let mut succ_debuginfos =
-                        self.basic_blocks[identical_succ].after_last_stmt_debuginfos.clone();
-                    self.basic_blocks[bb].after_last_stmt_debuginfos.append(&mut succ_debuginfos);
-                }
 
                 let mut inner_changed = true;
                 merged_blocks.clear();
@@ -236,28 +228,36 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
         // goto chains. We should probably benchmark different sizes.
         let mut terminators: SmallVec<[_; 1]> = Default::default();
         let mut current = *start;
+        // If each successor has only one predecessor, it's a trivial goto chain.
+        // We can move all debuginfos to the last basic block.
+        let mut trivial_goto_chain = true;
         while let Some(terminator) = self.take_terminator_if_simple_goto(current) {
             let Terminator { kind: TerminatorKind::Goto { target }, .. } = terminator else {
                 unreachable!();
             };
+            trivial_goto_chain &= self.pred_count[target] == 1;
             terminators.push((current, terminator));
             current = target;
         }
         let last = current;
         *changed |= *start != last;
         *start = last;
-        let mut succ = last;
         while let Some((current, mut terminator)) = terminators.pop() {
             let Terminator { kind: TerminatorKind::Goto { ref mut target }, .. } = terminator
             else {
                 unreachable!();
             };
-            if *target != last {
-                let mut succ_debuginfos =
-                    self.basic_blocks[succ].after_last_stmt_debuginfos.clone();
-                self.basic_blocks[current].after_last_stmt_debuginfos.extend(&mut succ_debuginfos);
+            if trivial_goto_chain {
+                let mut pred_debuginfos =
+                    std::mem::take(&mut self.basic_blocks[current].after_last_stmt_debuginfos);
+                let debuginfos = if let Some(stmt) = self.basic_blocks[last].statements.first_mut()
+                {
+                    &mut stmt.debuginfos
+                } else {
+                    &mut self.basic_blocks[last].after_last_stmt_debuginfos
+                };
+                debuginfos.prepend(&mut pred_debuginfos);
             }
-            succ = current;
             *changed |= *target != last;
             *target = last;
             debug!("collapsing goto chain from {:?} to {:?}", current, target);

tests/mir-opt/debuginfo/simplifycfg.preserve_debuginfo_3.SimplifyCfg-final.diff

@@ -0,0 +1,37 @@
+- // MIR for `preserve_debuginfo_3` before SimplifyCfg-final
++ // MIR for `preserve_debuginfo_3` after SimplifyCfg-final
+  
+  fn preserve_debuginfo_3(_1: &Foo, _2: bool) -> i32 {
+      debug foo_a => _3;
+      debug foo_b => _4;
+      debug foo_c => _5;
+      let mut _0: i32;
+      let mut _3: &i32;
+      let mut _4: &i64;
+      let mut _5: &i32;
+  
+      bb0: {
+-         switchInt(copy _2) -> [1: bb1, otherwise: bb2];
++         switchInt(copy _2) -> [1: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+-         // DBG: _3 = &((*_1).0: i32);
+-         goto -> bb3;
+-     }
+- 
+-     bb2: {
+          // DBG: _4 = &((*_1).1: i64);
+          _0 = copy ((*_1).2: i32);
+          return;
+      }
+  
+-     bb3: {
++     bb2: {
++         // DBG: _3 = &((*_1).0: i32);
+          // DBG: _5 = &((*_1).2: i32);
+          _0 = copy ((*_1).0: i32);
+          return;
+      }
+  }
+  

tests/mir-opt/debuginfo/simplifycfg.rs

@@ -119,6 +119,54 @@ pub fn preserve_debuginfo_2(foo: &Foo) -> i32 {
     }
 }
 
+// EMIT_MIR simplifycfg.preserve_debuginfo_3.SimplifyCfg-final.diff
+#[custom_mir(dialect = "runtime", phase = "post-cleanup")]
+pub fn preserve_debuginfo_3(foo: &Foo, c: bool) -> i32 {
+    // CHECK-LABEL: fn preserve_debuginfo_3
+    // CHECK: debug foo_a => [[foo_a:_[0-9]+]];
+    // CHECK: debug foo_b => [[foo_b:_[0-9]+]];
+    // CHECK: debug foo_c => [[foo_c:_[0-9]+]];
+    // CHECK: bb0: {
+    // CHECK-NEXT: switchInt(copy _2) -> [1: bb2, otherwise: bb1];
+    // CHECK: bb1: {
+    // CHECK-NEXT: DBG: [[foo_b]] = &((*_1).1: i64)
+    // CHECK-NEXT: _0 = copy ((*_1).2: i32);
+    // CHECK-NEXT: return;
+    // CHECK: bb2: {
+    // CHECK-NEXT: DBG: [[foo_a]] = &((*_1).0: i32)
+    // CHECK-NEXT: DBG: [[foo_c]] = &((*_1).2: i32)
+    // CHECK-NEXT: _0 = copy ((*_1).0: i32);
+    // CHECK-NEXT: return;
+    mir! {
+        let _foo_a: &i32;
+        let _foo_b: &i64;
+        let _foo_c: &i32;
+        debug foo_a => _foo_a;
+        debug foo_b => _foo_b;
+        debug foo_c => _foo_c;
+        {
+            match c {
+                true => tmp,
+                _ => ret,
+            }
+        }
+        tmp = {
+            _foo_a = &(*foo).a;
+            Goto(ret_1)
+        }
+        ret = {
+            _foo_b = &(*foo).b;
+            RET = (*foo).c;
+            Return()
+        }
+        ret_1 = {
+            _foo_c = &(*foo).c;
+            RET = (*foo).a;
+            Return()
+        }
+    }
+}
+
 // EMIT_MIR simplifycfg.preserve_debuginfo_identical_succs.SimplifyCfg-final.diff
 #[custom_mir(dialect = "runtime", phase = "post-cleanup")]
 pub fn preserve_debuginfo_identical_succs(foo: &Foo, c: bool) -> i32 {