feat: JIT (WIP)

src/builtins/mod.rs

@@ -1,7 +1,8 @@
 use hashbrown::HashMap;
 use std::rc::Rc;
 
-use crate::ty::internal::{AttrSet, Const, PrimOp, Value};
+use crate::ty::internal::{AttrSet, PrimOp, Value};
+use crate::ty::common::Const;
 use crate::vm::{Env, VM};
 
 pub fn env<'vm>(vm: &'vm VM) -> Env<'vm> {

src/bytecode.rs

@@ -2,7 +2,8 @@ use hashbrown::HashMap;
 
 use ecow::EcoString;
 
-use crate::ty::internal::{Const, Param};
+use crate::ty::internal::Param;
+use crate::ty::common::Const;
 
 type Slice<T> = Box<[T]>;
 

src/error.rs

@@ -10,6 +10,8 @@ pub enum Error {
     DowngradeError(String),
     #[error("error occurred during evaluation stage: {0}")]
     EvalError(String),
+    #[error("error occurred during JIT compile stage: {0}")]
+    CompileError(#[from] inkwell::builder::BuilderError),
     #[error("unknown error")]
     Unknown,
 }

src/ir.rs

@@ -5,7 +5,7 @@ use rnix::ast::{self, Expr};
 
 use crate::compile::*;
 use crate::error::*;
-use crate::ty::internal as i;
+use crate::ty::common as c;
 
 pub fn downgrade(expr: Expr) -> Result<Downgraded> {
     let mut ctx = DowngradeContext::new();
@@ -126,15 +126,15 @@ pub struct DynamicAttrPair(pub Ir, pub Ir);
 pub struct DowngradeContext {
     thunks: Vec<Ir>,
     funcs: Vec<Func>,
-    consts: Vec<i::Const>,
+    consts: Vec<c::Const>,
-    constmap: HashMap<i::Const, usize>,
+    constmap: HashMap<c::Const, usize>,
     symbols: Vec<EcoString>,
     symmap: HashMap<EcoString, usize>,
 }
 
 pub struct Downgraded {
     pub top_level: Ir,
-    pub consts: Box<[i::Const]>,
+    pub consts: Box<[c::Const]>,
     pub symbols: Vec<EcoString>,
     pub symmap: HashMap<EcoString, usize>,
     pub thunks: Box<[Ir]>,
@@ -158,7 +158,7 @@ impl DowngradeContext {
         LoadFunc { idx }
     }
 
-    fn new_const(&mut self, cnst: i::Const) -> Const {
+    fn new_const(&mut self, cnst: c::Const) -> Const {
         if let Some(&idx) = self.constmap.get(&cnst) {
             Const { idx }
         } else {

src/jit.rs

@@ -1,5 +1,3 @@
-use std::pin::Pin;
-
 use inkwell::builder::Builder;
 use inkwell::context::Context;
 use inkwell::execution_engine::ExecutionEngine;
@@ -8,9 +6,12 @@ use inkwell::types::{BasicType, BasicTypeEnum, FunctionType, StructType};
 use inkwell::values::{BasicValueEnum, FunctionValue, IntValue};
 use inkwell::{AddressSpace, OptimizationLevel};
 
+use crate::bytecode::OpCode;
 use crate::stack::Stack;
+use crate::ty::internal::{Func, Value};
+use crate::error::*;
 
-use super::STACK_SIZE;
+const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
 
 #[repr(usize)]
 pub enum ValueTag {
@@ -27,25 +28,31 @@ pub enum ValueTag {
 #[repr(C)]
 pub struct JITValue {
     tag: ValueTag,
-    data: u64,
+    data: JITValueData
 }
 
+#[repr(C)]
+pub union JITValueData {
+    int: i64,
+    float: f64,
+    boolean: bool,
+}
+
+pub type JITFunc = fn(usize, usize, JITValue) -> JITValue;
+
 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>> {
+    pub fn new(context: &'ctx Context) -> 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());
@@ -55,24 +62,38 @@ impl<'ctx> JITContext<'ctx> {
             false,
         );
 
-        Pin::new(Box::new(JITContext {
+        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) {}
+    fn new_bool(&self, b: bool) -> IntValue {
+        self.context.bool_type().const_int(b as u64, false)
+    }
+
+    pub fn compile_function(&self, func: Func) -> Result<()> {
+        let mut stack = Stack::<_, STACK_SIZE>::new();
+        let mut iter = func.func.opcodes.iter().copied();
+        let func_ = self.module.add_function("fn", self.func_type, None);
+        while let Some(opcode) = iter.next() {
+            self.single_op(opcode, &func_, &mut stack)?;
+        }
+        Ok(())
+    }
+
+    fn single_op<const CAP: usize>(&self, opcode: OpCode, func: &FunctionValue, stack: &mut Stack<BasicValueEnum, CAP>) -> Result<()> {
+        todo!()
+    }
 }

src/lib.rs

@@ -5,6 +5,7 @@ mod builtins;
 mod bytecode;
 mod stack;
 mod ty;
+mod jit;
 
 pub mod compile;
 pub mod error;

src/stack.rs

@@ -3,11 +3,6 @@ use std::ops::Deref;
 
 use crate::error::*;
 
-pub struct Stack<T, const CAP: usize> {
-    items: [MaybeUninit<T>; CAP],
-    top: usize,
-}
-
 macro_rules! into {
     ($e:expr) => {
         // SAFETY: This macro is used to transmute `MaybeUninit<Value<'vm>>` to `Value<'vm>`
@@ -18,6 +13,17 @@ macro_rules! into {
     };
 }
 
+pub struct Stack<T, const CAP: usize> {
+    items: [MaybeUninit<T>; CAP],
+    top: usize,
+}
+
+impl<T, const CAP: usize> Default for Stack<T, CAP> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 impl<T, const CAP: usize> Stack<T, CAP> {
     pub fn new() -> Self {
         Stack {

src/ty/common.rs

@@ -1,6 +1,8 @@
+use std::hash::Hash;
 use std::fmt::{Display, Formatter, Result as FmtResult};
 
-use derive_more::Constructor;
+use derive_more::{Constructor, IsVariant, Unwrap};
+use ecow::EcoString;
 
 #[derive(Clone, Debug, PartialEq, Constructor, Hash)]
 pub struct Catchable {
@@ -12,3 +14,92 @@ impl Display for Catchable {
         write!(f, "<error: {}>", self.msg)
     }
 }
+
+#[derive(Debug, Clone, IsVariant, Unwrap)]
+pub enum Const {
+    Bool(bool),
+    Int(i64),
+    Float(f64),
+    String(EcoString),
+    Null,
+}
+
+impl Hash for Const {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        use Const::*;
+        std::mem::discriminant(self).hash(state);
+        match self {
+            Int(x) => x.hash(state),
+            Float(x) => x.to_bits().hash(state),
+            Bool(x) => x.hash(state),
+            String(x) => x.hash(state),
+            Null => (),
+        }
+    }
+}
+
+impl Display for Const {
+    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
+        use Const::*;
+        match self {
+            Int(x) => write!(f, "{x}"),
+            Float(x) => write!(f, "{x}"),
+            Bool(x) => write!(f, "{x}"),
+            String(x) => write!(f, "{x:?}"),
+            Null => write!(f, "null"),
+        }
+    }
+}
+
+impl From<bool> for Const {
+    fn from(value: bool) -> Self {
+        Const::Bool(value)
+    }
+}
+
+impl From<i64> for Const {
+    fn from(value: i64) -> Self {
+        Const::Int(value)
+    }
+}
+
+impl From<f64> for Const {
+    fn from(value: f64) -> Self {
+        Const::Float(value)
+    }
+}
+
+impl From<EcoString> for Const {
+    fn from(value: EcoString) -> Self {
+        Const::String(value)
+    }
+}
+
+impl From<String> for Const {
+    fn from(value: String) -> Self {
+        Const::String(value.into())
+    }
+}
+
+impl From<&str> for Const {
+    fn from(value: &str) -> Self {
+        Const::String(value.into())
+    }
+}
+
+impl PartialEq for Const {
+    fn eq(&self, other: &Self) -> bool {
+        use Const::*;
+        match (self, other) {
+            (Bool(a), Bool(b)) => a == b,
+            (Int(a), Int(b)) => a == b,
+            (Float(a), Float(b)) => a == b,
+            (Int(a), Float(b)) => *a as f64 == *b,
+            (Float(a), Int(b)) => *b as f64 == *a,
+            (String(a), String(b)) => a == b,
+            _ => false,
+        }
+    }
+}
+
+impl Eq for Const {}

src/ty/internal/cnst.rs

@@ -1,79 +0,0 @@
-use std::hash::Hash;
-
-use derive_more::{IsVariant, Unwrap};
-use ecow::EcoString;
-
-#[derive(Debug, Clone, IsVariant, Unwrap)]
-pub enum Const {
-    Bool(bool),
-    Int(i64),
-    Float(f64),
-    String(EcoString),
-    Null,
-}
-
-impl Hash for Const {
-    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        use Const::*;
-        match self {
-            Int(x) => x.hash(state),
-            Float(x) => x.to_bits().hash(state),
-            Bool(x) => x.hash(state),
-            String(x) => x.hash(state),
-            x @ Null => x.hash(state),
-        }
-    }
-}
-
-impl From<bool> for Const {
-    fn from(value: bool) -> Self {
-        Const::Bool(value)
-    }
-}
-
-impl From<i64> for Const {
-    fn from(value: i64) -> Self {
-        Const::Int(value)
-    }
-}
-
-impl From<f64> for Const {
-    fn from(value: f64) -> Self {
-        Const::Float(value)
-    }
-}
-
-impl From<EcoString> for Const {
-    fn from(value: EcoString) -> Self {
-        Const::String(value)
-    }
-}
-
-impl From<String> for Const {
-    fn from(value: String) -> Self {
-        Const::String(value.into())
-    }
-}
-
-impl From<&str> for Const {
-    fn from(value: &str) -> Self {
-        Const::String(value.into())
-    }
-}
-
-impl PartialEq for Const {
-    fn eq(&self, other: &Self) -> bool {
-        use Const::*;
-        match (self, other) {
-            (Bool(a), Bool(b)) => a == b,
-            (Int(a), Int(b)) => a == b,
-            (Float(a), Float(b)) => a == b,
-            (Int(a), Float(b)) => *a as f64 == *b,
-            (Float(a), Int(b)) => *b as f64 == *a,
-            (String(a), String(b)) => a == b,
-            _ => false,
-        }
-    }
-}
-
-impl Eq for Const {}

src/ty/internal/func.rs

@@ -1,9 +1,12 @@
-use derive_more::Constructor;
+use std::cell::{Cell, OnceCell};
+
 use itertools::Itertools;
+use derive_more::Constructor;
 
 use crate::bytecode::Func as BFunc;
 use crate::error::Result;
 use crate::ir;
+use crate::jit::JITFunc;
 use crate::ty::internal::{Thunk, Value};
 use crate::vm::{Env, VM};
 
@@ -37,14 +40,12 @@ impl From<ir::Param> for Param {
     }
 }
 
-pub type JITFunc<'vm> =
-    unsafe extern "C" fn(vm: *mut VM<'_>, *mut Env<'vm>, *mut Value<'vm>) -> Value<'vm>;
-
 #[derive(Debug, Clone, Constructor)]
 pub struct Func<'vm> {
     pub func: &'vm BFunc,
     pub env: Env<'vm>,
-    pub compiled: Option<JITFunc<'vm>>,
+    pub compiled: OnceCell<JITFunc>,
+    pub count: Cell<usize>
 }
 
 impl<'vm> Func<'vm> {

src/ty/internal/mod.rs

@@ -6,7 +6,7 @@ use std::rc::Rc;
 
 use derive_more::{IsVariant, Unwrap};
 
-use super::common as c;
+use super::common::*;
 use super::public as p;
 
 use crate::bytecode::OpCodes;
@@ -14,14 +14,12 @@ use crate::error::*;
 use crate::vm::{Env, VM};
 
 mod attrset;
-mod cnst;
 mod func;
 mod list;
 mod primop;
 mod string;
 
 pub use attrset::*;
-pub use cnst::Const;
 pub use func::*;
 pub use list::List;
 pub use primop::*;
@@ -33,7 +31,7 @@ pub enum Value<'vm> {
     ThunkRef(&'vm Thunk<'vm>),
     AttrSet(Rc<AttrSet<'vm>>),
     List(Rc<List<'vm>>),
-    Catchable(c::Catchable),
+    Catchable(Catchable),
     PrimOp(Rc<PrimOp<'vm>>),
     PartialPrimOp(Rc<PartialPrimOp<'vm>>),
     Func(Rc<Func<'vm>>),
@@ -81,7 +79,7 @@ pub enum ValueAsRef<'v, 'vm: 'v> {
     Thunk(&'v Thunk<'vm>),
     AttrSet(&'v AttrSet<'vm>),
     List(&'v List<'vm>),
-    Catchable(&'v c::Catchable),
+    Catchable(&'v Catchable),
     PrimOp(&'v PrimOp<'vm>),
     PartialPrimOp(&'v PartialPrimOp<'vm>),
     Func(&'v Func<'vm>),
@@ -93,7 +91,7 @@ pub enum ValueAsMut<'v, 'vm: 'v> {
     Thunk(&'v Thunk<'vm>),
     AttrSet(&'v mut AttrSet<'vm>),
     List(&'v mut List<'vm>),
-    Catchable(&'v mut c::Catchable),
+    Catchable(&'v mut Catchable),
     PrimOp(&'v mut PrimOp<'vm>),
     PartialPrimOp(&'v mut PartialPrimOp<'vm>),
     Func(&'v Func<'vm>),

src/ty/public.rs

@@ -10,10 +10,6 @@ use rpds::VectorSync;
 
 use super::common::*;
 
-mod cnst;
-
-pub use cnst::Const;
-
 #[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Constructor)]
 pub struct Symbol(EcoString);
 
@@ -132,3 +128,4 @@ impl Display for Value {
         }
     }
 }
+

src/ty/public/cnst.rs

@@ -1,137 +0,0 @@
-use std::fmt::{Display, Formatter, Result as FmtResult};
-
-use derive_more::{IsVariant, Unwrap};
-use ecow::EcoString;
-
-use crate::error::Error;
-
-use super::super::internal as i;
-
-#[derive(Debug, Clone, IsVariant, Unwrap)]
-pub enum Const {
-    Bool(bool),
-    Int(i64),
-    Float(f64),
-    String(EcoString),
-    Null,
-}
-
-impl Display for Const {
-    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
-        use Const::*;
-        match self {
-            Bool(b) => write!(f, "{b}"),
-            Int(i) => write!(f, "{i}"),
-            Float(float) => write!(f, "{float}"),
-            String(s) => write!(f, "{s}"),
-            Null => write!(f, "null"),
-        }
-    }
-}
-
-impl From<i::Const> for Const {
-    fn from(value: i::Const) -> Self {
-        use i::Const::*;
-        match value {
-            Bool(bool) => Const::Bool(bool),
-            Int(int) => Const::Int(int),
-            Float(float) => Const::Float(float),
-            String(string) => Const::String(string),
-            Null => Const::Null,
-        }
-    }
-}
-
-impl From<bool> for Const {
-    fn from(value: bool) -> Self {
-        Const::Bool(value)
-    }
-}
-
-impl From<i64> for Const {
-    fn from(value: i64) -> Self {
-        Const::Int(value)
-    }
-}
-
-impl From<f64> for Const {
-    fn from(value: f64) -> Self {
-        Const::Float(value)
-    }
-}
-
-impl From<EcoString> for Const {
-    fn from(value: EcoString) -> Self {
-        Const::String(value)
-    }
-}
-
-impl From<String> for Const {
-    fn from(value: String) -> Self {
-        Const::String(value.into())
-    }
-}
-
-impl From<&str> for Const {
-    fn from(value: &str) -> Self {
-        Const::String(value.into())
-    }
-}
-
-impl<'a> TryFrom<&'a Const> for &'a bool {
-    type Error = Error;
-
-    fn try_from(value: &'a Const) -> Result<Self, Self::Error> {
-        match value {
-            Const::Bool(b) => Ok(b),
-            _ => panic!(),
-        }
-    }
-}
-impl<'a> TryFrom<&'a Const> for &'a i64 {
-    type Error = Error;
-
-    fn try_from(value: &'a Const) -> Result<Self, Self::Error> {
-        match value {
-            Const::Int(int) => Ok(int),
-            _ => panic!(),
-        }
-    }
-}
-
-impl<'a> TryFrom<&'a Const> for &'a f64 {
-    type Error = Error;
-
-    fn try_from(value: &'a Const) -> Result<Self, Self::Error> {
-        match value {
-            Const::Float(float) => Ok(float),
-            _ => panic!(),
-        }
-    }
-}
-
-impl<'a> TryFrom<&'a Const> for &'a str {
-    type Error = Error;
-
-    fn try_from(value: &'a Const) -> Result<Self, Self::Error> {
-        match value {
-            Const::String(string) => Ok(string),
-            _ => panic!(),
-        }
-    }
-}
-
-impl PartialEq for Const {
-    fn eq(&self, other: &Self) -> bool {
-        use Const::*;
-        match (self, other) {
-            (Bool(a), Bool(b)) => a == b,
-            (Int(a), Int(b)) => a == b,
-            (Float(a), Float(b)) => a == b,
-            (String(a), String(b)) => a == b,
-            _ => false,
-        }
-    }
-}
-
-impl Eq for Const {}

src/vm/mod.rs

@@ -1,37 +1,35 @@
 use hashbrown::{HashMap, HashSet};
-use std::cell::RefCell;
+use std::cell::{Cell, OnceCell, RefCell};
-use std::pin::Pin;
 
 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> {
+pub fn run(prog: Program, jit: JITContext<'_>) -> Result<p::Value> {
-    let vm = VM::new(
+    let vm = VM {
-        prog.thunks,
+        thunks: prog.thunks,
-        prog.funcs,
+        funcs: prog.funcs,
-        RefCell::new(prog.symbols),
+        symbols: RefCell::new(prog.symbols),
-        RefCell::new(prog.symmap),
+        symmap: RefCell::new(prog.symmap),
-        prog.consts,
+        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;