implement ForceMode

fix/src/runtime/vm.rs

@@ -32,9 +32,10 @@ impl From<Box<Error>> for VmError {
     }
 }
 
-#[derive(Collect, Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Collect, Clone, Copy, Debug, PartialEq, Eq, Default)]
 #[collect(require_static)]
 pub(super) enum ForceMode {
+    #[default]
     AsIs,
     Shallow,
     Deep,
@@ -1208,9 +1209,10 @@ impl Runtime {
     pub(super) fn run(
         &mut self,
         ip: InstructionPtr,
-        _mode: ForceMode,
+        mode: ForceMode,
     ) -> Result<crate::value::Value> {
         self.pc = ip.0;
+        self.force_mode = mode;
         self.arena.mutate_root(|mc, root| {
             if root.current_env.is_none() {
                 root.current_env = Some(Gc::new(mc, RefLock::new(Env::empty())));
@@ -1255,7 +1257,7 @@ impl Runtime {
         });
     }
 
-    pub(super) fn force_tos(&mut self) {
+    pub(super) fn force_tos(&mut self) -> Action {
         loop {
             let run = self.arena.mutate_root(|_mc, root| {
                 let thunk = root
@@ -1287,14 +1289,16 @@ impl Runtime {
                 }
             });
             if !run {
-                return;
+                return Action::Continue;
             }
             loop {
                 self.check_gc();
                 match self.execute_one() {
                     Action::Continue => (),
                     Action::Return => break,
-                    Action::Done(_) => unreachable!(),
+                    // FIXME: poison thunk
+                    Action::Done(err @ Err(_)) => return Action::Done(err),
+                    Action::Done(Ok(_)) => unreachable!(),
                 }
             }
             self.arena.mutate_root(|mc, root| {
@@ -1305,7 +1309,7 @@ impl Runtime {
                     *thunk.borrow_mut(mc) = ThunkState::Evaluated(val);
                 }
                 *thunk = val;
-            })
+            });
         }
     }
 
@@ -1331,21 +1335,227 @@ impl Runtime {
 
     pub(super) fn handle_return(&mut self) -> Action {
         self.force_tos();
-        self.arena.mutate_root(|_, root| {
+        let done= self.arena.mutate_root(|_, root| {
             let Some(frame) = root.frames.pop() else {
-                // FIXME: ForceMode
-                root.current_env = None;
-                return Action::Done(Ok(convert_value(
-                    root.stack.pop().expect("what the heck"),
-                    &self.strings,
-                )));
+                return true;
             };
-
             self.pc = frame.pc;
             root.current_env = Some(frame.env);
+            false
+        });
+        if !done {
+            return Action::Return;
+        }
+        match self.force_mode {
+            ForceMode::AsIs => (),
+            ForceMode::Shallow => {
+                if let done @ Action::Done(_) = self.force_tos_shallow() {
+                    return done
+                }
+            }
+            ForceMode::Deep => {
+                if let done @ Action::Done(_) = self.force_tos_shallow() {
+                    return done
+                }
+            }
+        }
+        let val = self.arena.mutate_root(|_, root| {
+            root.current_env = None;
+            convert_value(
+                root.stack.pop().expect("stack underflow"),
+                &self.strings,
+            )
+        });
+        Action::Done(Ok(val))
+    }
 
-            Action::Return
-        })
+    pub(super) fn force_tos_shallow(&mut self) -> Action {
+        if let err @ Action::Done(Err(_)) = self.force_tos() {
+            return err;
+        }
+
+        let (is_list, is_attrs) = self.arena.mutate_root(|_, root| {
+            let tos = *root.stack.tos().expect("stack underflow");
+            (tos.as_gc::<List<'_>>().is_some(), tos.as_gc::<AttrSet<'_>>().is_some())
+        });
+
+        if is_list {
+            let len = self.arena.mutate_root(|_, root| {
+                let list = root.pop_stack().as_gc::<List<'_>>().unwrap();
+                for &item in list.inner.iter() {
+                    root.temp_stack.push(item);
+                }
+                list.inner.len()
+            });
+
+            let eval_base = self.arena.mutate_root(|_, root| root.temp_stack.len());
+
+            for i in 0..len {
+                self.arena.mutate_root(|_, root| {
+                    let item = root.temp_stack[eval_base - len + i];
+                    root.push_stack(item);
+                });
+
+                if let err @ Action::Done(Err(_)) = self.force_tos() {
+                    self.arena.mutate_root(|_, root| root.temp_stack.truncate(eval_base - len));
+                    return err;
+                }
+
+                self.arena.mutate_root(|_, root| {
+                    let eval_item = root.pop_stack();
+                    root.temp_stack[eval_base - len + i] = eval_item;
+                });
+            }
+
+            self.arena.mutate_root(|mc, root| {
+                let items: SmallVec<[Value; 4]> = root.temp_stack[eval_base - len..eval_base].iter().copied().collect();
+                root.temp_stack.truncate(eval_base - len);
+                
+                // Reconstruct List
+                let new_list = Gc::new(mc, List { inner: items });
+                root.push_stack(Value::new_gc(new_list));
+            });
+            
+        } else if is_attrs {
+            let len = self.arena.mutate_root(|_, root| {
+                let attrs = root.pop_stack().as_gc::<AttrSet<'_>>().unwrap();
+                for &(key, item) in attrs.iter() {
+                    root.temp_stack.push(Value::new_inline(key));
+                    root.temp_stack.push(item);
+                }
+                attrs.len()
+            });
+
+            let eval_base = self.arena.mutate_root(|_, root| root.temp_stack.len());
+
+            for i in 0..len {
+                self.arena.mutate_root(|_, root| {
+                    let item = root.temp_stack[eval_base - len * 2 + i * 2 + 1];
+                    root.push_stack(item);
+                });
+
+                if let err @ Action::Done(Err(_)) = self.force_tos() {
+                    self.arena.mutate_root(|_, root| root.temp_stack.truncate(eval_base - len * 2));
+                    return err;
+                }
+
+                self.arena.mutate_root(|_, root| {
+                    let eval_item = root.pop_stack();
+                    root.temp_stack[eval_base - len * 2 + i * 2 + 1] = eval_item;
+                });
+            }
+
+            self.arena.mutate_root(|mc, root| {
+                let mut kv = SmallVec::with_capacity(len);
+                let mut i = eval_base - len * 2;
+                while i < eval_base {
+                    let key = root.temp_stack[i].as_inline::<StringId>().unwrap();
+                    let val = root.temp_stack[i + 1];
+                    kv.push((key, val));
+                    i += 2;
+                }
+                kv.sort_by_key(|(k, _)| *k);
+                root.temp_stack.truncate(eval_base - len * 2);
+                
+                let new_attrs = Gc::new(mc, unsafe { AttrSet::from_sorted_unchecked(kv) });
+                root.push_stack(Value::new_gc(new_attrs));
+            });
+        }
+
+        Action::Continue
+    }
+
+    pub(super) fn force_tos_deep(&mut self) -> Action {
+        if let err @ Action::Done(Err(_)) = self.force_tos() {
+            return err;
+        }
+
+        let (is_list, is_attrs) = self.arena.mutate_root(|_, root| {
+            let tos = *root.stack.tos().expect("stack underflow");
+            (tos.as_gc::<List<'_>>().is_some(), tos.as_gc::<AttrSet<'_>>().is_some())
+        });
+
+        if is_list {
+            let len = self.arena.mutate_root(|_, root| {
+                let list = root.pop_stack().as_gc::<List<'_>>().unwrap();
+                for &item in list.inner.iter() {
+                    root.temp_stack.push(item);
+                }
+                list.inner.len()
+            });
+
+            let eval_base = self.arena.mutate_root(|_, root| root.temp_stack.len());
+
+            for i in 0..len {
+                self.arena.mutate_root(|_, root| {
+                    let item = root.temp_stack[eval_base - len + i];
+                    root.push_stack(item);
+                });
+
+                if let err @ Action::Done(Err(_)) = self.force_tos_deep() {
+                    self.arena.mutate_root(|_, root| root.temp_stack.truncate(eval_base - len));
+                    return err;
+                }
+
+                self.arena.mutate_root(|_, root| {
+                    let eval_item = root.pop_stack();
+                    root.temp_stack[eval_base - len + i] = eval_item;
+                });
+            }
+
+            self.arena.mutate_root(|mc, root| {
+                let items: SmallVec<[Value; 4]> = root.temp_stack[eval_base - len..eval_base].iter().copied().collect();
+                root.temp_stack.truncate(eval_base - len);
+                let new_list = Gc::new(mc, List { inner: items });
+                root.push_stack(Value::new_gc(new_list));
+            });
+            
+        } else if is_attrs {
+            let len = self.arena.mutate_root(|_, root| {
+                let attrs = root.pop_stack().as_gc::<AttrSet<'_>>().unwrap();
+                for &(key, item) in attrs.iter() {
+                    root.temp_stack.push(Value::new_inline(key));
+                    root.temp_stack.push(item);
+                }
+                attrs.len()
+            });
+
+            let eval_base = self.arena.mutate_root(|_, root| root.temp_stack.len());
+
+            for i in 0..len {
+                self.arena.mutate_root(|_, root| {
+                    let item = root.temp_stack[eval_base - len * 2 + i * 2 + 1];
+                    root.push_stack(item);
+                });
+
+                if let err @ Action::Done(Err(_)) = self.force_tos_deep() {
+                    self.arena.mutate_root(|_, root| root.temp_stack.truncate(eval_base - len * 2));
+                    return err;
+                }
+
+                self.arena.mutate_root(|_, root| {
+                    let eval_item = root.pop_stack();
+                    root.temp_stack[eval_base - len * 2 + i * 2 + 1] = eval_item;
+                });
+            }
+
+            self.arena.mutate_root(|mc, root| {
+                let mut kv = SmallVec::with_capacity(len);
+                let mut i = eval_base - len * 2;
+                while i < eval_base {
+                    let key = root.temp_stack[i].as_inline::<StringId>().unwrap();
+                    let val = root.temp_stack[i + 1];
+                    kv.push((key, val));
+                    i += 2;
+                }
+                kv.sort_by_key(|(k, _)| *k);
+                root.temp_stack.truncate(eval_base - len * 2);
+                let new_attrs = Gc::new(mc, unsafe { AttrSet::from_sorted_unchecked(kv) });
+                root.push_stack(Value::new_gc(new_attrs));
+            });
+        }
+
+        Action::Continue
     }
 
     fn handle_vm_error(&mut self, e: VmError) -> Action {