refactor with

fix-codegen/src/disassembler.rs

@@ -98,7 +98,7 @@ impl<'a, Ctx: DisassemblerContext> Disassembler<'a, Ctx> {
                 self.read_u32();
             }
             Builtins => {}
-            ReplBinding | ScopedImportBinding | WithLookup => {
+            ReplBinding | ScopedImportBinding => {
                 self.read_u32();
             }
         }
@@ -426,16 +426,14 @@ impl<'a, Ctx: DisassemblerContext> Disassembler<'a, Ctx> {
                 let span_id = self.read_u32();
                 ("Assert", format!("text_id={} span={}", raw_idx, span_id))
             }
-            Op::PushWith => {
-                self.read_operand_data();
-                ("PushWith", String::new())
-            }
-            Op::PopWith => ("PopWith", String::new()),
-            Op::PrepareWith => ("PrepareWith", String::new()),
             Op::LookupWith => {
                 let idx = self.read_u32();
                 let name = self.ctx.resolve_string(idx);
-                ("LookupWith", format!("{:?}", name))
+                let n = self.read_u8();
+                for _ in 0..n {
+                    self.read_operand_data();
+                }
+                ("LookupWith", format!("sym={:?} n={}", name, n))
             }
 
             Op::LoadBuiltins => ("LoadBuiltins", String::new()),

fix-codegen/src/lib.rs

@@ -84,10 +84,7 @@ pub enum Op {
 
     Assert,
 
-    PushWith,
-    PopWith,
     LookupWith,
-    PrepareWith,
 
     LoadBuiltins,
     LoadBuiltin,
@@ -120,7 +117,6 @@ pub enum OperandType {
     Builtins,
     ReplBinding,
     ScopedImportBinding,
-    WithLookup,
 }
 
 pub enum Const {
@@ -152,7 +148,6 @@ pub enum InlineOperand {
     Builtins,
     ReplBinding(StringId),
     ScopedImportBinding(StringId),
-    WithLookup(StringId),
 }
 
 pub fn compile_bytecode(ir: RawIrRef<'_>, ctx: &mut impl BytecodeContext) -> InstructionPtr {
@@ -203,7 +198,6 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
             Builtins => InlineOperand::Builtins,
             ReplBinding(id) => InlineOperand::ReplBinding(id),
             ScopedImportBinding(id) => InlineOperand::ScopedImportBinding(id),
-            WithLookup(id) => InlineOperand::WithLookup(id),
         }
     }
 
@@ -240,10 +234,6 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
                 self.emit_u8(OperandType::ScopedImportBinding as u8);
                 self.emit_str_id(id);
             }
-            WithLookup(id) => {
-                self.emit_u8(OperandType::WithLookup as u8);
-                self.emit_str_id(id);
-            }
         }
     }
 
@@ -353,98 +343,6 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         }
     }
 
-    fn count_with_thunks(&self, ir: RawIrRef<'_>) -> usize {
-        match ir {
-            Ir::With { thunks, body, .. } => thunks.len() + self.count_with_thunks(*body),
-            Ir::TopLevel { thunks, body } => thunks.len() + self.count_with_thunks(*body),
-            Ir::If { cond, consq, alter } => {
-                self.count_with_thunks(*cond)
-                    + self.count_with_thunks(*consq)
-                    + self.count_with_thunks(*alter)
-            }
-            Ir::BinOp { lhs, rhs, .. } => {
-                self.count_with_thunks(*lhs) + self.count_with_thunks(*rhs)
-            }
-            Ir::UnOp { rhs, .. } => self.count_with_thunks(*rhs),
-            Ir::Call { func, .. } => self.count_with_thunks(*func),
-            Ir::Assert {
-                assertion, expr, ..
-            } => self.count_with_thunks(*assertion) + self.count_with_thunks(*expr),
-            Ir::Select { expr, .. } => self.count_with_thunks(*expr),
-            Ir::HasAttr { lhs, .. } => self.count_with_thunks(*lhs),
-            Ir::ConcatStrings { parts, .. } => {
-                parts.iter().map(|p| self.count_with_thunks(*p)).sum()
-            }
-            _ => 0,
-        }
-    }
-
-    fn collect_all_thunks<'ir>(
-        &self,
-        own_thunks: &[(ThunkId, RawIrRef<'ir>)],
-        body: RawIrRef<'ir>,
-    ) -> Vec<(ThunkId, RawIrRef<'ir>)> {
-        let mut all = Vec::from(own_thunks);
-        self.collect_with_thunks_recursive(body, &mut all);
-        let mut i = 0;
-        while i < all.len() {
-            let thunk_body = all[i].1;
-            self.collect_with_thunks_recursive(thunk_body, &mut all);
-            i += 1;
-        }
-        all
-    }
-
-    fn collect_with_thunks_recursive<'ir>(
-        &self,
-        ir: RawIrRef<'ir>,
-        out: &mut Vec<(ThunkId, RawIrRef<'ir>)>,
-    ) {
-        match ir {
-            Ir::With { thunks, body, .. } => {
-                for &(id, inner) in thunks.iter() {
-                    out.push((id, inner));
-                }
-                self.collect_with_thunks_recursive(*body, out);
-            }
-            Ir::TopLevel { thunks, body } => {
-                for &(id, inner) in thunks.iter() {
-                    out.push((id, inner));
-                }
-                self.collect_with_thunks_recursive(*body, out);
-            }
-            Ir::If { cond, consq, alter } => {
-                self.collect_with_thunks_recursive(*cond, out);
-                self.collect_with_thunks_recursive(*consq, out);
-                self.collect_with_thunks_recursive(*alter, out);
-            }
-            Ir::BinOp { lhs, rhs, .. } => {
-                self.collect_with_thunks_recursive(*lhs, out);
-                self.collect_with_thunks_recursive(*rhs, out);
-            }
-            Ir::UnOp { rhs, .. } => self.collect_with_thunks_recursive(*rhs, out),
-            Ir::Call { func, .. } => {
-                self.collect_with_thunks_recursive(*func, out);
-            }
-            Ir::Assert {
-                assertion, expr, ..
-            } => {
-                self.collect_with_thunks_recursive(*assertion, out);
-                self.collect_with_thunks_recursive(*expr, out);
-            }
-            Ir::Select { expr, .. } => {
-                self.collect_with_thunks_recursive(*expr, out);
-            }
-            Ir::HasAttr { lhs, .. } => self.collect_with_thunks_recursive(*lhs, out),
-            Ir::ConcatStrings { parts, .. } => {
-                for p in parts.iter() {
-                    self.collect_with_thunks_recursive(*p, out);
-                }
-            }
-            _ => (),
-        }
-    }
-
     fn push_scope(&mut self, has_arg: bool, thunk_ids: &[ThunkId]) {
         let depth = self.scope_stack.len().try_into().expect("scope too deep!");
         let thunk_base = if has_arg { 1u32 } else { 0u32 };
@@ -463,17 +361,11 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
     fn emit_toplevel(&mut self, ir: RawIrRef<'_>) {
         match ir {
             &Ir::TopLevel { body, ref thunks } => {
-                let with_thunk_count = self.count_with_thunks(body);
+                let thunk_ids: Vec<ThunkId> = thunks.iter().map(|&(id, _)| id).collect();
-                let total_slots = thunks.len() + with_thunk_count;
-
-                let all_thunks = self.collect_all_thunks(thunks, body);
-                let thunk_ids: Vec<ThunkId> = all_thunks.iter().map(|&(id, _)| id).collect();
-
                 self.push_scope(false, &thunk_ids);
-
+                if !thunks.is_empty() {
-                if total_slots > 0 {
                     self.emit_op(Op::AllocLocals);
-                    self.emit_u32(total_slots as u32);
+                    self.emit_u32(thunks.len().try_into().expect("too many thunks"));
                 }
 
                 self.emit_scope_thunks(thunks);
@@ -658,18 +550,20 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
                 self.emit_op(Op::LoadScopedBinding);
                 self.emit_str_id(name);
             }
-            &Ir::With {
+            Ir::WithLookup { sym, namespaces } => {
-                namespace,
+                // counter
-                body,
+                self.emit_expr(&Ir::Int(0));
-                ref thunks,
-            } => {
-                self.emit_with(namespace, body, thunks);
-            }
-            &Ir::WithLookup(name) => {
-                // TODO: specialize shallow with lookups
-                self.emit_op(Op::PrepareWith);
                 self.emit_op(Op::LookupWith);
-                self.emit_str_id(name);
+                self.emit_str_id(*sym);
+                self.emit_u8(
+                    namespaces
+                        .len()
+                        .try_into()
+                        .expect("too many `with` namespaces"),
+                );
+                for namespace in namespaces {
+                    self.emit_maybe_thunk(namespace);
+                }
             }
             &Ir::MaybeThunk(thunk) => {
                 use MaybeThunk::*;
@@ -723,12 +617,6 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
                         self.emit_op(Op::LoadScopedBinding);
                         self.emit_str_id(name);
                     }
-                    WithLookup(name) => {
-                        // TODO: specialize shallow with lookups
-                        self.emit_op(Op::PrepareWith);
-                        self.emit_op(Op::LookupWith);
-                        self.emit_str_id(name);
-                    }
                 }
             }
         }
@@ -822,11 +710,7 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         param: &Option<Param<'ir>>,
         body: RawIrRef<'ir>,
     ) {
-        let with_thunk_count = self.count_with_thunks(body);
+        let thunk_ids: Vec<ThunkId> = thunks.iter().map(|&(id, _)| id).collect();
-        let total_slots = thunks.len() + with_thunk_count;
-
-        let all_thunks = self.collect_all_thunks(thunks, body);
-        let thunk_ids: Vec<ThunkId> = all_thunks.iter().map(|&(id, _)| id).collect();
 
         let skip_patch = self.emit_jump_placeholder();
         let entry_point = self.ctx.get_code().len() as u32;
@@ -845,7 +729,7 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         {
             self.emit_op(Op::MakePatternClosure);
             self.emit_u32(entry_point);
-            self.emit_u32(total_slots as u32);
+            self.emit_u32(thunks.len().try_into().expect("too many thunks"));
             self.emit_u16(required.len() as u16);
             self.emit_u16(optional.len() as u16);
             self.emit_bool(*ellipsis);
@@ -864,7 +748,7 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         } else {
             self.emit_op(Op::MakeClosure);
             self.emit_u32(entry_point);
-            self.emit_u32(total_slots as u32);
+            self.emit_u32(thunks.len().try_into().expect("too many thunks"));
         }
     }
 
@@ -991,19 +875,6 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         self.emit_op(Op::HasAttrResolve);
     }
 
-    fn emit_with(
-        &mut self,
-        namespace: &MaybeThunk,
-        body: RawIrRef<'_>,
-        thunks: &[(ThunkId, RawIrRef<'_>)],
-    ) {
-        self.emit_op(Op::PushWith);
-        self.emit_maybe_thunk(namespace);
-        self.emit_scope_thunks(thunks);
-        self.emit_expr(body);
-        self.emit_op(Op::PopWith);
-    }
-
     fn emit_toplevel_inner(&mut self, body: RawIrRef<'_>, thunks: &[(ThunkId, RawIrRef<'_>)]) {
         self.emit_scope_thunks(thunks);
         self.emit_expr(body);

fix-ir/src/downgrade.rs

@@ -44,7 +44,7 @@ pub trait DowngradeContext<'id: 'ir, 'ir> {
 
     fn intern_string(&mut self, sym: impl AsRef<str>) -> StringId;
     fn resolve_sym(&self, id: StringId) -> Symbol<'_>;
-    fn lookup(&self, sym: StringId, span: TextRange) -> Result<GhostRoMaybeThunkRef<'id, 'ir>>;
+    fn lookup(&mut self, sym: StringId, span: TextRange) -> Result<GhostRoMaybeThunkRef<'id, 'ir>>;
 
     fn get_current_source(&self) -> Source;
 
@@ -54,7 +54,7 @@ pub trait DowngradeContext<'id: 'ir, 'ir> {
     fn with_let_scope<F, R>(&mut self, bindings: &[StringId], f: F) -> Result<R>
     where
         F: FnOnce(&mut Self) -> Result<(Vec<'ir, GhostRoMaybeThunkRef<'id, 'ir>>, R)>;
-    fn with_with_scope<F, R>(&mut self, f: F) -> R
+    fn with_with_scope<F, R>(&mut self, namespace: GhostRoMaybeThunkRef<'id, 'ir>, f: F) -> R
     where
         F: FnOnce(&mut Self) -> R;
     fn with_thunk_scope<F, R>(&mut self, f: F) -> (R, Vec<'ir, (ThunkId, GhostRoIrRef<'id, 'ir>)>)
@@ -448,15 +448,7 @@ impl<'id: 'ir, 'ir, Ctx: DowngradeContext<'id, 'ir>> Downgrade<'id, 'ir, Ctx> fo
         let namespace = ctx.maybe_thunk(namespace);
 
         let body_expr = self.body().require(ctx, span)?;
-        let (body, thunks) =
+        ctx.with_with_scope(namespace, |ctx| body_expr.downgrade(ctx))
-            ctx.with_thunk_scope(|ctx| ctx.with_with_scope(|ctx| body_expr.downgrade(ctx)));
-        let body = body?;
-
-        Ok(ctx.new_expr(Ir::With {
-            namespace,
-            body,
-            thunks,
-        }))
     }
 }
 

fix-ir/src/lib.rs

@@ -93,7 +93,6 @@ pub enum MaybeThunk {
     Builtins,
     ReplBinding(StringId),
     ScopedImportBinding(StringId),
-    WithLookup(StringId),
 }
 
 pub trait Ref<'ir> {
@@ -182,12 +181,10 @@ pub enum Ir<'ir, R: RefExt<'ir> + ?Sized + 'ir> {
         span: TextRange,
     },
 
-    With {
+    WithLookup {
-        namespace: R::MaybeThunkRef,
+        sym: StringId,
-        body: R::IrRef,
+        namespaces: Vec<'ir, R::MaybeThunkRef>,
-        thunks: Vec<'ir, (ThunkId, R::IrRef)>,
     },
-    WithLookup(StringId),
 
     // Function related
     Func {

fix-vm/src/bytecode_reader.rs

@@ -125,10 +125,6 @@ impl<'a> BytecodeReader<'a> {
                 let id = self.read_string_id();
                 OperandData::ScopedImportBinding(id)
             }
-            OperandType::WithLookup => {
-                let id = self.read_string_id();
-                OperandData::WithLookup(id)
-            }
         }
     }
 

fix-vm/src/dispatch_tailcall.rs

@@ -187,10 +187,7 @@ tail_fn!(op_resolve_path, (ctx));
 
 tail_fn!(op_assert, (ctx, reader, mc));
 
-tail_fn!(op_push_with, (ctx, reader, mc));
-tail_fn!(op_pop_with, ());
 tail_fn!(op_lookup_with, (ctx, reader, mc));
-tail_fn!(op_prepare_with, ());
 
 tail_fn!(op_load_builtins, ());
 tail_fn!(op_load_builtin, (reader));
@@ -281,10 +278,7 @@ table! {
 
     Assert                => op_assert,
 
-    PushWith              => op_push_with,
-    PopWith               => op_pop_with,
     LookupWith            => op_lookup_with,
-    PrepareWith           => op_prepare_with,
 
     LoadBuiltins          => op_load_builtins,
     LoadBuiltin           => op_load_builtin,

fix-vm/src/instructions/builtins.rs

@@ -46,10 +46,8 @@ impl<'gc> Vm<'gc> {
     fn return_from_primop(&mut self, reader: &mut BytecodeReader<'_>) -> Step {
         let Some(CallFrame {
             pc: ret_pc,
-            stack_depth: _,
             thunk: _,
             env,
-            with_env,
         }) = self.call_stack.pop()
         else {
             unreachable!()
@@ -57,7 +55,6 @@ impl<'gc> Vm<'gc> {
         reader.set_pc(ret_pc);
         self.call_depth -= 1;
         self.env = env;
-        self.with_env = with_env;
         Step::Continue(())
     }
 
@@ -372,6 +369,25 @@ impl<'gc> Vm<'gc> {
         self.finish_ok(ctx.convert_value(val.relax()))
     }
 
+    fn is_value_in_seen(&self, seen: Gc<'gc, List<'gc>>, val: Value<'gc>) -> bool {
+        if !is_container(val) {
+            return false;
+        }
+        let target = val.to_bits();
+        for &v in seen.inner.borrow().iter() {
+            if v.to_bits() == target {
+                return true;
+            }
+        }
+        false
+    }
+
+    fn add_value_to_seen(&self, seen: Gc<'gc, List<'gc>>, mc: &Mutation<'gc>, val: Value<'gc>) {
+        if is_container(val) {
+            seen.unlock(mc).borrow_mut().push(val);
+        }
+    }
+
     pub(crate) fn primop_call_pattern(
         &mut self,
         ctx: &mut impl VmRuntimeCtx,
@@ -421,25 +437,6 @@ impl<'gc> Vm<'gc> {
 
         Step::Continue(())
     }
-
-    fn is_value_in_seen(&self, seen: Gc<'gc, List<'gc>>, val: Value<'gc>) -> bool {
-        if !is_container(val) {
-            return false;
-        }
-        let target = val.to_bits();
-        for &v in seen.inner.borrow().iter() {
-            if v.to_bits() == target {
-                return true;
-            }
-        }
-        false
-    }
-
-    fn add_value_to_seen(&self, seen: Gc<'gc, List<'gc>>, mc: &Mutation<'gc>, val: Value<'gc>) {
-        if is_container(val) {
-            seen.unlock(mc).borrow_mut().push(val);
-        }
-    }
 }
 
 fn is_container(val: Value<'_>) -> bool {

fix-vm/src/instructions/calls.rs

@@ -29,10 +29,8 @@ impl<'gc> crate::Vm<'gc> {
                 self.push(arg);
                 self.call_stack.push(CallFrame {
                     pc: resume_pc,
-                    stack_depth: 0,
                     thunk: None,
                     env: self.env,
-                    with_env: self.with_env,
                 });
                 reader.set_pc(PrimOpPhase::CallPattern.ip() as usize);
                 return Step::Continue(());
@@ -44,10 +42,8 @@ impl<'gc> crate::Vm<'gc> {
             let new_env = Gc::new(mc, RefLock::new(Env::with_arg(arg, n_locals, env)));
             self.call_stack.push(CallFrame {
                 pc: resume_pc,
-                stack_depth: 0,
                 thunk: None,
                 env: self.env,
-                with_env: self.with_env,
             });
             reader.set_pc(ip as usize);
             self.env = new_env;
@@ -56,10 +52,8 @@ impl<'gc> crate::Vm<'gc> {
                 self.push(arg);
                 self.call_stack.push(CallFrame {
                     pc: resume_pc,
-                    stack_depth: 0,
                     thunk: None,
                     env: self.env,
-                    with_env: self.with_env,
                 });
                 reader.set_pc(primop.dispatch_ip as usize)
             } else {
@@ -78,10 +72,8 @@ impl<'gc> crate::Vm<'gc> {
                 self.push(arg);
                 self.call_stack.push(CallFrame {
                     pc: resume_pc,
-                    stack_depth: 0,
                     thunk: None,
                     env: self.env,
-                    with_env: self.with_env,
                 });
                 reader.set_pc(app.primop.dispatch_ip as usize)
             } else {
@@ -119,10 +111,8 @@ impl<'gc> crate::Vm<'gc> {
         let val = self.force_and_retry::<StrictValue>(reader, mc)?;
         let Some(CallFrame {
             pc: ret_pc,
-            stack_depth,
             thunk,
             env,
-            with_env,
         }) = self.call_stack.pop()
         else {
             match self.force_mode {
@@ -137,10 +127,8 @@ impl<'gc> crate::Vm<'gc> {
                     self.push(val.relax());
                     self.call_stack.push(CallFrame {
                         pc: PrimOpPhase::ForceResultDeepFinish.ip() as usize,
-                        stack_depth: 0,
                         thunk: None,
                         env: self.env,
-                        with_env: self.with_env,
                     });
                     self.call_depth += 1;
                     reader.set_pc(PrimOpPhase::DeepSeq.ip() as usize);
@@ -151,13 +139,11 @@ impl<'gc> crate::Vm<'gc> {
         reader.set_pc(ret_pc);
         if let Some(outer_thunk) = thunk {
             *outer_thunk.borrow_mut(mc) = ThunkState::Evaluated(val);
-            self.replace(stack_depth, val.relax());
         } else {
             self.call_depth -= 1;
             self.push(val.relax())
         }
         self.env = env;
-        self.with_env = with_env;
         Step::Continue(())
     }
 }

fix-vm/src/instructions/closures.rs

@@ -15,7 +15,6 @@ impl<'gc> crate::Vm<'gc> {
             RefLock::new(ThunkState::Pending {
                 ip: entry_point as usize,
                 env: self.env,
-                with_env: self.with_env,
             }),
         );
         self.push(Value::new_gc(thunk));

fix-vm/src/instructions/with_scope.rs

@@ -1,51 +1,11 @@
 use fix_common::Symbol;
 use fix_error::Error;
-use gc_arena::Gc;
+use smallvec::SmallVec;
 
 use crate::value::*;
-use crate::{BytecodeReader, CallFrame, Step, VmRuntimeCtx, WithEnv};
+use crate::{Break, BytecodeReader, CallFrame, Step, VmRuntimeCtx};
 
 impl<'gc> crate::Vm<'gc> {
-    #[inline(always)]
-    pub(crate) fn op_push_with(
-        &mut self,
-        ctx: &mut impl VmRuntimeCtx,
-        reader: &mut BytecodeReader<'_>,
-        mc: &gc_arena::Mutation<'gc>,
-    ) -> Step {
-        let env = reader.read_operand_data(ctx).resolve(mc, self);
-        let scope = Gc::new(
-            mc,
-            WithEnv {
-                env,
-                prev: self.with_env,
-            },
-        );
-        self.with_env = Some(scope);
-        Step::Continue(())
-    }
-
-    #[inline(always)]
-    pub(crate) fn op_pop_with(&mut self) -> Step {
-        let Some(scope) = self.with_env else {
-            unreachable!("no with_scope to pop");
-        };
-        self.with_env = scope.prev;
-        Step::Continue(())
-    }
-
-    #[inline(always)]
-    pub(crate) fn op_prepare_with(&mut self) -> Step {
-        self.call_stack.push(CallFrame {
-            pc: usize::MAX,
-            stack_depth: 0,
-            thunk: None,
-            env: self.env,
-            with_env: self.with_env,
-        });
-        Step::Continue(())
-    }
-
     #[inline(always)]
     pub(crate) fn op_lookup_with(
         &mut self,
@@ -53,31 +13,62 @@ impl<'gc> crate::Vm<'gc> {
         reader: &mut BytecodeReader<'_>,
         mc: &gc_arena::Mutation<'gc>,
     ) -> Step {
-        let name = reader.read_string_id();
+        #[allow(clippy::unwrap_used)]
+        let counter = self.peek_forced(0).as_inline::<i32>().unwrap();
 
-        let Some(&WithEnv { env, prev }) = self.with_env.as_deref() else {
+        let name = reader.read_string_id();
-            let Some(CallFrame { with_env, .. }) = self.call_stack.pop() else {
+        let n = reader.read_u8();
-                unreachable!()
+        let mut namespaces = SmallVec::<[_; 2]>::new();
+        for _ in 0..n {
+            namespaces.push(reader.read_operand_data(ctx).resolve(mc, self));
+        }
+
+        let resume_pc = reader.inst_start_pc();
+        let namespace = match namespaces[counter as usize].restrict() {
+            Ok(val) => val,
+            Err(thunk) => {
+                let mut state = thunk.borrow_mut(mc);
+                match *state {
+                    ThunkState::Pending { ip, env } => {
+                        *state = ThunkState::Blackhole;
+                        self.call_stack.push(CallFrame {
+                            thunk: Some(thunk),
+                            pc: resume_pc,
+                            env: self.env,
+                        });
+                        self.env = env;
+                        reader.set_pc(ip);
+                        return Step::Break(Break::Force);
+                    }
+                    ThunkState::Evaluated(v) => v,
+                    ThunkState::Apply { func, arg } => {
+                        self.call_stack.push(CallFrame {
+                            thunk: Some(thunk),
+                            pc: resume_pc,
+                            env: self.env,
+                        });
+                        self.push(func);
+                        return self.call(reader, mc, arg, resume_pc);
+                    }
+                    ThunkState::Blackhole => {
+                        return self.finish_err(Error::eval_error("infinite recursion encountered"));
+                    }
+                }
+            }
         };
-            self.with_env = with_env;
+
+        if let Some(val) = namespace.as_gc::<AttrSet>().and_then(|attrs| attrs.lookup(name)) {
+            self.replace(0, val);
+        } else if counter + 1 == n as i32 {
             return self.finish_err(Error::eval_error(format!(
                 "undefined variable '{}'",
                 Symbol::from(ctx.resolve_string(name))
             )));
-        };
+        } else {
-        self.push(env);
+            self.replace(0, Value::new_inline(counter + 1));
-        let env = self.force_and_retry::<Gc<AttrSet>>(reader, mc)?;
+            reader.set_pc(resume_pc);
-        let Some(val) = env.lookup(name) else {
+        }
-            reader.set_pc(reader.inst_start_pc());
-            self.with_env = prev;
-            return Step::Continue(());
-        };
 
-        self.push(val);
-        let Some(CallFrame { with_env, .. }) = self.call_stack.pop() else {
-            unreachable!()
-        };
-        self.with_env = with_env;
         Step::Continue(())
     }
 }

fix-vm/src/lib.rs

@@ -26,9 +26,9 @@ mod value;
 pub use value::StaticValue;
 use value::*;
 mod helpers;
-pub(crate) mod instructions;
+mod instructions;
-pub(crate) use bytecode_reader::BytecodeReader;
+use bytecode_reader::BytecodeReader;
-pub(crate) use forced::Forced;
+use forced::Forced;
 use helpers::*;
 
 type VmResult<T> = std::result::Result<T, VmError>;
@@ -83,7 +83,7 @@ pub trait VmCode {
     ) -> fix_error::Result<InstructionPtr>;
 }
 
-pub(crate) trait VmRuntimeCtxExt: VmRuntimeCtx {
+trait VmRuntimeCtxExt: VmRuntimeCtx {
     fn get_string<'a, 'gc: 'a>(&'a self, val: StrictValue<'gc>) -> Option<&'a str>;
     fn get_string_id<'a, 'gc: 'a>(
         &'a mut self,
@@ -193,39 +193,38 @@ impl<T: VmRuntimeCtx> ConvertValueWithSeen for T {
 }
 
 #[repr(u8)]
-pub(crate) enum Break {
+enum Break {
     Force,
     Done,
 }
 
-pub(crate) type Step = std::ops::ControlFlow<Break>;
+type Step = std::ops::ControlFlow<Break>;
 
 #[derive(Collect)]
 #[collect(no_drop)]
 pub struct Vm<'gc> {
-    pub(crate) stack: Vec<Value<'gc>>,
+    stack: Vec<Value<'gc>>,
-    pub(crate) call_stack: Vec<CallFrame<'gc>>,
+    call_stack: Vec<CallFrame<'gc>>,
-    pub(crate) call_depth: usize,
+    call_depth: usize,
     #[allow(dead_code)]
     #[collect(require_static)]
-    pub(crate) error_context: Vec<ErrorFrame>,
+    error_context: Vec<ErrorFrame>,
 
-    pub(crate) env: GcEnv<'gc>,
+    env: GcEnv<'gc>,
-    pub(crate) with_env: Option<GcWithEnv<'gc>>,
 
-    pub(crate) import_cache: HashMap<PathBuf, Value<'gc>>,
+    import_cache: HashMap<PathBuf, Value<'gc>>,
 
-    pub(crate) builtins: Value<'gc>,
+    builtins: Value<'gc>,
-    pub(crate) empty_list: Value<'gc>,
+    empty_list: Value<'gc>,
-    pub(crate) empty_attrs: Value<'gc>,
+    empty_attrs: Value<'gc>,
 
-    pub(crate) force_mode: ForceMode,
+    force_mode: ForceMode,
 
     #[collect(require_static)]
-    pub(crate) result: Option<Result<fix_common::Value>>,
+    result: Option<Result<fix_common::Value>>,
 }
 
-pub(crate) enum OperandData {
+enum OperandData {
     Const(StaticValue),
     BigInt(i64),
     Local { layer: u8, idx: u32 },
@@ -233,11 +232,10 @@ pub(crate) enum OperandData {
     Builtins,
     ReplBinding(StringId),
     ScopedImportBinding(StringId),
-    WithLookup(StringId),
 }
 
 impl OperandData {
-    pub(crate) fn resolve<'gc>(&self, mc: &Mutation<'gc>, root: &Vm<'gc>) -> Value<'gc> {
+    fn resolve<'gc>(&self, mc: &Mutation<'gc>, root: &Vm<'gc>) -> Value<'gc> {
         use OperandData::*;
         match *self {
             Const(sv) => sv.into(),
@@ -260,7 +258,6 @@ impl OperandData {
             Builtins => root.builtins,
             ReplBinding(_id) => todo!(),
             ScopedImportBinding(_id) => todo!(),
-            WithLookup(_id) => todo!(),
         }
     }
 }
@@ -332,7 +329,6 @@ impl<'gc> Vm<'gc> {
             error_context: Vec::with_capacity(1024),
 
             env: Gc::new(mc, RefLock::new(Env::empty())),
-            with_env: None,
 
             import_cache: HashMap::new(),
 
@@ -347,19 +343,19 @@ impl<'gc> Vm<'gc> {
     }
 
     #[inline(always)]
-    pub(crate) fn finish_ok(&mut self, val: fix_common::Value) -> Step {
+    fn finish_ok(&mut self, val: fix_common::Value) -> Step {
         self.result = Some(Ok(val));
         Step::Break(Break::Done)
     }
 
     #[inline(always)]
-    pub(crate) fn finish_err(&mut self, err: Box<Error>) -> Step {
+    fn finish_err(&mut self, err: Box<Error>) -> Step {
         self.result = Some(Err(err));
         Step::Break(Break::Done)
     }
 
     #[inline(always)]
-    pub(crate) fn finish_type_err(&mut self, expected: NixType, got: NixType) -> Step {
+    fn finish_type_err(&mut self, expected: NixType, got: NixType) -> Step {
         self.result = Some(Err(Error::eval_error(format!(
             "expected {expected}, got {got}"
         ))));
@@ -367,24 +363,24 @@ impl<'gc> Vm<'gc> {
     }
 
     #[inline(always)]
-    pub(crate) fn finish_vm_err(&mut self, err: VmError) -> Step {
+    fn finish_vm_err(&mut self, err: VmError) -> Step {
         self.finish_err(err.into_error())
     }
 
     #[inline(always)]
-    pub(crate) fn push(&mut self, val: Value<'gc>) {
+    fn push(&mut self, val: Value<'gc>) {
         self.stack.push(val);
     }
 
     #[inline(always)]
     #[must_use]
-    pub(crate) fn pop(&mut self) -> Value<'gc> {
+    fn pop(&mut self) -> Value<'gc> {
         self.stack.pop().expect("stack underflow")
     }
 
     #[inline(always)]
     #[must_use]
-    pub(crate) fn peek(&mut self, depth: usize) -> Value<'gc> {
+    fn peek(&mut self, depth: usize) -> Value<'gc> {
         *self
             .stack
             .get(self.stack.len() - depth - 1)
@@ -393,7 +389,7 @@ impl<'gc> Vm<'gc> {
 
     #[inline(always)]
     #[must_use]
-    pub(crate) fn peek_forced(&mut self, depth: usize) -> StrictValue<'gc> {
+    fn peek_forced(&mut self, depth: usize) -> StrictValue<'gc> {
         self.stack
             .get(self.stack.len() - depth - 1)
             .expect("stack underflow")
@@ -402,7 +398,7 @@ impl<'gc> Vm<'gc> {
     }
 
     #[inline(always)]
-    pub(crate) fn replace(&mut self, depth: usize, val: Value<'gc>) {
+    fn replace(&mut self, depth: usize, val: Value<'gc>) {
         let len = self.stack.len();
         *self
             .stack
@@ -412,7 +408,7 @@ impl<'gc> Vm<'gc> {
 
     #[inline(always)]
     #[cfg_attr(debug_assertions, track_caller)]
-    pub(crate) fn pop_forced(&mut self) -> StrictValue<'gc> {
+    fn pop_forced(&mut self) -> StrictValue<'gc> {
         self.stack
             .pop()
             .expect("stack underflow")
@@ -447,7 +443,7 @@ impl<'gc> Vm<'gc> {
     /// * The return value must be propagated with `?` so that
     ///   `Break::Force` correctly unwinds to the dispatch loop.
     #[inline(always)]
-    pub(crate) fn force_and_retry<T: Forced<'gc>>(
+    fn force_and_retry<T: Forced<'gc>>(
         &mut self,
         reader: &mut BytecodeReader<'_>,
         mc: &Mutation<'gc>,
@@ -458,7 +454,7 @@ impl<'gc> Vm<'gc> {
     /// Same as [`force_and_retry`](Self::force_and_retry) but allows
     /// specifying a custom resume PC.
     #[inline(always)]
-    pub(crate) fn force_and_retry_pc<T: Forced<'gc>>(
+    fn force_and_retry_pc<T: Forced<'gc>>(
         &mut self,
         reader: &mut BytecodeReader<'_>,
         mc: &Mutation<'gc>,
@@ -470,7 +466,7 @@ impl<'gc> Vm<'gc> {
 
     #[inline(always)]
     #[allow(unused)]
-    pub(crate) fn force_slot(
+    fn force_slot(
         &mut self,
         depth: usize,
         reader: &mut BytecodeReader<'_>,
@@ -480,7 +476,7 @@ impl<'gc> Vm<'gc> {
     }
 
     #[inline(always)]
-    pub(crate) fn force_slot_to_pc(
+    fn force_slot_to_pc(
         &mut self,
         depth: usize,
         reader: &mut BytecodeReader<'_>,
@@ -492,17 +488,14 @@ impl<'gc> Vm<'gc> {
         };
         let mut state = thunk.borrow_mut(mc);
         match *state {
-            ThunkState::Pending { ip, env, with_env } => {
+            ThunkState::Pending { ip, env } => {
                 *state = ThunkState::Blackhole;
                 self.call_stack.push(CallFrame {
                     thunk: Some(thunk),
-                    stack_depth: depth,
                     pc: resume_pc,
                     env: self.env,
-                    with_env: self.with_env,
                 });
                 self.env = env;
-                self.with_env = with_env;
                 reader.set_pc(ip);
                 Step::Break(Break::Force)
             }
@@ -513,10 +506,8 @@ impl<'gc> Vm<'gc> {
             ThunkState::Apply { func, arg } => {
                 self.call_stack.push(CallFrame {
                     thunk: Some(thunk),
-                    stack_depth: depth,
                     pc: resume_pc,
                     env: self.env,
-                    with_env: self.with_env,
                 });
                 self.push(func);
                 self.call(reader, mc, arg, resume_pc)
@@ -536,20 +527,18 @@ struct ErrorFrame {
 
 #[derive(Collect, Debug)]
 #[collect(no_drop)]
-pub(crate) struct CallFrame<'gc> {
+struct CallFrame<'gc> {
-    pub(crate) pc: usize,
+    pc: usize,
-    pub(crate) stack_depth: usize,
+    thunk: Option<Gc<'gc, Thunk<'gc>>>,
-    pub(crate) thunk: Option<Gc<'gc, Thunk<'gc>>>,
+    env: Gc<'gc, RefLock<Env<'gc>>>,
-    pub(crate) env: Gc<'gc, RefLock<Env<'gc>>>,
-    pub(crate) with_env: Option<Gc<'gc, WithEnv<'gc>>>,
 }
 
-pub(crate) enum Action {
+enum Action {
     Continue { pc: usize },
     Done(Result<fix_common::Value>),
 }
 
-pub(crate) enum NixNum {
+enum NixNum {
     Int(i64),
     Float(f64),
 }
@@ -699,10 +688,7 @@ impl<'gc> Vm<'gc> {
 
                 Assert => self.op_assert(ctx, &mut reader, mc),
 
-                PushWith => self.op_push_with(ctx, &mut reader, mc),
-                PopWith => self.op_pop_with(),
                 LookupWith => self.op_lookup_with(ctx, &mut reader, mc),
-                PrepareWith => self.op_prepare_with(),
 
                 LoadBuiltins => self.op_load_builtins(),
                 LoadBuiltin => self.op_load_builtin(&mut reader),

fix-vm/src/value.rs

@@ -527,7 +527,6 @@ pub(crate) enum ThunkState<'gc> {
     Pending {
         ip: usize,
         env: GcEnv<'gc>,
-        with_env: Option<GcWithEnv<'gc>>,
     },
     Apply {
         func: Value<'gc>,

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
     }