feat(jit): with

src/env.rs

@@ -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();
     }
 

src/eval/jit/compile.rs

@@ -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)
     }
 }
 

src/eval/jit/helpers.rs

@@ -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()) }
 }

src/eval/jit/mod.rs

@@ -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,
         }
     }

src/eval/mod.rs

@@ -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();
+        env.lookup_arg(self.level).ok()
-        Ok(result.unwrap())
     }
 }
 
@@ -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),

src/ir/ctx.rs

@@ -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);
+                    .or_insert(len)
-                len
             }
             Index::Func(idx) => {
                 let len = self.thunk_deps.len();
-                self.func_deps[idx].entry(dep).or_insert(len);
+                *self.func_deps[idx].entry(dep).or_insert(len)
-                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(
+                            core::ptr::write(thunk, Ir::Const(super::Const { val: Const::Null }));
-                                &mut thunk.0,
-                                Ir::Const(super::Const { val: Const::Null }),
-                            );
                             return Err(err);
                         }
                     }