From dedf84a1a9cea1a6b3acf8b8db3adcc474ba214b Mon Sep 17 00:00:00 2001 From: imxyy_soope_ Date: Thu, 17 Jul 2025 15:57:02 +0800 Subject: [PATCH] feat(jit): attrs & list --- src/eval/jit/compile.rs | 44 +++++++--- src/eval/jit/helpers.rs | 37 +++++++- src/eval/jit/mod.rs | 185 ++++++++++++++++++++++++++++++++++++++-- 3 files changed, 247 insertions(+), 19 deletions(-) diff --git a/src/eval/jit/compile.rs b/src/eval/jit/compile.rs index 2272729..6d3c30f 100644 --- a/src/eval/jit/compile.rs +++ b/src/eval/jit/compile.rs @@ -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); diff --git a/src/eval/jit/helpers.rs b/src/eval/jit/helpers.rs index cb28eaf..564f75a 100644 --- a/src/eval/jit/helpers.rs +++ b/src/eval/jit/helpers.rs @@ -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, +) { + unsafe { + ret.write(Value::List(Vec::from_raw_parts(ptr, len, len).into())); + } +} + +pub unsafe extern "C" fn helper_create_attrs(ret: &mut MaybeUninit>) { + ret.write(HashMap::new()); +} + +pub unsafe extern "C" fn helper_push_attr( + attrs: &mut HashMap, + sym_ptr: *const u8, + sym_len: usize, + val: NonNull, +) { + 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>, ret: &mut MaybeUninit) { + 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::(len).unwrap()) } } diff --git a/src/eval/jit/mod.rs b/src/eval/jit/mod.rs index edefa0d..3740b1d 100644 --- a/src/eval/jit/mod.rs +++ b/src/eval/jit/mod.rs @@ -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, +} + +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, + strings: HashSet, } 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 + } } } }