mir: Simplify the handling of debuginfos

Author: dianqk

compiler/rustc_codegen_ssa/src/mir/statement.rs

@@ -1,4 +1,4 @@
-use rustc_middle::mir::{self, NonDivergingIntrinsic, StmtDebugInfo};
+use rustc_middle::mir::{self, NonDivergingIntrinsic, StmtDebugInfo, StmtDebugInfos};
 use rustc_middle::{bug, span_bug};
 use tracing::instrument;
 
@@ -165,9 +165,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     pub(crate) fn codegen_stmt_debuginfos(
         &mut self,
         bx: &mut Bx,
-        debuginfos: &[StmtDebugInfo<'tcx>],
+        debuginfos: &StmtDebugInfos<'tcx>,
     ) {
-        for debuginfo in debuginfos {
+        for debuginfo in debuginfos.iter() {
             self.codegen_stmt_debuginfo(bx, debuginfo);
         }
     }

compiler/rustc_middle/src/mir/mod.rs

@@ -4,8 +4,8 @@
 
 use std::borrow::Cow;
 use std::fmt::{self, Debug, Formatter};
+use std::iter;
 use std::ops::{Index, IndexMut};
-use std::{iter, mem};
 
 pub use basic_blocks::{BasicBlocks, SwitchTargetValue};
 use either::Either;
@@ -1344,7 +1344,7 @@ pub struct BasicBlockData<'tcx> {
 
     /// All debuginfos happen before the statement.
     /// Put debuginfos here when the last statement is eliminated.
-    pub after_last_stmt_debuginfos: Vec<StmtDebugInfo<'tcx>>,
+    pub after_last_stmt_debuginfos: StmtDebugInfos<'tcx>,
 
     /// Terminator for this block.
     ///
@@ -1375,7 +1375,7 @@ impl<'tcx> BasicBlockData<'tcx> {
     ) -> BasicBlockData<'tcx> {
         BasicBlockData {
             statements,
-            after_last_stmt_debuginfos: Vec::new(),
+            after_last_stmt_debuginfos: StmtDebugInfos::default(),
             terminator,
             is_cleanup,
         }
@@ -1418,41 +1418,27 @@ impl<'tcx> BasicBlockData<'tcx> {
     where
         F: FnMut(&Statement<'tcx>) -> bool,
     {
-        // Fast path without debugging information
-        if self.statements.iter().all(|stmt| stmt.debuginfos.is_empty()) {
-            self.statements.retain_mut(|stmt| f(stmt));
-            return;
-        }
-        let mut debuginfos = Vec::new();
+        let mut debuginfos = StmtDebugInfos::default();
         self.statements.retain_mut(|stmt| {
             let retain = f(stmt);
             if retain {
-                if !debuginfos.is_empty() {
-                    if !stmt.debuginfos.is_empty() {
-                        debuginfos.append(&mut stmt.debuginfos);
-                    }
-                    mem::swap(&mut debuginfos, &mut stmt.debuginfos);
-                }
+                stmt.debuginfos.prepend(&mut debuginfos);
             } else {
-                if !stmt.debuginfos.is_empty() {
-                    debuginfos.append(&mut stmt.debuginfos);
-                }
+                debuginfos.append(&mut stmt.debuginfos);
             }
             retain
         });
-        if !debuginfos.is_empty() {
-            self.after_last_stmt_debuginfos.append(&mut debuginfos);
-        }
+        self.after_last_stmt_debuginfos.prepend(&mut debuginfos);
     }
 
     pub fn strip_nops(&mut self) {
         self.retain_statements(|stmt| !matches!(stmt.kind, StatementKind::Nop))
     }
 
     pub fn drop_debuginfo(&mut self) {
-        self.after_last_stmt_debuginfos = Vec::new();
+        self.after_last_stmt_debuginfos.drop_debuginfo();
         for stmt in self.statements.iter_mut() {
-            stmt.debuginfos = Vec::new();
+            stmt.debuginfos.drop_debuginfo();
         }
     }
 }

compiler/rustc_middle/src/mir/statement.rs

@@ -1,5 +1,7 @@
 //! Functionality for statements, operands, places, and things that appear in them.
 
+use std::ops;
+
 use tracing::{debug, instrument};
 
 use super::interpret::GlobalAlloc;
@@ -15,7 +17,7 @@ use crate::ty::CoroutineArgsExt;
 pub struct Statement<'tcx> {
     pub source_info: SourceInfo,
     pub kind: StatementKind<'tcx>,
-    pub debuginfos: Vec<StmtDebugInfo<'tcx>>,
+    pub debuginfos: StmtDebugInfos<'tcx>,
 }
 
 impl<'tcx> Statement<'tcx> {
@@ -35,7 +37,7 @@ impl<'tcx> Statement<'tcx> {
     }
 
     pub fn new(source_info: SourceInfo, kind: StatementKind<'tcx>) -> Self {
-        Statement { source_info, kind, debuginfos: Vec::new() }
+        Statement { source_info, kind, debuginfos: StmtDebugInfos::default() }
     }
 }
 
@@ -961,6 +963,76 @@ impl RawPtrKind {
     }
 }
 
+#[derive(Default, Debug, Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub struct StmtDebugInfos<'tcx>(Vec<StmtDebugInfo<'tcx>>);
+
+impl<'tcx> StmtDebugInfos<'tcx> {
+    pub fn push(&mut self, debuginfo: StmtDebugInfo<'tcx>) {
+        self.0.push(debuginfo);
+    }
+
+    pub fn drop_debuginfo(&mut self) {
+        self.0.clear();
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    pub fn prepend(&mut self, debuginfos: &mut Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        if self.is_empty() {
+            std::mem::swap(self, debuginfos);
+            return;
+        };
+        debuginfos.0.append(self);
+        std::mem::swap(debuginfos, self);
+    }
+
+    pub fn append(&mut self, debuginfos: &mut Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        if self.is_empty() {
+            std::mem::swap(self, debuginfos);
+            return;
+        };
+        self.0.append(debuginfos);
+    }
+
+    pub fn extend(&mut self, debuginfos: &Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        self.0.extend_from_slice(debuginfos);
+    }
+
+    pub fn retain<F>(&mut self, f: F)
+    where
+        F: FnMut(&StmtDebugInfo<'tcx>) -> bool,
+    {
+        self.0.retain(f);
+    }
+}
+
+impl<'tcx> ops::Deref for StmtDebugInfos<'tcx> {
+    type Target = Vec<StmtDebugInfo<'tcx>>;
+
+    #[inline]
+    fn deref(&self) -> &Vec<StmtDebugInfo<'tcx>> {
+        &self.0
+    }
+}
+
+impl<'tcx> ops::DerefMut for StmtDebugInfos<'tcx> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Vec<StmtDebugInfo<'tcx>> {
+        &mut self.0
+    }
+}
+
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
 pub enum StmtDebugInfo<'tcx> {
     AssignRef(Local, Option<Place<'tcx>>),

compiler/rustc_mir_transform/src/simplify.rs

@@ -158,20 +158,16 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                 let mut terminator =
                     self.basic_blocks[bb].terminator.take().expect("invalid terminator state");
 
-                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;
-                    });
+                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.
-                    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 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;
@@ -193,12 +189,11 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
                         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);
+                            stmt.debuginfos.prepend(&mut parent_bb_last_debuginfos);
                         }
                         statements.append(&mut self.basic_blocks[from].statements);
-                        parent_bb_last_debuginfos
-                            .append(&mut self.basic_blocks[from].after_last_stmt_debuginfos);
+                        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;
@@ -236,12 +231,7 @@ 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,
-        pred_debuginfos: &mut Vec<StmtDebugInfo<'tcx>>,
-    ) {
+    fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
         // 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();
@@ -254,20 +244,20 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
             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 {
-                self.basic_blocks[current]
-                    .after_last_stmt_debuginfos
-                    .extend_from_slice(pred_debuginfos);
+                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);
             }
-            pred_debuginfos
-                .extend_from_slice(&self.basic_blocks[current].after_last_stmt_debuginfos);
-            *changed |= *target != last;
+            succ = current;
             *target = last;
             debug!("collapsing goto chain from {:?} to {:?}", current, target);