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

feat(jit): attrs & list

Browse Source
This commit is contained in:
imxyy_soope_
2025-07-17 15:57:02 +08:00
parent 2909483afb
commit dedf84a1a9
3 changed files with 247 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/eval/jit/compile.rs
View File

@@ -14,13 +14,38 @@ pub trait JITCompile {
impl JITCompile for Attrs {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
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 JITCompile for List {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
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())
}
}
@@ -281,6 +306,7 @@ impl JITCompile for BinOp {
ctx.builder.seal_block(eq_block);
ctx.builder.seal_block(neq_block);
ctx.builder.switch_to_block(exit_block);
ctx.free_slot(rhs);
lhs
}
@@ -343,8 +369,7 @@ impl JITCompile for If {
let exit_block = ctx.builder.create_block();
let error_block = ctx.builder.create_block();
let judge_block = ctx.builder.create_block();
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = ctx.builder.create_sized_stack_slot(slot);
let slot = ctx.alloca();
let is_bool = ctx
.builder
@@ -393,8 +418,7 @@ impl JITCompile for If {
impl JITCompile for LoadFunc {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = ctx.builder.create_sized_stack_slot(slot);
let slot = ctx.alloca();
let tag = ctx.builder.ins().iconst(types::I64, Value::FUNC as i64);
let val = ctx.builder.ins().iconst(types::I64, self.idx as i64);
ctx.builder.ins().stack_store(tag, slot, 0);
@@ -459,8 +483,7 @@ impl JITCompile for ConcatStrings {
impl JITCompile for Const {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
use c::Const::*;
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = ctx.builder.create_sized_stack_slot(slot);
let slot = ctx.alloca();
match self.val {
Bool(x) => {
let tag = ctx.builder.ins().iconst(types::I64, Value::BOOL as i64);
@@ -491,7 +514,7 @@ impl JITCompile for Const {
impl JITCompile for Str {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
ctx.create_string(&self.val)
}
}
@@ -515,8 +538,7 @@ impl JITCompile for LetVar {
impl JITCompile for Thunk {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = ctx.builder.create_sized_stack_slot(slot);
let slot = ctx.alloca();
let tag = ctx.builder.ins().iconst(types::I64, Value::THUNK as i64);
let val = ctx.builder.ins().iconst(types::I64, self.idx as i64);
ctx.builder.ins().stack_store(tag, slot, 0);

37
src/eval/jit/helpers.rs
View File

@@ -3,6 +3,9 @@ use std::alloc::Layout;
use std::alloc::alloc;
use std::mem::MaybeUninit;
use std::ptr::NonNull;
use std::rc::Rc;
use hashbrown::HashMap;
use crate::env::Env;
use crate::eval::Engine;
@@ -94,11 +97,43 @@ pub unsafe extern "C" fn helper_create_string(
) {
unsafe {
ret.write(Value::String(
str::from_utf8_unchecked(slice::from_raw_parts(ptr, len)).to_string(),
str::from_utf8_unchecked(slice::from_raw_parts(ptr, len)).to_owned(),
));
}
}
pub unsafe extern "C" fn helper_create_list(
ptr: *mut Value,
len: usize,
ret: &mut MaybeUninit<Value>,
) {
unsafe {
ret.write(Value::List(Vec::from_raw_parts(ptr, len, len).into()));
}
}
pub unsafe extern "C" fn helper_create_attrs(ret: &mut MaybeUninit<HashMap<String, Value>>) {
ret.write(HashMap::new());
}
pub unsafe extern "C" fn helper_push_attr(
attrs: &mut HashMap<String, Value>,
sym_ptr: *const u8,
sym_len: usize,
val: NonNull<Value>,
) {
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(attrs: NonNull<HashMap<String, Value>>, ret: &mut MaybeUninit<Value>) {
ret.write(Value::AttrSet(Rc::new(unsafe { attrs.read() }.into())));
}
pub unsafe extern "C" fn helper_alloc_array(len: usize) -> *mut u8 {
unsafe { alloc(Layout::array::<Value>(len).unwrap()) }
}

185
src/eval/jit/mod.rs
View File

@@ -1,8 +1,11 @@
use std::ops::Deref;
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::HashSet;
use crate::engine::Engine;
use crate::env::Env;
@@ -15,16 +18,49 @@ mod helpers;
pub use compile::JITCompile;
use helpers::*;
pub type JITFunc = unsafe extern "C" fn(*const Engine, *const Env, *mut Value);
type F = unsafe extern "C" fn(*const Engine, *const Env, *mut Value);
pub struct JITFunc {
func: F,
strings: HashSet<String>,
}
impl Deref for JITFunc {
type Target = F;
fn deref(&self) -> &Self::Target {
&self.func
}
}
pub struct JITContext<'comp, 'ctx> {
pub compiler: &'comp mut JITCompiler,
pub builder: FunctionBuilder<'ctx>,
free_slots: Vec<StackSlot>,
strings: HashSet<String>,
}
impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
fn new(compiler: &'comp mut JITCompiler, builder: FunctionBuilder<'ctx>) -> Self {
Self { compiler, builder }
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 {
@@ -41,6 +77,9 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
}
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()
@@ -53,8 +92,7 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
.compiler
.module
.declare_func_in_func(self.compiler.create_string, self.builder.func);
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = self.builder.create_sized_stack_slot(slot);
let slot = self.alloca();
let ret = self
.builder
.ins()
@@ -63,6 +101,77 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
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 dbg(&mut self, slot: StackSlot) {
let ptr = self
.builder
@@ -101,8 +210,7 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
}
fn lookup_arg(&mut self, env: ir::Value, idx: usize) -> StackSlot {
let slot = StackSlotData::new(StackSlotKind::ExplicitSlot, 32, 3);
let slot = self.builder.create_sized_stack_slot(slot);
let slot = self.alloca();
let lookup_arg = self
.compiler
.module
@@ -235,6 +343,10 @@ pub struct JITCompiler {
alloc_array: FuncId,
create_string: FuncId,
create_list: FuncId,
create_attrs: FuncId,
push_attr: FuncId,
finalize_attrs: FuncId,
dbg: FuncId,
}
@@ -264,6 +376,10 @@ impl JITCompiler {
builder.symbol("helper_alloc_array", helper_alloc_array as _);
builder.symbol("helper_create_string", helper_create_string as _);
builder.symbol("helper_create_list", helper_create_list as _);
builder.symbol("helper_create_attrs", helper_create_attrs as _);
builder.symbol("helper_push_attr", helper_push_attr as _);
builder.symbol("helper_finalize_attrs", helper_finalize_attrs as _);
builder.symbol("helper_dbg", helper_dbg as _);
let mut module = JITModule::new(builder);
@@ -407,6 +523,52 @@ impl JITCompiler {
.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 dbg_sig = module.make_signature();
dbg_sig.params.push(AbiParam {
value_type: ptr_type,
@@ -440,6 +602,10 @@ impl JITCompiler {
alloc_array,
create_string,
create_list,
create_attrs,
push_attr,
finalize_attrs,
dbg,
}
}
@@ -480,6 +646,8 @@ impl JITCompiler {
ctx.builder.seal_all_blocks();
ctx.builder.finalize();
let strings = ctx.strings;
println!("{:?}", ir);
println!("{}", func.display());
self.ctx.func = func;
@@ -489,7 +657,10 @@ impl JITCompiler {
let _ = self.builder_ctx.insert(builder_ctx);
unsafe {
std::mem::transmute::<*const u8, JITFunc>(self.module.get_finalized_function(func_id))
JITFunc {
func: std::mem::transmute::<*const u8, F>(self.module.get_finalized_function(func_id)),
strings
}
}
}
}