split SelectDefault -> SelectStatic & Jump...

fix-codegen/src/disassembler.rs

@@ -255,9 +255,7 @@ impl<'a, Ctx: DisassemblerContext> Disassembler<'a, Ctx> {
                 ("MakePatternClosure", arg_str)
             }
 
-            Op::Call => {
+            Op::Call => ("Call", String::new()),
-                ("Call", String::new())
-            },
 
             Op::MakeAttrs => {
                 let count = self.read_u32();
@@ -265,17 +263,24 @@ impl<'a, Ctx: DisassemblerContext> Disassembler<'a, Ctx> {
             }
             Op::MakeEmptyAttrs => ("MakeEmptyAttrs", String::new()),
 
-            Op::Select => {
+            Op::SelectStatic => {
-                let path_len = self.read_u16();
-                let span_id = self.read_u32();
-                ("Select", format!("path_len={} span={}", path_len, span_id))
-            }
-            Op::SelectDefault => {
-                let path_len = self.read_u16();
                 let span_id = self.read_u32();
+                let key_id = self.read_u32();
                 (
-                    "SelectDefault",
+                    "SelectStatic",
-                    format!("path_len={} span={}", path_len, span_id),
+                    format!("key={} span={}", self.ctx.resolve_string(key_id), span_id),
+                )
+            }
+            Op::SelectDynamic => {
+                let span_id = self.read_u32();
+                ("SelectDynamic", format!("span={}", span_id))
+            }
+            Op::JumpIfSelectSucceeded => {
+                let offset = self.read_i32();
+                let target = (current_pc as isize + 1 + 4 + offset as isize) as usize;
+                (
+                    "JumpIfSelectSucceeded",
+                    format!("-> {:04x} offset={}", target, offset),
                 )
             }
             Op::HasAttr => {

fix-codegen/src/lib.rs

@@ -45,8 +45,9 @@ pub enum Op {
 
     MakeAttrs,
     MakeEmptyAttrs,
-    Select,
+    SelectStatic,
-    SelectDefault,
+    SelectDynamic,
+    JumpIfSelectSucceeded,
     HasAttr,
 
     MakeList,
@@ -114,13 +115,6 @@ pub enum OperandType {
     BigInt,
 }
 
-#[repr(u8)]
-#[derive(Debug, Clone, Copy, TryFromPrimitive)]
-pub enum AttrKeyType {
-    Static,
-    Dynamic,
-}
-
 pub enum Const {
     Smi(i32),
     Float(f64),
@@ -130,6 +124,13 @@ pub enum Const {
     Null,
 }
 
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, TryFromPrimitive)]
+pub enum AttrKeyType {
+    Static,
+    Dynamic,
+}
+
 pub enum InlineOperand {
     Const(Const),
     Local { layer: u16, local: u32 },
@@ -626,6 +627,7 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
                 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);
@@ -819,34 +821,32 @@ impl<'a, Ctx: BytecodeContext> BytecodeEmitter<'a, Ctx> {
         span: TextRange,
     ) {
         self.emit_expr(expr);
-        for attr in attrpath.iter() {
+
-            if let Attr::Dynamic(expr, _) = *attr {
-                self.emit_expr(expr);
-            }
-        }
-        if let Some(default) = default {
-            self.emit_expr(default);
-            let span_id = self.ctx.register_span(span);
-            self.emit_op(Op::SelectDefault);
-            self.emit_u16(attrpath.len() as u16);
-            self.emit_u32(span_id);
-        } else {
-            let span_id = self.ctx.register_span(span);
-            self.emit_op(Op::Select);
-            self.emit_u16(attrpath.len() as u16);
-            self.emit_u32(span_id);
-        }
         for attr in attrpath.iter() {
             match *attr {
                 Attr::Str(sym, _) => {
-                    self.emit_u8(AttrKeyType::Static as u8);
+                    let span_id = self.ctx.register_span(span);
+                    self.emit_op(Op::SelectStatic);
+                    self.emit_u32(span_id);
                     self.emit_str_id(sym);
                 }
-                Attr::Dynamic(_, _) => {
+                Attr::Dynamic(key_expr, _) => {
-                    self.emit_u8(AttrKeyType::Dynamic as u8);
+                    self.emit_expr(key_expr);
+                    let span_id = self.ctx.register_span(span);
+                    self.emit_op(Op::SelectDynamic);
+                    self.emit_u32(span_id);
                 }
             }
         }
+
+        if let Some(default) = default {
+            self.emit_op(Op::JumpIfSelectSucceeded);
+            let placeholder = self.emit_i32_placeholder();
+            let before: i32 = self.ctx.get_code().len().try_into().unwrap();
+            self.emit_expr(default);
+            let after: i32 = self.ctx.get_code().len().try_into().unwrap();
+            self.patch_i32(placeholder, after - before);
+        }
     }
 
     fn emit_has_attr(&mut self, lhs: RawIrRef<'_>, rhs: &[Attr<RawIrRef<'_>>]) {

fix-vm/src/lib.rs

@@ -215,11 +215,13 @@ impl<'gc> Vm<'gc> {
     }
 
     #[inline(always)]
+    #[must_use]
     fn pop_stack(&mut self) -> Value<'gc> {
         self.stack.pop().expect("stack underflow")
     }
 
     #[inline(always)]
+    #[must_use]
     fn peek_stack(&mut self, depth: usize) -> Value<'gc> {
         *self
             .stack
@@ -264,7 +266,7 @@ struct CallFrame<'gc> {
 }
 
 pub(crate) enum Action {
-    Continue,
+    Continue { pc: usize },
     Done(Result<fix_common::Value>),
 }
 
@@ -308,11 +310,6 @@ struct AttrEntry {
     val: OperandData,
 }
 
-enum SelectKeyData {
-    Static(StringId),
-    Dynamic,
-}
-
 macro_rules! try_vm {
     ($self:ident; $expr:expr) => {
         match $expr {
@@ -335,8 +332,9 @@ impl Vm<'_> {
 
         let mut pc = ip.0;
         loop {
-            match arena.mutate_root(|mc, root| root.execute_batch(&mut ctx, &mut pc, mc)) {
+            match arena.mutate_root(|mc, root| root.execute_batch(&mut ctx, pc, mc)) {
-                Action::Continue => {
+                Action::Continue { pc: new_pc } => {
+                    pc = new_pc;
                     if arena.metrics().allocation_debt() > COLLECTOR_GRANULARITY {
                         if arena.collection_phase() == CollectionPhase::Sweeping {
                             arena.collect_debt();
@@ -356,7 +354,7 @@ impl<'gc> Vm<'gc> {
     fn execute_batch(
         &mut self,
         ctx: &mut impl VmContext,
-        pc: &mut usize,
+        mut pc: usize,
         mc: &Mutation<'gc>,
     ) -> Action {
         use fix_codegen::Op::{self, *};
@@ -397,21 +395,8 @@ impl<'gc> Vm<'gc> {
                     }
                 }
             }};
-            (SelectKeyData, $n:expr) => {{
-                let mut keys = SmallVec::<[SelectKeyData; 4]>::with_capacity($n as usize);
-                for _ in 0..$n {
-                    let tag = read!(u8);
-                    let Ok(ty) = AttrKeyType::try_from_primitive(tag)
-                        .map_err(|err| panic!("unknown key tag: {:#04x}", err.number));
-                    match ty {
-                        AttrKeyType::Static => keys.push(SelectKeyData::Static(read!(StringId))),
-                        AttrKeyType::Dynamic => keys.push(SelectKeyData::Dynamic),
-                    };
-                }
-                keys
-            }};
             ($type:ty) => {
-                <$type>::from_le_bytes(read_array(ctx.bytecode(), pc))
+                <$type>::from_le_bytes(read_array(ctx.bytecode(), &mut pc))
             };
         }
 
@@ -433,7 +418,7 @@ impl<'gc> Vm<'gc> {
                                     with_env: self.with_env,
                                 });
 
-                                *pc = ip;
+                                pc = ip;
                                 self.env = env;
                                 self.with_env = with_env;
                                 *state = ThunkState::Blackhole;
@@ -454,18 +439,18 @@ impl<'gc> Vm<'gc> {
             }
 
             if fuel == 0 {
-                return Action::Continue;
+                return Action::Continue { pc };
             }
             fuel -= 1;
 
             // Save PC for Instruction Retry
-            let inst_start_pc = *pc;
+            let inst_start_pc = pc;
-            let byte = ctx.bytecode()[*pc];
+            let byte = ctx.bytecode()[pc];
             if !likely_stable::likely((0..Op::Illegal as u8).contains(&byte)) {
                 panic!("unknown opcode: {byte:#04x}")
             }
             let op = unsafe { std::mem::transmute::<u8, Op>(byte) };
-            *pc += 1;
+            pc += 1;
 
             match op {
                 PushSmi => {
@@ -606,13 +591,13 @@ impl<'gc> Vm<'gc> {
                             let new_env =
                                 Gc::new(mc, RefLock::new(Env::with_arg(arg, n_locals, env)));
                             self.call_stack.push(CallFrame {
-                                pc: *pc,
+                                pc,
                                 stack_depth: 0,
                                 thunk: None,
                                 env: self.env,
                                 with_env: self.with_env,
                             });
-                            *pc = ip as usize;
+                            pc = ip as usize;
                             self.env = new_env;
                         }
                     } else {
@@ -656,148 +641,81 @@ impl<'gc> Vm<'gc> {
                     self.push_stack(self.empty_attrs);
                 }
 
-                Select => {
+                SelectStatic => {
-                    let n = read!(u16) as usize;
                     let _span_id = read!(u32);
+                    let key = read!(StringId);
 
-                    let keys = read!(SelectKeyData, n);
+                    try_force!(0, inst_start_pc);
-                    let dyn_count = keys
-                        .iter()
-                        .filter(|k| matches!(k, SelectKeyData::Dynamic))
-                        .count();
 
-                    // Force target (at depth `dyn_count`) and all dynamic keys on top of it.
+                    let attrs = self.peek_stack(0).restrict().expect("forced");
-                    for i in 0..=dyn_count {
+                    let Some(attrset) = attrs.as_gc::<AttrSet>() else {
-                        try_force!(i, inst_start_pc);
+                        return Action::Done(Err(Error::eval_error("value is not a set while a set was expected")));
-                    }
-
-                    // Stack Layout: [..., target, dyn1, dyn2, ..., dyn_m]
-                    let target_idx = self.stack.len() - dyn_count - 1;
-                    let mut current_dyn_key_idx = target_idx + 1;
-
-                    let mut current_val = self.stack[target_idx].restrict().expect("forced");
-                    let mut result_val = None;
-                    let mut error = None;
-
-                    for (i, key) in keys.iter().enumerate() {
-                        let key_sid = match key {
-                            SelectKeyData::Static(sid) => *sid,
-                            SelectKeyData::Dynamic => {
-                                let v = self.stack[current_dyn_key_idx]
-                                    .restrict()
-                                    .expect("dynamic key must be forced");
-                                current_dyn_key_idx += 1;
-                                if let Some(sid) = v.as_inline::<StringId>() {
-                                    sid
-                                } else if let Some(ns) = v.as_gc::<NixString>() {
-                                    ctx.intern_string(ns.as_str())
-                                } else {
-                                    panic!("dynamic select key must be a string")
-                                }
-                            }
                     };
 
-                        let Some(attrset) = current_val.as_gc::<AttrSet>() else {
+                    match attrset.lookup(key) {
-                            error = Some(vm_err("value is not a set while a set was expected"));
+                        Some(v) => {
+                            self.replace_stack(0, v);
+                        }
+                        None => {
+                            loop {
+                                let byte = ctx.bytecode()[pc];
+                                if byte == SelectStatic as u8 {
+                                    pc += 1 + 4 + 4;
+                                } else if byte == SelectDynamic as u8 {
+                                    pc += 1 + 4;
+                                } else if byte == JumpIfSelectSucceeded as u8 {
+                                    pc += 1 + 4;
+                                    let _ = self.pop_stack();
                                     break;
+                                } else {
+                                    let name = ctx.resolve_string(key);
+                                    return Action::Done(Err(Error::eval_error(format!("attribute '{name}' missing"))));
+                                }
+                            }
+                        }
+                    }
+                }
+                SelectDynamic => {
+                    let _span_id = read!(u32);
+
+                    try_force!(0, inst_start_pc);
+                    try_force!(1, inst_start_pc);
+
+                    let key_val = self.stack[self.stack.len() - 1]
+                        .restrict()
+                        .expect("dynamic key must be forced");
+                    let key_sid = if let Some(sid) = key_val.as_inline::<StringId>() {
+                        sid
+                    } else if let Some(ns) = key_val.as_gc::<NixString>() {
+                        ctx.intern_string(ns.as_str())
+                    } else {
+                        return self.handle_vm_error(vm_err("dynamic select key must be a string"));
+                    };
+
+                    let attrset_val = self.stack[self.stack.len() - 2].restrict().expect("forced");
+                    let Some(attrset) = attrset_val.as_gc::<AttrSet>() else {
+                        return self.handle_vm_error(vm_err(
+                            "value is not a set while a set was expected",
+                        ));
                     };
 
                     match attrset.lookup(key_sid) {
                         Some(v) => {
-                                if i < n - 1 {
+                            self.stack.truncate(self.stack.len() - 2);
-                                    // FIXME: Proper async force hook inside select chain for nested thunks
+                            self.push_stack(v);
-                                    current_val = v.restrict().unwrap_or_else(|_| {
-                                        panic!("intermediate select values must be forced")
-                                    });
-                                } else {
-                                    result_val = Some(v);
-                                }
                         }
                         None => {
                             let name = ctx.resolve_string(key_sid);
-                                error = Some(vm_err(format!("attribute '{name}' missing")));
+                            return self
-                                break;
+                                .handle_vm_error(vm_err(format!("attribute '{name}' missing")));
                         }
                     }
                 }
+                JumpIfSelectSucceeded => {
+                    let offset = read!(i32);
+                    pc = ((pc as isize) + (offset as isize)) as usize;
+                }
 
-                    // Clean up the target and all dynamic keys
-                    self.stack.truncate(target_idx);
-
-                    if let Some(e) = error {
-                        return self.handle_vm_error(e);
-                    }
-                    if let Some(v) = result_val {
-                        self.push_stack(v);
-                    }
-                }
-                SelectDefault => {
-                    let n = read!(u16) as usize;
-                    let _span_id = read!(u32);
-
-                    let keys = read!(SelectKeyData, n);
-                    let dyn_count = keys
-                        .iter()
-                        .filter(|k| matches!(k, SelectKeyData::Dynamic))
-                        .count();
-
-                    // Stack layout: [..., target, default_val, dyn1, dyn2, ..., dyn_m]
-                    for i in 0..=dyn_count + 1 {
-                        try_force!(i, inst_start_pc);
-                    }
-
-                    let target_idx = self.stack.len() - dyn_count - 2;
-                    let default_idx = target_idx + 1;
-                    let mut current_dyn_key_idx = default_idx + 1;
-
-                    let mut current_val = self.stack[target_idx].restrict().expect("forced");
-                    let mut result_val = None;
-                    let mut use_default = false;
-
-                    for (i, key) in keys.iter().enumerate() {
-                        let key_sid = match key {
-                            SelectKeyData::Static(sid) => *sid,
-                            SelectKeyData::Dynamic => {
-                                let v = self.stack[current_dyn_key_idx]
-                                    .restrict()
-                                    .expect("dynamic key must be forced");
-                                current_dyn_key_idx += 1;
-                                if let Some(sid) = v.as_inline::<StringId>() {
-                                    sid
-                                } else if let Some(ns) = v.as_gc::<NixString>() {
-                                    ctx.intern_string(ns.as_str())
-                                } else {
-                                    panic!("dynamic select key must be a string")
-                                }
-                            }
-                        };
-
-                        if let Some(attrset) = current_val.as_gc::<AttrSet>()
-                            && let Some(v) = attrset.lookup(key_sid)
-                        {
-                            if i < n - 1 {
-                                current_val = v.restrict().unwrap_or_else(|_| {
-                                    panic!("intermediate select values must be forced")
-                                });
-                            } else {
-                                result_val = Some(v);
-                            }
-                            continue;
-                        }
-
-                        use_default = true;
-                        break;
-                    }
-
-                    let def = self.stack[default_idx];
-                    self.stack.truncate(target_idx);
-
-                    if use_default {
-                        self.push_stack(def);
-                    } else if let Some(v) = result_val {
-                        self.push_stack(v);
-                    }
-                }
                 HasAttr => {
                     let _n = read!(u16) as usize;
                     todo!("HasAttr");
@@ -946,7 +864,7 @@ impl<'gc> Vm<'gc> {
                     try_force!(0, inst_start_pc);
                     let cond = self.pop_stack();
                     if cond.as_inline::<bool>() == Some(false) {
-                        *pc = ((*pc as isize) + (offset as isize)) as usize;
+                        pc = ((pc as isize) + (offset as isize)) as usize;
                     }
                 }
                 JumpIfTrue => {
@@ -954,12 +872,12 @@ impl<'gc> Vm<'gc> {
                     try_force!(0, inst_start_pc);
                     let cond = self.pop_stack();
                     if cond.as_inline::<bool>() == Some(true) {
-                        *pc = ((*pc as isize) + (offset as isize)) as usize;
+                        pc = ((pc as isize) + (offset as isize)) as usize;
                     }
                 }
                 Jump => {
                     let offset = read!(i32);
-                    *pc = ((*pc as isize) + (offset as isize)) as usize;
+                    pc = ((pc as isize) + (offset as isize)) as usize;
                 }
 
                 ConcatStrings => {
@@ -1027,7 +945,7 @@ impl<'gc> Vm<'gc> {
 
                     let env = self.pop_stack().as_gc::<AttrSet>().unwrap();
                     let Some(val) = env.lookup(name) else {
-                        *pc = inst_start_pc;
+                        pc = inst_start_pc;
                         self.with_env = prev;
                         continue 'dispatch;
                     };
@@ -1066,7 +984,7 @@ impl<'gc> Vm<'gc> {
                 }
 
                 Return => {
-                    if let Some(result) = self.handle_return(pc, ctx, mc) {
+                    if let Some(result) = self.handle_return(&mut pc, ctx, mc) {
                         return Action::Done(result);
                     }
                 }