feat: WIP

TODO.md

@@ -8,11 +8,6 @@
   - [ ] graph update
   - [ ] stack storage
 
-阻拦变量解析的是 `with import` 而不是一般的 `with`
-解析阶段可以解决除涉及非字符串常量拼接路径导入之外的问题
-
-静态分析图与求值期图分离，动态生成求值期图
-
 依赖类型：
 1. 强依赖 (x!)
   - builtins.toString x => x!
@@ -35,121 +30,3 @@ e.g.
 3. HIR -> LIR (resolve var, build graph)
 4. LIR -> Value
 
-```mermaid
-flowchart TD
-  Start([Context::eval]) --> Eval[EvalContext::eval]
-  Eval --> AddNode[add node to graph]
-  AddNode --> CheckType{check expression type}
-
-  CheckType -->|AttrSet| A_EvalKeys[eval keys]
-  A_EvalKeys --> A_Construct[construct attrset]
-  A_Construct --> A_Return[return attrset]
-  A_Return --> End
-  A_EvalKeys -.->|eval| Eval
-
-  CheckType -->|List| L_Construct[construct list]
-  L_Construct --> L_Return[return list]
-  L_Return --> End
-
-  CheckType -->|HasAttr| HA_ForceAttrSet[force attrset]
-  HA_ForceAttrSet --> HA_ForceAttrPath[force attrpath]
-  HA_ForceAttrPath --> HA_Return[return bool]
-  HA_Return --> End
-  HA_ForceAttrSet -.->|eval| Eval
-  HA_ForceAttrPath -.->|eval| Eval
-
-  CheckType -->|BinOp| B_ForceLeft[force left operand]
-  B_ForceLeft --> B_ForceRight[force right operand]
-  B_ForceRight --> B_Apply[apply operator]
-  B_Apply --> B_Return[return value]
-  B_Return --> End
-  B_ForceLeft -.->|eval| Eval
-  B_ForceRight -.->|eval| Eval
-
-  CheckType -->|UnOp| U_ForceOperand[force operand]
-  U_ForceOperand --> U_Apply[apply operator]
-  U_Apply --> U_Return[return value]
-  U_Return --> End
-  U_ForceOperand -.->|eval| Eval
-
-  CheckType -->|Select| S_ForceAttrSet[force attrset]
-  S_ForceAttrSet --> S_ForceAttrPath[force attrpath]
-  S_ForceAttrPath --> S_CallEval[call eval on value]
-  S_CallEval --> S_Return[return value]
-  S_Return --> End
-  S_ForceAttrSet -.->|eval| Eval
-  S_ForceAttrPath -.->|eval| Eval
-  S_CallEval -.-> Eval
-
-  CheckType -->|If| I_ForceCond[force condition]
-  I_ForceCond --> I_Cond{condition}
-  I_Cond -->|true| I_ForceConsq[force consequence]
-  I_Cond -->|false| I_ForceAlter[force alternative]
-  I_ForceConsq --> I_CallEvalConsq[call eval on consequence]
-  I_ForceAlter --> I_CallEvalAlter[call eval on alternative]
-  I_CallEvalConsq --> I_ReturnConsq[return value]
-  I_CallEvalAlter --> I_ReturnAlter[return value]
-  I_ReturnConsq --> End
-  I_ReturnAlter --> End
-  I_ForceCond -.->|eval| Eval
-  I_ForceConsq -.->|eval| Eval
-  I_ForceAlter -.->|eval| Eval
-  I_CallEvalConsq -.-> Eval
-  I_CallEvalAlter -.-> Eval
-
-  CheckType -->|Call| C_RegArg[register argument]
-  C_RegArg --> C_ForceBody[force body]
-  C_ForceBody --> C_CallEval[call eval on body]
-  C_CallEval --> C_Return[return value]
-  C_Return --> End
-  C_ForceBody -.->|eval| Eval
-  C_CallEval -.-> Eval
-
-  CheckType -->|With| W_EnterWith[enter with]
-  W_EnterWith --> W_ForceBody[force body]
-  W_ForceBody --> W_CallEval[call eval on body]
-  W_CallEval --> W_Return[return value]
-  W_Return --> End
-  W_ForceBody -.->|eval| Eval
-  W_CallEval -.-> Eval
-
-  CheckType -->|Assert| As_ForceCond[force condition]
-  As_ForceCond --> As_Cond{condition}
-  As_Cond -->|true| As_ForceBody[force body]
-  As_Cond -->|false| As_Throw[throw Catchable]
-  As_ForceBody --> As_CallEval[call eval on body]
-  As_CallEval --> End
-  As_ForceCond -.->|eval| Eval
-  As_ForceBody -.->|eval| Eval
-  As_CallEval -.-> Eval
-
-  CheckType -->|ConcatStrings| CS_ForceParts[force string parts]
-  CS_ForceParts --> CS_Construct[construct string]
-  CS_Construct --> End
-  CS_ForceParts -.->|eval| Eval
-
-  CheckType -->|Const| Co_Return[return constant]
-  Co_Return --> End
-
-  CheckType -->|Str| St_Return[return string]
-  St_Return --> End
-
-  CheckType -->|Var| V_Lookup[lookup var]
-  V_Lookup --> V_Return[return value]
-  V_Return --> End
-
-  CheckType -->|Arg| Ar_Lookup[lookup arg]
-  Ar_Lookup --> End
-
-  CheckType -->|Func| F_Construct[construct function]
-  F_Construct --> End
-
-  CheckType -->|StrictRef| SR_Eval[eval referenced expr]
-  SR_Eval --> Eval
-
-  CheckType -->|LazyRef| LR_Resolve[resolve dynamic variable lookups of the referenced expr]
-  LR_Resolve --> LR_Contruct[construct lazy reference]
-  LR_Contruct --> End
-
-  End([return result])
-```

docs/eval/eval.mmd

@@ -0,0 +1,96 @@
+flowchart TD
+  Start([EvalCtx::eval]) --> Eval[EvalContext::eval]
+  Eval --> AddNode[为本即将求值表达式分配 ValueId，添加至图中]
+  AddNode --> CheckType{检查表达式类型}
+
+  CheckType -->|AttrSet| A_EvalKeys[强制求值所有键]
+  A_EvalKeys --> A_Construct[构造 AttrSet]
+  A_Construct --> End
+  A_EvalKeys -.->|eval| Eval
+
+  CheckType -->|List| L_Construct[构造列表]
+  L_Construct --> End
+
+  CheckType -->|HasAttr| HA_ForceAttrSet[强制求值 AttrSet]
+  HA_ForceAttrSet --> HA_ForceAttrPath[强制求值 AttrPath]
+  HA_ForceAttrPath --> HA_Result[求出结果]
+  HA_Result --> End
+  HA_ForceAttrSet -.->|eval| Eval
+  HA_ForceAttrPath -.->|eval| Eval
+
+  CheckType -->|BinOp| B_ForceLeft[强制求值左操作数]
+  B_ForceLeft --> B_ForceRight[强制求值右操作数]
+  B_ForceRight --> B_Apply[应用操作符]
+  B_Apply --> End
+  B_ForceLeft -.->|eval| Eval
+  B_ForceRight -.->|eval| Eval
+
+  CheckType -->|UnOp| U_ForceOperand[强制求值操作数]
+  U_ForceOperand --> U_Apply[应用操作符]
+  U_Apply --> End
+  U_ForceOperand -.->|eval| Eval
+
+  CheckType -->|Select| S_ForceAttrSet[强制求值 AttrSet]
+  S_ForceAttrSet --> S_ForceAttrPath[强制求值 AttrPath]
+  S_ForceAttrPath --> S_Result[获取结果]
+  S_Result --> End
+  S_ForceAttrSet -.->|eval| Eval
+  S_ForceAttrPath -.->|eval| Eval
+
+  CheckType -->|If| I_ForceCond[强制求值条件]
+  I_ForceCond --> I_Cond{判断条件}
+  I_Cond -->|true| I_Consq[返回 then]
+  I_Cond -->|false| I_Alter[返回 else]
+  I_Consq --> End
+  I_Alter --> End
+  I_ForceCond -.->|eval| Eval
+
+  CheckType -->|Call| C_RegArg[注册函数参数]
+  C_RegArg --> C_ForceBody[强制求值函数体]
+  C_ForceBody --> End
+  C_ForceBody -.->|eval| Eval
+
+  CheckType -->|With| W_ForceNameSpace[强制求值命名空间]
+  W_ForceNameSpace --> W_EnterWith[进入命名空间]
+  W_EnterWith --> W_ForceBody[强制求值表达式]
+  W_ForceBody --> W_ExitWith[退出命名空间]
+  W_ExitWith --> End
+  W_ForceNameSpace -.->|eval| Eval
+  W_ForceBody -.->|eval| Eval
+
+  CheckType -->|Assert| As_ForceCond[强制求值断言]
+  As_ForceCond --> As_Cond{判断断言}
+  As_Cond -->|true| As_ForceBody[强制求值表达式]
+  As_Cond -->|false| As_Throw[抛出 Catchable]
+  As_ForceBody --> End
+  As_ForceCond -.->|eval| Eval
+  As_ForceBody -.->|eval| Eval
+
+  CheckType -->|ConcatStrings| CS_ForceParts[强制求值所有字符串片段]
+  CS_ForceParts --> CS_Construct[构造字符串]
+  CS_Construct --> End
+  CS_ForceParts -.->|eval| Eval
+
+  CheckType -->|Const| Co_Construct[构造常量]
+  Co_Construct --> End
+
+  CheckType -->|Str| St_Construct[构造字符串]
+  St_Construct --> End
+
+  CheckType -->|Var| V_Lookup[动态查找变量]
+  V_Lookup --> End
+
+  CheckType -->|Arg| Ar_Lookup[查找参数]
+  Ar_Lookup --> End
+
+  CheckType -->|Func| F_Construct[构造函数]
+  F_Construct --> End
+
+  CheckType -->|StrictRef| SR_Eval[求值被引用表达式]
+  SR_Eval --> Eval
+
+  CheckType -->|LazyRef| LR_Resolve[解析动态变量引用]
+  LR_Resolve --> LR_Contruct[构造惰性引用]
+  LR_Contruct --> End
+
+  End([返回结果])

docs/resolve/resolve.mmd

@@ -0,0 +1,92 @@
+flowchart TD
+  Start([ResolveCtx::resolve_root]) --> Resolve[ResolveContext::resolve]
+  Resolve --> CheckType{检查表达式类型}
+
+  CheckType -->|AttrSet| A_ResolveKeys[解析所有键]
+  A_ResolveKeys --> A_ResolveValues[解析所有值]
+  A_ResolveValues --> End
+  A_ResolveKeys -.->|resolve| Resolve
+  A_ResolveValues -.->|resolve| Resolve
+
+  CheckType -->|List| L_ResolveElements[解析所有元素]
+  L_ResolveElements --> End
+  L_ResolveElements -.->|resolve| Resolve
+
+  CheckType -->|HasAttr| HA_ResolveAttrSet[解析 AttrSet]
+  HA_ResolveAttrSet --> HA_ResolveAttrPath[解析 AttrPath]
+  HA_ResolveAttrPath --> End
+  HA_ResolveAttrSet -.->|resolve| Resolve
+  HA_ResolveAttrPath -.->|resolve| Resolve
+
+  CheckType -->|BinOp| B_ResolveLeft[解析左操作数]
+  B_ResolveLeft --> B_ResolveRight[解析右操作数]
+  B_ResolveRight --> End
+  B_ResolveLeft -.->|resolve| Resolve
+  B_ResolveRight -.->|resolve| Resolve
+
+  CheckType -->|UnOp| U_ResolveOperand[解析操作数]
+  U_ResolveOperand --> End
+  U_ResolveOperand -.->|resolve| Resolve
+
+  CheckType -->|Select| S_ResolveAttrSet[解析 AttrSet]
+  S_ResolveAttrSet --> S_ResolveAttrPath[解析 AttrPath]
+  S_ResolveAttrPath --> End
+  S_ResolveAttrSet -.->|resolve| Resolve
+  S_ResolveAttrPath -.->|resolve| Resolve
+
+  CheckType -->|If| I_ResolveCond[解析条件]
+  I_ResolveCond --> I_ResolveConsq[解析 then 分支]
+  I_ResolveConsq --> I_ResolveAlter[解析 else 分支]
+  I_ResolveAlter --> End
+  I_ResolveCond -.->|resolve| Resolve
+  I_ResolveConsq -.->|resolve| Resolve
+  I_ResolveAlter -.->|resolve| Resolve
+
+  CheckType -->|Call| C_ResolveFunc[解析函数]
+  C_ResolveFunc --> C_ResolveArg[解析参数]
+  C_ResolveArg --> End
+  C_ResolveFunc -.->|resolve| Resolve
+  C_ResolveArg -.->|resolve| Resolve
+
+  CheckType -->|Let| Le_EnterLet[进入 Let 环境]
+  Le_EnterLet --> Le_ResolveValues[解析所有键]
+  Le_ResolveValues --> Le_ResolveBody[解析表达式]
+  Le_ResolveBody --> Le_ExitLet[退出 Let 环境]
+  Le_ExitLet --> End
+  Le_ResolveValues -.->|resolve| Resolve
+  Le_ResolveBody -.->|resolve| Resolve
+
+  CheckType -->|With| W_ResolveNameSpace[解析命名空间]
+  W_ResolveNameSpace --> W_EnterWith[进入命名空间]
+  W_EnterWith --> W_ResolveBody[解析表达式]
+  W_ResolveBody --> W_ExitWith[退出命名空间]
+  W_ExitWith --> End
+  W_ResolveNameSpace -.->|resolve| Resolve
+  W_ResolveBody -.->|resolve| Resolve
+
+  CheckType -->|Assert| As_ResolveCond[解析断言]
+  As_ResolveCond --> As_ResolveBody[解析表达式]
+  As_ResolveBody --> End
+  As_ResolveCond -.->|resolve| Resolve
+  As_ResolveBody -.->|resolve| Resolve
+
+  CheckType -->|ConcatStrings| CS_ResolveParts[解析所有字符串片段]
+  CS_ResolveParts --> End
+  CS_ResolveParts -.->|resolve| Resolve
+
+  CheckType -->|Const| End
+
+  CheckType -->|Str| End
+
+  CheckType -->|Var| V_Lookup[查找变量]
+  V_Lookup --> End
+
+  CheckType -->|Arg| Ar_Lookup[查找参数]
+  Ar_Lookup --> End
+
+  CheckType -->|Func| F_EnterEnv[进入函数环境]
+  F_EnterEnv --> F_ResolveBody[解析函数体]
+  F_ResolveBody --> F_ExitEnv[退出函数环境]
+  F_ExitEnv --> End
+
+  End([返回结果])

evaluator/nixjit_context/src/downgrade.rs

@@ -22,10 +22,15 @@ impl<'ctx, 'bump> DowngradeCtx<'ctx, 'bump> {
 
 impl DowngradeCtx<'_, '_> {
     fn get_ir(&self, id: ExprId) -> &RefCell<Hir> {
+        // SAFETY: The `ExprId` is guaranteed to be valid and correspond to an expression
+        // allocated within this context, making the raw index access safe.
         let idx = unsafe { id.raw() } - self.ctx.lirs.len() - self.ctx.hirs.len();
         if cfg!(debug_assertions) {
             self.irs.get(idx).unwrap()
         } else {
+            // SAFETY: The index calculation is guarded by the logic that creates `ExprId`s,
+            // ensuring it's always within the bounds of the `irs` vector in release builds.
+            // The debug build's `unwrap()` serves as a runtime check for this invariant.
             unsafe { self.irs.get_unchecked(idx) }
         }
     }

evaluator/nixjit_context/src/eval.rs

@@ -16,7 +16,6 @@ use petgraph::prelude::DiGraph;
 
 use super::Context;
 
-#[derive(Default)]
 struct ValueCache(
     #[cfg(debug_assertions)]
     Option<Value>,
@@ -24,6 +23,19 @@ struct ValueCache(
     MaybeUninit<Value>
 );
 
+impl Default for ValueCache {
+    fn default() -> Self {
+        #[cfg(debug_assertions)]
+        {
+            Self(None)
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            Self(MaybeUninit::uninit())
+        }
+    }
+}
+
 impl ValueCache {
     fn insert(&mut self, val: Value) {
         #[cfg(debug_assertions)]
@@ -39,8 +51,10 @@ impl ValueCache {
 impl Drop for ValueCache {
     fn drop(&mut self) {
         #[cfg(not(debug_assertions))]
+        unsafe {
             self.0.assume_init_drop();
         }
+    }
 }
 
 pub struct EvalCtx<'ctx, 'bump> {

evaluator/nixjit_context/src/lib.rs

@@ -1,15 +1,10 @@
 use std::cell::Cell;
-use std::marker::PhantomPinned;
-use std::ops::{Deref, DerefMut};
 use std::ptr::NonNull;
 
 use bumpalo::{Bump, boxed::Box};
 use hashbrown::HashMap;
 use itertools::Itertools;
-use petgraph::{
+use petgraph::graphmap::DiGraphMap;
-    dot::{Config, Dot},
-    graphmap::DiGraphMap,
-};
 
 use nixjit_builtins::{
     Builtins, BuiltinsContext,
@@ -19,7 +14,7 @@ use nixjit_error::{Error, Result};
 use nixjit_eval::{Args, EvalContext, Value};
 use nixjit_hir::{DowngradeContext, Hir};
 use nixjit_ir::{AttrSet, ExprId, Param, PrimOpId, StackIdx};
-use nixjit_lir::{Lir, ResolveContext};
+use nixjit_lir::Lir;
 
 use crate::downgrade::DowngradeCtx;
 use crate::eval::EvalCtx;
@@ -29,40 +24,6 @@ mod downgrade;
 mod eval;
 mod resolve;
 
-#[derive(Debug)]
-struct Pin<'bump, T> {
-    ptr: Box<'bump, T>,
-    _marker: PhantomPinned,
-}
-
-impl<T> Pin<'_, T> {
-    fn into_inner(self) -> T {
-        Box::into_inner(self.ptr)
-    }
-}
-
-impl<T> Deref for Pin<'_, T> {
-    type Target = T;
-    fn deref(&self) -> &Self::Target {
-        self.ptr.as_ref()
-    }
-}
-
-impl<T> DerefMut for Pin<'_, T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.ptr.as_mut()
-    }
-}
-
-impl<'bump, T> Pin<'bump, T> {
-    fn new_in(x: T, bump: &'bump Bump) -> Self {
-        Self {
-            ptr: Box::new_in(x, bump),
-            _marker: PhantomPinned,
-        }
-    }
-}
-
 /// The main evaluation context.
 ///
 /// This struct orchestrates the entire Nix expression evaluation process,
@@ -70,7 +31,7 @@ impl<'bump, T> Pin<'bump, T> {
 pub struct Context<'bump> {
     ir_count: usize,
     hirs: Vec<Hir>,
-    lirs: Vec<Pin<'bump, Lir>>,
+    lirs: Vec<Box<'bump, Lir>>,
     /// Maps a function's body `ExprId` to its parameter definition.
     funcs: HashMap<ExprId, Param>,
 
@@ -112,17 +73,17 @@ impl<'bump> Context<'bump> {
             .chain(scoped.iter().map(|&(_, arity, f)| (arity, f)))
             .collect_array()
             .unwrap();
-        let lirs = (0..global.len()).map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx) }))
+        let lirs = (0..global.len())
+            .map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx) }))
             .chain(
                 (0..scoped.len())
                     .map(|idx| Lir::PrimOp(unsafe { PrimOpId::from_raw(idx + GLOBAL_LEN) })),
             )
             .chain(core::iter::once(Lir::AttrSet(AttrSet {
-                stcs: global.into_iter().enumerate().map(|(idx, (name, ..))| {
+                stcs: global
-                        (name.to_string(), unsafe {
+                    .into_iter()
-                            ExprId::from_raw(idx)
+                    .enumerate()
-                        })
+                    .map(|(idx, (name, ..))| (name.to_string(), unsafe { ExprId::from_raw(idx) }))
-                    })
                     .chain(scoped.into_iter().enumerate().map(|(idx, (name, ..))| {
                         (name.to_string(), unsafe {
                             ExprId::from_raw(idx + GLOBAL_LEN)
@@ -137,10 +98,10 @@ impl<'bump> Context<'bump> {
             .chain(core::iter::once(Lir::Thunk(unsafe {
                 ExprId::from_raw(GLOBAL_LEN + SCOPED_LEN)
             })))
-            .map(|lir| Pin::new_in(lir, bump))
+            .map(|lir| Box::new_in(lir, bump))
-            .collect();
+            .collect_vec();
         Self {
-            ir_count: 0,
+            ir_count: lirs.len(),
             hirs: Vec::new(),
             lirs,
             funcs: HashMap::new(),
@@ -159,8 +120,8 @@ impl<'bump> Context<'bump> {
         DowngradeCtx::new(self)
     }
 
-    pub fn resolve_ctx<'a>(&'a mut self) -> ResolveCtx<'a, 'bump> {
+    pub fn resolve_ctx<'a>(&'a mut self, root: ExprId) -> ResolveCtx<'a, 'bump> {
-        ResolveCtx::new(self)
+        ResolveCtx::new(self, root)
     }
 
     pub fn eval_ctx<'a>(&'a mut self) -> EvalCtx<'a, 'bump> {
@@ -183,27 +144,19 @@ impl<'bump> Context<'bump> {
         let root = self
             .downgrade_ctx()
             .downgrade_root(root.tree().expr().unwrap())?;
-        self.resolve_ctx().resolve_root(root)?;
+        self.resolve_ctx(root).resolve_root()?;
-        println!(
-            "{:?}",
-            Dot::with_config(&self.graph, &[Config::EdgeNoLabel])
-        );
-        for (idx, ir) in self.lirs.iter().enumerate() {
-            println!("{:?} {:#?}", unsafe { ExprId::from_raw(idx) }, &ir);
-        }
         Ok(self.eval_ctx().eval_root(root)?.to_public())
     }
 
     pub fn add_binding(&mut self, ident: &str, expr: &str) -> Result<()> {
         let root = rnix::Root::parse(expr);
-        if !root.errors().is_empty() {
+        let root_expr = root
-            return Err(Error::parse_error(
+            .ok()
-                root.errors().iter().map(|err| err.to_string()).join("; "),
+            .map_err(|err| Error::parse_error(err.to_string()))?
-            ));
+            .expr()
-        }
+            .unwrap();
-        let root_expr = root.tree().expr().unwrap();
         let expr_id = self.downgrade_ctx().downgrade_root(root_expr)?;
-        self.resolve_ctx().resolve_root(expr_id)?;
+        self.resolve_ctx(expr_id).resolve_root()?;
         unsafe { self.repl_scope.as_mut() }.insert(ident.to_string(), expr_id);
         Ok(())
     }

evaluator/nixjit_context/src/resolve.rs

@@ -1,15 +1,16 @@
 use std::cell::{Cell, RefCell};
 
+use bumpalo::boxed::Box;
 use derive_more::Unwrap;
 use hashbrown::HashMap;
 
 use nixjit_error::Result;
 use nixjit_hir::Hir;
-use nixjit_ir::{Const, ExprId, Param, PrimOpId, StackIdx};
+use nixjit_ir::{Const, ExprId, Param, StackIdx};
 use nixjit_lir::{Lir, LookupResult, Resolve, ResolveContext};
 use replace_with::replace_with_and_return;
 
-use super::{Context, Pin};
+use super::Context;
 
 #[derive(Clone)]
 enum Scope<'ctx> {
@@ -30,18 +31,35 @@ enum Ir {
     Lir(Lir),
 }
 
+struct Closure {
+    id: ExprId,
+    arg: ExprId,
+    deps: Cell<usize>
+}
+
+impl Closure {
+    fn new(id: ExprId, arg: ExprId) -> Self {
+        Self {
+            id,
+            arg,
+            deps: Cell::new(0)
+        }
+    }
+}
+
 pub struct ResolveCtx<'ctx, 'bump> {
     ctx: &'ctx mut Context<'bump>,
-    irs: Vec<Pin<'bump, RefCell<Ir>>>,
+    irs: Vec<Box<'bump, RefCell<Ir>>>,
+    root: ExprId,
+    root_deps: Cell<usize>,
+    closures: Vec<Closure>,
     scopes: Vec<Scope<'ctx>>,
     has_with: bool,
     with_used: bool,
-    closures: Vec<(ExprId, Option<ExprId>, Cell<usize>)>,
-    current_expr: Option<ExprId>,
 }
 
 impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
-    pub fn new(ctx: &'ctx mut Context<'bump>) -> Self {
+    pub fn new(ctx: &'ctx mut Context<'bump>, root: ExprId) -> Self {
         Self {
             scopes: vec![
                 Scope::Builtins(unsafe { ctx.global_scope.as_ref() }),
@@ -52,14 +70,30 @@ impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
             irs: core::mem::take(&mut ctx.hirs)
                 .into_iter()
                 .map(|hir| Ir::Hir(hir).into())
-                .map(|ir| Pin::new_in(ir, ctx.bump))
+                .map(|ir| Box::new_in(ir, ctx.bump))
                 .collect(),
             ctx,
+            root,
+            root_deps: Cell::new(0),
             closures: Vec::new(),
-            current_expr: None,
         }
     }
 
+    pub fn resolve_root(mut self) -> Result<()> {
+        let ret = self.resolve(self.root);
+        if ret.is_ok() {
+            self.ctx.lirs.extend(
+                self.irs
+                    .into_iter()
+                    .map(Box::into_inner)
+                    .map(RefCell::into_inner)
+                    .map(Ir::unwrap_lir)
+                    .map(|lir| Box::new_in(lir, self.ctx.bump)),
+            );
+        }
+        ret
+    }
+
     fn get_ir(&self, id: ExprId) -> &RefCell<Ir> {
         let idx = unsafe { id.raw() } - self.ctx.lirs.len();
         if cfg!(debug_assertions) {
@@ -71,14 +105,13 @@ impl<'ctx, 'bump> ResolveCtx<'ctx, 'bump> {
 
     fn new_lir(&mut self, lir: Lir) -> ExprId {
         self.irs
-            .push(Pin::new_in(RefCell::new(Ir::Lir(lir)), self.ctx.bump));
+            .push(Box::new_in(RefCell::new(Ir::Lir(lir)), self.ctx.bump));
         self.ctx.alloc_id()
     }
 }
 
 impl ResolveContext for ResolveCtx<'_, '_> {
     fn resolve(&mut self, expr: ExprId) -> Result<()> {
-        let prev_expr = self.current_expr.replace(expr);
         let result = unsafe {
             let ctx = &mut *(self as *mut Self);
             let ir = self.get_ir(expr);
@@ -103,26 +136,9 @@ impl ResolveContext for ResolveCtx<'_, '_> {
                 },
             )
         };
-        self.current_expr = prev_expr;
         result
     }
 
-    fn resolve_root(mut self, expr: ExprId) -> Result<()> {
-        self.closures.push((expr, None, Cell::new(0)));
-        let ret = self.resolve(expr);
-        if ret.is_ok() {
-            self.ctx.lirs.extend(
-                self.irs
-                    .into_iter()
-                    .map(Pin::into_inner)
-                    .map(RefCell::into_inner)
-                    .map(Ir::unwrap_lir)
-                    .map(|lir| crate::Pin::new_in(lir, self.ctx.bump)),
-            );
-        }
-        ret
-    }
-
     fn lookup(&mut self, name: &str) -> LookupResult {
         let mut closure_depth = 0;
         // Then search from outer to inner scopes for dependencies
@@ -135,11 +151,8 @@ impl ResolveContext for ResolveCtx<'_, '_> {
                 }
                 Scope::Let(scope) | &Scope::Repl(scope) => {
                     if let Some(&dep) = scope.get(name) {
-                        let (expr, _, deps) = self
+                        let (expr, deps) = self.closures.last().map_or_else(|| (self.root, &self.root_deps), |closure| (closure.id, &closure.deps));
-                            .closures
+                        let idx = self.ctx.add_dep(expr, dep, deps);
-                            .last()
-                            .unwrap();
-                        let idx = self.ctx.add_dep(*expr, dep, deps);
                         return LookupResult::Stack(idx);
                     }
                 }
@@ -149,12 +162,10 @@ impl ResolveContext for ResolveCtx<'_, '_> {
                         // We need to find the corresponding parameter expression to create dependency
                         // For now, we need to handle this case by creating a dependency to the parameter
                         let mut iter = self.closures.iter().rev().take(closure_depth + 1).rev();
-                        let Some((func, Some(arg), count)) = iter.next() else {
+                        let Closure { id: func, arg, deps: count } = iter.next().unwrap();
-                            unreachable!()
-                        };
                         let mut cur = self.ctx.add_dep(*func, *arg, count);
-                        for (func, _, count) in iter {
+                        for Closure { id: func, deps: count, .. } in iter {
-                            self.irs.push(Pin::new_in(
+                            self.irs.push(Box::new_in(
                                 RefCell::new(Ir::Lir(Lir::StackRef(cur))),
                                 self.ctx.bump,
                             ));
@@ -176,10 +187,8 @@ impl ResolveContext for ResolveCtx<'_, '_> {
     }
 
     fn lookup_arg(&mut self) -> StackIdx {
-        let Some((func, Some(arg), count)) = self.closures.last() else {
+        let Closure { id: func, arg, deps } = self.closures.last().unwrap();
-            unreachable!()
+        self.ctx.add_dep(*func, *arg, deps)
-        };
-        self.ctx.add_dep(*func, *arg, count)
     }
 
     fn new_func(&mut self, body: ExprId, param: Param) {
@@ -214,7 +223,7 @@ impl ResolveContext for ResolveCtx<'_, '_> {
         f: impl FnOnce(&mut Self) -> T,
     ) -> T {
         let arg = self.new_lir(Lir::Arg(nixjit_ir::Arg));
-        self.closures.push((func, Some(arg), Cell::new(0)));
+        self.closures.push(Closure::new(func, arg));
         self.scopes.push(Scope::Arg(ident));
         let res = f(self);
         self.scopes.pop();

evaluator/nixjit_lir/src/lib.rs

@@ -71,8 +71,6 @@ pub trait ResolveContext {
     /// Triggers the resolution of a given expression.
     fn resolve(&mut self, expr: ExprId) -> Result<()>;
 
-    fn resolve_root(self, expr: ExprId) -> Result<()>;
-
     /// Looks up a variable by name in the current scope.
     fn lookup(&mut self, name: &str) -> LookupResult;
 
@@ -140,8 +138,10 @@ impl<Ctx: ResolveContext> Resolve<Ctx> for AttrSet {
         for (_, &v) in self.stcs.iter() {
             ctx.resolve(v)?;
         }
-        for &(k, v) in self.dyns.iter() {
+        for &(k, _) in self.dyns.iter() {
             ctx.resolve(k)?;
+        }
+        for &(_, v) in self.dyns.iter() {
             ctx.resolve(v)?;
         }
         Ok(self.to_lir())