refactor: simplify vm creation API

Author: fffonion

pd-edge/src/bin/pd-edge-console.rs

@@ -337,7 +337,7 @@ async fn run_loaded_program_once(state: &SharedState) -> Result<(), Box<dyn std:
         state.rate_limiter.clone(),
     )));
     let async_ops = new_shared_vm_async_ops(); 
-    let mut vm = Vm::with_locals_shared(loaded.program.clone(), loaded.local_count);
+    let mut vm = Vm::new_shared(loaded.program.clone());
     vm.set_async_bridge(Box::new(VmAsyncOpBridge::new(async_ops.clone())));
 
     register_host_module(&mut vm, context, async_ops.clone())?;

pd-edge/src/runtime.rs

@@ -10,7 +10,7 @@ use std::{
 use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 use tracing::{info, warn};
-use vm::{Program, decode_program, infer_local_count, validate_program};
+use vm::{Program, decode_program, validate_program};
 
 use crate::{
     HOST_FUNCTION_COUNT,
@@ -40,7 +40,6 @@ pub struct SharedState {
 #[derive(Clone)]
 pub struct LoadedProgram {
     pub program: Arc<Program>,
-    pub local_count: usize,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -466,28 +465,12 @@ pub async fn apply_program_from_bytes(state: &SharedState, bytes: &[u8]) -> Prog
         };
     }
 
-    let local_count = match infer_local_count(&program) {
-        Ok(local_count) => local_count,
-        Err(err) => {
-            state.record_program_apply_failure();
-            let message = format!("invalid bytecode: {err}");
-            warn!("{} local inference error: {err}", category_program());
-            return ProgramApplyReport {
-                applied: false,
-                constants: None,
-                code_bytes: None,
-                local_count: None,
-                message: Some(message),
-            };
-        }
-    };
-
+    let local_count = program.local_count;
     let const_count = program.constants.len();
     let code_len = program.code.len();
     let mut guard = state.active_program.write().await;
     *guard = Some(Arc::new(LoadedProgram {
         program: Arc::new(program),
-        local_count,
     }));
     state.record_program_apply_success();
     info!(

pd-edge/src/runtime/vm_runner.rs

@@ -34,14 +34,12 @@ pub async fn execute_vm_with_context(
     debug: VmDebugInvocation,
     register_host_modules: HostModuleRegistrar,
 ) -> Result<VmExecutionOutcome, VmExecutionError> {
-    let local_count = program.local_count;
     let program = program.program.clone();
     let async_ops = new_shared_vm_async_ops();
 
     let task = tokio::task::spawn_blocking(move || {
         run_vm_blocking(
             program,
-            local_count,
             vm_context,
             debug_session,
             debug,
@@ -59,14 +57,13 @@ pub async fn execute_vm_with_context(
 
 fn run_vm_blocking(
     program: std::sync::Arc<vm::Program>,
-    local_count: usize,
     vm_context: SharedProxyVmContext,
     debug_session: SharedDebugSession,
     debug: VmDebugInvocation,
     async_ops: SharedVmAsyncOps,
     register_host_modules: HostModuleRegistrar,
 ) -> Result<VmExecutionOutcome, VmExecutionError> {
-    let mut vm = Vm::with_locals_shared(program, local_count);
+    let mut vm = Vm::new_shared(program);
     vm.set_async_bridge(Box::new(VmAsyncOpBridge::new(async_ops.clone())));
     register_host_modules(&mut vm, vm_context.clone(), async_ops)
         .map_err(VmExecutionError::HostRegistration)?;

pd-edge/tests/compile_tests.rs

@@ -100,7 +100,7 @@ fn compile_edge_source_file_prefers_local_module_over_host_namespace_fallback()
         "local runtime module should win over host namespace fallback"
     );
 
-    let mut vm = Vm::with_locals(compiled.program, compiled.locals);
+    let mut vm = Vm::new(compiled.program);
     let status = vm.run().expect("vm should run");
     assert_eq!(status, VmStatus::Halted);
     assert_eq!(vm.stack(), &[Value::Int(42)]);
@@ -144,7 +144,7 @@ fn compile_edge_source_file_with_options_can_override_runtime_module() {
         "runtime module override should replace host import fallback"
     );
 
-    let mut vm = Vm::with_locals(compiled.program, compiled.locals);
+    let mut vm = Vm::new(compiled.program);
     let status = vm.run().expect("vm should run");
     assert_eq!(status, VmStatus::Halted);
     assert_eq!(vm.stack(), &[Value::Int(42)]);

pd-vm/src/bin/pd-vm-run.rs

@@ -97,7 +97,7 @@ fn run_main() -> Result<(), Box<dyn std::error::Error>> {
         return Ok(());
     }
     let recording_program = cli.record_path.as_ref().map(|_| compiled.program.clone());
-    let mut vm = Vm::with_locals(compiled.program, compiled.locals);
+    let mut vm = Vm::new(compiled.program);
     if let Some(fuel) = cli.fuel {
         vm.set_fuel(fuel);
     }
@@ -525,7 +525,7 @@ fn run_repl() -> Result<(), Box<dyn std::error::Error>> {
                         continue;
                     }
                 };
-                let mut vm = Vm::with_locals(compiled.program, compiled.locals);
+                let mut vm = Vm::new(compiled.program);
                 if let Err(err) = register_functions(&mut vm, &compiled.functions) {
                     println!("{err}");
                     continue;

pd-vm/src/bytecode.rs

@@ -19,15 +19,18 @@ pub struct HostImport {
 pub struct Program {
     pub constants: Vec<Value>,
     pub code: Vec<u8>,
+    pub local_count: usize,
     pub imports: Vec<HostImport>,
     pub debug: Option<crate::debug_info::DebugInfo>,
 }
 
 impl Program {
     pub fn new(constants: Vec<Value>, code: Vec<u8>) -> Self {
+        let local_count = infer_local_count_from_code(&code);
         Self {
             constants,
             code,
+            local_count,
             imports: Vec::new(),
             debug: None,
         }
@@ -38,9 +41,11 @@ impl Program {
         code: Vec<u8>,
         debug: Option<crate::debug_info::DebugInfo>,
     ) -> Self {
+        let local_count = infer_local_count_from_code(&code);
         Self {
             constants,
             code,
+            local_count,
             imports: Vec::new(),
             debug,
         }
@@ -52,13 +57,72 @@ impl Program {
         imports: Vec<HostImport>,
         debug: Option<crate::debug_info::DebugInfo>,
     ) -> Self {
+        let local_count = infer_local_count_from_code(&code);
         Self {
             constants,
             code,
+            local_count,
             imports,
             debug,
         }
     }
+
+    pub fn with_local_count(mut self, local_count: usize) -> Self {
+        self.local_count = local_count;
+        self
+    }
+}
+
+fn infer_local_count_from_code(code: &[u8]) -> usize {
+    let mut ip = 0usize;
+    let mut max_local_index: Option<u8> = None;
+
+    while let Some(&opcode) = code.get(ip) {
+        ip += 1;
+        match opcode {
+            x if x == OpCode::Nop as u8
+                || x == OpCode::Ret as u8
+                || x == OpCode::Add as u8
+                || x == OpCode::Sub as u8
+                || x == OpCode::Mul as u8
+                || x == OpCode::Div as u8
+                || x == OpCode::Neg as u8
+                || x == OpCode::Ceq as u8
+                || x == OpCode::Clt as u8
+                || x == OpCode::Cgt as u8
+                || x == OpCode::Pop as u8
+                || x == OpCode::Dup as u8
+                || x == OpCode::Shl as u8
+                || x == OpCode::Shr as u8
+                || x == OpCode::Mod as u8
+                || x == OpCode::And as u8
+                || x == OpCode::Or as u8 => {}
+            x if x == OpCode::Ldc as u8
+                || x == OpCode::Br as u8
+                || x == OpCode::Brfalse as u8 => {
+                    if ip + 4 > code.len() {
+                        break;
+                    }
+                    ip += 4;
+                }
+            x if x == OpCode::Ldloc as u8 || x == OpCode::Stloc as u8 => {
+                let Some(&index) = code.get(ip) else {
+                    break;
+                };
+                ip += 1;
+                max_local_index = Some(max_local_index.map_or(index, |prev| prev.max(index)));
+            }
+            x if x == OpCode::Call as u8 => {
+                if ip + 3 > code.len() {
+                    break;
+                }
+                ip += 3;
+            }
+            _ => break,
+        }
+    }
+
+    max_local_index.map_or(0, |index| index as usize + 1)
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]

pd-vm/src/compiler/mod.rs

@@ -236,7 +236,7 @@ struct LocalDebugRange {
 impl CompiledProgram {
     #[cfg(feature = "runtime")]
     pub fn into_vm(self) -> Vm {
-        Vm::with_locals(self.program, self.locals)
+        Vm::new(self.program)
     }
 }
 
@@ -332,6 +332,7 @@ fn compile_parsed_output(
     let mut program = compiler
         .compile_program(&stmts)
         .map_err(SourceError::Compile)?;
+    program.local_count = locals;
     program.imports = runtime_import_functions
         .iter()
         .map(|func| HostImport {

pd-vm/src/debugger.rs

@@ -583,7 +583,8 @@ impl DebugCommandBridge {
             state.attached = true;
             state.current_line = current_line(vm);
             state.pending_request = None;
-            state.pending_response = None;
+            // Keep any in-flight response visible to execute() to avoid
+            // races when the debugger resumes and immediately re-attaches.
             self.inner.changed.notify_all();
         }
 
@@ -1867,7 +1868,7 @@ mod tests {
                 }],
             }),
         );
-        let mut vm = Vm::with_locals(program, 1);
+        let mut vm = Vm::new(program.with_local_count(1));
         let status = vm.run().expect("vm should run");
         assert_eq!(status, crate::vm::VmStatus::Halted);
         vm
@@ -1889,7 +1890,7 @@ mod tests {
                 }],
             }),
         );
-        Vm::with_locals(program, 1)
+        Vm::new(program.with_local_count(1))
     }
 
     fn vm_with_scoped_named_locals() -> Vm {
@@ -1925,7 +1926,7 @@ mod tests {
                 ],
             }),
         );
-        let mut vm = Vm::with_locals(program, 1);
+        let mut vm = Vm::new(program.with_local_count(1));
         let status = vm.run().expect("vm should run");
         assert_eq!(status, crate::vm::VmStatus::Halted);
         vm
@@ -2560,7 +2561,7 @@ mod tests {
         debugger.stop_on_entry();
 
         let join = std::thread::spawn(move || {
-            let mut vm = Vm::with_locals(program, 3);
+            let mut vm = Vm::new(program.with_local_count(3));
             vm.run_with_debugger(&mut debugger)
                 .expect("debugged vm run should succeed")
         });

pd-vm/src/vm/jit/native/aarch64.rs

@@ -2775,7 +2775,7 @@ mod tests {
 
     #[test]
     fn stloc_step_bridge_moves_owned_values_safely() {
-        let mut vm = Vm::with_locals(Program::new(Vec::new(), Vec::new()), 1);
+        let mut vm = Vm::new(Program::new(Vec::new(), Vec::new()).with_local_count(1));
         vm.locals[0] = Value::String("old".to_string());
         vm.stack.push(Value::String("new".to_string()));
 
@@ -2812,7 +2812,7 @@ mod tests {
 
     #[test]
     fn ldloc_step_bridge_clones_owned_values() {
-        let mut vm = Vm::with_locals(Program::new(Vec::new(), Vec::new()), 1);
+        let mut vm = Vm::new(Program::new(Vec::new(), Vec::new()).with_local_count(1));
         vm.locals[0] = Value::Map(vec![(Value::Int(1), Value::Int(2))]);
 
         let status =
@@ -2959,9 +2959,8 @@ mod tests {
             executions: 0,
         };
 
-        let mut vm = Vm::with_locals(
-            Program::new(vec![Value::Int(200), Value::Int(1)], vec![0; 128]),
-            2,
+        let mut vm = Vm::new(
+            Program::new(vec![Value::Int(200), Value::Int(1)], vec![0; 128]).with_local_count(2),
         );
         vm.locals[0] = Value::Int(0);
         vm.locals[1] = Value::Int(0);

pd-vm/src/vm/jit/native/x86_64.rs

@@ -3616,7 +3616,7 @@ mod tests {
 
     #[test]
     fn stloc_step_bridge_moves_owned_values_safely() {
-        let mut vm = Vm::with_locals(Program::new(Vec::new(), Vec::new()), 1);
+        let mut vm = Vm::new(Program::new(Vec::new(), Vec::new()).with_local_count(1));
         vm.locals[0] = Value::String("old".to_string());
         vm.stack.push(Value::String("new".to_string()));