Generate MIR pass names for profiling on the fly and pass the body DefId as argument

Author: Zoxc

compiler/rustc_middle/src/mir/mod.rs

@@ -25,6 +25,7 @@ use rustc_target::abi::{FieldIdx, Size, VariantIdx};
 
 use polonius_engine::Atom;
 pub use rustc_ast::Mutability;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::graph::dominators::Dominators;
 use rustc_index::{Idx, IndexSlice, IndexVec};
@@ -35,6 +36,8 @@ use rustc_span::{Span, DUMMY_SP};
 use either::Either;
 
 use std::borrow::Cow;
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
 use std::fmt::{self, Debug, Display, Formatter, Write};
 use std::ops::{Index, IndexMut};
 use std::{iter, mem};
@@ -97,6 +100,36 @@ impl<'tcx> HasLocalDecls<'tcx> for Body<'tcx> {
     }
 }
 
+thread_local! {
+    static PASS_NAMES: RefCell<FxHashMap<&'static str, &'static str>> = {
+        RefCell::new(FxHashMap::default())
+    };
+}
+
+/// Converts a MIR pass name into a snake case form to match the profiling naming style.
+fn to_profiler_name(type_name: &'static str) -> &'static str {
+    PASS_NAMES.with(|names| match names.borrow_mut().entry(type_name) {
+        Entry::Occupied(e) => *e.get(),
+        Entry::Vacant(e) => {
+            let snake_case: String = type_name
+                .chars()
+                .flat_map(|c| {
+                    if c.is_ascii_uppercase() {
+                        vec!['_', c.to_ascii_lowercase()]
+                    } else if c == '-' {
+                        vec!['_']
+                    } else {
+                        vec![c]
+                    }
+                })
+                .collect();
+            let result = &*String::leak(format!("mir_pass{}", snake_case));
+            e.insert(result);
+            result
+        }
+    })
+}
+
 /// A streamlined trait that you can implement to create a pass; the
 /// pass will be named after the type, and it will consist of a main
 /// loop that goes over each available MIR and applies `run_pass`.
@@ -106,6 +139,10 @@ pub trait MirPass<'tcx> {
         if let Some((_, tail)) = name.rsplit_once(':') { tail } else { name }
     }
 
+    fn profiler_name(&self) -> &'static str {
+        to_profiler_name(self.name())
+    }
+
     /// Returns `true` if this pass is enabled with the current combination of compiler flags.
     fn is_enabled(&self, _sess: &Session) -> bool {
         true

compiler/rustc_mir_transform/src/pass_manager.rs

@@ -94,6 +94,8 @@ fn run_passes_inner<'tcx>(
     let overridden_passes = &tcx.sess.opts.unstable_opts.mir_enable_passes;
     trace!(?overridden_passes);
 
+    let prof_arg = tcx.sess.prof.enabled().then(|| format!("{:?}", body.source.def_id()));
+
     if !body.should_skip() {
         for pass in passes {
             let name = pass.name();
@@ -121,10 +123,14 @@ fn run_passes_inner<'tcx>(
                 validate_body(tcx, body, format!("before pass {name}"));
             }
 
-            tcx.sess
-                .prof
-                .generic_activity_with_arg("mir_pass", name)
-                .run(|| pass.run_pass(tcx, body));
+            if let Some(prof_arg) = &prof_arg {
+                tcx.sess
+                    .prof
+                    .generic_activity_with_arg(pass.profiler_name(), &**prof_arg)
+                    .run(|| pass.run_pass(tcx, body));
+            } else {
+                pass.run_pass(tcx, body);
+            }
 
             if dump_enabled {
                 dump_mir_for_pass(tcx, body, &name, true);