Initialize the atomic mutex in a constructor for proc macros

Author: bjorn3

src/atomic_shim.rs

@@ -10,7 +10,7 @@ use crate::prelude::*;
 pub static mut __cg_clif_global_atomic_mutex: libc::pthread_mutex_t = libc::PTHREAD_MUTEX_INITIALIZER;
 
 pub(crate) fn init_global_lock(module: &mut Module<impl Backend>, bcx: &mut FunctionBuilder<'_>) {
-    if std::env::var("CG_CLIF_JIT").is_ok () {
+    if std::env::var("CG_CLIF_JIT").is_ok() {
         // When using JIT, dylibs won't find the __cg_clif_global_atomic_mutex data object defined here,
         // so instead define it in the cg_clif dylib.
 
@@ -47,6 +47,39 @@ pub(crate) fn init_global_lock(module: &mut Module<impl Backend>, bcx: &mut Func
     bcx.ins().call(pthread_mutex_init, &[atomic_mutex, nullptr]);
 }
 
+pub(crate) fn init_global_lock_constructor(
+    module: &mut Module<impl Backend>,
+    constructor_name: &str
+) -> FuncId {
+    let sig = Signature::new(CallConv::SystemV);
+    let init_func_id = module
+        .declare_function(constructor_name, Linkage::Export, &sig)
+        .unwrap();
+
+    let mut ctx = Context::new();
+    ctx.func = Function::with_name_signature(ExternalName::user(0, 0), sig);
+    {
+        let mut func_ctx = FunctionBuilderContext::new();
+        let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
+
+        let block = bcx.create_block();
+        bcx.switch_to_block(block);
+
+        crate::atomic_shim::init_global_lock(module, &mut bcx);
+
+        bcx.ins().return_(&[]);
+        bcx.seal_all_blocks();
+        bcx.finalize();
+    }
+    module.define_function(
+        init_func_id,
+        &mut ctx,
+        &mut cranelift_codegen::binemit::NullTrapSink {},
+    ).unwrap();
+
+    init_func_id
+}
+
 pub(crate) fn lock_global_lock(fx: &mut FunctionCx<'_, '_, impl Backend>) {
     let atomic_mutex = fx.module.declare_data(
         "__cg_clif_global_atomic_mutex",

src/backend.rs

@@ -97,6 +97,32 @@ impl WriteDebugInfo for ObjectProduct {
     }
 }
 
+// FIXME remove once atomic instructions are implemented in Cranelift.
+pub(crate) trait AddConstructor {
+    fn add_constructor(&mut self, func_id: FuncId);
+}
+
+impl AddConstructor for ObjectProduct {
+    fn add_constructor(&mut self, func_id: FuncId) {
+        let symbol = self.function_symbol(func_id);
+        let segment = self.object.segment_name(object::write::StandardSegment::Data);
+        let init_array_section = self.object.add_section(segment.to_vec(), b".init_array".to_vec(), SectionKind::Data);
+        self.object.append_section_data(
+            init_array_section,
+            &std::iter::repeat(0).take(8 /*FIXME pointer size*/).collect::<Vec<u8>>(),
+            8,
+        );
+        self.object.add_relocation(init_array_section, object::write::Relocation {
+            offset: 0,
+            size: 64, // FIXME pointer size
+            kind: RelocationKind::Absolute,
+            encoding: RelocationEncoding::Generic,
+            symbol,
+            addend: 0,
+        }).unwrap();
+    }
+}
+
 pub(crate) trait Emit {
     fn emit(self) -> Vec<u8>;
 }
@@ -140,7 +166,7 @@ pub(crate) fn with_object(sess: &Session, name: &str, f: impl FnOnce(&mut Object
     metadata_object.write().unwrap()
 }
 
-pub(crate) type Backend = impl cranelift_module::Backend<Product: Emit + WriteDebugInfo>;
+pub(crate) type Backend = impl cranelift_module::Backend<Product: AddConstructor + Emit + WriteDebugInfo>;
 
 pub(crate) fn make_module(sess: &Session, name: String) -> Module<Backend> {
     let module: Module<ObjectBackend> = Module::new(

src/driver/aot.rs

@@ -11,7 +11,7 @@ use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 
 use crate::prelude::*;
 
-use crate::backend::{Emit, WriteDebugInfo};
+use crate::backend::{AddConstructor, Emit, WriteDebugInfo};
 
 fn new_module(tcx: TyCtxt<'_>, name: String) -> Module<crate::backend::Backend> {
     let module = crate::backend::make_module(tcx.sess, name);
@@ -35,8 +35,9 @@ fn emit_module<B: Backend>(
     mut module: Module<B>,
     debug: Option<DebugContext<'_>>,
     unwind_context: UnwindContext<'_>,
+    map_product: impl FnOnce(B::Product) -> B::Product,
 ) -> ModuleCodegenResult
-    where B::Product: Emit + WriteDebugInfo,
+    where B::Product: AddConstructor + Emit + WriteDebugInfo,
 {
     module.finalize_definitions();
     let mut product = module.finish();
@@ -47,6 +48,8 @@ fn emit_module<B: Backend>(
 
     unwind_context.emit(&mut product);
 
+    let product = map_product(product);
+
     let tmp_file = tcx
         .output_filenames(LOCAL_CRATE)
         .temp_path(OutputType::Object, Some(&name));
@@ -110,7 +113,23 @@ fn module_codegen(tcx: TyCtxt<'_>, cgu_name: rustc_span::Symbol) -> ModuleCodege
     let cgu = tcx.codegen_unit(cgu_name);
     let mono_items = cgu.items_in_deterministic_order(tcx);
 
-    let module = new_module(tcx, cgu_name.as_str().to_string());
+    let mut module = new_module(tcx, cgu_name.as_str().to_string());
+
+    // Initialize the global atomic mutex using a constructor for proc-macros.
+    // FIXME implement atomic instructions in Cranelift.
+    let mut init_atomics_mutex_from_constructor = None;
+    if tcx.sess.crate_types().contains(&rustc_session::config::CrateType::ProcMacro) {
+        if mono_items.iter().any(|(mono_item, _)| {
+            match mono_item {
+                rustc_middle::mir::mono::MonoItem::Static(def_id) => {
+                    tcx.symbol_name(Instance::mono(tcx, *def_id)).name.as_str().contains("__rustc_proc_macro_decls_")
+                }
+                _ => false,
+            }
+        }) {
+            init_atomics_mutex_from_constructor = Some(crate::atomic_shim::init_global_lock_constructor(&mut module, &format!("{}_init_atomics_mutex", cgu_name.as_str())));
+        }
+    }
 
     let mut cx = crate::CodegenCx::new(tcx, module, tcx.sess.opts.debuginfo != DebugInfo::None);
     super::codegen_mono_items(&mut cx, mono_items);
@@ -124,6 +143,13 @@ fn module_codegen(tcx: TyCtxt<'_>, cgu_name: rustc_span::Symbol) -> ModuleCodege
         module,
         debug,
         unwind_context,
+        |mut product| {
+            if let Some(func_id) = init_atomics_mutex_from_constructor {
+                product.add_constructor(func_id);
+            }
+
+            product
+        }
     );
 
     codegen_global_asm(tcx, &cgu.name().as_str(), &global_asm);
@@ -196,6 +222,7 @@ pub(super) fn run_aot(
             allocator_module,
             None,
             allocator_unwind_context,
+            |product| product,
         );
         if let Some((id, product)) = work_product {
             work_products.insert(id, product);