imxyy1soope1/nixjit
1
0
Fork 0
Code Issues Pull Requests Actions Projects 1 Releases Activity

feat(jit): fib!

Browse Source
This commit is contained in:
imxyy_soope_
2025-05-19 19:29:25 +08:00
parent 6d26716412
commit 9e172bf013
10 changed files with 363 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files

218
src/jit/mod.rs
View File

@@ -1,18 +1,19 @@
use std::rc::Rc;
use inkwell::OptimizationLevel;
use inkwell::basic_block::BasicBlock;
use inkwell::builder::Builder;
use inkwell::context::Context;
use inkwell::execution_engine::{ExecutionEngine, JitFunction};
use inkwell::module::Module;
use inkwell::values::{BasicValueEnum, PointerValue};
use inkwell::values::{BasicValueEnum, FunctionValue, PointerValue};
use crate::bytecode::{Func, OpCode, UnOp};
use crate::error::*;
use crate::stack::Stack;
use crate::ty::common::Const;
use crate::ty::internal::{Thunk, Value};
use crate::vm::{Env, VM};
use crate::vm::{LetEnv, VM};
mod helpers;
@@ -48,17 +49,19 @@ pub struct JITValue {
pub union JITValueData {
int: i64,
float: f64,
boolean: bool,
bool: bool,
ptr: *const (),
}
impl<'jit: 'vm, 'vm> Into<Value<'jit, 'vm>> for JITValue {
fn into(self) -> Value<'jit, 'vm> {
impl<'jit: 'vm, 'vm> From<JITValue> for Value<'jit, 'vm> {
fn from(value: JITValue) -> Self {
use ValueTag::*;
match self.tag {
Int => Value::Const(Const::Int(unsafe { self.data.int })),
match value.tag {
Int => Value::Const(Const::Int(unsafe { value.data.int })),
Null => Value::Const(Const::Null),
_ => todo!("not implemented for {:?}", self.tag),
Function => Value::Func(unsafe { Rc::from_raw(value.data.ptr as *const _) }),
Thunk => Value::Thunk(unsafe { Rc::from_raw(value.data.ptr as *const _) }),
_ => todo!("not implemented for {:?}", value.tag),
}
}
}
@@ -70,13 +73,25 @@ impl From<Value<'_, '_>> for JITValue {
tag: ValueTag::Int,
data: JITValueData { int },
},
Value::Func(func) => JITValue {
tag: ValueTag::Function,
data: JITValueData {
ptr: Rc::into_raw(func) as *const _
}
},
Value::Thunk(thunk) => JITValue {
tag: ValueTag::Thunk,
data: JITValueData {
ptr: Rc::into_raw(thunk) as *const _
}
},
_ => todo!(),
}
}
}
pub type JITFunc<'jit, 'vm> =
unsafe extern "C" fn(*const VM<'jit>, *const Env<'jit, 'vm>) -> JITValue;
unsafe extern "C" fn(*const VM<'jit>, *const LetEnv<'jit, 'vm>) -> JITValue;
pub struct JITContext<'ctx> {
context: &'ctx Context,
@@ -125,17 +140,25 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
let func_ = self
.module
.add_function("nixjit_function", self.helpers.func_type, None);
let entry = self.context.append_basic_block(func_, "entry");
self.builder.position_at_end(entry);
let env = func_.get_nth_param(1).unwrap().into_pointer_value();
while let Some(opcode) = iter.next() {
self.single_op(opcode, vm, env, &mut stack)?;
}
let entry = self.context.append_basic_block(func_, "entry");
self.build_expr(
&mut iter,
vm,
env,
&mut stack,
func_,
entry,
func.opcodes.len(),
)?;
assert_eq!(stack.len(), 1);
let value = stack.pop();
let exit = self.context.append_basic_block(func_, "exit");
self.builder.build_unconditional_branch(exit)?;
self.builder.position_at_end(exit);
self.builder.build_return(Some(&value))?;
if func_.verify(false) {
func_.print_to_stderr();
if func_.verify(true) {
unsafe {
let name = func_.get_name().to_str().unwrap();
let addr = self.execution_engine.get_function(name).unwrap();
@@ -146,13 +169,82 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
}
}
fn build_expr<const CAP: usize>(
&self,
iter: &mut impl Iterator<Item = OpCode>,
vm: &'vm VM<'_>,
env: PointerValue<'ctx>,
stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
func: FunctionValue<'ctx>,
bb: BasicBlock<'ctx>,
mut length: usize,
) -> Result<usize> {
self.builder.position_at_end(bb);
while length > 1 {
let opcode = iter.next().unwrap();
let br = self.single_op(opcode, vm, env, stack)?;
length -= 1;
if br > 0 {
let consq = self.context.append_basic_block(func, "consq");
let alter = self.context.append_basic_block(func, "alter");
let cont = self.context.append_basic_block(func, "cont");
let cond = self
.builder
.build_alloca(self.helpers.value_type, "cond_alloca")?;
let result = self
.builder
.build_alloca(self.helpers.value_type, "result_alloca")?;
self.builder.build_store(cond, stack.pop())?;
self.builder.build_conditional_branch(
self.builder
.build_load(
self.context.bool_type(),
self.builder.build_struct_gep(
self.helpers.value_type,
cond,
1,
"gep_cond",
)?,
"load_cond",
)?
.into_int_value(),
consq,
alter,
)?;
length -= br;
let br = self.build_expr(iter, vm, env, stack, func, consq, br)?;
self.builder.build_store(result, stack.pop())?;
self.builder.build_unconditional_branch(cont)?;
length -= br;
self.build_expr(iter, vm, env, stack, func, alter, br)?;
self.builder.build_store(result, stack.pop())?;
self.builder.build_unconditional_branch(cont)?;
self.builder.position_at_end(cont);
stack.push(self.builder.build_load(
self.helpers.value_type,
result,
"load_result",
)?)?;
}
}
if length > 0 {
self.single_op(iter.next().unwrap(), vm, env, stack)
} else {
Ok(0)
}
}
#[inline(always)]
fn single_op<const CAP: usize>(
&self,
opcode: OpCode,
vm: &'vm VM<'_>,
env: PointerValue<'ctx>,
stack: &mut Stack<BasicValueEnum<'ctx>, CAP>,
) -> Result<()> {
) -> Result<usize> {
match opcode {
OpCode::Const { idx } => {
use Const::*;
@@ -186,54 +278,14 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
.try_as_basic_value()
.left()
.unwrap(),
_ => todo!(),
Not => self
.builder
.build_direct_call(self.helpers.not, &[rhs.into(), env.into()], "call_neg")?
.try_as_basic_value()
.left()
.unwrap(),
})?
}
OpCode::Func { idx } => {
let func = vm.get_func(idx);
let jit_func_ptr = self.compile_function(func, vm)?;
let jit_value = self
.helpers
.value_type
.const_named_struct(&[
self.helpers
.int_type
.const_int(ValueTag::Function as _, false)
.into(),
self.helpers
.ptr_int_type
.const_int(unsafe { jit_func_ptr.as_raw() } as _, false)
.into(),
])
.into();
stack.push(jit_value)?;
}
OpCode::Call { arity } => {
// Assuming arity is 1 for the test case
assert_eq!(arity, 1);
let arg = stack.pop();
let func_value = stack.pop();
let func_ptr = self
.builder
.build_extract_value(func_value.into_struct_value(), 1, "func_ptr")?
.into_pointer_value();
let result = self
.builder
.build_direct_call(
self.helpers.call,
&[
self.new_ptr(vm).into(),
env.into(),
func_ptr.into(),
arg.into(),
],
"call_func",
)?
.try_as_basic_value()
.left()
.unwrap();
stack.push(result)?;
}
OpCode::BinOp { op } => {
use crate::bytecode::BinOp;
let rhs = stack.pop();
@@ -269,9 +321,22 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
let result = self
.builder
.build_direct_call(
self.helpers.add,
self.helpers.eq,
&[lhs.into(), rhs.into()],
"call_add",
"call_eq",
)?
.try_as_basic_value()
.left()
.unwrap();
stack.push(result)?;
}
BinOp::Or => {
let result = self
.builder
.build_direct_call(
self.helpers.or,
&[lhs.into(), rhs.into()],
"call_or",
)?
.try_as_basic_value()
.left()
@@ -299,8 +364,29 @@ impl<'vm, 'ctx: 'vm> JITContext<'ctx> {
.left()
.unwrap(),
)?,
OpCode::Call { arity } => {
// TODO:
assert_eq!(arity, 1);
let mut args = Vec::with_capacity(arity);
for _ in 0..arity {
args.insert(0, stack.pop());
}
let func = self.builder
.build_direct_call(self.helpers.force, &[stack.pop().into(), self.new_ptr(vm).into()], "force")?
.try_as_basic_value()
.left()
.unwrap();
let ret = self.builder
.build_direct_call(self.helpers.call, &[func.into(), args[0].into(), self.new_ptr(vm).into()], "call")?
.try_as_basic_value()
.left()
.unwrap();
stack.push(ret)?;
}
OpCode::JmpIfFalse { step } => return Ok(step),
OpCode::Jmp { step } => return Ok(step),
_ => todo!("{opcode:?} not implemented"),
}
Ok(())
Ok(0)
}
}