refactor with

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 Some(CallFrame { with_env, .. }) = self.call_stack.pop() else {
-                unreachable!()
-            };
-            self.with_env = with_env;
+        let name = reader.read_string_id();
+        let n = reader.read_u8();
+        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"));
+                    }
+                }
+            }
+        };
+
+        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))
             )));
-        };
-        self.push(env);
-        let env = self.force_and_retry::<Gc<AttrSet>>(reader, mc)?;
-        let Some(val) = env.lookup(name) else {
-            reader.set_pc(reader.inst_start_pc());
-            self.with_env = prev;
-            return Step::Continue(());
-        };
+        } else {
+            self.replace(0, Value::new_inline(counter + 1));
+            reader.set_pc(resume_pc);
+        }
 
-        self.push(val);
-        let Some(CallFrame { with_env, .. }) = self.call_stack.pop() else {
-            unreachable!()
-        };
-        self.with_env = with_env;
         Step::Continue(())
     }
 }