feat: JIT (WIP)

src/jit/mod.rs

@@ -0,0 +1,206 @@
+use std::rc::Rc;
+
+use inkwell::builder::Builder;
+use inkwell::context::Context;
+use inkwell::execution_engine::ExecutionEngine;
+use inkwell::module::Module;
+use inkwell::types::{FunctionType, StructType};
+use inkwell::values::BasicValueEnum;
+use inkwell::{AddressSpace, OptimizationLevel};
+
+use crate::bytecode::OpCode;
+use crate::error::*;
+use crate::stack::Stack;
+use crate::ty::common::Const;
+use crate::ty::internal::{Func, Thunk, Value};
+use crate::vm::VM;
+
+#[cfg(test)]
+mod test;
+
+const STACK_SIZE: usize = 8 * 1024 / size_of::<Value>();
+
+#[repr(u64)]
+pub enum ValueTag {
+    Int,
+    Float,
+    String,
+    Bool,
+    AttrSet,
+    List,
+    Function,
+    Thunk,
+    Path,
+}
+
+#[repr(C)]
+pub struct JITValue {
+    tag: ValueTag,
+    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>,
+
+    value_type: StructType<'ctx>,
+    func_type: FunctionType<'ctx>,
+}
+
+impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
+    pub fn new(context: &'ctx Context) -> Self {
+        let module = context.create_module("nixjit");
+
+        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,
+        );
+
+        JITContext {
+            execution_engine: module
+                .create_jit_execution_engine(OptimizationLevel::Default)
+                .unwrap(),
+            builder: context.create_builder(),
+            context,
+            module,
+
+            value_type,
+            func_type,
+        }
+    }
+
+    fn new_int(&self, int: i64) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Int as _, false)
+                    .into(),
+                self.context.i64_type().const_int(int as _, false).into(),
+            ])
+            .into()
+    }
+
+    fn new_float(&self, float: f64) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Float as _, false)
+                    .into(),
+                self.context.f64_type().const_float(float).into(),
+            ])
+            .into()
+    }
+
+    fn new_bool(&self, bool: bool) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Bool as _, false)
+                    .into(),
+                self.context.bool_type().const_int(bool as _, false).into(),
+            ])
+            .into()
+    }
+
+    fn new_null(&self) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Float as _, false)
+                    .into(),
+                self.context.i64_type().const_zero().into(),
+            ])
+            .into()
+    }
+
+    fn const_string(&self, string: *const u8) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Float as _, false)
+                    .into(),
+                self.context
+                    .ptr_sized_int_type(self.execution_engine.get_target_data(), None)
+                    .const_int(string as _, false)
+                    .into(),
+            ])
+            .into()
+    }
+
+    fn new_thunk(&self, thunk: *const Thunk) -> BasicValueEnum<'ctx> {
+        self.value_type
+            .const_named_struct(&[
+                self.context
+                    .i64_type()
+                    .const_int(ValueTag::Thunk as _, false)
+                    .into(),
+                self.context
+                    .ptr_sized_int_type(self.execution_engine.get_target_data(), None)
+                    .const_int(thunk as _, false)
+                    .into(),
+            ])
+            .into()
+    }
+
+    pub fn compile_function(&self, func: &Func, vm: &'vm VM<'_>) -> Result<JITFunc> {
+        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);
+        let entry = self.context.append_basic_block(func_, "entry");
+        self.builder.position_at_end(entry);
+        while let Some(opcode) = iter.next() {
+            self.single_op(opcode, vm, &mut stack)?;
+        }
+        assert_eq!(stack.len(), 1);
+        self.builder.build_return(Some(&stack.pop()))?;
+        if func_.verify(false) {
+            unsafe {
+                Ok(std::mem::transmute(self.execution_engine.get_function_address(func_.get_name().to_str().unwrap()).unwrap()))
+            }
+        } else {
+            todo!()
+        }
+    }
+
+    fn single_op<const CAP: usize>(
+        &self,
+        opcode: OpCode,
+        vm: &'vm VM<'_>,
+        stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
+    ) -> Result<()> {
+        match opcode {
+            OpCode::Const { idx } => {
+                use Const::*;
+                match vm.get_const(idx) {
+                    Int(int) => stack.push(self.new_int(int))?,
+                    Float(float) => stack.push(self.new_float(float))?,
+                    Bool(bool) => stack.push(self.new_bool(bool))?,
+                    String(string) => stack.push(self.const_string(string.as_ptr()))?,
+                    Null => stack.push(self.new_null())?,
+                }
+            }
+            OpCode::LoadThunk { idx } => stack.push(self.new_thunk(Rc::new(Thunk::new(vm.get_thunk(idx))).as_ref() as _))?,
+            _ => todo!()
+        }
+        Ok(())
+    }
+}