-
-
Save ZhangHanDong/68f257098e2b64065b498b26499b1f32 to your computer and use it in GitHub Desktop.
An almost legit Rust JIT compiler
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
mod jit { | |
use libc::{c_void, dlclose, dlopen, dlsym, RTLD_NOW}; | |
use std::ffi::CString; | |
use std::fs::File; | |
use std::io::prelude::*; | |
use std::io::SeekFrom; | |
use std::process::Command; | |
const SOURCE_PATH: &'static str = "/tmp/jit.rs"; | |
const LIB_PATH: &'static str = "/tmp/librsjit.so"; | |
const FUN_NAME: &'static str = "calc"; | |
/// A JIT engine using rustc backed by a single source file. | |
pub struct JitEngine { | |
file: File, | |
} | |
impl JitEngine { | |
pub fn new() -> Self { | |
let file = File::create(SOURCE_PATH).expect("Could not create file"); | |
Self { file } | |
} | |
pub fn compile(&mut self, expression: &str) -> Fun { | |
// Reset the source file | |
self.file.set_len(0).unwrap(); | |
self.file.seek(SeekFrom::Start(0)).unwrap(); | |
// Write the rust program | |
self.file | |
.write_all( | |
format!( | |
" | |
#[no_mangle] | |
pub extern fn calc(a: i64, b: i64) -> i64 {{ | |
{} | |
}}", | |
expression | |
) | |
.as_bytes(), | |
) | |
.unwrap(); | |
// Compile the sources | |
Command::new("rustc") | |
.args(&["--crate-type=dylib", SOURCE_PATH, "-o"]) | |
.arg(LIB_PATH) | |
.status() | |
.unwrap(); | |
unsafe { Fun::new(LIB_PATH, FUN_NAME) } | |
} | |
} | |
/// A function from a library dynamically linked. | |
pub struct Fun { | |
fun: fn(a: i64, b: i64) -> i64, | |
handle: *mut c_void, | |
} | |
impl Fun { | |
unsafe fn new(lib_path: &str, fun_name: &str) -> Self { | |
// Load the library | |
let filename = CString::new(lib_path).unwrap(); | |
let handle = dlopen(filename.as_ptr(), RTLD_NOW); | |
if handle.is_null() { | |
panic!("Failed to resolve dlopen") | |
} | |
// Look for the function in the library | |
let fun_name = CString::new(fun_name).unwrap(); | |
let fun = dlsym(handle, fun_name.as_ptr()); | |
if fun.is_null() { | |
panic!("Failed to resolve '{}'", &fun_name.to_str().unwrap()); | |
} | |
// dlsym returns a C 'void*', cast it to a function pointer | |
let fun = std::mem::transmute::<*mut c_void, fn(i64, i64) -> i64>(fun); | |
Self { fun, handle } | |
} | |
pub fn call(&self, a: i64, b: i64) -> i64 { | |
(self.fun)(a, b) | |
} | |
} | |
impl Drop for Fun { | |
fn drop(&mut self) { | |
unsafe { | |
let ret = dlclose(self.handle); | |
if ret != 0 { | |
panic!("Error while closing lib"); | |
} | |
} | |
} | |
} | |
} | |
use std::io; | |
fn main() { | |
let mut jit = jit::JitEngine::new(); | |
loop { | |
println!("Value for a:"); | |
let a = read_value(); | |
println!("Value for b:"); | |
let b = read_value(); | |
println!("Expression:"); | |
let expression = read_expression(); | |
let fun = jit.compile(&expression); | |
let result = fun.call(a, b); | |
println!("{}\n", result); | |
} | |
} | |
fn read_value() -> i64 { | |
let mut buffer = String::new(); | |
io::stdin().read_line(&mut buffer).unwrap(); | |
let v = i64::from_str_radix(&buffer[..buffer.len() - 1], 10).unwrap(); | |
v | |
} | |
fn read_expression() -> String { | |
let mut buffer = String::new(); | |
io::stdin().read_line(&mut buffer).unwrap(); | |
buffer | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment