src/driver/jit.rs

   1 use std::ffi::CString;
   2 use std::os::raw::{c_char, c_int};
   3
   4 use rustc_codegen_ssa::CrateInfo;
   5
   6 use crate::prelude::*;
   7
   8 pub(super) fn run_jit(tcx: TyCtxt<'_>) -> ! {
   9     use cranelift_simplejit::{SimpleJITBackend, SimpleJITBuilder};
  10
  11     // Rustc opens us without the RTLD_GLOBAL flag, so __cg_clif_global_atomic_mutex will not be
  12     // exported. We fix this by opening ourself again as global.
  13     // FIXME remove once atomic_shim is gone
  14     let cg_dylib = std::ffi::OsString::from(&tcx.sess.opts.debugging_opts.codegen_backend.as_ref().unwrap());
  15     std::mem::forget(libloading::os::unix::Library::open(Some(cg_dylib), libc::RTLD_NOW | libc::RTLD_GLOBAL).unwrap());
  16
  17
  18     let imported_symbols = load_imported_symbols_for_jit(tcx);
  19
  20     let mut jit_builder = SimpleJITBuilder::with_isa(
  21         crate::build_isa(tcx.sess, false),
  22         cranelift_module::default_libcall_names(),
  23     );
  24     jit_builder.symbols(imported_symbols);
  25     let mut jit_module: Module<SimpleJITBackend> = Module::new(jit_builder);
  26     assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
  27
  28     let sig = Signature {
  29         params: vec![
  30             AbiParam::new(jit_module.target_config().pointer_type()),
  31             AbiParam::new(jit_module.target_config().pointer_type()),
  32         ],
  33         returns: vec![AbiParam::new(
  34             jit_module.target_config().pointer_type(), /*isize*/
  35         )],
  36         call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
  37     };
  38     let main_func_id = jit_module
  39         .declare_function("main", Linkage::Import, &sig)
  40         .unwrap();
  41
  42     if !tcx.sess.opts.output_types.should_codegen() {
  43         tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
  44     }
  45
  46     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
  47     let mono_items = cgus
  48         .iter()
  49         .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
  50         .flatten()
  51         .collect::<FxHashMap<_, (_, _)>>()
  52         .into_iter()
  53         .collect::<Vec<(_, (_, _))>>();
  54
  55     let mut cx = crate::CodegenCx::new(tcx, jit_module, false);
  56
  57     let (mut jit_module, _debug, mut unwind_context) = super::time(tcx, "codegen mono items", || {
  58         let mut global_asm = Vec::new();
  59         super::codegen_mono_items(&mut cx, &mut global_asm, mono_items);
  60         for hir_id in global_asm {
  61             let item = tcx.hir().expect_item(hir_id);
  62             tcx.sess.span_err(item.span, "Global asm is not supported in JIT mode");
  63         }
  64         tcx.sess.time("finalize CodegenCx", || cx.finalize())
  65     });
  66     crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);
  67     crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);
  68
  69     jit_module.finalize_definitions();
  70
  71     let _unwind_register_guard = unsafe { unwind_context.register_jit(&mut jit_module) };
  72
  73     tcx.sess.abort_if_errors();
  74
  75     let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
  76
  77     println!("Rustc codegen cranelift will JIT run the executable, because the CG_CLIF_JIT env var is set");
  78
  79     let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
  80         unsafe { ::std::mem::transmute(finalized_main) };
  81
  82     let args = ::std::env::var("CG_CLIF_JIT_ARGS").unwrap_or_else(|_| String::new());
  83     let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
  84         .chain(args.split(" "))
  85         .map(|arg| CString::new(arg).unwrap())
  86         .collect::<Vec<_>>();
  87     let argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();
  88     // TODO: Rust doesn't care, but POSIX argv has a NULL sentinel at the end
  89
  90     let ret = f(args.len() as c_int, argv.as_ptr());
  91
  92     jit_module.finish();
  93     std::process::exit(ret);
  94 }
  95
  96 fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
  97     use rustc_middle::middle::dependency_format::Linkage;
  98
  99     let mut dylib_paths = Vec::new();
 100
 101     let crate_info = CrateInfo::new(tcx);
 102     let formats = tcx.dependency_formats(LOCAL_CRATE);
 103     let data = &formats
 104         .iter()
 105         .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
 106         .unwrap()
 107         .1;
 108     for &(cnum, _) in &crate_info.used_crates_dynamic {
 109         let src = &crate_info.used_crate_source[&cnum];
 110         match data[cnum.as_usize() - 1] {
 111             Linkage::NotLinked | Linkage::IncludedFromDylib => {}
 112             Linkage::Static => {
 113                 let name = tcx.crate_name(cnum);
 114                 let mut err = tcx
 115                     .sess
 116                     .struct_err(&format!("Can't load static lib {}", name.as_str()));
 117                 err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
 118                 err.emit();
 119             }
 120             Linkage::Dynamic => {
 121                 dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
 122             }
 123         }
 124     }
 125
 126     let mut imported_symbols = Vec::new();
 127     for path in dylib_paths {
 128         use object::Object;
 129         let lib = libloading::Library::new(&path).unwrap();
 130         let obj = std::fs::read(path).unwrap();
 131         let obj = object::File::parse(&obj).unwrap();
 132         imported_symbols.extend(obj.dynamic_symbols().filter_map(|(_idx, symbol)| {
 133             let name = symbol.name().unwrap().to_string();
 134             if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
 135                 return None;
 136             }
 137             let dlsym_name = if cfg!(target_os = "macos") {
 138                 // On macOS `dlsym` expects the name without leading `_`.
 139                 assert!(name.starts_with("_"), "{:?}", name);
 140                 &name[1..]
 141             } else {
 142                 &name
 143             };
 144             let symbol: libloading::Symbol<'_, *const u8> =
 145                 unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
 146             Some((name, *symbol))
 147         }));
 148         std::mem::forget(lib)
 149     }
 150
 151     tcx.sess.abort_if_errors();
 152
 153     imported_symbols
 154 }