feat: JIT (WIP)

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 rhs = self.rhs.compile(ctx, func);
+        let lhs = self.lhs.compile(ctx, func, values);
+        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(
-                        res,
-                        ctx.helpers.new_value(Int, val.into())
-                    )
+                    .build_store(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(
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
+                    .build_store(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(
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
+                    .build_store(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(
-                        res,
-                        ctx.helpers.new_value(Float, val.into())
-                    )
+                    .build_store(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(
-                        res,
-                        ctx.helpers.new_value(Bool, val.into())
-                    )
+                    .build_store(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> {
-        ctx.helpers.new_value(ValueTag::Function, ctx.helpers.const_int(self.idx as i64).into())
+    fn compile<'gc>(
+        &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> {
-        todo!()
+    fn compile<'gc>(
+        &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> {
-        todo!()
+    fn compile<'gc>(
+        &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> {
-        todo!()
+    fn compile<'gc>(
+        &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> {
-        todo!()
+    fn compile<'gc>(
+        &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> {
-        todo!()
+    fn compile<'gc>(
+        &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> {
-        ctx.helpers.new_value(ValueTag::Thunk, ctx.helpers.const_int(self.idx as i64).into())
+    fn compile<'gc>(
+        &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!()
     }
 }