Implement global_asm! using an external assembler

Fixes #1061

5 files changed, 121 insertions, 18 deletions

diff --git a/example/mini_core.rs b/example/mini_core.rs
index a8db81fa06d..94336154748 100644
--- a/example/mini_core.rs
+++ b/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"]
diff --git a/example/mini_core_hello_world.rs b/example/mini_core_hello_world.rs
index 12c6638dd92..d83fb2aece9 100644
--- a/example/mini_core_hello_world.rs
+++ b/example/mini_core_hello_world.rs
@@ -284,6 +284,26 @@ fn main() {
 
     #[cfg(not(jit))]
     test_tls();
+
+    #[cfg(not(jit))]
+    unsafe {
+        global_asm_test();
+    }
+}
+
+#[cfg(not(jit))]
+extern "C" {
+    fn global_asm_test();
+}
+
+#[cfg(not(jit))]
+global_asm! {
+    "
+    .global global_asm_test
+    global_asm_test:
+    // comment that would normally be removed by LLVM
+    ret
+    "
 }
 
 #[repr(C)]
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index 9b05a146020..028de9b1564 100644
--- a/src/driver/aot.rs
+++ b/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,73 @@ 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;
+    }
+
+    // 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("as")
+        .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("ld")
+        .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() {
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index 93535cd16de..138792084b2 100644
--- a/src/driver/jit.rs
+++ b/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);
diff --git a/src/driver/mod.rs b/src/driver/mod.rs
index 112741b6191..8872b6f1da4 100644
--- a/src/driver/mod.rs
+++ b/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);
         }
     }
 }