ForceMode (WIP)

fix/src/codegen.rs

@@ -7,7 +7,7 @@ use rnix::TextRange;
 use string_interner::Symbol as _;
 
 use crate::ir::{ArgId, Attr, BinOpKind, Ir, Param, RawIrRef, StringId, ThunkId, UnOpKind};
-use crate::runtime::{BUILTINS, BuiltinId};
+use crate::runtime::BUILTINS;
 use crate::runtime::value::{Null, PrimOp, StaticValue};
 
 pub struct InstructionPtr(pub(crate) usize);

fix/src/runtime.rs

@@ -47,6 +47,7 @@ pub struct Runtime {
     fuel: usize,
     error_contexts: Stack<8192, ErrorFrame>,
     arena: Arena<Rootable![GcRoot<'_>]>,
+    force_mode: ForceMode,
 }
 
 impl Runtime {
@@ -81,6 +82,7 @@ impl Runtime {
             fuel: Self::DEFAULT_FUEL_AMOUNT,
             error_contexts: Stack::new(),
             arena,
+            force_mode: ForceMode::default(),
         })
     }
 

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,
@@ -578,13 +579,181 @@ impl Runtime {
                 let n = self.read_u16() as usize;
                 let _span_id = self.read_u32();
                 let keys = self.read_attr_keys(n);
-                todo!("implement Select (force + lookup)");
+
+                // Move dynamic key values from stack to temp_stack
+                let dyn_count = keys
+                    .iter()
+                    .filter(|k| matches!(k, SelectKeyData::Dynamic))
+                    .count();
+                let temp_base = self.arena.mutate_root(|_, root| {
+                    let base = root.temp_stack.len();
+                    for _ in 0..dyn_count {
+                        let v = root.pop_stack();
+                        root.temp_stack.push(v);
+                    }
+                    base
+                });
+
+                // Force the expr (now TOS)
+                self.force_tos();
+
+                for (i, key) in keys.iter().enumerate() {
+                    let key_sid = match key {
+                        &SelectKeyData::Static(sid) => sid,
+                        SelectKeyData::Dynamic => {
+                            self.arena.mutate_root(|_, root| {
+                                let v =
+                                    root.temp_stack.pop().expect("missing dynamic key");
+                                root.push_stack(v);
+                            });
+                            self.force_tos();
+                            let key_data: std::result::Result<StringId, String> =
+                                self.arena.mutate_root(|_, root| {
+                                    let v = root.pop_stack();
+                                    if let Some(sid) = v.as_inline::<StringId>() {
+                                        Ok(sid)
+                                    } else if let Some(ns) = v.as_gc::<NixString>() {
+                                        Err(ns.as_str().to_owned())
+                                    } else {
+                                        panic!("dynamic select key must be a string")
+                                    }
+                                });
+                            match key_data {
+                                Ok(sid) => sid,
+                                Err(s) => StringId(self.strings.get_or_intern(&s)),
+                            }
+                        }
+                    };
+
+                    let result = self.arena.mutate_root(|_, root| {
+                        let val = root.pop_stack();
+                        let Some(attrset) = val.as_gc::<AttrSet<'_>>() else {
+                            return Err(vm_err(
+                                "value is not a set while a set was expected",
+                            ));
+                        };
+                        match attrset.lookup(key_sid) {
+                            Some(v) => {
+                                root.push_stack(v);
+                                Ok(true)
+                            }
+                            None => Ok(false),
+                        }
+                    });
+
+                    match result {
+                        Err(e) => {
+                            self.arena
+                                .mutate_root(|_, root| root.temp_stack.truncate(temp_base));
+                            return Runtime::handle_vm_error(self, e);
+                        }
+                        Ok(false) => {
+                            self.arena
+                                .mutate_root(|_, root| root.temp_stack.truncate(temp_base));
+                            let name =
+                                self.strings.resolve(key_sid.0).unwrap_or("«unknown»");
+                            return Runtime::handle_vm_error(
+                                self,
+                                vm_err(format!("attribute '{name}' missing")),
+                            );
+                        }
+                        Ok(true) => {
+                            if i < n - 1 {
+                                self.force_tos();
+                            }
+                        }
+                    }
+                }
+
+                self.arena
+                    .mutate_root(|_, root| root.temp_stack.truncate(temp_base));
             }
             SelectDefault => {
                 let n = self.read_u16() as usize;
                 let _span_id = self.read_u32();
                 let keys = self.read_attr_keys(n);
-                todo!("implement SelectDefault (force + lookup with default)");
+
+                let dyn_count = keys
+                    .iter()
+                    .filter(|k| matches!(k, SelectKeyData::Dynamic))
+                    .count();
+                let temp_base = self.arena.mutate_root(|_, root| {
+                    let base = root.temp_stack.len();
+                    // Default value is on top of the stack (pushed last by codegen)
+                    let default_val = root.pop_stack();
+                    root.temp_stack.push(default_val);
+                    // Then dynamic keys
+                    for _ in 0..dyn_count {
+                        let v = root.pop_stack();
+                        root.temp_stack.push(v);
+                    }
+                    base
+                });
+
+                // Force the expr (now TOS)
+                self.force_tos();
+
+                let mut use_default = false;
+                for (i, key) in keys.iter().enumerate() {
+                    let key_sid = match key {
+                        &SelectKeyData::Static(sid) => sid,
+                        SelectKeyData::Dynamic => {
+                            self.arena.mutate_root(|_, root| {
+                                let v =
+                                    root.temp_stack.pop().expect("missing dynamic key");
+                                root.push_stack(v);
+                            });
+                            self.force_tos();
+                            let key_data: std::result::Result<StringId, String> =
+                                self.arena.mutate_root(|_, root| {
+                                    let v = root.pop_stack();
+                                    if let Some(sid) = v.as_inline::<StringId>() {
+                                        Ok(sid)
+                                    } else if let Some(ns) = v.as_gc::<NixString>() {
+                                        Err(ns.as_str().to_owned())
+                                    } else {
+                                        panic!("dynamic select key must be a string")
+                                    }
+                                });
+                            match key_data {
+                                Ok(sid) => sid,
+                                Err(s) => StringId(self.strings.get_or_intern(&s)),
+                            }
+                        }
+                    };
+
+                    let found = self.arena.mutate_root(|_, root| {
+                        let val = root.pop_stack();
+                        if let Some(attrset) = val.as_gc::<AttrSet<'_>>() {
+                            if let Some(v) = attrset.lookup(key_sid) {
+                                root.push_stack(v);
+                                return true;
+                            }
+                        }
+                        // Not a set or key missing → use default
+                        false
+                    });
+
+                    if !found {
+                        use_default = true;
+                        break;
+                    }
+
+                    if i < n - 1 {
+                        self.force_tos();
+                    }
+                }
+
+                if use_default {
+                    self.arena.mutate_root(|_, root| {
+                        let default_val = root.temp_stack[temp_base];
+                        root.temp_stack.truncate(temp_base);
+                        root.push_stack(default_val);
+                    });
+                } else {
+                    self.arena
+                        .mutate_root(|_, root| root.temp_stack.truncate(temp_base));
+                }
             }
             HasAttr => {
                 let n = self.read_u16() as usize;
@@ -1040,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())));
@@ -1087,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
@@ -1119,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| {
@@ -1137,7 +1309,7 @@ impl Runtime {
                     *thunk.borrow_mut(mc) = ThunkState::Evaluated(val);
                 }
                 *thunk = val;
-            })
+            });
         }
     }
 
@@ -1163,21 +1335,48 @@ 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 {
+        // FIXME: shallow
+        self.force_tos()
+    }
+
+    pub(super) fn force_tos_deep(&mut self) -> Action {
+        // FIXME: deep
+        self.force_tos()
     }
 
     fn handle_vm_error(&mut self, e: VmError) -> Action {