feat: JIT (WIP)

src/vm/jit.rs

@@ -1,78 +0,0 @@
-use std::pin::Pin;
-
-use inkwell::builder::Builder;
-use inkwell::context::Context;
-use inkwell::execution_engine::ExecutionEngine;
-use inkwell::module::Module;
-use inkwell::types::{BasicType, BasicTypeEnum, FunctionType, StructType};
-use inkwell::values::{BasicValueEnum, FunctionValue, IntValue};
-use inkwell::{AddressSpace, OptimizationLevel};
-
-use crate::stack::Stack;
-
-use super::STACK_SIZE;
-
-#[repr(usize)]
-pub enum ValueTag {
-    Int,
-    String,
-    Bool,
-    AttrSet,
-    List,
-    Function,
-    Thunk,
-    Path,
-}
-
-#[repr(C)]
-pub struct JITValue {
-    tag: ValueTag,
-    data: u64,
-}
-
-pub struct JITContext<'ctx> {
-    context: &'ctx Context,
-    module: Module<'ctx>,
-    builder: Builder<'ctx>,
-    execution_engine: ExecutionEngine<'ctx>,
-    stack: Stack<BasicValueEnum<'ctx>, STACK_SIZE>,
-    cur_func: Option<FunctionValue<'ctx>>,
-
-    value_type: StructType<'ctx>,
-    func_type: FunctionType<'ctx>,
-}
-
-impl<'ctx> JITContext<'ctx> {
-    pub fn new(context: &'ctx Context) -> Pin<Box<Self>> {
-        let module = context.create_module("nixjit");
-        let stack = Stack::new();
-
-        let int_type = context.i64_type();
-        let pointer_type = context.ptr_type(AddressSpace::default());
-        let value_type = context.struct_type(&[int_type.into(), int_type.into()], false);
-        let func_type = value_type.fn_type(
-            &[pointer_type.into(), pointer_type.into(), value_type.into()],
-            false,
-        );
-
-        Pin::new(Box::new(JITContext {
-            execution_engine: module
-                .create_jit_execution_engine(OptimizationLevel::Default)
-                .unwrap(),
-            builder: context.create_builder(),
-            context,
-            module,
-            stack,
-            cur_func: None,
-
-            value_type,
-            func_type,
-        }))
-    }
-
-    fn new_int(&self, int: i64) -> IntValue {
-        self.context.i64_type().const_int(int as u64, false)
-    }
-
-    pub fn start_trace(&mut self) {}
-}

src/vm/mod.rs

@@ -1,37 +1,35 @@
 use hashbrown::{HashMap, HashSet};
-use std::cell::RefCell;
-use std::pin::Pin;
+use std::cell::{Cell, OnceCell, RefCell};
 
 use crate::builtins::env;
 use crate::bytecode::{BinOp, Func as F, OpCode, OpCodes, Program, UnOp};
 use crate::error::*;
 use crate::ty::internal::*;
+use crate::ty::common::Const;
 use crate::ty::public::{self as p, Symbol};
-
 use crate::stack::Stack;
+use crate::jit::JITContext;
 
 use derive_more::Constructor;
 use ecow::EcoString;
 pub use env::Env;
-pub use jit::JITContext;
 
 mod env;
-mod jit;
 
 #[cfg(test)]
 mod test;
 
-pub const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
+const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
 
-pub fn run(prog: Program, jit: Pin<Box<JITContext<'_>>>) -> Result<p::Value> {
-    let vm = VM::new(
-        prog.thunks,
-        prog.funcs,
-        RefCell::new(prog.symbols),
-        RefCell::new(prog.symmap),
-        prog.consts,
+pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
+    let vm = VM {
+        thunks: prog.thunks,
+        funcs: prog.funcs,
+        symbols: RefCell::new(prog.symbols),
+        symmap: RefCell::new(prog.symmap),
+        consts: prog.consts,
         jit,
-    );
+    };
     let env = env(&vm);
     let mut seen = HashSet::new();
     let value = vm
@@ -47,7 +45,7 @@ pub struct VM<'jit> {
     symbols: RefCell<Vec<EcoString>>,
     symmap: RefCell<HashMap<EcoString, usize>>,
     consts: Box<[Const]>,
-    jit: Pin<Box<JITContext<'jit>>>,
+    jit: JITContext<'jit>,
 }
 
 impl<'vm, 'jit: 'vm> VM<'jit> {
@@ -137,7 +135,7 @@ impl<'vm, 'jit: 'vm> VM<'jit> {
             }
             OpCode::Func { idx } => {
                 let func = self.get_func(idx);
-                stack.push(Value::Func(Func::new(func, env.clone(), None).into()))?;
+                stack.push(Value::Func(Func::new(func, env.clone(), OnceCell::new(), Cell::new(0)).into()))?;
             }
             OpCode::UnOp { op } => {
                 use UnOp::*;

src/vm/test.rs

@@ -10,8 +10,8 @@ use rpds::vector_sync;
 
 use crate::compile::compile;
 use crate::ir::downgrade;
-use crate::ty::public::Symbol;
 use crate::ty::public::*;
+use crate::ty::common::Const;
 use crate::vm::JITContext;
 
 use super::run;