Merge pull request #1089 from bjorn3/custom_driver

Add custom rustc driver that uses cg_clif

14 files changed, 191 insertions, 50 deletions

diff --git a/Readme.md b/Readme.md
index c4e55975541..680ff877656 100644
--- a/Readme.md
+++ b/Readme.md
@@ -36,7 +36,24 @@ If you compiled cg_clif in debug mode (aka you didn't pass `--release` to `./tes
 > You should prefer using the Cargo method.
 
 ```bash
-$ rustc +$(cat $cg_clif_dir/rust-toolchain) -Cpanic=abort -Zcodegen-backend=$cg_clif_dir/target/release/librustc_codegen_cranelift.so --sysroot $cg_clif_dir/build_sysroot/sysroot my_crate.rs
+$ $cg_clif_dir/target/release/cg_clif my_crate.rs
+```
+
+### Jit mode
+
+In jit mode cg_clif will immediately execute your code without creating an executable file.
+
+> This requires all dependencies to be available as dynamic library.
+> The jit mode will probably need cargo integration to make this possible.
+
+```bash
+$ $cg_clif_dir/cargo.sh jit
+```
+
+or
+
+```bash
+$ $cg_clif_dir/target/release/cg_clif --jit my_crate.rs
 ```
 
 ### Shell
@@ -45,7 +62,7 @@ These are a few functions that allow you to easily run rust code from the shell
 
 ```bash
 function jit_naked() {
-    echo "$@" | CG_CLIF_JIT=1 rustc -Zcodegen-backend=$cg_clif_dir/target/release/librustc_codegen_cranelift.so --sysroot $cg_clif_dir/build_sysroot/sysroot - -Cprefer-dynamic
+    echo "$@" | $cg_clif_dir/target/release/cg_clif - --jit
 }
 
 function jit() {
diff --git a/build_sysroot/build_sysroot.sh b/build_sysroot/build_sysroot.sh
index 0df05f57563..2f09453e0f7 100755
--- a/build_sysroot/build_sysroot.sh
+++ b/build_sysroot/build_sysroot.sh
@@ -14,8 +14,13 @@ popd >/dev/null
 rm -r target/*/{debug,release}/{build,deps,examples,libsysroot*,native} 2>/dev/null || true
 rm -r sysroot/ 2>/dev/null || true
 
+# Use rustc with cg_clif as hotpluggable backend instead of the custom cg_clif driver so that
+# build scripts are still compiled using cg_llvm.
+export RUSTC=rustc
+export RUSTFLAGS=$RUSTFLAGS" -Ztrim-diagnostic-paths=no -Zcodegen-backend=$(pwd)/../target/"$CHANNEL"/librustc_codegen_cranelift."$dylib_ext" --sysroot $(pwd)/sysroot"
+
 # Build libs
-export RUSTFLAGS="$RUSTFLAGS -Z force-unstable-if-unmarked"
+export RUSTFLAGS="$RUSTFLAGS -Zforce-unstable-if-unmarked -Cpanic=abort"
 if [[ "$1" == "--release" ]]; then
     sysroot_channel='release'
     # FIXME Enable incremental again once rust-lang/rust#74946 is fixed
diff --git a/cargo.sh b/cargo.sh
index ba4b7d6c942..cebc3e67363 100755
--- a/cargo.sh
+++ b/cargo.sh
@@ -12,12 +12,11 @@ TOOLCHAIN=$(cat rust-toolchain)
 
 popd >/dev/null
 
-if [[ $(rustc -V) != $(rustc +${TOOLCHAIN} -V) ]]; then
-    echo "rustc_codegen_cranelift is build for $(rustc +${TOOLCHAIN} -V) but the default rustc version is $(rustc -V)."
-    echo "Using $(rustc +${TOOLCHAIN} -V)."
-fi
-
 cmd=$1
 shift
 
+if [[ "$cmd" = "jit" ]]; then
+cargo +${TOOLCHAIN} rustc $@ -- --jit
+else
 cargo +${TOOLCHAIN} $cmd $@
+fi
diff --git a/docs/env_vars.md b/docs/env_vars.md
index 3707909a978..07b75622a58 100644
--- a/docs/env_vars.md
+++ b/docs/env_vars.md
@@ -1,8 +1,6 @@
 # List of env vars recognized by cg_clif
 
 <dl>
-    <dt>CG_CLIF_JIT</dt>
-    <dd>Enable JIT mode to immediately run a program instead of writing an executable file.</dd>
     <dt>CG_CLIF_JIT_ARGS</dt>
     <dd>When JIT mode is enable pass these arguments to the program.</dd>
     <dt>CG_CLIF_INCR_CACHE_DISABLED</dt>
diff --git a/scripts/config.sh b/scripts/config.sh
index 728667ec367..acd048893e9 100644
--- a/scripts/config.sh
+++ b/scripts/config.sh
@@ -32,15 +32,18 @@ if [[ "$HOST_TRIPLE" != "$TARGET_TRIPLE" ]]; then
    fi
 fi
 
-export RUSTFLAGS=$linker' -Ztrim-diagnostic-paths=no -Cpanic=abort -Cdebuginfo=2 -Zpanic-abort-tests -Zcodegen-backend='$(pwd)'/target/'$CHANNEL'/librustc_codegen_cranelift.'$dylib_ext' --sysroot '$(pwd)'/build_sysroot/sysroot'
-export RUSTDOCFLAGS=$RUSTFLAGS
+export RUSTC=$(pwd)/"target/"$CHANNEL"/cg_clif"
+export RUSTFLAGS=$linker
+export RUSTDOCFLAGS=$linker' -Ztrim-diagnostic-paths=no -Cpanic=abort -Zpanic-abort-tests '\
+'-Zcodegen-backend='$(pwd)'/target/'$CHANNEL'/librustc_codegen_cranelift.'$dylib_ext' --sysroot '$(pwd)'/build_sysroot/sysroot'
 
 # FIXME remove once the atomic shim is gone
 if [[ `uname` == 'Darwin' ]]; then
    export RUSTFLAGS="$RUSTFLAGS -Clink-arg=-undefined -Clink-arg=dynamic_lookup"
 fi
 
-export LD_LIBRARY_PATH="$(pwd)/target/out:$(pwd)/build_sysroot/sysroot/lib/rustlib/$TARGET_TRIPLE/lib"
+export LD_LIBRARY_PATH="$(pwd)/target/out:$(pwd)/build_sysroot/sysroot/lib/rustlib/"$TARGET_TRIPLE"/lib:\
+$(pwd)/target/"$CHANNEL":$(rustc --print sysroot)/lib"
 export DYLD_LIBRARY_PATH=$LD_LIBRARY_PATH
 
 export CG_CLIF_DISPLAY_CG_TIME=1
diff --git a/scripts/filter_profile.rs b/scripts/filter_profile.rs
index 75f2f61b21a..c70c3ec47f3 100755
--- a/scripts/filter_profile.rs
+++ b/scripts/filter_profile.rs
@@ -5,7 +5,7 @@ CHANNEL="release"
 pushd $(dirname "$0")/../
 source scripts/config.sh
 popd
-CG_CLIF_JIT=1 PROFILE=$1 OUTPUT=$2 exec rustc $RUSTFLAGS $0 --crate-type bin -Cprefer-dynamic
+PROFILE=$1 OUTPUT=$2 exec $RUSTC $RUSTFLAGS --jit $0
 #*/
 
 //! This program filters away uninteresting samples and trims uninteresting frames for stackcollapse
diff --git a/src/atomic_shim.rs b/src/atomic_shim.rs
index 2f9f9e3a4f7..f29e269c01f 100644
--- a/src/atomic_shim.rs
+++ b/src/atomic_shim.rs
@@ -10,10 +10,14 @@ 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() {
+pub(crate) fn init_global_lock(
+    module: &mut Module<impl Backend>,
+    bcx: &mut FunctionBuilder<'_>,
+    use_jit: bool,
+) {
+    if use_jit {
         // 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.
+        // so instead we define it in the cg_clif dylib.
 
         return;
     }
@@ -80,7 +84,7 @@ pub(crate) fn init_global_lock_constructor(
         let block = bcx.create_block();
         bcx.switch_to_block(block);
 
-        crate::atomic_shim::init_global_lock(module, &mut bcx);
+        crate::atomic_shim::init_global_lock(module, &mut bcx, false);
 
         bcx.ins().return_(&[]);
         bcx.seal_all_blocks();
diff --git a/src/bin/cg_clif.rs b/src/bin/cg_clif.rs
new file mode 100644
index 00000000000..8a39e9b0590
--- /dev/null
+++ b/src/bin/cg_clif.rs
@@ -0,0 +1,94 @@
+#![feature(rustc_private)]
+
+extern crate rustc_data_structures;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate rustc_session;
+extern crate rustc_target;
+
+use rustc_data_structures::profiling::print_time_passes_entry;
+use rustc_interface::interface;
+use rustc_session::config::ErrorOutputType;
+use rustc_session::early_error;
+use rustc_target::spec::PanicStrategy;
+
+#[derive(Default)]
+pub struct TimePassesCallbacks {
+    time_passes: bool,
+}
+
+impl rustc_driver::Callbacks for TimePassesCallbacks {
+    fn config(&mut self, config: &mut interface::Config) {
+        // If a --prints=... option has been given, we don't print the "total"
+        // time because it will mess up the --prints output. See #64339.
+        self.time_passes = config.opts.prints.is_empty()
+            && (config.opts.debugging_opts.time_passes || config.opts.debugging_opts.time);
+
+        // FIXME workaround for an ICE
+        config.opts.debugging_opts.trim_diagnostic_paths = false;
+
+        config.opts.cg.panic = Some(PanicStrategy::Abort);
+        config.opts.debugging_opts.panic_abort_tests = true;
+        config.opts.maybe_sysroot = Some(
+            std::env::current_exe()
+                .unwrap()
+                .parent()
+                .unwrap()
+                .parent()
+                .unwrap()
+                .parent()
+                .unwrap()
+                .join("build_sysroot")
+                .join("sysroot"),
+        );
+    }
+}
+
+fn main() {
+    let start = std::time::Instant::now();
+    rustc_driver::init_rustc_env_logger();
+    let mut callbacks = TimePassesCallbacks::default();
+    rustc_driver::install_ice_hook();
+    let exit_code = rustc_driver::catch_with_exit_code(|| {
+        let mut use_jit = false;
+
+        let mut args = std::env::args_os()
+            .enumerate()
+            .map(|(i, arg)| {
+                arg.into_string().unwrap_or_else(|arg| {
+                    early_error(
+                        ErrorOutputType::default(),
+                        &format!("Argument {} is not valid Unicode: {:?}", i, arg),
+                    )
+                })
+            })
+            .filter(|arg| {
+                if arg == "--jit" {
+                    use_jit = true;
+                    false
+                } else {
+                    true
+                }
+            })
+            .collect::<Vec<_>>();
+        if use_jit {
+            args.push("-Cprefer-dynamic".to_string());
+        }
+        rustc_driver::run_compiler(
+            &args,
+            &mut callbacks,
+            None,
+            None,
+            Some(Box::new(move |_| {
+                Box::new(rustc_codegen_cranelift::CraneliftCodegenBackend {
+                    config: rustc_codegen_cranelift::BackendConfig {
+                        use_jit,
+                    }
+                })
+            })),
+        )
+    });
+    // The extra `\t` is necessary to align this label with the others.
+    print_time_passes_entry(callbacks.time_passes, "\ttotal", start.elapsed());
+    std::process::exit(exit_code)
+}
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index 2cf136ceb67..fc01398c40d 100644
--- a/src/driver/aot.rs
+++ b/src/driver/aot.rs
@@ -150,7 +150,7 @@ fn module_codegen(tcx: TyCtxt<'_>, cgu_name: rustc_span::Symbol) -> ModuleCodege
     super::codegen_mono_items(&mut cx, mono_items);
     let (mut module, global_asm, debug, mut unwind_context) =
         tcx.sess.time("finalize CodegenCx", || cx.finalize());
-    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut unwind_context);
+    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut unwind_context, false);
 
     let codegen_result = emit_module(
         tcx,
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index edeeb7670e9..5a899afb97c 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -11,21 +11,23 @@ use crate::prelude::*;
 pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
     use cranelift_simplejit::{SimpleJITBackend, SimpleJITBuilder};
 
-    // Rustc opens us without the RTLD_GLOBAL flag, so __cg_clif_global_atomic_mutex will not be
-    // exported. We fix this by opening ourself again as global.
-    // FIXME remove once atomic_shim is gone
-    let cg_dylib = std::ffi::OsString::from(
-        &tcx.sess
-            .opts
-            .debugging_opts
-            .codegen_backend
-            .as_ref()
-            .unwrap(),
-    );
-    std::mem::forget(
-        libloading::os::unix::Library::open(Some(cg_dylib), libc::RTLD_NOW | libc::RTLD_GLOBAL)
-            .unwrap(),
-    );
+    #[cfg(unix)]
+    unsafe {
+        // When not using our custom driver rustc will open us without the RTLD_GLOBAL flag, so
+        // __cg_clif_global_atomic_mutex will not be exported. We fix this by opening ourself again
+        // as global.
+        // FIXME remove once atomic_shim is gone
+
+        let mut dl_info: libc::Dl_info = std::mem::zeroed();
+        assert_ne!(
+            libc::dladdr(run_jit as *const libc::c_void, &mut dl_info),
+            0
+        );
+        assert_ne!(
+            libc::dlopen(dl_info.dli_fname, libc::RTLD_NOW | libc::RTLD_GLOBAL),
+            std::ptr::null_mut(),
+        );
+    }
 
     let imported_symbols = load_imported_symbols_for_jit(tcx);
 
@@ -74,7 +76,7 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
     if !global_asm.is_empty() {
         tcx.sess.fatal("Global asm is not supported in JIT mode");
     }
-    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);
+    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context, true);
     crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);
 
     jit_module.finalize_definitions();
@@ -85,7 +87,7 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
 
     let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
 
-    println!("Rustc codegen cranelift will JIT run the executable, because the CG_CLIF_JIT env var is set");
+    println!("Rustc codegen cranelift will JIT run the executable, because --jit was passed");
 
     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
         unsafe { ::std::mem::transmute(finalized_main) };
diff --git a/src/driver/mod.rs b/src/driver/mod.rs
index 0224f929ac7..881221c0105 100644
--- a/src/driver/mod.rs
+++ b/src/driver/mod.rs
@@ -16,15 +16,19 @@ pub(crate) fn codegen_crate(
     tcx: TyCtxt<'_>,
     metadata: EncodedMetadata,
     need_metadata_module: bool,
+    config: crate::BackendConfig,
 ) -> Box<dyn Any> {
     tcx.sess.abort_if_errors();
 
-    if std::env::var("CG_CLIF_JIT").is_ok()
-        && tcx
+    if config.use_jit {
+        let is_executable = tcx
             .sess
             .crate_types()
-            .contains(&rustc_session::config::CrateType::Executable)
-    {
+            .contains(&rustc_session::config::CrateType::Executable);
+        if !is_executable {
+            tcx.sess.fatal("can't jit non-executable crate");
+        }
+
         #[cfg(feature = "jit")]
         let _: ! = jit::run_jit(tcx);
 
diff --git a/src/lib.rs b/src/lib.rs
index 64f4e490aed..52845b11da0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -181,7 +181,14 @@ impl<'tcx, B: Backend + 'static> CodegenCx<'tcx, B> {
     }
 }
 
-struct CraneliftCodegenBackend;
+#[derive(Copy, Clone, Debug)]
+pub struct BackendConfig {
+    pub use_jit: bool,
+}
+
+pub struct CraneliftCodegenBackend {
+    pub config: BackendConfig,
+}
 
 impl CodegenBackend for CraneliftCodegenBackend {
     fn init(&self, sess: &Session) {
@@ -223,7 +230,7 @@ impl CodegenBackend for CraneliftCodegenBackend {
         metadata: EncodedMetadata,
         need_metadata_module: bool,
     ) -> Box<dyn Any> {
-        let res = driver::codegen_crate(tcx, metadata, need_metadata_module);
+        let res = driver::codegen_crate(tcx, metadata, need_metadata_module, self.config);
 
         rustc_symbol_mangling::test::report_symbol_names(tcx);
 
@@ -345,5 +352,9 @@ fn build_isa(sess: &Session, enable_pic: bool) -> Box<dyn isa::TargetIsa + 'stat
 /// This is the entrypoint for a hot plugged rustc_codegen_cranelift
 #[no_mangle]
 pub fn __rustc_codegen_backend() -> Box<dyn CodegenBackend> {
-    Box::new(CraneliftCodegenBackend)
+    Box::new(CraneliftCodegenBackend {
+        config: BackendConfig {
+            use_jit: false,
+        }
+    })
 }
diff --git a/src/main_shim.rs b/src/main_shim.rs
index 520a54b8e53..c4b21dcecd3 100644
--- a/src/main_shim.rs
+++ b/src/main_shim.rs
@@ -9,6 +9,7 @@ pub(crate) fn maybe_create_entry_wrapper(
     tcx: TyCtxt<'_>,
     module: &mut Module<impl Backend + 'static>,
     unwind_context: &mut UnwindContext<'_>,
+    use_jit: bool,
 ) {
     let (main_def_id, use_start_lang_item) = match tcx.entry_fn(LOCAL_CRATE) {
         Some((def_id, entry_ty)) => (
@@ -32,6 +33,7 @@ pub(crate) fn maybe_create_entry_wrapper(
         unwind_context,
         main_def_id,
         use_start_lang_item,
+        use_jit,
     );
 
     fn create_entry_fn(
@@ -40,6 +42,7 @@ pub(crate) fn maybe_create_entry_wrapper(
         unwind_context: &mut UnwindContext<'_>,
         rust_main_def_id: DefId,
         use_start_lang_item: bool,
+        use_jit: bool,
     ) {
         let main_ret_ty = tcx.fn_sig(rust_main_def_id).output();
         // Given that `main()` has no arguments,
@@ -83,7 +86,7 @@ pub(crate) fn maybe_create_entry_wrapper(
             let arg_argc = bcx.append_block_param(block, m.target_config().pointer_type());
             let arg_argv = bcx.append_block_param(block, m.target_config().pointer_type());
 
-            crate::atomic_shim::init_global_lock(m, &mut bcx);
+            crate::atomic_shim::init_global_lock(m, &mut bcx, use_jit);
 
             let main_func_ref = m.declare_func_in_func(main_func_id, &mut bcx.func);
 
diff --git a/test.sh b/test.sh
index 3028f5955ed..0d2b5c6a392 100755
--- a/test.sh
+++ b/test.sh
@@ -2,18 +2,19 @@
 set -e
 
 # Build cg_clif
+export RUSTFLAGS="-Zrun_dsymutil=no"
 if [[ "$1" == "--release" ]]; then
     export CHANNEL='release'
-    cargo rustc --release -- -Zrun_dsymutil=no
+    cargo build --release
 else
     export CHANNEL='debug'
-    cargo rustc -- -Zrun_dsymutil=no
+    cargo build --bin cg_clif
 fi
 
 # Config
 source scripts/config.sh
 export CG_CLIF_INCR_CACHE_DISABLED=1
-RUSTC="rustc $RUSTFLAGS -L crate=target/out --out-dir target/out"
+RUSTC=$RUSTC" "$RUSTFLAGS" -L crate=target/out --out-dir target/out -Cdebuginfo=2"
 
 # Cleanup
 rm -r target/out || true
@@ -28,7 +29,7 @@ $RUSTC example/example.rs --crate-type lib --target $TARGET_TRIPLE
 
 if [[ "$JIT_SUPPORTED" = "1" ]]; then
     echo "[JIT] mini_core_hello_world"
-    CG_CLIF_JIT=1 CG_CLIF_JIT_ARGS="abc bcd" $RUSTC --crate-type bin -Cprefer-dynamic example/mini_core_hello_world.rs --cfg jit --target $HOST_TRIPLE
+    CG_CLIF_JIT_ARGS="abc bcd" $RUSTC --jit example/mini_core_hello_world.rs --cfg jit --target $HOST_TRIPLE
 else
     echo "[JIT] mini_core_hello_world (skipped)"
 fi
@@ -51,7 +52,7 @@ $RUN_WRAPPER ./target/out/alloc_example
 
 if [[ "$JIT_SUPPORTED" = "1" ]]; then
     echo "[JIT] std_example"
-    CG_CLIF_JIT=1 $RUSTC --crate-type bin -Cprefer-dynamic example/std_example.rs --target $HOST_TRIPLE
+    $RUSTC --jit example/std_example.rs --target $HOST_TRIPLE
 else
     echo "[JIT] std_example (skipped)"
 fi
@@ -86,7 +87,7 @@ pushd simple-raytracer
 if [[ "$HOST_TRIPLE" = "$TARGET_TRIPLE" ]]; then
     echo "[BENCH COMPILE] ebobby/simple-raytracer"
     hyperfine --runs ${RUN_RUNS:-10} --warmup 1 --prepare "cargo clean" \
-    "RUSTFLAGS='' cargo build" \
+    "RUSTC=rustc RUSTFLAGS='' cargo build" \
     "../cargo.sh build"
 
     echo "[BENCH RUN] ebobby/simple-raytracer"