Merge pull request rust-lang#1062 from bjorn3/global_asm

Implement global_asm! using an external assembler

Author: bjorn3

example/mini_core.rs

@@ -562,6 +562,10 @@ pub macro line() { /* compiler built-in */ }
 #[rustc_macro_transparency = "semitransparent"]
 pub macro cfg() { /* compiler built-in */ }
 
+#[rustc_builtin_macro]
+#[rustc_macro_transparency = "semitransparent"]
+pub macro global_asm() { /* compiler built-in */ }
+
 pub static A_STATIC: u8 = 42;
 
 #[lang = "panic_location"]

example/mini_core_hello_world.rs

@@ -284,6 +284,26 @@ fn main() {
 
     #[cfg(not(jit))]
     test_tls();
+
+    #[cfg(all(not(jit), target_os = "linux"))]
+    unsafe {
+        global_asm_test();
+    }
+}
+
+#[cfg(all(not(jit), target_os = "linux"))]
+extern "C" {
+    fn global_asm_test();
+}
+
+#[cfg(all(not(jit), target_os = "linux"))]
+global_asm! {
+    "
+    .global global_asm_test
+    global_asm_test:
+    // comment that would normally be removed by LLVM
+    ret
+    "
 }
 
 #[repr(C)]

src/archive.rs

@@ -220,8 +220,10 @@ impl<'a> ArchiveBuilder<'a> for ArArchiveBuilder<'a> {
         std::mem::drop(builder);
 
         if self.update_symbols {
+            let ranlib = crate::toolchain::get_toolchain_binary(self.config.sess, "ranlib");
+
             // Run ranlib to be able to link the archive
-            let status = std::process::Command::new("ranlib")
+            let status = std::process::Command::new(ranlib)
                 .arg(self.config.dst)
                 .status()
                 .expect("Couldn't run ranlib");

src/driver/aot.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::middle::cstore::EncodedMetadata;
 use rustc_middle::mir::mono::CodegenUnit;
@@ -110,19 +112,33 @@ fn module_codegen(tcx: TyCtxt<'_>, cgu_name: rustc_span::Symbol) -> ModuleCodege
 
     let module = new_module(tcx, cgu_name.as_str().to_string());
 
+    let mut global_asm = Vec::new();
     let mut cx = crate::CodegenCx::new(tcx, module, tcx.sess.opts.debuginfo != DebugInfo::None);
-    super::codegen_mono_items(&mut cx, mono_items);
+    super::codegen_mono_items(&mut cx, &mut global_asm, mono_items);
     let (mut module, debug, mut unwind_context) = tcx.sess.time("finalize CodegenCx", || cx.finalize());
     crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut unwind_context);
 
-    emit_module(
+    let global_asm = global_asm.into_iter().map(|hir_id| {
+        let item = tcx.hir().expect_item(hir_id);
+        if let rustc_hir::ItemKind::GlobalAsm(rustc_hir::GlobalAsm { asm }) = item.kind {
+            asm.as_str().to_string()
+        } else {
+            bug!("Expected GlobalAsm found {:?}", item);
+        }
+    }).collect::<Vec<String>>().join("\n");
+
+    let codegen_result = emit_module(
         tcx,
         cgu.name().as_str().to_string(),
         ModuleKind::Regular,
         module,
         debug,
         unwind_context,
-    )
+    );
+
+    codegen_global_asm(tcx, &cgu.name().as_str(), &global_asm);
+
+    codegen_result
 }
 
 pub(super) fn run_aot(
@@ -253,6 +269,85 @@ pub(super) fn run_aot(
     }, work_products))
 }
 
+fn codegen_global_asm(tcx: TyCtxt<'_>, cgu_name: &str, global_asm: &str) {
+    use std::io::Write;
+    use std::process::{Command, Stdio};
+
+    if global_asm.is_empty() {
+        return;
+    }
+
+    if tcx.sess.target.target.options.is_like_osx || tcx.sess.target.target.options.is_like_windows {
+        if global_asm.contains("__rust_probestack") {
+            return;
+        }
+
+        // FIXME fix linker error on macOS
+        tcx.sess.fatal("global_asm! is not yet supported on macOS and Windows");
+    }
+
+    let assembler = crate::toolchain::get_toolchain_binary(tcx.sess, "as");
+    let linker = crate::toolchain::get_toolchain_binary(tcx.sess, "ld");
+
+    // Remove all LLVM style comments
+    let global_asm = global_asm.lines().map(|line| {
+        if let Some(index) = line.find("//") {
+            &line[0..index]
+        } else {
+            line
+        }
+    }).collect::<Vec<_>>().join("\n");
+
+    let output_object_file = tcx
+        .output_filenames(LOCAL_CRATE)
+        .temp_path(OutputType::Object, Some(cgu_name));
+
+    // Assemble `global_asm`
+    let global_asm_object_file = add_file_stem_postfix(output_object_file.clone(), ".asm");
+    let mut child = Command::new(assembler)
+        .arg("-o").arg(&global_asm_object_file)
+        .stdin(Stdio::piped())
+        .spawn()
+        .expect("Failed to spawn `as`.");
+    child.stdin.take().unwrap().write_all(global_asm.as_bytes()).unwrap();
+    let status = child.wait().expect("Failed to wait for `as`.");
+    if !status.success() {
+        tcx.sess.fatal(&format!("Failed to assemble `{}`", global_asm));
+    }
+
+    // Link the global asm and main object file together
+    let main_object_file = add_file_stem_postfix(output_object_file.clone(), ".main");
+    std::fs::rename(&output_object_file, &main_object_file).unwrap();
+    let status = Command::new(linker)
+        .arg("-r") // Create a new object file
+        .arg("-o").arg(output_object_file)
+        .arg(&main_object_file)
+        .arg(&global_asm_object_file)
+        .status()
+        .unwrap();
+    if !status.success() {
+        tcx.sess.fatal(&format!(
+            "Failed to link `{}` and `{}` together",
+            main_object_file.display(),
+            global_asm_object_file.display(),
+        ));
+    }
+
+    std::fs::remove_file(global_asm_object_file).unwrap();
+    std::fs::remove_file(main_object_file).unwrap();
+}
+
+fn add_file_stem_postfix(mut path: PathBuf, postfix: &str) -> PathBuf {
+    let mut new_filename = path.file_stem().unwrap().to_owned();
+    new_filename.push(postfix);
+    if let Some(extension) = path.extension() {
+        new_filename.push(".");
+        new_filename.push(extension);
+    }
+    path.set_file_name(new_filename);
+    path
+}
+
 // Adapted from https://github.com/rust-lang/rust/blob/303d8aff6092709edd4dbd35b1c88e9aa40bf6d8/src/librustc_codegen_ssa/base.rs#L922-L953
 fn determine_cgu_reuse<'tcx>(tcx: TyCtxt<'tcx>, cgu: &CodegenUnit<'tcx>) -> CguReuse {
     if !tcx.dep_graph.is_fully_enabled() {

src/driver/jit.rs

@@ -55,7 +55,12 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
     let mut cx = crate::CodegenCx::new(tcx, jit_module, false);
 
     let (mut jit_module, _debug, mut unwind_context) = super::time(tcx, "codegen mono items", || {
-        super::codegen_mono_items(&mut cx, mono_items);
+        let mut global_asm = Vec::new();
+        super::codegen_mono_items(&mut cx, &mut global_asm, mono_items);
+        for hir_id in global_asm {
+            let item = tcx.hir().expect_item(hir_id);
+            tcx.sess.span_err(item.span, "Global asm is not supported in JIT mode");
+        }
         tcx.sess.time("finalize CodegenCx", || cx.finalize())
     });
     crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);

src/driver/mod.rs

@@ -1,5 +1,6 @@
 use std::any::Any;
 
+use rustc_hir::HirId;
 use rustc_middle::middle::cstore::EncodedMetadata;
 use rustc_middle::mir::mono::{Linkage as RLinkage, MonoItem, Visibility};
 
@@ -31,6 +32,7 @@ pub(crate) fn codegen_crate(
 
 fn codegen_mono_items<'tcx>(
     cx: &mut crate::CodegenCx<'tcx, impl Backend + 'static>,
+    global_asm: &mut Vec<HirId>,
     mono_items: Vec<(MonoItem<'tcx>, (RLinkage, Visibility))>,
 ) {
     cx.tcx.sess.time("predefine functions", || {
@@ -49,12 +51,13 @@ fn codegen_mono_items<'tcx>(
 
     for (mono_item, (linkage, visibility)) in mono_items {
         let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
-        trans_mono_item(cx, mono_item, linkage);
+        trans_mono_item(cx, global_asm, mono_item, linkage);
     }
 }
 
 fn trans_mono_item<'tcx, B: Backend + 'static>(
     cx: &mut crate::CodegenCx<'tcx, B>,
+    global_asm: &mut Vec<HirId>,
     mono_item: MonoItem<'tcx>,
     linkage: Linkage,
 ) {
@@ -91,19 +94,7 @@ fn trans_mono_item<'tcx, B: Backend + 'static>(
             crate::constant::codegen_static(&mut cx.constants_cx, def_id);
         }
         MonoItem::GlobalAsm(hir_id) => {
-            let item = tcx.hir().expect_item(hir_id);
-            if let rustc_hir::ItemKind::GlobalAsm(rustc_hir::GlobalAsm { asm }) = item.kind {
-                // FIXME implement global asm using an external assembler
-                if asm.as_str().contains("__rust_probestack") {
-                    return;
-                } else {
-                    tcx
-                        .sess
-                        .fatal(&format!("Unimplemented global asm mono item \"{}\"", asm));
-                }
-            } else {
-                bug!("Expected GlobalAsm found {:?}", item);
-            }
+            global_asm.push(hir_id);
         }
     }
 }

src/lib.rs

@@ -64,6 +64,7 @@ mod optimize;
 mod pointer;
 mod pretty_clif;
 mod target_features_whitelist;
+mod toolchain;
 mod trap;
 mod unsize;
 mod value_and_place;

src/toolchain.rs

@@ -0,0 +1,115 @@
+use std::path::PathBuf;
+
+use rustc_middle::bug;
+use rustc_session::Session;
+use rustc_target::spec::LinkerFlavor;
+
+/// Tries to infer the path of a binary for the target toolchain from the linker name.
+pub(crate) fn get_toolchain_binary(sess: &Session, tool: &str) -> PathBuf {
+    let (mut linker, _linker_flavor) = linker_and_flavor(sess);
+    let linker_file_name = linker.file_name().and_then(|name| name.to_str()).unwrap_or_else(|| {
+        sess.fatal("couldn't extract file name from specified linker")
+    });
+
+    if linker_file_name == "ld.lld" {
+        if tool != "ld" {
+            linker.set_file_name(tool)
+        }
+    } else {
+        let tool_file_name = linker_file_name
+            .replace("ld", tool)
+            .replace("gcc", tool)
+            .replace("clang", tool)
+            .replace("cc", tool);
+
+        linker.set_file_name(tool_file_name)
+    }
+
+    linker
+}
+
+// Adapted from https://github.com/rust-lang/rust/blob/5db778affee7c6600c8e7a177c48282dab3f6292/src/librustc_codegen_ssa/back/link.rs#L848-L931
+fn linker_and_flavor(sess: &Session) -> (PathBuf, LinkerFlavor) {
+    fn infer_from(
+        sess: &Session,
+        linker: Option<PathBuf>,
+        flavor: Option<LinkerFlavor>,
+    ) -> Option<(PathBuf, LinkerFlavor)> {
+        match (linker, flavor) {
+            (Some(linker), Some(flavor)) => Some((linker, flavor)),
+            // only the linker flavor is known; use the default linker for the selected flavor
+            (None, Some(flavor)) => Some((
+                PathBuf::from(match flavor {
+                    LinkerFlavor::Em => {
+                        if cfg!(windows) {
+                            "emcc.bat"
+                        } else {
+                            "emcc"
+                        }
+                    }
+                    LinkerFlavor::Gcc => {
+                        if cfg!(any(target_os = "solaris", target_os = "illumos")) {
+                            // On historical Solaris systems, "cc" may have
+                            // been Sun Studio, which is not flag-compatible
+                            // with "gcc".  This history casts a long shadow,
+                            // and many modern illumos distributions today
+                            // ship GCC as "gcc" without also making it
+                            // available as "cc".
+                            "gcc"
+                        } else {
+                            "cc"
+                        }
+                    }
+                    LinkerFlavor::Ld => "ld",
+                    LinkerFlavor::Msvc => "link.exe",
+                    LinkerFlavor::Lld(_) => "lld",
+                    LinkerFlavor::PtxLinker => "rust-ptx-linker",
+                }),
+                flavor,
+            )),
+            (Some(linker), None) => {
+                let stem = linker.file_stem().and_then(|stem| stem.to_str()).unwrap_or_else(|| {
+                    sess.fatal("couldn't extract file stem from specified linker")
+                });
+
+                let flavor = if stem == "emcc" {
+                    LinkerFlavor::Em
+                } else if stem == "gcc"
+                    || stem.ends_with("-gcc")
+                    || stem == "clang"
+                    || stem.ends_with("-clang")
+                {
+                    LinkerFlavor::Gcc
+                } else if stem == "ld" || stem == "ld.lld" || stem.ends_with("-ld") {
+                    LinkerFlavor::Ld
+                } else if stem == "link" || stem == "lld-link" {
+                    LinkerFlavor::Msvc
+                } else if stem == "lld" || stem == "rust-lld" {
+                    LinkerFlavor::Lld(sess.target.target.options.lld_flavor)
+                } else {
+                    // fall back to the value in the target spec
+                    sess.target.target.linker_flavor
+                };
+
+                Some((linker, flavor))
+            }
+            (None, None) => None,
+        }
+    }
+
+    // linker and linker flavor specified via command line have precedence over what the target
+    // specification specifies
+    if let Some(ret) = infer_from(sess, sess.opts.cg.linker.clone(), sess.opts.cg.linker_flavor) {
+        return ret;
+    }
+
+    if let Some(ret) = infer_from(
+        sess,
+        sess.target.target.options.linker.clone().map(PathBuf::from),
+        Some(sess.target.target.linker_flavor),
+    ) {
+        return ret;
+    }
+
+    bug!("Not enough information provided to determine how to invoke the linker");
+}