feat(vm): preserve non-yielding registry bindings

Author: fffonion

src/vm/host.rs

@@ -105,6 +105,7 @@ enum RegistryEntryKind {
     StackStatic(StaticHostStackFunction),
     ArgsFactory(Arc<HostArgsFactory>),
     ArgsStatic(StaticHostArgsFunction),
+    ArgsStaticNonYielding(StaticHostArgsFunction),
 }
 
 #[derive(Clone)]
@@ -308,6 +309,32 @@ impl HostFunctionRegistry {
         self.invalidate_plan_cache();
     }
 
+    pub fn register_static_non_yielding_args(
+        &mut self,
+        name: impl Into<String>,
+        arity: u8,
+        function: StaticHostArgsFunction,
+    ) {
+        let name = name.into();
+        if let Some(&slot) = self.by_name.get(&name)
+            && let Some(entry) = Arc::make_mut(&mut self.entries).get_mut(slot as usize)
+        {
+            entry.arity = arity;
+            entry.kind = RegistryEntryKind::ArgsStaticNonYielding(function);
+            self.invalidate_plan_cache();
+            return;
+        }
+
+        let entries = Arc::make_mut(&mut self.entries);
+        let slot = entries.len() as u16;
+        entries.push(RegistryEntry {
+            arity,
+            kind: RegistryEntryKind::ArgsStaticNonYielding(function),
+        });
+        Arc::make_mut(&mut self.by_name).insert(name, slot);
+        self.invalidate_plan_cache();
+    }
+
     pub fn bind_vm_cached(&self, vm: &mut Vm) -> VmResult<()> {
         let plan = self.prepare_shared_plan(&vm.program.imports)?;
         self.bind_vm_with_plan(vm, &plan)
@@ -415,6 +442,9 @@ impl HostFunctionRegistry {
                 RegistryEntryKind::ArgsStatic(function) => {
                     vm.register_static_args_function(*function);
                 }
+                RegistryEntryKind::ArgsStaticNonYielding(function) => {
+                    vm.register_static_non_yielding_args_function(*function);
+                }
             }
         }
         vm.install_resolved_calls(plan.resolved_calls.clone())?;

tests/compiler/compiler_common_tests.rs

@@ -577,6 +577,37 @@ fn host_function_registry_caches_static_args_function_pointer_plan_across_vms()
     assert_eq!(vm2.stack(), &[Value::Int(6)]);
 }
 
+#[test]
+fn host_function_registry_caches_static_non_yielding_args_function_pointer_plan_across_vms() {
+    let source = include_str!("../../examples/example.rss");
+    let compiled = compile_source(source).expect("compile should succeed");
+
+    let mut registry = HostFunctionRegistry::new();
+    registry.register_static_non_yielding_args("print", 1, |args| {
+        Ok(CallOutcome::Return(args.to_vec().into()))
+    });
+    registry.register_static_non_yielding_args("add_one", 1, static_add_one_args);
+    let plan = registry
+        .prepare_plan(&compiled.program.imports)
+        .expect("plan should build");
+
+    let mut vm1 = Vm::new(compiled.program.clone());
+    registry
+        .bind_vm_with_plan(&mut vm1, &plan)
+        .expect("cached static non-yielding args host binding should succeed");
+    let status1 = vm1.run().expect("vm should run");
+    assert_eq!(status1, VmStatus::Halted);
+    assert_eq!(vm1.stack(), &[Value::Int(6)]);
+
+    let mut vm2 = Vm::new(compiled.program);
+    registry
+        .bind_vm_with_plan(&mut vm2, &plan)
+        .expect("cached static non-yielding args host binding should succeed");
+    let status2 = vm2.run().expect("vm should run");
+    assert_eq!(status2, VmStatus::Halted);
+    assert_eq!(vm2.stack(), &[Value::Int(6)]);
+}
+
 #[test]
 fn break_and_continue_outside_loop_are_rejected() {
     let break_err = match compile_source("break;") {