Expose run_jit to allow for custom behaviour in JIT flow

src/base.rs

@@ -10,6 +10,7 @@ use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_index::IndexVec;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::InlineAsmMacro;
+use rustc_middle::mir::mono::{MonoItem, MonoItemData};
 use rustc_middle::ty::TypeVisitableExt;
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::layout::{FnAbiOf, HasTypingEnv};
@@ -22,6 +23,34 @@ use crate::inline_asm::codegen_naked_asm;
 use crate::prelude::*;
 use crate::pretty_clif::CommentWriter;
 
+pub(crate) struct MeasuremeProfiler(pub 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)
+    }
+}
+
 pub(crate) struct CodegenedFunction {
     symbol_name: String,
     func_id: FuncId,
@@ -30,6 +59,47 @@ pub(crate) struct CodegenedFunction {
     func_debug_cx: Option<FunctionDebugContext>,
 }
 
+pub(crate) fn predefine_mono_items<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    module: &mut dyn Module,
+    mono_items: &[(MonoItem<'tcx>, MonoItemData)],
+) {
+    tcx.prof.generic_activity("predefine functions").run(|| {
+        let is_compiler_builtins = tcx.is_compiler_builtins(LOCAL_CRATE);
+        for &(mono_item, data) in mono_items {
+            match mono_item {
+                MonoItem::Fn(instance) => {
+                    let name = tcx.symbol_name(instance).name;
+                    let _inst_guard = crate::PrintOnPanic(|| format!("{:?} {}", instance, name));
+                    let sig =
+                        get_function_sig(tcx, module.target_config().default_call_conv, instance);
+                    let linkage = crate::linkage::get_clif_linkage(
+                        mono_item,
+                        data.linkage,
+                        data.visibility,
+                        is_compiler_builtins,
+                    );
+                    let is_naked = tcx
+                        .codegen_fn_attrs(instance.def_id())
+                        .flags
+                        .contains(CodegenFnAttrFlags::NAKED);
+                    module
+                        .declare_function(
+                            name,
+                            // Naked functions are defined in a separate object
+                            // file from the codegen unit rustc expects them to
+                            // be defined in.
+                            if is_naked { Linkage::Import } else { linkage },
+                            &sig,
+                        )
+                        .unwrap();
+                }
+                MonoItem::Static(_) | MonoItem::GlobalAsm(_) => {}
+            }
+        }
+    });
+}
+
 pub(crate) fn codegen_fn<'tcx>(
     tcx: TyCtxt<'tcx>,
     cx: &mut crate::CodegenCx,

src/driver/aot.rs

@@ -27,7 +27,7 @@ use rustc_session::Session;
 use rustc_session::config::{DebugInfo, OutFileName, OutputFilenames, OutputType};
 
 use crate::CodegenCx;
-use crate::base::CodegenedFunction;
+use crate::base::{CodegenedFunction, MeasuremeProfiler, predefine_mono_items};
 use crate::concurrency_limiter::{ConcurrencyLimiter, ConcurrencyLimiterToken};
 use crate::debuginfo::TypeDebugContext;
 use crate::global_asm::GlobalAsmConfig;
@@ -511,7 +511,7 @@ fn codegen_cgu_content(
         cgu_name,
     );
     let mut type_dbg = TypeDebugContext::default();
-    super::predefine_mono_items(tcx, module, &mono_items);
+    predefine_mono_items(tcx, module, &mono_items);
     let mut codegened_functions = vec![];
     for (mono_item, _) in mono_items {
         match mono_item {
@@ -563,7 +563,7 @@ fn module_codegen(
 
     OngoingModuleCodegen::Async(std::thread::spawn(move || {
         profiler.clone().generic_activity_with_arg("compile functions", &*cgu_name).run(|| {
-            cranelift_codegen::timing::set_thread_profiler(Box::new(super::MeasuremeProfiler(
+            cranelift_codegen::timing::set_thread_profiler(Box::new(MeasuremeProfiler(
                 profiler.clone(),
             )));