feat(jit): with

This commit is contained in:
2025-07-18 09:53:41 +08:00
parent 74e819c678
commit b556f1ea2d
6 changed files with 101 additions and 20 deletions

View File

@@ -107,7 +107,7 @@ impl Env {
self.args.extend(args);
}
pub fn pop_with(&mut self) {
pub fn exit_with(&mut self) {
self.with.pop();
}

View File

@@ -463,7 +463,12 @@ impl JITCompile for Let {
impl JITCompile for With {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
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
}
}
@@ -519,7 +524,7 @@ impl JITCompile for Str {
impl JITCompile for Var {
fn compile(&self, ctx: &mut JITContext, engine: ir::Value, env: ir::Value) -> StackSlot {
todo!()
ctx.lookup(env, &self.sym)
}
}

View File

@@ -137,6 +137,19 @@ pub unsafe extern "C" fn helper_finalize_attrs(
ret.write(Value::AttrSet(Rc::new(unsafe { attrs.read() }.into())));
}
pub unsafe extern "C" fn helper_enter_with(
env: &mut Env,
namespace: NonNull<Value>,
) {
env.enter_with(unsafe { namespace.read() }.unwrap_attr_set().into_inner());
}
pub unsafe extern "C" fn helper_exit_with(
env: &mut Env,
) {
env.exit_with();
}
pub unsafe extern "C" fn helper_alloc_array(len: usize) -> *mut u8 {
unsafe { alloc(Layout::array::<Value>(len).unwrap()) }
}

View File

@@ -170,6 +170,26 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
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
@@ -207,6 +227,29 @@ impl<'comp, 'ctx> JITContext<'comp, 'ctx> {
.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_arg(&mut self, env: ir::Value, idx: usize) -> StackSlot {
let slot = self.alloca();
let lookup_arg = self
@@ -345,6 +388,8 @@ pub struct JITCompiler {
create_attrs: FuncId,
push_attr: FuncId,
finalize_attrs: FuncId,
enter_with: FuncId,
exit_with: FuncId,
dbg: FuncId,
}
@@ -384,6 +429,8 @@ impl JITCompiler {
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_enter_with", helper_enter_with as _);
builder.symbol("helper_exit_with", helper_exit_with as _);
builder.symbol("helper_dbg", helper_dbg as _);
let mut module = JITModule::new(builder);
@@ -577,6 +624,28 @@ impl JITCompiler {
)
.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,
@@ -614,6 +683,8 @@ impl JITCompiler {
create_attrs,
push_attr,
finalize_attrs,
enter_with,
exit_with,
dbg,
}
}

View File

@@ -230,7 +230,7 @@ impl Evaluate for ir::With {
// TODO: Error Handling
env.enter_with(namespace.unwrap_attr_set().into_inner());
let ret = self.expr.eval(engine, env);
env.pop_with();
env.exit_with();
ret
}
}
@@ -284,7 +284,6 @@ impl Evaluate for ir::Const {
impl Evaluate for ir::Var {
fn eval(&self, _: &mut Engine, env: &mut Env) -> Result<Value> {
env.lookup_with(&self.sym).ok_or_else(|| {
Error::EvalError(format!(
"variable {} not found",
@@ -296,8 +295,7 @@ impl Evaluate for ir::Var {
impl Evaluate for ir::Arg {
fn eval(&self, _: &mut Engine, env: &mut Env) -> Result<Value> {
let result = env.lookup_arg(self.level).clone().ok();
Ok(result.unwrap())
env.lookup_arg(self.level).ok()
}
}
@@ -315,7 +313,6 @@ impl Evaluate for ir::Thunk {
impl Evaluate for ir::MaybeThunk {
fn eval(&self, engine: &mut Engine, env: &mut Env) -> Result<Value> {
match self {
ir::MaybeThunk::Const(cnst) => cnst.eval(engine, env),
ir::MaybeThunk::String(string) => string.eval(engine, env),

View File

@@ -175,15 +175,13 @@ impl DowngradeContext {
match this {
Index::Thunk(idx) => {
let len = self.thunk_deps.len();
self.thunk_deps[idx]
*self.thunk_deps[idx]
.entry(dep.unwrap_thunk())
.or_insert(len);
len
.or_insert(len)
}
Index::Func(idx) => {
let len = self.thunk_deps.len();
self.func_deps[idx].entry(dep).or_insert(len);
len
*self.func_deps[idx].entry(dep).or_insert(len)
}
}
}
@@ -230,16 +228,13 @@ impl DowngradeContext {
let self_ptr = self as *mut Self;
self.thunks.get_mut(idx).map_or_else(
|| unreachable!(),
|thunk| {
|(thunk, _)| {
unsafe {
let (old, _) = core::ptr::read(thunk);
let old = core::ptr::read(thunk);
match old.resolve(Index::Thunk(idx), self_ptr.as_mut().unwrap(), env) {
Ok(ok) => core::ptr::write(&mut thunk.0, ok),
Ok(ok) => core::ptr::write(thunk, ok),
Err(err) => {
core::ptr::write(
&mut thunk.0,
Ir::Const(super::Const { val: Const::Null }),
);
core::ptr::write(thunk, Ir::Const(super::Const { val: Const::Null }));
return Err(err);
}
}