feat: JIT (WIP)

src/bin/repl.rs

@@ -1,9 +1,10 @@
+use inkwell::context::Context;
 use rustyline::error::ReadlineError;
 use rustyline::{DefaultEditor, Result};
 
 use nixjit::compile::compile;
 use nixjit::ir::downgrade;
-use nixjit::vm::run;
+use nixjit::vm::{run, JITContext};
 
 macro_rules! unwrap {
     ($e:expr) => {
@@ -31,7 +32,9 @@ fn main() -> Result<()> {
                     rnix::Root::parse(expr.as_str()).tree().expr().unwrap()
                 ));
                 let prog = compile(downgraded);
-                println!("{}", unwrap!(run(prog)));
+                let ctx = Context::create();
+                let jit = JITContext::new(&ctx);
+                println!("{}", unwrap!(run(prog, jit)));
             }
             Err(ReadlineError::Interrupted) => {
                 println!("CTRL-C");

src/builtins/mod.rs

@@ -1,4 +1,4 @@
-use std::cell::{OnceCell, RefCell};
+use std::cell::RefCell;
 use std::rc::Rc;
 
 use crate::ty::common::Symbol;

src/ty/internal/attrset.rs

@@ -10,6 +10,7 @@ use super::super::common::Symbol;
 use super::super::public as p;
 use super::{ToPublic, Value};
 
+#[repr(C)]
 #[derive(Debug, Constructor, Clone, PartialEq)]
 pub struct AttrSet<'vm> {
     data: HashTrieMap<Symbol, Value<'vm>>,
@@ -82,7 +83,7 @@ impl<'vm> AttrSet<'vm> {
         &self.data
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
         let mut map: Vec<_> = self
             .data
             .into_iter()
@@ -169,7 +170,7 @@ impl<'vm> RecAttrSet<'vm> {
         RecAttrSet { data }
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
         let mut map: Vec<_> = self
             .data
             .map

src/ty/internal/func.rs

@@ -8,7 +8,7 @@ use crate::bytecode::OpCodes;
 use crate::error::Result;
 use crate::ir;
 use crate::ty::internal::Value;
-use crate::vm::{CapturedEnv, VM};
+use crate::vm::{CapturedEnv, Env, VM};
 
 #[derive(Debug, Clone)]
 pub enum Param {
@@ -40,15 +40,18 @@ impl From<ir::Param> for Param {
     }
 }
 
+pub type JITFunc<'vm> = unsafe extern "C" fn(vm: *mut VM<'vm, '_>, *mut Env<'vm>, *mut Value<'vm>) -> Value<'vm>;
+
 #[derive(Debug, Clone)]
 pub struct Func<'vm> {
     pub env: OnceCell<CapturedEnv<'vm>>,
     pub param: Param,
     pub opcodes: OpCodes,
+    pub compiled: Option<JITFunc<'vm>>
 }
 
 impl<'vm> Func<'vm> {
-    pub fn call(&'vm self, vm: &VM<'vm>, arg: Value<'vm>) -> Result<Value<'vm>> {
+    pub fn call(&'vm self, vm: &VM<'vm, '_>, arg: Value<'vm>) -> Result<Value<'vm>> {
         use Param::*;
 
         let env = self.env.get().unwrap().clone().released();

src/ty/internal/list.rs

@@ -30,7 +30,7 @@ impl<'vm> List<'vm> {
         self
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<()> {
+    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<()> {
         let mut vec: Vec<_> = self.data.iter().cloned().collect();
         vec.iter_mut()
             .map(|v| v.force_deep(vm).map(|_| ()))

src/ty/internal/mod.rs

@@ -127,7 +127,7 @@ impl<'vm> Value<'vm> {
         }
     }
 
-    pub fn call(self, vm: &VM<'vm>, args: Vec<Value<'vm>>) -> Result<Value<'vm>> {
+    pub fn call(self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Result<Value<'vm>> {
         use Value::*;
         Ok(match self {
             PrimOp(func) => func.call(vm, args),
@@ -365,7 +365,7 @@ impl<'vm> Value<'vm> {
         self
     }
 
-    pub fn force(&mut self, vm: &VM<'vm>) -> Result<&mut Self> {
+    pub fn force(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> {
         if let Value::Thunk(thunk) = self {
             let value = thunk.force(vm)?;
             *self = value
@@ -376,7 +376,7 @@ impl<'vm> Value<'vm> {
         Ok(self)
     }
 
-    pub fn force_deep(&mut self, vm: &VM<'vm>) -> Result<&mut Self> {
+    pub fn force_deep(&mut self, vm: &VM<'vm, '_>) -> Result<&mut Self> {
         match self {
             Value::Thunk(thunk) => {
                 let mut value = thunk.force(vm)?;
@@ -447,7 +447,7 @@ impl<'vm> Thunk<'vm> {
         *self.env.borrow_mut() = Some(env);
     }
 
-    pub fn force(&self, vm: &VM<'vm>) -> Result<Value<'vm>> {
+    pub fn force(&self, vm: &VM<'vm, '_>) -> Result<Value<'vm>> {
         match &*self.thunk.borrow() {
             _Thunk::Value(value) => return Ok(value.as_ref().clone()),
             _Thunk::SuspendedFrom(from) => {

src/ty/internal/primop.rs

@@ -8,7 +8,7 @@ use super::Value;
 pub struct PrimOp<'vm> {
     pub name: &'static str,
     arity: u8,
-    func: fn(&VM<'vm>, Vec<Value<'vm>>) -> Value<'vm>,
+    func: fn(&VM<'vm, '_>, Vec<Value<'vm>>) -> Value<'vm>,
 }
 
 impl PartialEq for PrimOp<'_> {
@@ -18,7 +18,7 @@ impl PartialEq for PrimOp<'_> {
 }
 
 impl<'vm> PrimOp<'vm> {
-    pub fn call(self, vm: &VM<'vm>, args: Vec<Value<'vm>>) -> Value<'vm> {
+    pub fn call(self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Value<'vm> {
         if (args.len() as u8) < self.arity {
             Value::PartialPrimOp(PartialPrimOp {
                 name: self.name,
@@ -39,7 +39,7 @@ pub struct PartialPrimOp<'vm> {
     pub name: &'static str,
     arity: u8,
     args: Vec<Value<'vm>>,
-    func: fn(&VM<'vm>, Vec<Value<'vm>>) -> Value<'vm>,
+    func: fn(&VM<'vm, '_>, Vec<Value<'vm>>) -> Value<'vm>,
 }
 
 impl PartialEq for PartialPrimOp<'_> {
@@ -49,7 +49,7 @@ impl PartialEq for PartialPrimOp<'_> {
 }
 
 impl<'vm> PartialPrimOp<'vm> {
-    pub fn call(mut self, vm: &VM<'vm>, args: Vec<Value<'vm>>) -> Value<'vm> {
+    pub fn call(mut self, vm: &VM<'vm, '_>, args: Vec<Value<'vm>>) -> Value<'vm> {
         let len = args.len() as u8;
         self.args.extend(args);
         if len < self.arity {

src/vm/mod.rs

@@ -11,6 +11,7 @@ use crate::ty::public as p;
 use stack::{STACK_SIZE, Stack};
 
 pub use env::{CapturedEnv, Env};
+pub use jit::JITContext;
 
 mod env;
 mod jit;
@@ -19,7 +20,7 @@ mod stack;
 #[cfg(test)]
 mod test;
 
-pub fn run(prog: Program) -> Result<p::Value> {
+pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
     let vm = VM::new(
         prog.thunks,
         prog.funcs
@@ -28,8 +29,10 @@ pub fn run(prog: Program) -> Result<p::Value> {
                 env: OnceCell::new(),
                 param: f.param,
                 opcodes: f.opcodes,
+                compiled: None,
             })
             .collect(),
+        jit
     );
     let env = env();
     let temp = vm.eval(prog.top_level, env)?;
@@ -37,18 +40,19 @@ pub fn run(prog: Program) -> Result<p::Value> {
     Ok(temp)
 }
 
-pub struct VM<'vm> {
+pub struct VM<'vm, 'jit> {
     thunks: Box<[Thunk<'vm>]>,
     funcs: Box<[Func<'vm>]>,
+    jit: JITContext<'jit>,
 }
 
-impl<'vm> VM<'vm> {
-    fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>) -> Self {
+impl<'vm, 'jit> VM<'vm, 'jit> {
+    fn new(thunks: Box<[OpCodes]>, funcs: Box<[Func<'vm>]>, jit: JITContext<'jit>) -> Self {
         let thunks = thunks
             .into_iter()
             .map(|opcodes| Thunk::new(opcodes))
             .collect();
-        VM { thunks, funcs }
+        VM { thunks, funcs, jit }
     }
 
     pub fn get_thunk(&self, idx: usize) -> &'vm Thunk<'vm> {

src/vm/test.rs

@@ -1,5 +1,6 @@
 extern crate test;
 
+use inkwell::context::Context;
 use test::{Bencher, black_box};
 
 use ecow::EcoString;
@@ -9,6 +10,7 @@ use crate::compile::compile;
 use crate::ir::downgrade;
 use crate::ty::common::Symbol;
 use crate::ty::public::*;
+use crate::vm::JITContext;
 
 use super::run;
 
@@ -17,7 +19,9 @@ fn test_expr(expr: &str, expected: Value) {
     let downgraded = downgrade(rnix::Root::parse(expr).tree().expr().unwrap()).unwrap();
     let prog = compile(downgraded);
     dbg!(&prog);
-    assert_eq!(run(prog).unwrap(), expected);
+    let ctx = Context::create();
+    let jit = JITContext::new(&ctx);
+    assert_eq!(run(prog, jit).unwrap(), expected);
 }
 
 macro_rules! thunk {