refactor: reduce coupling

This commit is contained in:
2025-07-28 21:37:27 +08:00
parent 78e3c5a26e
commit 7afb2a7b1c
53 changed files with 2964 additions and 3444 deletions

View File

@@ -0,0 +1,19 @@
[package]
name = "nixjit_jit"
version = "0.1.0"
edition = "2024"
[dependencies]
hashbrown = "0.15"
cranelift = "0.122"
cranelift-module = "0.122"
cranelift-jit = "0.122"
cranelift-native = "0.122"
nixjit_error = { path = "../nixjit_error" }
nixjit_eval = { path = "../nixjit_eval" }
nixjit_hir = { path = "../nixjit_hir" }
nixjit_ir = { path = "../nixjit_ir" }
nixjit_lir = { path = "../nixjit_lir" }
nixjit_value = { path = "../nixjit_value" }

View File

@@ -0,0 +1,537 @@
use cranelift::codegen::ir::{self, StackSlot};
use cranelift::prelude::*;
use nixjit_eval::{EvalContext, Value};
use nixjit_ir::*;
use nixjit_lir::Lir;
use super::{Context, JITContext};
pub trait JITCompile<Ctx: JITContext> {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot;
}
impl<Ctx: JITContext> JITCompile<Ctx> for ExprId {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Lir {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for AttrSet {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let attrs = ctx.create_attrs();
for (k, v) in self.stcs.iter() {
let v = v.compile(ctx, engine, env);
ctx.push_attr(attrs, k, v);
}
ctx.finalize_attrs(attrs)
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for List {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let array = ctx.alloc_array(self.items.len());
for (i, item) in self.items.iter().enumerate() {
let item = item.compile(ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, item, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, item, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, item, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, item, 24);
ctx.builder
.ins()
.store(MemFlags::new(), tag, array, i as i32 * 32);
ctx.builder
.ins()
.store(MemFlags::new(), val0, array, i as i32 * 32 + 8);
ctx.builder
.ins()
.store(MemFlags::new(), val1, array, i as i32 * 32 + 16);
ctx.builder
.ins()
.store(MemFlags::new(), val2, array, i as i32 * 32 + 24);
}
ctx.create_list(array, self.items.len())
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for HasAttr {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for BinOp {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use BinOpKind::*;
let lhs = self.lhs.compile(ctx, engine, env);
let rhs = self.rhs.compile(ctx, engine, env);
let lhs_tag = ctx.get_tag(lhs);
let rhs_tag = ctx.get_tag(rhs);
let eq = ctx.builder.ins().icmp(IntCC::Equal, lhs_tag, rhs_tag);
let eq_block = ctx.builder.create_block();
let neq_block = ctx.builder.create_block();
let exit_block = ctx.builder.create_block();
ctx.builder.ins().brif(eq, eq_block, [], neq_block, []);
match self.kind {
Add => {
ctx.builder.switch_to_block(eq_block);
let default_block = ctx.builder.create_block();
let int_block = ctx.builder.create_block();
let float_block = ctx.builder.create_block();
let float_check_block = ctx.builder.create_block();
let is_int =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::INT as i64);
ctx.builder
.ins()
.brif(is_int, int_block, [], float_check_block, []);
ctx.builder.switch_to_block(int_block);
let lhs_value = ctx.get_small_value(types::I64, lhs);
let rhs_value = ctx.get_small_value(types::I64, rhs);
let result = ctx.builder.ins().iadd(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: Non-float
ctx.builder.switch_to_block(float_check_block);
let is_float =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::FLOAT as i64);
ctx.builder
.ins()
.brif(is_float, float_block, [], default_block, []);
ctx.builder.switch_to_block(float_block);
let lhs_value = ctx.get_small_value(types::F64, lhs);
let rhs_value = ctx.get_small_value(types::F64, rhs);
let result = ctx.builder.ins().fadd(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: finish this
ctx.builder.switch_to_block(default_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(neq_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.seal_block(default_block);
ctx.builder.seal_block(int_block);
ctx.builder.seal_block(float_check_block);
ctx.builder.seal_block(float_block);
}
Sub => {
ctx.builder.switch_to_block(eq_block);
let default_block = ctx.builder.create_block();
let int_block = ctx.builder.create_block();
let float_block = ctx.builder.create_block();
let float_check_block = ctx.builder.create_block();
let is_int =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::INT as i64);
ctx.builder
.ins()
.brif(is_int, int_block, [], float_check_block, []);
ctx.builder.switch_to_block(int_block);
let lhs_value = ctx.get_small_value(types::I64, lhs);
let rhs_value = ctx.get_small_value(types::I64, rhs);
let result = ctx.builder.ins().isub(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: Non-float
ctx.builder.switch_to_block(float_check_block);
let is_float =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::FLOAT as i64);
ctx.builder
.ins()
.brif(is_float, float_block, [], default_block, []);
ctx.builder.switch_to_block(float_block);
let lhs_value = ctx.get_small_value(types::F64, lhs);
let rhs_value = ctx.get_small_value(types::F64, rhs);
let result = ctx.builder.ins().fsub(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: finish this
ctx.builder.switch_to_block(default_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(neq_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.seal_block(default_block);
ctx.builder.seal_block(int_block);
ctx.builder.seal_block(float_check_block);
ctx.builder.seal_block(float_block);
}
Div => {
ctx.builder.switch_to_block(eq_block);
let default_block = ctx.builder.create_block();
let int_block = ctx.builder.create_block();
let float_block = ctx.builder.create_block();
let float_check_block = ctx.builder.create_block();
let is_int =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::INT as i64);
ctx.builder
.ins()
.brif(is_int, int_block, [], float_check_block, []);
ctx.builder.switch_to_block(int_block);
let lhs_value = ctx.get_small_value(types::I64, lhs);
let rhs_value = ctx.get_small_value(types::I64, rhs);
let result = ctx.builder.ins().sdiv(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: Non-float
ctx.builder.switch_to_block(float_check_block);
let is_float =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::FLOAT as i64);
ctx.builder
.ins()
.brif(is_float, float_block, [], default_block, []);
ctx.builder.switch_to_block(float_block);
let lhs_value = ctx.get_small_value(types::F64, lhs);
let rhs_value = ctx.get_small_value(types::F64, rhs);
let result = ctx.builder.ins().fdiv(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, &[]);
// FIXME: finish this
ctx.builder.switch_to_block(default_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(neq_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.seal_block(default_block);
ctx.builder.seal_block(int_block);
ctx.builder.seal_block(float_check_block);
ctx.builder.seal_block(float_block);
}
And => {
ctx.builder.switch_to_block(eq_block);
let bool_block = ctx.builder.create_block();
let non_bool_block = ctx.builder.create_block();
let is_bool =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::BOOL as i64);
ctx.builder
.ins()
.brif(is_bool, bool_block, [], non_bool_block, []);
ctx.builder.switch_to_block(bool_block);
let lhs_value = ctx.get_small_value(types::I64, lhs);
let rhs_value = ctx.get_small_value(types::I64, rhs);
let result = ctx.builder.ins().band(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, []);
ctx.builder.switch_to_block(non_bool_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(neq_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.seal_block(bool_block);
ctx.builder.seal_block(non_bool_block);
}
Or => {
ctx.builder.switch_to_block(eq_block);
let bool_block = ctx.builder.create_block();
let non_bool_block = ctx.builder.create_block();
let is_bool =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, lhs_tag, Value::<Ctx>::BOOL as i64);
ctx.builder
.ins()
.brif(is_bool, bool_block, [], non_bool_block, []);
ctx.builder.switch_to_block(bool_block);
let lhs_value = ctx.get_small_value(types::I64, lhs);
let rhs_value = ctx.get_small_value(types::I64, rhs);
let result = ctx.builder.ins().bor(lhs_value, rhs_value);
ctx.builder.ins().stack_store(lhs_tag, lhs, 0);
ctx.builder.ins().stack_store(result, lhs, 8);
ctx.builder.ins().jump(exit_block, []);
ctx.builder.switch_to_block(non_bool_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(neq_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.seal_block(bool_block);
ctx.builder.seal_block(non_bool_block);
}
Eq => {
ctx.builder.switch_to_block(eq_block);
ctx.eq(lhs, rhs);
ctx.builder.ins().jump(exit_block, []);
ctx.builder.switch_to_block(neq_block);
ctx.eq(lhs, rhs);
ctx.builder.ins().jump(exit_block, []);
}
_ => todo!(),
}
ctx.builder.seal_block(exit_block);
ctx.builder.seal_block(eq_block);
ctx.builder.seal_block(neq_block);
ctx.builder.switch_to_block(exit_block);
ctx.free_slot(rhs);
lhs
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for UnOp {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Attr {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use Attr::*;
match self {
Str(string) => ctx.create_string(string),
Dynamic(ir) => ir.compile(ctx, engine, env),
}
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Select {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let val = self.expr.compile(ctx, engine, env);
let attrpath = ctx.alloc_array(self.attrpath.len());
for (i, attr) in self.attrpath.iter().enumerate() {
let arg = attr.compile(ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, arg, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, arg, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, arg, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, arg, 24);
ctx.builder
.ins()
.store(MemFlags::new(), tag, attrpath, i as i32 * 32);
ctx.builder
.ins()
.store(MemFlags::new(), val0, attrpath, i as i32 * 32 + 8);
ctx.builder
.ins()
.store(MemFlags::new(), val1, attrpath, i as i32 * 32 + 16);
ctx.builder
.ins()
.store(MemFlags::new(), val2, attrpath, i as i32 * 32 + 24);
}
ctx.select(val, attrpath, self.attrpath.len(), engine, env);
val
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for If {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let cond = self.cond.compile(ctx, engine, env);
let cond_type = ctx.builder.ins().stack_load(types::I64, cond, 0);
let cond_value = ctx.builder.ins().stack_load(types::I64, cond, 8);
let true_block = ctx.builder.create_block();
let false_block = ctx.builder.create_block();
let exit_block = ctx.builder.create_block();
let error_block = ctx.builder.create_block();
let judge_block = ctx.builder.create_block();
let slot = ctx.alloca();
let is_bool =
ctx.builder
.ins()
.icmp_imm(IntCC::Equal, cond_type, Value::<Ctx>::BOOL as i64);
ctx.builder
.ins()
.brif(is_bool, judge_block, [], error_block, []);
ctx.builder.switch_to_block(judge_block);
ctx.builder
.ins()
.brif(cond_value, true_block, [], false_block, []);
ctx.builder.switch_to_block(true_block);
let ret = self.consq.compile(ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, ret, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, ret, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, ret, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, ret, 24);
ctx.builder.ins().stack_store(tag, slot, 0);
ctx.builder.ins().stack_store(val0, slot, 8);
ctx.builder.ins().stack_store(val1, slot, 16);
ctx.builder.ins().stack_store(val2, slot, 24);
ctx.builder.ins().jump(exit_block, []);
ctx.builder.switch_to_block(false_block);
let ret = self.alter.compile(ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, ret, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, ret, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, ret, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, ret, 24);
ctx.builder.ins().stack_store(tag, slot, 0);
ctx.builder.ins().stack_store(val0, slot, 8);
ctx.builder.ins().stack_store(val1, slot, 16);
ctx.builder.ins().stack_store(val2, slot, 24);
ctx.builder.ins().jump(exit_block, []);
ctx.builder.switch_to_block(error_block);
ctx.builder.ins().trap(TrapCode::unwrap_user(1));
ctx.builder.switch_to_block(exit_block);
slot
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Call {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let func = self.func.compile(ctx, engine, env);
let args = ctx.alloc_array(self.args.len());
for (i, arg) in self.args.iter().enumerate() {
let arg = arg.compile(ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, arg, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, arg, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, arg, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, arg, 24);
ctx.builder
.ins()
.store(MemFlags::new(), tag, args, i as i32 * 32);
ctx.builder
.ins()
.store(MemFlags::new(), val0, args, i as i32 * 32 + 8);
ctx.builder
.ins()
.store(MemFlags::new(), val1, args, i as i32 * 32 + 16);
ctx.builder
.ins()
.store(MemFlags::new(), val2, args, i as i32 * 32 + 24);
}
ctx.call(func, args, self.args.len(), engine, env);
func
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for With {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
let namespace = self.namespace.compile(ctx, engine, env);
ctx.enter_with(env, namespace);
let ret = self.expr.compile(ctx, engine, env);
ctx.exit_with(env);
ctx.free_slot(namespace);
ret
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Assert {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for ConcatStrings {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Const {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
use nixjit_value::Const::*;
let slot = ctx.alloca();
match self.val {
Bool(x) => {
let tag = ctx
.builder
.ins()
.iconst(types::I64, Value::<Ctx>::BOOL as i64);
let val = ctx.builder.ins().iconst(types::I64, x as i64);
ctx.builder.ins().stack_store(tag, slot, 0);
ctx.builder.ins().stack_store(val, slot, 8);
}
Int(x) => {
let tag = ctx
.builder
.ins()
.iconst(types::I64, Value::<Ctx>::INT as i64);
let val = ctx.builder.ins().iconst(types::I64, x);
ctx.builder.ins().stack_store(tag, slot, 0);
ctx.builder.ins().stack_store(val, slot, 8);
}
Float(x) => {
let tag = ctx
.builder
.ins()
.iconst(types::I64, Value::<Ctx>::FLOAT as i64);
let val = ctx.builder.ins().f64const(x);
ctx.builder.ins().stack_store(tag, slot, 0);
ctx.builder.ins().stack_store(val, slot, 8);
}
Null => {
let tag = ctx
.builder
.ins()
.iconst(types::I64, Value::<Ctx>::NULL as i64);
ctx.builder.ins().stack_store(tag, slot, 0);
}
}
slot
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Str {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
ctx.create_string(&self.val)
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Var {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
ctx.lookup(env, &self.sym)
}
}
impl<Ctx: JITContext> JITCompile<Ctx> for Path {
fn compile(&self, ctx: &mut Context<Ctx>, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
}
}

View File

@@ -0,0 +1,165 @@
use core::{slice, str};
use std::alloc::Layout;
use std::alloc::alloc;
use std::mem::MaybeUninit;
use std::ptr::NonNull;
use hashbrown::HashMap;
use nixjit_eval::{AttrSet, EvalContext, List, Value};
use super::JITContext;
pub extern "C" fn helper_call<Ctx: JITContext>(
func: &mut Value<Ctx>,
args_ptr: *mut Value<Ctx>,
args_len: usize,
ctx: &mut Ctx,
) {
// TODO: Error Handling
let args = core::ptr::slice_from_raw_parts_mut(args_ptr, args_len);
let args = unsafe { Box::from_raw(args) };
func.call(args.into_iter().collect(), ctx).unwrap();
}
pub extern "C" fn helper_lookup_stack<Ctx: JITContext>(
ctx: &Ctx,
offset: usize,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
ret.write(ctx.lookup_stack(offset).clone());
}
pub extern "C" fn helper_lookup_arg<Ctx: JITContext>(
ctx: &Ctx,
offset: usize,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
ret.write(ctx.lookup_arg(offset).clone());
}
pub extern "C" fn helper_lookup<Ctx: JITContext>(
ctx: &Ctx,
sym_ptr: *const u8,
sym_len: usize,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
// TODO: Error Handling
unsafe {
ret.write(
ctx.lookup_with(str::from_utf8_unchecked(slice::from_raw_parts(
sym_ptr, sym_len,
)))
.unwrap()
.clone(),
);
}
}
pub extern "C" fn helper_select<Ctx: JITContext>(
val: &mut Value<Ctx>,
path_ptr: *mut Value<Ctx>,
path_len: usize,
) {
let path = core::ptr::slice_from_raw_parts_mut(path_ptr, path_len);
let path = unsafe { Box::from_raw(path) };
val.select(path.into_iter().map(|mut val| {
val.coerce_to_string();
Ok(val.unwrap_string())
}))
.unwrap();
}
pub extern "C" fn helper_select_with_default<Ctx: JITContext>(
val: &mut Value<Ctx>,
path_ptr: *mut Value<Ctx>,
path_len: usize,
default: NonNull<Value<Ctx>>,
) {
let path = core::ptr::slice_from_raw_parts_mut(path_ptr, path_len);
let path = unsafe { Box::from_raw(path) };
val.select_with_default(
path.into_iter().map(|mut val| {
val.coerce_to_string();
Ok(val.unwrap_string())
}),
unsafe { default.read() },
)
.unwrap();
}
pub extern "C" fn helper_eq<Ctx: JITContext>(lhs: &mut Value<Ctx>, rhs: &Value<Ctx>) {
lhs.eq(rhs);
}
pub unsafe extern "C" fn helper_create_string<Ctx: JITContext>(
ptr: *const u8,
len: usize,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
unsafe {
ret.write(Value::String(
str::from_utf8_unchecked(slice::from_raw_parts(ptr, len)).to_owned(),
));
}
}
pub unsafe extern "C" fn helper_create_list<Ctx: JITContext>(
ptr: *mut Value<Ctx>,
len: usize,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
unsafe {
ret.write(Value::List(
List::from(Vec::from_raw_parts(ptr, len, len)).into(),
));
}
}
pub unsafe extern "C" fn helper_create_attrs<Ctx: JITContext>(
ret: &mut MaybeUninit<HashMap<String, Value<Ctx>>>,
) {
ret.write(HashMap::new());
}
pub unsafe extern "C" fn helper_push_attr<Ctx: JITContext>(
attrs: &mut HashMap<String, Value<Ctx>>,
sym_ptr: *const u8,
sym_len: usize,
val: NonNull<Value<Ctx>>,
) {
unsafe {
attrs.insert(
str::from_utf8_unchecked(slice::from_raw_parts(sym_ptr, sym_len)).to_owned(),
val.read(),
);
}
}
pub unsafe extern "C" fn helper_finalize_attrs<Ctx: JITContext>(
attrs: NonNull<HashMap<String, Value<Ctx>>>,
ret: &mut MaybeUninit<Value<Ctx>>,
) {
ret.write(Value::AttrSet(
AttrSet::from(unsafe { attrs.read() }).into(),
));
}
pub unsafe extern "C" fn helper_enter_with<Ctx: JITContext>(
ctx: &mut Ctx,
namespace: NonNull<Value<Ctx>>,
) {
ctx.enter_with(unsafe { namespace.read() }.unwrap_attr_set().into_inner());
}
pub unsafe extern "C" fn helper_exit_with<Ctx: JITContext>(ctx: &mut Ctx) {
ctx.exit_with();
}
pub unsafe extern "C" fn helper_alloc_array<Ctx: JITContext>(len: usize) -> *mut u8 {
unsafe { alloc(Layout::array::<Value<Ctx>>(len).unwrap()) }
}
pub extern "C" fn helper_dbg<Ctx: JITContext>(value: &Value<Ctx>) {
println!("{value:?}")
}

View File

@@ -0,0 +1,762 @@
use std::marker::PhantomData;
use std::ops::Deref;
use std::rc::Rc;
use cranelift::codegen::ir::Function;
use cranelift::codegen::ir::{self, ArgumentExtension, ArgumentPurpose, StackSlot};
use cranelift::prelude::*;
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{FuncId, Linkage, Module};
use hashbrown::{HashMap, HashSet};
use nixjit_eval::{EvalContext, Value};
use nixjit_lir::Lir;
mod compile;
mod helpers;
pub use compile::JITCompile;
use helpers::*;
pub trait JITContext: EvalContext + Sized {
fn lookup_stack(&self, offset: usize) -> &Value<Self>;
fn lookup_arg(&self, offset: usize) -> &Value<Self>;
fn enter_with(&mut self, namespace: Rc<HashMap<String, Value<Self>>>);
fn exit_with(&mut self);
}
type F<Ctx: JITContext> = unsafe extern "C" fn(*const Ctx, *mut Value<Ctx>);
pub struct JITFunc<Ctx: JITContext> {
func: F<Ctx>,
strings: HashSet<String>,
}
impl<Ctx: JITContext> Deref for JITFunc<Ctx> {
type Target = F<Ctx>;
fn deref(&self) -> &Self::Target {
&self.func
}
}
struct Context<'comp, 'ctx, Ctx: JITContext> {
pub compiler: &'comp mut JITCompiler<Ctx>,
pub builder: FunctionBuilder<'ctx>,
free_slots: Vec<StackSlot>,
strings: HashSet<String>,
}
impl<'comp, 'ctx, Ctx: JITContext> Context<'comp, 'ctx, Ctx> {
fn new(compiler: &'comp mut JITCompiler<Ctx>, builder: FunctionBuilder<'ctx>) -> Self {
Self {
compiler,
builder,
free_slots: Vec::new(),
strings: HashSet::new(),
}
}
fn alloca(&mut self) -> StackSlot {
self.free_slots.pop().map_or_else(
|| {
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
self.builder.create_sized_stack_slot(slot)
},
|x| x,
)
}
fn free_slot(&mut self, slot: StackSlot) {
self.free_slots.push(slot);
}
fn alloc_array(&mut self, len: usize) -> ir::Value {
let len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, len as i64);
let alloc_array = self
.compiler
.module
.declare_func_in_func(self.compiler.alloc_array, self.builder.func);
let inst = self.builder.ins().call(alloc_array, &[len]);
self.builder.inst_results(inst)[0]
}
fn create_string(&mut self, string: &str) -> StackSlot {
let string = self
.strings
.get_or_insert_with(string, |_| string.to_owned());
let ptr = self
.builder
.ins()
.iconst(self.compiler.ptr_type, string.as_ptr() as i64);
let len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, string.len() as i64);
let create_string = self
.compiler
.module
.declare_func_in_func(self.compiler.create_string, self.builder.func);
let slot = self.alloca();
let ret = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(create_string, &[ptr, len, ret]);
slot
}
fn create_list(&mut self, ptr: ir::Value, len: usize) -> StackSlot {
let len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, len as i64);
let create_list = self
.compiler
.module
.declare_func_in_func(self.compiler.create_list, self.builder.func);
let slot = self.alloca();
let ret = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(create_list, &[ptr, len, ret]);
slot
}
fn create_attrs(&mut self) -> StackSlot {
let create_attrs = self
.compiler
.module
.declare_func_in_func(self.compiler.create_attrs, self.builder.func);
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 40, 3);
let slot = self.builder.create_sized_stack_slot(slot);
let ret = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(create_attrs, &[ret]);
slot
}
fn push_attr(&mut self, attrs: StackSlot, sym: &str, val: StackSlot) {
self.free_slot(attrs);
self.free_slot(val);
let attrs = self.builder.ins().stack_addr(types::I64, attrs, 0);
let val = self.builder.ins().stack_addr(types::I64, val, 0);
let sym = self.strings.get_or_insert_with(sym, |_| sym.to_owned());
let ptr = self
.builder
.ins()
.iconst(self.compiler.ptr_type, sym.as_ptr() as i64);
let len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, sym.len() as i64);
let push_attr = self
.compiler
.module
.declare_func_in_func(self.compiler.push_attr, self.builder.func);
self.builder.ins().call(push_attr, &[attrs, ptr, len, val]);
}
fn finalize_attrs(&mut self, attrs: StackSlot) -> StackSlot {
let attrs = self.builder.ins().stack_addr(types::I64, attrs, 0);
let finalize_attrs = self
.compiler
.module
.declare_func_in_func(self.compiler.finalize_attrs, self.builder.func);
let slot = self.alloca();
let ret = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(finalize_attrs, &[attrs, ret]);
slot
}
fn enter_with(&mut self, env: ir::Value, namespace: StackSlot) {
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, namespace, 0);
let enter_with = self
.compiler
.module
.declare_func_in_func(self.compiler.enter_with, self.builder.func);
self.builder.ins().call(enter_with, &[env, ptr]);
}
fn exit_with(&mut self, env: ir::Value) {
let exit_with = self
.compiler
.module
.declare_func_in_func(self.compiler.exit_with, self.builder.func);
self.builder.ins().call(exit_with, &[env]);
}
fn dbg(&mut self, slot: StackSlot) {
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
let dbg = self
.compiler
.module
.declare_func_in_func(self.compiler.dbg, self.builder.func);
self.builder.ins().call(dbg, &[ptr]);
}
fn call(
&mut self,
func: StackSlot,
args_ptr: ir::Value,
args_len: usize,
engine: ir::Value,
env: ir::Value,
) {
let args_len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, args_len as i64);
let call = self
.compiler
.module
.declare_func_in_func(self.compiler.call, self.builder.func);
let func = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, func, 0);
self.builder
.ins()
.call(call, &[func, args_ptr, args_len, engine, env]);
}
fn lookup(&mut self, env: ir::Value, sym: &str) -> StackSlot {
let sym = self.strings.get_or_insert_with(sym, |_| sym.to_owned());
let ptr = self
.builder
.ins()
.iconst(self.compiler.ptr_type, sym.as_ptr() as i64);
let len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, sym.len() as i64);
let lookup = self
.compiler
.module
.declare_func_in_func(self.compiler.lookup, self.builder.func);
let slot = self.alloca();
let ret = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(lookup, &[env, ptr, len, ret]);
slot
}
fn lookup_stack(&mut self, env: ir::Value, idx: usize) -> StackSlot {
let slot = self.alloca();
let lookup_stack = self
.compiler
.module
.declare_func_in_func(self.compiler.lookup_stack, self.builder.func);
let idx = self
.builder
.ins()
.iconst(self.compiler.ptr_type, idx as i64);
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(lookup_stack, &[env, idx, ptr]);
slot
}
fn lookup_arg(&mut self, env: ir::Value, idx: usize) -> StackSlot {
let slot = self.alloca();
let lookup_arg = self
.compiler
.module
.declare_func_in_func(self.compiler.lookup_arg, self.builder.func);
let idx = self
.builder
.ins()
.iconst(self.compiler.ptr_type, idx as i64);
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder.ins().call(lookup_arg, &[env, idx, ptr]);
slot
}
fn select(
&mut self,
slot: StackSlot,
path_ptr: ir::Value,
path_len: usize,
engine: ir::Value,
env: ir::Value,
) {
let select = self
.compiler
.module
.declare_func_in_func(self.compiler.select, self.builder.func);
let path_len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, path_len as i64);
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
self.builder
.ins()
.call(select, &[ptr, path_ptr, path_len, engine, env]);
}
fn select_with_default(
&mut self,
slot: StackSlot,
path_ptr: ir::Value,
path_len: usize,
default: StackSlot,
engine: ir::Value,
env: ir::Value,
) {
let select_with_default = self
.compiler
.module
.declare_func_in_func(self.compiler.select_with_default, self.builder.func);
let path_len = self
.builder
.ins()
.iconst(self.compiler.ptr_type, path_len as i64);
let ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, slot, 0);
let default_ptr = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, default, 0);
self.builder.ins().call(
select_with_default,
&[ptr, path_ptr, path_len, default_ptr, engine, env],
);
}
pub fn eq(&mut self, lhs: StackSlot, rhs: StackSlot) {
let lhs = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, lhs, 0);
let rhs = self
.builder
.ins()
.stack_addr(self.compiler.ptr_type, rhs, 0);
let eq = self
.compiler
.module
.declare_func_in_func(self.compiler.eq, self.builder.func);
self.builder.ins().call(eq, &[lhs, rhs]);
}
pub fn get_tag(&mut self, slot: StackSlot) -> ir::Value {
self.builder.ins().stack_load(types::I64, slot, 0)
}
pub fn get_small_value(&mut self, ty: Type, slot: StackSlot) -> ir::Value {
self.builder.ins().stack_load(ty, slot, 8)
}
}
pub struct JITCompiler<Ctx: JITContext> {
ctx: codegen::Context,
module: JITModule,
builder_ctx: Option<FunctionBuilderContext>,
_marker: PhantomData<Ctx>,
int_type: Type,
float_type: Type,
bool_type: Type,
ptr_type: Type,
value_type: Type,
func_sig: Signature,
call: FuncId,
lookup_stack: FuncId,
lookup_arg: FuncId,
lookup: FuncId,
select: FuncId,
select_with_default: FuncId,
eq: FuncId,
alloc_array: FuncId,
create_string: FuncId,
create_list: FuncId,
create_attrs: FuncId,
push_attr: FuncId,
finalize_attrs: FuncId,
enter_with: FuncId,
exit_with: FuncId,
dbg: FuncId,
}
impl<Ctx: JITContext> Default for JITCompiler<Ctx> {
fn default() -> Self {
Self::new()
}
}
impl<Ctx: JITContext> JITCompiler<Ctx> {
pub fn new() -> Self {
let mut flag_builder = settings::builder();
flag_builder.set("use_colocated_libcalls", "false").unwrap();
flag_builder.set("is_pic", "false").unwrap();
let isa_builder = cranelift_native::builder().unwrap_or_else(|msg| {
panic!("host machine is not supported: {msg}");
});
let isa = isa_builder
.finish(settings::Flags::new(flag_builder))
.unwrap();
let mut builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
builder.symbol("helper_call", helper_call::<Ctx> as _);
builder.symbol("helper_lookup_stack", helper_lookup_stack::<Ctx> as _);
builder.symbol("helper_lookup_arg", helper_lookup_arg::<Ctx> as _);
builder.symbol("helper_lookup", helper_lookup::<Ctx> as _);
builder.symbol("helper_select", helper_select::<Ctx> as _);
builder.symbol(
"helper_select_with_default",
helper_select_with_default::<Ctx> as _,
);
builder.symbol("helper_eq", helper_eq::<Ctx> as _);
builder.symbol("helper_alloc_array", helper_alloc_array::<Ctx> as _);
builder.symbol("helper_create_string", helper_create_string::<Ctx> as _);
builder.symbol("helper_create_list", helper_create_list::<Ctx> as _);
builder.symbol("helper_create_attrs", helper_create_attrs::<Ctx> as _);
builder.symbol("helper_push_attr", helper_push_attr::<Ctx> as _);
builder.symbol("helper_finalize_attrs", helper_finalize_attrs::<Ctx> as _);
builder.symbol("helper_enter_with", helper_enter_with::<Ctx> as _);
builder.symbol("helper_exit_with", helper_exit_with::<Ctx> as _);
builder.symbol("helper_dbg", helper_dbg::<Ctx> as _);
let mut module = JITModule::new(builder);
let ctx = module.make_context();
let int_type = types::I64;
let float_type = types::F64;
let bool_type = types::I8;
let ptr_type = module.target_config().pointer_type();
let value_type = types::I128;
// fn(*const Context, *const Env, *mut Value)
let mut func_sig = module.make_signature();
func_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 3],
);
// fn(func: &mut Value, args_ptr: *mut Value, args_len: usize, engine: &mut Context, env:
// &mut Env)
let mut call_sig = module.make_signature();
call_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 5],
);
let call = module
.declare_function("helper_call", Linkage::Import, &call_sig)
.unwrap();
let mut lookup_stack_sig = module.make_signature();
lookup_stack_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 3],
);
let lookup_stack = module
.declare_function("helper_lookup_stack", Linkage::Import, &lookup_stack_sig)
.unwrap();
// fn(env: &Env, level: usize, ret: &mut MaybeUninit<Value>)
let mut lookup_arg_sig = module.make_signature();
lookup_arg_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 3],
);
let lookup_arg = module
.declare_function("helper_lookup_arg", Linkage::Import, &lookup_arg_sig)
.unwrap();
// fn(env: &Env, sym_ptr: *const u8, sym_len: usize, ret: &mut MaybeUninit<Value>)
let mut lookup_sig = module.make_signature();
lookup_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 4],
);
let lookup = module
.declare_function("helper_lookup", Linkage::Import, &lookup_sig)
.unwrap();
// fn(val: &mut Value, path_ptr: *mut Value, path_len: usize, engine: &mut Context, env: &mut Env)
let mut select_sig = module.make_signature();
select_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 5],
);
let select = module
.declare_function("helper_select", Linkage::Import, &select_sig)
.unwrap();
// fn(val: &mut Value, path_ptr: *mut Value, path_len: usize, default: NonNull<Value>, engine: &mut Context, env: &mut Env)
let mut select_with_default_sig = module.make_signature();
select_with_default_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 6],
);
let select_with_default = module
.declare_function(
"helper_select_with_default",
Linkage::Import,
&select_with_default_sig,
)
.unwrap();
let mut eq_sig = module.make_signature();
eq_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 2],
);
let eq = module
.declare_function("helper_eq", Linkage::Import, &eq_sig)
.unwrap();
let mut alloc_array_sig = module.make_signature();
alloc_array_sig.params.push(AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
});
alloc_array_sig.returns.push(AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
});
let alloc_array = module
.declare_function("helper_alloc_array", Linkage::Import, &alloc_array_sig)
.unwrap();
let mut create_string_sig = module.make_signature();
create_string_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 3],
);
let create_string = module
.declare_function("helper_create_string", Linkage::Import, &create_string_sig)
.unwrap();
let mut create_list_sig = module.make_signature();
create_list_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 3],
);
let create_list = module
.declare_function("helper_create_list", Linkage::Import, &create_list_sig)
.unwrap();
let mut create_attrs_sig = module.make_signature();
create_attrs_sig.params.push(AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
});
let create_attrs = module
.declare_function("helper_create_attrs", Linkage::Import, &create_attrs_sig)
.unwrap();
let mut push_attr_sig = module.make_signature();
push_attr_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 4],
);
let push_attr = module
.declare_function("helper_push_attr", Linkage::Import, &push_attr_sig)
.unwrap();
let mut finalize_attrs_sig = module.make_signature();
finalize_attrs_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 2],
);
let finalize_attrs = module
.declare_function(
"helper_finalize_attrs",
Linkage::Import,
&finalize_attrs_sig,
)
.unwrap();
let mut enter_with_sig = module.make_signature();
enter_with_sig.params.extend(
[AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
}; 2],
);
let enter_with = module
.declare_function("helper_enter_with", Linkage::Import, &enter_with_sig)
.unwrap();
let mut exit_with_sig = module.make_signature();
exit_with_sig.params.push(AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
});
let exit_with = module
.declare_function("helper_exit_with", Linkage::Import, &exit_with_sig)
.unwrap();
let mut dbg_sig = module.make_signature();
dbg_sig.params.push(AbiParam {
value_type: ptr_type,
purpose: ArgumentPurpose::Normal,
extension: ArgumentExtension::None,
});
let dbg = module
.declare_function("helper_dbg", Linkage::Import, &dbg_sig)
.unwrap();
Self {
builder_ctx: None,
_marker: PhantomData,
ctx,
module,
int_type,
float_type,
bool_type,
ptr_type,
value_type,
func_sig,
call,
lookup_stack,
lookup_arg,
lookup,
select,
select_with_default,
eq,
alloc_array,
create_string,
create_list,
create_attrs,
push_attr,
finalize_attrs,
enter_with,
exit_with,
dbg,
}
}
pub fn compile(&mut self, ir: &Lir, id: usize) -> JITFunc<Ctx> {
let func_id = self
.module
.declare_function(
format!("nixjit_thunk{id}").as_str(),
Linkage::Local,
&self.func_sig,
)
.unwrap();
let mut func = Function::new();
func.signature = self.func_sig.clone();
let mut builder_ctx = self.builder_ctx.take().unwrap_or_default();
let mut ctx = Context::new(self, FunctionBuilder::new(&mut func, &mut builder_ctx));
let entry = ctx.builder.create_block();
ctx.builder.append_block_params_for_function_params(entry);
ctx.builder.switch_to_block(entry);
let params = ctx.builder.block_params(entry);
let engine = params[0];
let env = params[1];
let ret = params[2];
let res = ir.compile(&mut ctx, engine, env);
let tag = ctx.builder.ins().stack_load(types::I64, res, 0);
let val0 = ctx.builder.ins().stack_load(types::I64, res, 8);
let val1 = ctx.builder.ins().stack_load(types::I64, res, 16);
let val2 = ctx.builder.ins().stack_load(types::I64, res, 24);
ctx.builder.ins().store(MemFlags::new(), tag, ret, 0);
ctx.builder.ins().store(MemFlags::new(), val0, ret, 8);
ctx.builder.ins().store(MemFlags::new(), val1, ret, 16);
ctx.builder.ins().store(MemFlags::new(), val2, ret, 24);
ctx.builder.ins().return_(&[]);
ctx.builder.seal_all_blocks();
ctx.builder.finalize();
let strings = ctx.strings;
if cfg!(debug_assertions) {
println!("{ir:#?}");
println!("{}", func.display());
}
self.ctx.func = func;
self.module.define_function(func_id, &mut self.ctx).unwrap();
self.module.finalize_definitions().unwrap();
self.ctx.clear();
let _ = self.builder_ctx.insert(builder_ctx);
unsafe {
JITFunc {
func: std::mem::transmute::<*const u8, F<Ctx>>(
self.module.get_finalized_function(func_id),
),
strings,
}
}
}
}