feat: JIT (WIP)

src/engine/mod.rs

@@ -1,10 +1,13 @@
+use std::cell::OnceCell;
 use std::rc::Rc;
 
 use hashbrown::HashSet;
+use inkwell::context::Context;
 use priority_queue::PriorityQueue;
 
 use crate::env::Env;
 use crate::error::Result;
+use crate::eval::jit::{JITContext, JITFunc};
 use crate::eval::Evaluate;
 use crate::ir::{Dep, Downgraded, Ir, SccNode};
 use crate::ty::internal as i;
@@ -16,29 +19,36 @@ mod test;
 type ThunkIdx = usize;
 type EnvIdx = usize;
 
-pub struct Engine {
+pub struct Engine<'ctx: 'exec, 'exec> {
     pub thunks: Box<[Ir]>,
     pub func_offset: usize,
     pub func_deps: Vec<HashSet<Dep>>,
+    jit: &'exec JITContext<'ctx>,
+    compiled: Box<[OnceCell<JITFunc<'ctx, 'exec>>]>,
     tasks: PriorityQueue<CompileTask, usize>,
 }
 
 pub fn eval(downgraded: Downgraded) -> Result<Value> {
+    let ctx = Context::create();
+    let jit = JITContext::new(&ctx);
     let mut engine = Engine::new(
         downgraded.thunks,
         downgraded.func_offset,
         downgraded.func_deps,
+        &jit
     );
     engine.eval(downgraded.graph)
 }
 
-impl Engine {
+impl<'ctx, 'exec> Engine<'ctx, 'exec> {
-    pub fn new(thunks: Box<[Ir]>, func_offset: usize, func_deps: Vec<HashSet<Dep>>) -> Self {
+    pub fn new(thunks: Box<[Ir]>, func_offset: usize, func_deps: Vec<HashSet<Dep>>, jit: &'exec JITContext<'ctx>) -> Self {
         Self {
+            compiled: (0..thunks.len()).map(|_| OnceCell::new()).collect(),
             tasks: PriorityQueue::new(),
             thunks,
             func_offset,
             func_deps,
+            jit,
         }
     }
 
@@ -69,8 +79,8 @@ impl Engine {
     }
 
     pub fn eval_thunk(&mut self, idx: usize, env: &mut Env) -> Result<i::Value> {
-        let self_mut = unsafe { &mut *(self as *mut Self) };
+        let func =  self.compiled[idx].get_or_init(|| self.jit.compile(&self.thunks[idx]));
-        self.thunks[idx].eval(self_mut, env)
+        Ok(unsafe { func(self as *const Engine, env as *const Env).into() })
     }
 
     pub fn eval_func_deps(&mut self, idx: usize, env: &mut Env) -> Result<()> {

src/env.rs

@@ -25,7 +25,10 @@ impl Env {
         }
     }
 
-    pub fn with_new_cache<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> (T, HashMap<usize, Value>) {
+    pub fn with_new_cache<T>(
+        &mut self,
+        f: impl FnOnce(&mut Self) -> T,
+    ) -> (T, HashMap<usize, Value>) {
         self.cache.push(HashMap::new());
         let ret = f(self);
         (ret, self.cache.pop().unwrap())
@@ -76,7 +79,7 @@ impl Env {
         self.args[self.args.len() - level - 1].clone()
     }
 
-    pub fn lookup_with(&self, symbol: &EcoString) -> Option<Value> {
+    pub fn lookup_with(&self, symbol: &str) -> Option<Value> {
         for with in self.with.iter().rev() {
             if let Some(ret) = with.get(symbol) {
                 return Some(ret.clone());

src/eval/jit/compile.rs

@@ -1,39 +1,66 @@
+use std::{alloc::Layout, ffi::CStr};
+
 use inkwell::values::{FunctionValue, StructValue};
 
-use crate::ir::*;
 use crate::ty::common as c;
 use crate::ty::internal::Value;
+use crate::{eval::jit::JITValue, ir::*};
 
 use super::{JITContext, ValueTag};
 
 pub trait JITCompile {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc>;
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc>;
 }
 
 impl JITCompile for Attrs {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for List {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for HasAttr {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for BinOp {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         use BinOpKind::*;
         use ValueTag::*;
-        let lhs = self.lhs.compile(ctx, func);
+        let lhs = self.lhs.compile(ctx, func, values);
-        let rhs = self.rhs.compile(ctx, func);
+        let rhs = self.rhs.compile(ctx, func, values);
         let lhs_tag = ctx.get_tag(lhs);
         let rhs_tag = ctx.get_tag(rhs);
         let tag = ctx
@@ -86,10 +113,7 @@ impl JITCompile for BinOp {
                     .build_int_add(ctx.get_int(lhs), ctx.get_int(rhs), "add")
                     .unwrap();
                 ctx.builder
-                    .build_store(
+                    .build_store(res, ctx.helpers.new_value(Int, val.into()))
-                        res,
-                        ctx.helpers.new_value(Int, val.into())
-                    )
                     .unwrap();
                 ctx.builder.position_at_end(int_float);
                 let val = ctx
@@ -107,10 +131,7 @@ impl JITCompile for BinOp {
                     )
                     .unwrap();
                 ctx.builder
-                    .build_store(
+                    .build_store(res, ctx.helpers.new_value(Float, val.into()))
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
                     .unwrap();
                 ctx.builder.position_at_end(float_int);
                 let val = ctx
@@ -128,10 +149,7 @@ impl JITCompile for BinOp {
                     )
                     .unwrap();
                 ctx.builder
-                    .build_store(
+                    .build_store(res, ctx.helpers.new_value(Float, val.into()))
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
                     .unwrap();
                 ctx.builder.position_at_end(int_int);
                 let val = ctx
@@ -139,10 +157,7 @@ impl JITCompile for BinOp {
                     .build_float_add(ctx.get_float(lhs), ctx.get_float(rhs), "add")
                     .unwrap();
                 ctx.builder
-                    .build_store(
+                    .build_store(res, ctx.helpers.new_value(Float, val.into()))
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
                     .unwrap();
                 ctx.builder.position_at_end(fallback);
             }
@@ -165,10 +180,7 @@ impl JITCompile for BinOp {
                     .build_or(ctx.get_bool(lhs), ctx.get_bool(rhs), "or")
                     .unwrap();
                 ctx.builder
-                    .build_store(
+                    .build_store(res, ctx.helpers.new_value(Bool, val.into()))
-                        res,
-                        ctx.helpers.new_value(Bool, val.into())
-                    )
                     .unwrap();
                 ctx.builder.position_at_end(fallback);
             }
@@ -184,9 +196,14 @@ impl JITCompile for BinOp {
 }
 
 impl JITCompile for UnOp {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!();
-        let rhs = self.rhs.compile(ctx, func);
+        let rhs = self.rhs.compile(ctx, func, values);
         let tag = ctx.get_tag(rhs);
         let fallback = ctx.context.append_basic_block(func, "fallback");
         let ret = ctx.context.append_basic_block(func, "fallback");
@@ -206,55 +223,143 @@ impl JITCompile for UnOp {
 }
 
 impl JITCompile for Select {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for If {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for LoadFunc {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, _: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        ctx.helpers.new_value(ValueTag::Function, ctx.helpers.const_int(self.idx as i64).into())
+        &self,
+        ctx: &JITContext<'gc>,
+        _: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        ctx.helpers.new_value(
+            ValueTag::Function,
+            ctx.helpers.const_int(self.idx as i64).into(),
+        )
     }
 }
 
 impl JITCompile for Call {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        todo!()
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        let args = ctx
+            .builder
+            .build_call(
+                ctx.helpers.alloc_array,
+                &[ctx.helpers.const_ptr_int(self.args.len()).into()],
+                "alloc_args",
+            )
+            .unwrap()
+            .try_as_basic_value()
+            .unwrap_left()
+            .into_pointer_value();
+        for (i, arg) in self.args.iter().enumerate() {
+            ctx.builder
+                .build_store(
+                    unsafe {
+                        args.const_in_bounds_gep(
+                            ctx.helpers.value_type,
+                            &[ctx.helpers.const_ptr_int(i)],
+                        )
+                    },
+                    arg.compile(ctx, func, values),
+                )
+                .unwrap();
+        }
+        ctx.builder
+            .build_call(
+                ctx.helpers.call,
+                &[
+                    self.func.compile(ctx, func, values).into(),
+                    args.into(),
+                    ctx.helpers.const_ptr_int(self.args.len()).into(),
+                    func.get_first_param().unwrap().into(),
+                    func.get_last_param().unwrap().into(),
+                ],
+                "call",
+            )
+            .unwrap()
+            .try_as_basic_value()
+            .unwrap_left()
+            .try_into()
+            .unwrap()
     }
 }
 
 impl JITCompile for Let {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        todo!()
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        unreachable!()
     }
 }
 
 impl JITCompile for With {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for Assert {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for ConcatStrings {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for Const {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         use c::Const::*;
         match self.val {
             Bool(x) => ctx.helpers.new_bool(x),
@@ -266,37 +371,97 @@ impl JITCompile for Const {
 }
 
 impl JITCompile for String {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }
 
 impl JITCompile for Var {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        todo!()
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        let env = func.get_nth_param(1).unwrap();
+        let ptr = self.sym.as_ptr();
+        let len = self.sym.len();
+        // values.push(Value::String(self.sym.clone()));
+        ctx.builder
+            .build_direct_call(
+                ctx.helpers.lookup,
+                &[
+                    env.into(),
+                    ctx.const_ptr(ptr as *const ()).into(),
+                    ctx.helpers.const_ptr_int(len).into(),
+                ],
+                "lookup",
+            )
+            .unwrap()
+            .try_as_basic_value()
+            .unwrap_left()
+            .try_into()
+            .unwrap()
     }
 }
 
 impl JITCompile for Arg {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        todo!()
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        let env = func.get_last_param().unwrap();
+        let ret = ctx.builder.build_alloca(ctx.helpers.value_type, "alloca_ret").unwrap();
+        ctx.builder
+            .build_direct_call(
+                ctx.helpers.lookup_arg,
+                &[env.into(), ctx.helpers.const_ptr_int(self.level).into(), ret.into()],
+                "lookup_arg",
+            )
+            .unwrap();
+        ctx.builder.build_load(ctx.helpers.value_type,ret, "load_ret").unwrap().into_struct_value()
     }
 }
 
 impl JITCompile for LetVar {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        todo!()
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        unreachable!()
     }
 }
 
 impl JITCompile for Thunk {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, _: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
-        ctx.helpers.new_value(ValueTag::Thunk, ctx.helpers.const_int(self.idx as i64).into())
+        &self,
+        ctx: &JITContext<'gc>,
+        _: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
+        ctx.helpers.new_value(
+            ValueTag::Thunk,
+            ctx.helpers.const_int(self.idx as i64).into(),
+        )
     }
 }
 
 impl JITCompile for Path {
-    fn compile<'gc>(&self, ctx: &JITContext<'gc>, func: FunctionValue<'gc>) -> StructValue<'gc> {
+    fn compile<'gc>(
+        &self,
+        ctx: &JITContext<'gc>,
+        func: FunctionValue<'gc>,
+        values: &mut Vec<Value>,
+    ) -> StructValue<'gc> {
         todo!()
     }
 }

src/eval/jit/helpers.rs

@@ -1,4 +1,7 @@
+use std::alloc::Layout;
+use std::ffi::CStr;
 use std::ptr::NonNull;
+use std::{slice, str};
 
 use inkwell::AddressSpace;
 use inkwell::context::Context;
@@ -33,39 +36,36 @@ pub struct Helpers<'ctx> {
     pub eq: FunctionValue<'ctx>,
     pub or: FunctionValue<'ctx>,
     pub call: FunctionValue<'ctx>,
-    pub arg: FunctionValue<'ctx>,
+    pub lookup_arg: FunctionValue<'ctx>,
-    pub lookup_let: FunctionValue<'ctx>,
     pub lookup: FunctionValue<'ctx>,
     pub force: FunctionValue<'ctx>,
+
+    pub alloc_array: FunctionValue<'ctx>,
 }
 
 impl<'ctx> Helpers<'ctx> {
     pub fn new(
-        context: &'ctx Context,
+        ctx: &'ctx Context,
         module: &Module<'ctx>,
         execution_engine: &ExecutionEngine<'ctx>,
     ) -> Self {
-        let int_type = context.i64_type();
+        let int_type = ctx.i64_type();
-        let float_type = context.f64_type();
+        let float_type = ctx.f64_type();
-        let bool_type = context.bool_type();
+        let bool_type = ctx.bool_type();
-        let ptr_int_type = context.ptr_sized_int_type(execution_engine.get_target_data(), None);
+        let ptr_int_type = ctx.ptr_sized_int_type(execution_engine.get_target_data(), None);
-        let ptr_type = context.ptr_type(AddressSpace::default());
+        let ptr_type = ctx.ptr_type(AddressSpace::default());
-        let value_type = context.struct_type(&[int_type.into(), int_type.into(), int_type.into()], false);
+        let value_type =
-        let func_type = value_type.fn_type(&[ptr_type.into()], false);
+            ctx.struct_type(&[int_type.into(), int_type.into(), int_type.into()], false);
+        let func_type = value_type.fn_type(&[ptr_type.into(), ptr_type.into()], false);
 
         let new_thunk = module.add_function(
             "new_thunk",
             value_type.fn_type(&[ptr_type.into(), ptr_type.into()], false),
             None,
         );
-        let debug = module.add_function(
-            "debug",
-            context.void_type().fn_type(&[value_type.into()], false),
-            None,
-        );
         let capture_env = module.add_function(
             "capture_env",
-            context
+            ctx
                 .void_type()
                 .fn_type(&[value_type.into(), ptr_type.into()], false),
             None,
@@ -106,6 +106,7 @@ impl<'ctx> Helpers<'ctx> {
                 &[
                     value_type.into(),
                     ptr_type.into(),
+                    ptr_int_type.into(),
                     ptr_type.into(),
                     ptr_type.into(),
                 ],
@@ -113,22 +114,25 @@ impl<'ctx> Helpers<'ctx> {
             ),
             None,
         );
-        let arg = module.add_function(
+        let debug = module.add_function(
-            "arg",
+            "debug",
-            value_type.fn_type(&[ptr_int_type.into(), ptr_type.into()], false),
+            ctx
+                .void_type()
+                .fn_type(&[ptr_type.into(), int_type.into()], false),
             None,
         );
-        let lookup_let = module.add_function(
+        let lookup_arg = module.add_function(
-            "lookup_let",
+            "lookup_arg",
-            value_type.fn_type(
+            // value_type.fn_type(&[ptr_type.into(), int_type.into()], false),
-                &[ptr_int_type.into(), ptr_int_type.into(), ptr_type.into()],
+            ctx.void_type().fn_type(&[ptr_type.into(), int_type.into(), ptr_type.into()], false),
-                false,
-            ),
             None,
         );
         let lookup = module.add_function(
             "lookup",
-            value_type.fn_type(&[ptr_int_type.into(), ptr_type.into()], false),
+            value_type.fn_type(
+                &[ptr_int_type.into(), ptr_type.into(), ptr_int_type.into()],
+                false,
+            ),
             None,
         );
         let force = module.add_function(
@@ -139,6 +143,11 @@ impl<'ctx> Helpers<'ctx> {
             ),
             None,
         );
+        let alloc_array = module.add_function(
+            "alloc_array",
+            value_type.fn_type(&[ptr_int_type.into()], false),
+            None,
+        );
 
         execution_engine.add_global_mapping(&new_thunk, helper_new_thunk as _);
         execution_engine.add_global_mapping(&debug, helper_debug as _);
@@ -150,10 +159,10 @@ impl<'ctx> Helpers<'ctx> {
         execution_engine.add_global_mapping(&eq, helper_eq as _);
         execution_engine.add_global_mapping(&or, helper_or as _);
         execution_engine.add_global_mapping(&call, helper_call as _);
-        execution_engine.add_global_mapping(&arg, helper_arg as _);
+        execution_engine.add_global_mapping(&lookup_arg, helper_lookup_arg as _);
-        execution_engine.add_global_mapping(&lookup_let, helper_lookup_let as _);
         execution_engine.add_global_mapping(&lookup, helper_lookup as _);
         execution_engine.add_global_mapping(&force, helper_force as _);
+        execution_engine.add_global_mapping(&alloc_array, helper_alloc_array as _);
 
         Helpers {
             int_type,
@@ -175,10 +184,11 @@ impl<'ctx> Helpers<'ctx> {
             eq,
             or,
             call,
-            arg,
+            lookup_arg,
-            lookup_let,
             lookup,
             force,
+
+            alloc_array,
         }
     }
 
@@ -186,10 +196,14 @@ impl<'ctx> Helpers<'ctx> {
         self.value_type.const_named_struct(&[
             self.const_int(tag as i64).into(),
             data,
-            self.int_type.const_zero().into()
+            self.int_type.const_zero().into(),
         ])
     }
 
+    pub fn const_ptr_int(&self, int: usize) -> IntValue<'ctx> {
+        self.ptr_int_type.const_int(int as _, false)
+    }
+
     pub fn const_int(&self, int: i64) -> IntValue<'ctx> {
         self.int_type.const_int(int as _, false)
     }
@@ -219,10 +233,6 @@ impl<'ctx> Helpers<'ctx> {
     }
 }
 
-extern "C" fn helper_debug(value: JITValue) {
-    dbg!(value.tag);
-}
-
 extern "C" fn helper_capture_env(thunk: JITValue, env: *const Env) {
     todo!()
 }
@@ -329,28 +339,47 @@ extern "C" fn helper_or(lhs: JITValue, rhs: JITValue) -> JITValue {
 
 extern "C" fn helper_call(
     func: JITValue,
-    args: Box<[JITValue]>,
+    args: *mut JITValue,
-    engine: NonNull<Engine>,
+    len: usize,
-    env: NonNull<Env>,
+    mut engine: NonNull<Engine>,
+    mut env: NonNull<Env>,
 ) -> JITValue {
-    let func = Value::from(func);
+    let mut func = Value::from(func);
+    // TODO: Error Handling
+    let args = core::ptr::slice_from_raw_parts_mut(args, len);
+    let args = unsafe { Box::from_raw(args) };
+    func.call(
+        args.into_iter().map(Value::from).collect(),
+        unsafe { engine.as_mut() },
+        unsafe { env.as_mut() },
+    )
+    .unwrap();
     todo!()
 }
 
-extern "C" fn helper_arg(idx: usize, env: *const Env) -> JITValue {
+extern "C" fn helper_debug(env: NonNull<Env>, level: u64) {
-    let env = unsafe { env.as_ref() }.unwrap();
+    dbg!(env, level);
-    let val: JITValue = env.lookup_arg(idx).clone().into();
+    dbg!(unsafe { env.as_ref() }.lookup_arg(level as usize));
+}
+
+extern "C" fn helper_lookup_arg(env_ptr: NonNull<Env>, level: u64, ret: NonNull<JITValue>) {
+    dbg!(env_ptr, level);
+    let env_ref = unsafe { env_ptr.as_ref() };
+    let val: JITValue = env_ref.lookup_arg(level as usize).into();
+    unsafe { ret.write(val) }
+}
+
+extern "C" fn helper_lookup(env: NonNull<Env>, ptr: *const u8, len: usize) -> JITValue {
+    let env = unsafe { env.as_ref() };
+    // TODO: Error Handling
+    let val: JITValue = env
+        .lookup_with(unsafe { str::from_utf8_unchecked(slice::from_raw_parts(ptr, len)) })
+        .unwrap()
+        .clone()
+        .into();
     val
 }
 
-extern "C" fn helper_lookup_let(level: usize, idx: usize, env: *const Env) -> JITValue {
-    todo!()
-}
-
-extern "C" fn helper_lookup(sym: usize, env: *const Env) -> JITValue {
-    todo!()
-}
-
 extern "C" fn helper_force(
     thunk: JITValue,
     vm: NonNull<Engine>,
@@ -365,3 +394,7 @@ extern "C" fn helper_force(
 extern "C" fn helper_new_thunk(opcodes: *const ()) -> JITValue {
     todo!()
 }
+
+unsafe extern "C" fn helper_alloc_array(len: usize) -> *mut u8 {
+    unsafe { std::alloc::alloc(Layout::array::<JITValue>(len).unwrap()) }
+}

src/eval/jit/mod.rs

@@ -7,8 +7,11 @@ use inkwell::builder::Builder;
 use inkwell::context::Context;
 use inkwell::execution_engine::ExecutionEngine;
 use inkwell::module::Module;
-use inkwell::values::{AnyValue, BasicMetadataValueEnum, FloatValue, IntValue, StructValue};
+use inkwell::values::{
+    AnyValue, BasicMetadataValueEnum, FloatValue, IntValue, PointerValue, StructValue,
+};
 
+use crate::engine::Engine;
 use crate::env::Env;
 use crate::ir::Ir;
 use crate::ty::internal::Value;
@@ -48,14 +51,14 @@ pub union JITValueData {
     float: f64,
     bool: bool,
     ptr: *const (),
-    slice: Slice
+    slice: Slice,
 }
 
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct Slice {
     ptr: *const (),
-    len: usize
+    len: usize,
 }
 
 impl From<JITValue> for Value {
@@ -78,35 +81,39 @@ impl From<Value> for JITValue {
                 tag: ValueTag::Int,
                 data: JITValueData { int },
             },
-            Value::List(list) => JITValue { tag: ValueTag::List, data: JITValueData { slice: Slice { ptr: list.as_ptr() as *const (), len: list.len() } } },
+            Value::List(list) => JITValue {
+                tag: ValueTag::List,
+                data: JITValueData {
+                    slice: Slice {
+                        ptr: list.as_ptr() as *const (),
+                        len: list.len(),
+                    },
+                },
+            },
             Value::Func(idx) => JITValue {
                 tag: ValueTag::Function,
-                data: JITValueData {
+                data: JITValueData { int: idx as i64 },
-                    int: idx as i64
-                },
             },
             Value::Thunk(idx) => JITValue {
                 tag: ValueTag::Thunk,
-                data: JITValueData {
+                data: JITValueData { int: idx as i64 },
-                    int: idx as i64
-                },
             },
             _ => todo!(),
         }
     }
 }
 
-pub struct JITFunc<'ctx>(F, PhantomData<&'ctx mut ()>);
+pub struct JITFunc<'ctx, 'exec>(F<'ctx, 'exec>, PhantomData<&'exec mut ()>);
-type F = unsafe extern "C" fn(*const Env) -> JITValue;
+type F<'ctx, 'exec> = unsafe extern "C" fn(*const Engine<'ctx, 'exec>, *const Env) -> JITValue;
 
-impl From<F> for JITFunc<'_> {
+impl<'ctx, 'exec> From<F<'ctx, 'exec>> for JITFunc<'ctx, 'exec> {
-    fn from(value: F) -> Self {
+    fn from(value: F<'ctx, 'exec>) -> Self {
         Self(value, PhantomData)
     }
 }
 
-impl Deref for JITFunc<'_> {
+impl<'ctx: 'exec, 'exec> Deref for JITFunc<'ctx, 'exec> {
-    type Target = F;
+    type Target = F<'ctx, 'exec>;
     fn deref(&self) -> &Self::Target {
         &self.0
     }
@@ -143,8 +150,28 @@ impl<'ctx> JITContext<'ctx> {
         }
     }
 
-    pub fn compile(&self, ir: Ir) -> JITFunc<'ctx> {
+    pub fn compile<'exec>(&'exec self, ir: &Ir) -> JITFunc<'ctx, 'exec> {
-        todo!()
+        let func = self
+            .module
+            .add_function("nixjit_func", self.helpers.func_type, None);
+        let entry = self.context.append_basic_block(func, "entry");
+        self.builder.position_at_end(entry);
+        // TODO:
+        let ret = ir.compile(self, func, &mut Vec::new());
+        self.builder.build_return(Some(&ret)).unwrap();
+        if func.verify(true) {
+            unsafe {
+                JITFunc(
+                    self.execution_engine
+                        .get_function(func.get_name().to_str().unwrap())
+                        .unwrap()
+                        .into_raw(),
+                    PhantomData,
+                )
+            }
+        } else {
+            todo!()
+        }
     }
 
     pub fn get_float(&self, val: StructValue<'ctx>) -> FloatValue<'ctx> {
@@ -219,11 +246,13 @@ impl<'ctx> JITContext<'ctx> {
             .into_int_value()
     }
 
-    pub fn call(&self, args: &[BasicMetadataValueEnum<'ctx>]) -> StructValue<'ctx> {
+    pub fn const_ptr(&self, ptr: *const ()) -> PointerValue<'ctx> {
         self.builder
-            .build_call(self.helpers.call, args, "call")
+            .build_int_to_ptr(
+                self.helpers.int_type.const_int(ptr as _, false),
+                self.helpers.ptr_type,
+                "ptrconv",
+            )
             .unwrap()
-            .as_any_value_enum()
-            .into_struct_value()
     }
 }

src/ir/mod.rs

@@ -90,9 +90,9 @@ macro_rules! ir {
         }
 
         impl JITCompile for Ir {
-            fn compile<'ctx>(&self, ctx: &JITContext<'ctx>, func: FunctionValue<'ctx>) -> StructValue<'ctx>{
+            fn compile<'ctx>(&self, ctx: &JITContext<'ctx>, func: FunctionValue<'ctx>, values: &mut Vec<Value>) -> StructValue<'ctx>{
                 match self {
-                    $(Ir::$ty(ir) => ir.compile(ctx, func),)*
+                    $(Ir::$ty(ir) => ir.compile(ctx, func, values),)*
                 }
             }
         }