simplifycfg: Preserve debuginfos when merging bbs

Author: dianqk

compiler/rustc_middle/src/mir/terminator.rs

@@ -444,6 +444,21 @@ impl<'tcx> Terminator<'tcx> {
         self.kind.successors()
     }
 
+    /// Return `Some` if all successors are identical.
+    #[inline]
+    pub fn identical_successor(&self) -> Option<BasicBlock> {
+        let mut successors = self.successors();
+        let Some(first_succ) = successors.next() else {
+            return None;
+        };
+        while let Some(succ) = successors.next() {
+            if first_succ != succ {
+                return None;
+            }
+        }
+        Some(first_succ)
+    }
+
     #[inline]
     pub fn successors_mut<'a>(&'a mut self, f: impl FnMut(&'a mut BasicBlock)) {
         self.kind.successors_mut(f)

compiler/rustc_mir_transform/src/simplify.rs

@@ -143,7 +143,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
         // statements itself to avoid moving the (relatively) large statements twice.
         // We do not push the statements directly into the target block (`bb`) as that is slower
         // due to additional reallocations
-        let mut merged_blocks = Vec::new();
+        let mut merged_blocks: Vec<BasicBlock> = Vec::new();
         let mut outer_changed = false;
         loop {
             let mut changed = false;
@@ -158,8 +158,21 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 let mut terminator =
                     self.basic_blocks[bb].terminator.take().expect("invalid terminator state");
 
-                terminator
-                    .successors_mut(|successor| self.collapse_goto_chain(successor, &mut changed));
+                let mut debuginfos = Vec::new();
+                if let Some(mut unique_succ) = terminator.identical_successor() {
+                    self.collapse_goto_chain(&mut unique_succ, &mut changed, &mut debuginfos);
+                    terminator.successors_mut(|successor| {
+                        *successor = unique_succ;
+                    });
+                    // Add debugging information from the goto chain only when all successors are identical,
+                    // otherwise, we may provide misleading debugging information within a branch.
+                    self.basic_blocks[bb].after_last_stmt_debuginfos.append(&mut debuginfos);
+                } else {
+                    terminator.successors_mut(|successor| {
+                        debuginfos.clear();
+                        self.collapse_goto_chain(successor, &mut changed, &mut debuginfos)
+                    });
+                }
 
                 let mut inner_changed = true;
                 merged_blocks.clear();
@@ -176,10 +189,19 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 if statements_to_merge > 0 {
                     let mut statements = std::mem::take(&mut self.basic_blocks[bb].statements);
                     statements.reserve(statements_to_merge);
+                    let mut parent_bb_last_debuginfos =
+                        std::mem::take(&mut self.basic_blocks[bb].after_last_stmt_debuginfos);
                     for &from in &merged_blocks {
+                        if let Some(stmt) = self.basic_blocks[from].statements.first_mut() {
+                            parent_bb_last_debuginfos.append(&mut stmt.debuginfos);
+                            std::mem::swap(&mut parent_bb_last_debuginfos, &mut stmt.debuginfos);
+                        }
                         statements.append(&mut self.basic_blocks[from].statements);
+                        parent_bb_last_debuginfos
+                            .append(&mut self.basic_blocks[from].after_last_stmt_debuginfos);
                     }
                     self.basic_blocks[bb].statements = statements;
+                    self.basic_blocks[bb].after_last_stmt_debuginfos = parent_bb_last_debuginfos;
                 }
 
                 self.basic_blocks[bb].terminator = Some(terminator);
@@ -214,7 +236,12 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
     }
 
     /// Collapse a goto chain starting from `start`
-    fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
+    fn collapse_goto_chain(
+        &mut self,
+        start: &mut BasicBlock,
+        changed: &mut bool,
+        pred_debuginfos: &mut Vec<StmtDebugInfo<'tcx>>,
+    ) {
         // Using `SmallVec` here, because in some logs on libcore oli-obk saw many single-element
         // goto chains. We should probably benchmark different sizes.
         let mut terminators: SmallVec<[_; 1]> = Default::default();
@@ -233,6 +260,13 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
             else {
                 unreachable!();
             };
+            if *target != last {
+                self.basic_blocks[current]
+                    .after_last_stmt_debuginfos
+                    .extend_from_slice(pred_debuginfos);
+            }
+            pred_debuginfos
+                .extend_from_slice(&self.basic_blocks[current].after_last_stmt_debuginfos);
             *changed |= *target != last;
             *target = last;
             debug!("collapsing goto chain from {:?} to {:?}", current, target);

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

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

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

@@ -0,0 +1,30 @@
+- // MIR for `preserve_debuginfo_1` before SimplifyCfg-final
++ // MIR for `preserve_debuginfo_1` after SimplifyCfg-final
+  
+  fn preserve_debuginfo_1(_1: &Foo, _2: &mut 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: {
+-         goto -> bb1;
+-     }
+- 
+-     bb1: {
+          (*_2) = const true;
+          // DBG: _3 = &((*_1).0: i32)
+-         goto -> bb2;
+-     }
+- 
+-     bb2: {
+          // DBG: _4 = &((*_1).1: i64)
+          _0 = copy ((*_1).2: i32);
+          // DBG: _5 = &((*_1).2: i32)
+          return;
+      }
+  }
+  

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

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

tests/mir-opt/debuginfo/simplifycfg.rs

@@ -0,0 +1,126 @@
+//@ test-mir-pass: SimplifyCfg-final
+//@ compile-flags: -Zmir-enable-passes=+DeadStoreElimination-initial,+SimplifyConstCondition-final
+
+#![feature(core_intrinsics, custom_mir)]
+#![crate_type = "lib"]
+
+use std::intrinsics::mir::*;
+
+pub struct Foo {
+    a: i32,
+    b: i64,
+    c: i32,
+}
+
+// EMIT_MIR simplifycfg.drop_debuginfo.SimplifyCfg-final.diff
+#[custom_mir(dialect = "runtime")]
+pub fn drop_debuginfo(foo: &Foo, c: bool) -> i32 {
+    // CHECK-LABEL: fn drop_debuginfo
+    // CHECK: debug foo_b => [[foo_b:_[0-9]+]];
+    // CHECK: bb0: {
+    // CHECK-NEXT: DBG: [[foo_b]] = &((*_1).1: i64)
+    // CHECK-NEXT: _0 = copy ((*_1).2: i32);
+    // CHECK-NEXT: return;
+    mir! {
+        let _foo_a: &i32;
+        let _foo_b: &i64;
+        debug foo_a => _foo_a;
+        debug foo_b => _foo_b;
+        {
+            match c {
+                true => tmp,
+                _ => ret,
+            }
+        }
+        tmp = {
+            // Because we don't know if `c` is always true, we must drop this debuginfo.
+            _foo_a = &(*foo).a;
+            Goto(ret)
+        }
+        ret = {
+            _foo_b = &(*foo).b;
+            RET = (*foo).c;
+            Return()
+        }
+    }
+}
+
+// EMIT_MIR simplifycfg.preserve_debuginfo_1.SimplifyCfg-final.diff
+#[custom_mir(dialect = "runtime")]
+pub fn preserve_debuginfo_1(foo: &Foo, v: &mut bool) -> i32 {
+    // CHECK-LABEL: fn preserve_debuginfo_1
+    // 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: (*_2) = const true;
+    // CHECK-NEXT: DBG: [[foo_a]] = &((*_1).0: i32)
+    // CHECK-NEXT: DBG: [[foo_b]] = &((*_1).1: i64)
+    // CHECK-NEXT: _0 = copy ((*_1).2: i32);
+    // CHECK-NEXT: DBG: [[foo_c]] = &((*_1).2: 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 true {
+                true => tmp,
+                _ => ret,
+            }
+        }
+        tmp = {
+            *v = true;
+            _foo_a = &(*foo).a;
+            Goto(ret)
+        }
+        ret = {
+            _foo_b = &(*foo).b;
+            RET = (*foo).c;
+            _foo_c = &(*foo).c;
+            Return()
+        }
+    }
+}
+
+// EMIT_MIR simplifycfg.preserve_debuginfo_2.SimplifyCfg-final.diff
+#[custom_mir(dialect = "runtime")]
+pub fn preserve_debuginfo_2(foo: &Foo) -> i32 {
+    // CHECK-LABEL: fn preserve_debuginfo_2
+    // 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: DBG: [[foo_a]] = &((*_1).0: i32)
+    // CHECK-NEXT: DBG: [[foo_b]] = &((*_1).1: i64)
+    // CHECK-NEXT: _0 = copy ((*_1).2: i32);
+    // CHECK-NEXT: DBG: [[foo_c]] = &((*_1).2: 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 true {
+                true => tmp,
+                _ => ret,
+            }
+        }
+        tmp = {
+            _foo_a = &(*foo).a;
+            Goto(ret)
+        }
+        ret = {
+            _foo_b = &(*foo).b;
+            RET = (*foo).c;
+            _foo_c = &(*foo).c;
+            Return()
+        }
+    }
+}