refactor with

fix/src/lib.rs

@@ -315,7 +315,7 @@ struct DowngradeCtx<'ctx, 'id, 'ir, R: VmRuntimeCtx> {
     runtime: &'ctx mut R,
     source: Source,
     scopes: Vec<Scope<'ctx, 'id, 'ir>>,
-    with_scope_count: u32,
+    with_stack: Vec<GhostRoMaybeThunkRef<'id, 'ir>>,
     arg_count: u32,
     thunk_count: &'ctx mut usize,
     thunk_scopes: Vec<ThunkScope<'id, 'ir>>,
@@ -341,7 +341,7 @@ impl<'ctx, 'id, 'ir, R: VmRuntimeCtx> DowngradeCtx<'ctx, 'id, 'ir, R> {
                 .collect(),
             thunk_count,
             arg_count: 0,
-            with_scope_count: 0,
+            with_stack: Vec::new(),
             thunk_scopes: vec![ThunkScope::new_in(bump)],
         }
     }
@@ -392,7 +392,7 @@ impl<'ctx: 'ir, 'id, 'ir, R: VmRuntimeCtx> DowngradeContext<'id, 'ir>
     }
 
     fn lookup(
-        &self,
+        &mut self,
         sym: StringId,
         span: rnix::TextRange,
     ) -> Result<GhostRoMaybeThunkRef<'id, 'ir>> {
@@ -438,10 +438,22 @@ impl<'ctx: 'ir, 'id, 'ir, R: VmRuntimeCtx> DowngradeContext<'id, 'ir>
             }
         }
 
-        if self.with_scope_count > 0 {
+        if !self.with_stack.is_empty() {
+            let id = ThunkId(*self.thunk_count);
+            *self.thunk_count = self.thunk_count.checked_add(1).expect("thunk id overflow");
+            let mut namespaces =
+                bumpalo::collections::Vec::with_capacity_in(self.with_stack.len(), self.bump);
+            namespaces.extend(self.with_stack.iter().rev().copied());
+            let body = self
+                .bump
+                .alloc(GhostCell::new(Ir::WithLookup { sym, namespaces }).into());
+            self.thunk_scopes
+                .last_mut()
+                .expect("no active thunk scope")
+                .add_binding(id, body);
             Ok(self
                 .bump
-                .alloc(GhostCell::new(MaybeThunk::WithLookup(sym)).into()))
+                .alloc(GhostCell::new(MaybeThunk::Thunk(id)).into()))
         } else {
             Err(Error::downgrade_error(
                 format!("'{}' not found", self.resolve_sym(sym)),
@@ -506,13 +518,13 @@ impl<'ctx: 'ir, 'id, 'ir, R: VmRuntimeCtx> DowngradeContext<'id, 'ir>
         f(guard.as_ctx())
     }
 
-    fn with_with_scope<F, Ret>(&mut self, f: F) -> Ret
+    fn with_with_scope<F, Ret>(&mut self, namespace: GhostRoMaybeThunkRef<'id, 'ir>, f: F) -> Ret
     where
         F: FnOnce(&mut Self) -> Ret,
     {
-        self.with_scope_count += 1;
+        self.with_stack.push(namespace);
         let ret = f(self);
-        self.with_scope_count -= 1;
+        self.with_stack.pop();
         ret
     }