simplifycfg: Preserve debuginfos when merging bbs

Author: dianqk

compiler/rustc_middle/src/mir/terminator.rs

@@ -444,6 +444,14 @@ 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 first_succ = successors.next()?;
+        if successors.all(|succ| first_succ == succ) { Some(first_succ) } else { None }
+    }
+
     #[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

@@ -144,7 +144,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;
@@ -159,8 +159,9 @@ 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));
+                terminator.successors_mut(|successor| {
+                    self.collapse_goto_chain(successor, &mut changed);
+                });
 
                 let mut inner_changed = true;
                 merged_blocks.clear();
@@ -177,10 +178,18 @@ 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() {
+                            stmt.debuginfos.prepend(&mut parent_bb_last_debuginfos);
+                        }
                         statements.append(&mut self.basic_blocks[from].statements);
+                        parent_bb_last_debuginfos =
+                            std::mem::take(&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);
@@ -220,10 +229,14 @@ 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;
         }
@@ -235,6 +248,17 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
             else {
                 unreachable!();
             };
+            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);
+            }
             *changed |= *target != last;
             *target = last;
             debug!("collapsing goto chain from {:?} to {:?}", current, target);

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

@@ -14,11 +14,13 @@
 - 
 -     bb1: {
 -         // DBG: _3 = &((*_1).0: i32);
+-         nop;
 -         goto -> bb2;
 -     }
 - 
 -     bb2: {
           // DBG: _4 = &((*_1).1: i64);
+-         nop;
           _0 = copy ((*_1).2: i32);
           return;
       }

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

@@ -17,13 +17,16 @@
 -     bb1: {
           (*_2) = const true;
           // DBG: _3 = &((*_1).0: i32);
+-         nop;
 -         goto -> bb2;
 -     }
 - 
 -     bb2: {
           // DBG: _4 = &((*_1).1: i64);
+-         nop;
           _0 = copy ((*_1).2: i32);
           // DBG: _5 = &((*_1).2: i32);
+-         nop;
           return;
       }
   }

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

@@ -16,13 +16,16 @@
 - 
 -     bb1: {
           // DBG: _2 = &((*_1).0: i32);
+-         nop;
 -         goto -> bb2;
 -     }
 - 
 -     bb2: {
           // DBG: _3 = &((*_1).1: i64);
+-         nop;
           _0 = copy ((*_1).2: i32);
           // DBG: _4 = &((*_1).2: i32);
+-         nop;
           return;
       }
   }

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

@@ -0,0 +1,40 @@
+- // 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);
+-         nop;
+-         goto -> bb3;
+-     }
+- 
+-     bb2: {
+          // DBG: _4 = &((*_1).1: i64);
+-         nop;
+          _0 = copy ((*_1).2: i32);
+          return;
+      }
+  
+-     bb3: {
++     bb2: {
++         // DBG: _3 = &((*_1).0: i32);
+          // DBG: _5 = &((*_1).2: i32);
+-         nop;
+          _0 = copy ((*_1).0: i32);
+          return;
+      }
+  }
+  

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

@@ -0,0 +1,32 @@
+- // MIR for `preserve_debuginfo_identical_succs` before SimplifyCfg-final
++ // MIR for `preserve_debuginfo_identical_succs` after SimplifyCfg-final
+  
+  fn preserve_debuginfo_identical_succs(_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: bb1];
+-     }
+- 
+-     bb1: {
+          // DBG: _3 = &((*_1).0: i32);
+-         nop;
+-         goto -> bb2;
+-     }
+- 
+-     bb2: {
+          // DBG: _4 = &((*_1).1: i64);
+-         nop;
+          _0 = copy ((*_1).2: i32);
+          // DBG: _5 = &((*_1).2: i32);
+-         nop;
+          return;
+      }
+  }
+