Integrate better with Cranelift's profiling infrastructure

4 files changed, 37 insertions, 21 deletions

diff --git a/src/base.rs b/src/base.rs
index 292f24263ac..a259a4f30b2 100644
--- a/src/base.rs
+++ b/src/base.rs
@@ -172,27 +172,6 @@ pub(crate) fn compile_fn(
     cx.profiler.generic_activity("define function").run(|| {
         context.want_disasm = cx.should_write_ir;
         module.define_function(codegened_func.func_id, context).unwrap();
-
-        if cx.profiler.enabled() {
-            let mut recording_args = false;
-            cx.profiler
-                .generic_activity_with_arg_recorder(
-                    "define function (clif pass timings)",
-                    |recorder| {
-                        let pass_times = cranelift_codegen::timing::take_current();
-                        // Replace newlines with | as measureme doesn't allow control characters like
-                        // newlines inside strings.
-                        recorder.record_arg(format!("{}", pass_times).replace('\n', " | "));
-                        recording_args = true;
-                    },
-                )
-                .run(|| {
-                    if recording_args {
-                        // Wait a tiny bit to ensure chrome's profiler doesn't hide the event
-                        std::thread::sleep(std::time::Duration::from_nanos(2))
-                    }
-                });
-        }
     });
 
     if cx.should_write_ir {
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index 762976c81f8..0e6c6ad95aa 100644
--- a/src/driver/aot.rs
+++ b/src/driver/aot.rs
@@ -324,6 +324,10 @@ fn module_codegen(
     OngoingModuleCodegen::Async(std::thread::spawn(move || {
         cx.profiler.clone().verbose_generic_activity_with_arg("compile functions", &*cgu_name).run(
             || {
+                cranelift_codegen::timing::set_thread_profiler(Box::new(super::MeasuremeProfiler(
+                    cx.profiler.clone(),
+                )));
+
                 let mut cached_context = Context::new();
                 for codegened_func in codegened_functions {
                     crate::base::compile_fn(
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index f6a48e3257b..3118105a4e2 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -224,6 +224,10 @@ pub(crate) fn codegen_and_compile_fn<'tcx>(
     module: &mut dyn Module,
     instance: Instance<'tcx>,
 ) {
+    cranelift_codegen::timing::set_thread_profiler(Box::new(super::MeasuremeProfiler(
+        cx.profiler.clone(),
+    )));
+
     tcx.prof.generic_activity("codegen and compile fn").run(|| {
         let _inst_guard =
             crate::PrintOnPanic(|| format!("{:?} {}", instance, tcx.symbol_name(instance).name));
diff --git a/src/driver/mod.rs b/src/driver/mod.rs
index d09d3a52975..5c52c9c18ad 100644
--- a/src/driver/mod.rs
+++ b/src/driver/mod.rs
@@ -4,6 +4,7 @@
 //! [`codegen_fn`]: crate::base::codegen_fn
 //! [`codegen_static`]: crate::constant::codegen_static
 
+use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_middle::mir::mono::{Linkage as RLinkage, MonoItem, Visibility};
 
 use crate::prelude::*;
@@ -39,3 +40,31 @@ fn predefine_mono_items<'tcx>(
         }
     });
 }
+
+struct MeasuremeProfiler(SelfProfilerRef);
+
+struct TimingGuard {
+    profiler: std::mem::ManuallyDrop<SelfProfilerRef>,
+    inner: Option<rustc_data_structures::profiling::TimingGuard<'static>>,
+}
+
+impl Drop for TimingGuard {
+    fn drop(&mut self) {
+        self.inner.take();
+        unsafe {
+            std::mem::ManuallyDrop::drop(&mut self.profiler);
+        }
+    }
+}
+
+impl cranelift_codegen::timing::Profiler for MeasuremeProfiler {
+    fn start_pass(&self, pass: cranelift_codegen::timing::Pass) -> Box<dyn std::any::Any> {
+        let mut timing_guard =
+            TimingGuard { profiler: std::mem::ManuallyDrop::new(self.0.clone()), inner: None };
+        timing_guard.inner = Some(
+            unsafe { &*(&*timing_guard.profiler as &SelfProfilerRef as *const SelfProfilerRef) }
+                .generic_activity(pass.description()),
+        );
+        Box::new(timing_guard)
+    }
+}