src/driver/jit.rs

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