Use libffi to call jitted functions.

skinnyBat · skinnyBat · commit c3f18a306e7a · 2020-08-28T19:52:04.000+12:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/jit/Cargo.toml b/jit/Cargo.toml
@@ -12,5 +12,6 @@ cranelift = "0.66.0"
 cranelift-module = "0.66.0"
 cranelift-simplejit = "0.66.0"
 num-traits = "0.2"
+libffi = "0.9.0"
 rustpython-bytecode = { path = "../bytecode", version = "0.1.1" }
 thiserror = "1.0"
diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
@@ -10,6 +10,16 @@ pub struct JitSig {
     pub ret: Option<JitType>,
 }
 
+impl JitSig {
+    pub fn to_cif(&self) -> libffi::middle::Cif {
+        let ret = match self.ret {
+            Some(ref ty) => ty.to_libffi(),
+            None => libffi::middle::Type::void(),
+        };
+        libffi::middle::Cif::new(Vec::new(), ret)
+    }
+}
+
 #[derive(Clone, PartialEq)]
 pub enum JitType {
     Int,
@@ -23,6 +33,13 @@ impl JitType {
             Self::Float => types::F64,
         }
     }
+
+    fn to_libffi(&self) -> libffi::middle::Type {
+        match self {
+            Self::Int => libffi::middle::Type::i64(),
+            Self::Float => libffi::middle::Type::f64(),
+        }
+    }
 }
 
 #[derive(Clone)]
@@ -132,8 +149,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 Ok(())
             }
             Instruction::BinaryOperation { op, .. } => {
-                let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
+                // the rhs is popped off first
                 let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
+                let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
                 match (a.ty, b.ty) {
                     (JitType::Int, JitType::Int) => match op {
                         BinaryOperator::Add => {
diff --git a/jit/src/lib.rs b/jit/src/lib.rs
@@ -1,5 +1,4 @@
 use std::fmt;
-use std::mem;
 
 use cranelift::prelude::*;
 use cranelift_module::{Backend, FuncId, Linkage, Module, ModuleError};
@@ -98,20 +97,13 @@ pub struct CompiledCode {
 
 impl CompiledCode {
     pub fn invoke(&self) -> Option<AbiValue> {
-        match self.sig.ret {
-            Some(JitType::Int) => {
-                let func = unsafe { mem::transmute::<_, fn() -> i64>(self.code) };
-                Some(AbiValue::Int(func()))
-            }
-            Some(JitType::Float) => {
-                let func = unsafe { mem::transmute::<_, fn() -> f64>(self.code) };
-                Some(AbiValue::Float(func()))
-            }
-            None => {
-                let func = unsafe { mem::transmute::<_, fn()>(self.code) };
-                func();
-                None
-            }
+        let cif = self.sig.to_cif();
+        unsafe {
+            let value = cif.call::<UnTypedAbiValue>(
+                libffi::middle::CodePtr::from_ptr(self.code as *const _),
+                &[],
+            );
+            self.sig.ret.as_ref().map(|ty| value.to_typed(ty))
         }
     }
 }
@@ -121,6 +113,21 @@ pub enum AbiValue {
     Int(i64),
 }
 
+union UnTypedAbiValue {
+    float: f64,
+    int: i64,
+    _void: (),
+}
+
+impl UnTypedAbiValue {
+    unsafe fn to_typed(&self, ty: &JitType) -> AbiValue {
+        match ty {
+            JitType::Int => AbiValue::Int(self.int),
+            JitType::Float => AbiValue::Float(self.float),
+        }
+    }
+}
+
 unsafe impl Send for CompiledCode {}
 unsafe impl Sync for CompiledCode {}
 
diff --git a/tests/snippets/jit.py b/tests/snippets/jit.py
@@ -3,6 +3,7 @@ def foo():
     a = 5
     return 10 + a
 
+
 def bar():
     a = 1e6
     return a / 5.0
@@ -12,9 +13,11 @@ def tests():
     assert foo() == 15
     assert bar() == 2e5
 
+
 tests()
 
 if hasattr(foo, "__jit__"):
     print("Has jit")
     foo.__jit__()
+    bar.__jit__()
     tests()
