Preserve VMBC local count in embedded decoder

Author: fffonion

pd-vm-nostd/src/program.rs

@@ -62,6 +62,11 @@ impl Program {
         }
     }
 
+    pub(crate) fn with_local_count(mut self, local_count: usize) -> Self {
+        self.local_count = local_count;
+        self
+    }
+
     pub fn constants(&self) -> &[Value] {
         &self.constants
     }

pd-vm-nostd/src/vmbc.rs

@@ -52,13 +52,17 @@ pub fn decode_program(bytes: &[u8]) -> Result<Program, WireError> {
         });
     }
 
-    skip_type_map(&mut cursor)?;
+    let encoded_local_count = skip_type_map(&mut cursor)?;
     skip_debug_info(&mut cursor)?;
     if !cursor.is_empty() {
         return Err(WireError::TrailingBytes);
     }
 
-    Ok(Program::new(constants, code, imports))
+    let program = Program::new(constants, code, imports);
+    Ok(match encoded_local_count {
+        Some(local_count) => program.with_local_count(local_count),
+        None => program,
+    })
 }
 
 fn reserve<T>(items: &mut Vec<T>, field: &'static str, count: usize) -> Result<(), WireError> {
@@ -71,9 +75,9 @@ fn read_value_type(raw: u8) -> Result<ValueType, WireError> {
     ValueType::try_from(raw).map_err(|()| WireError::InvalidValueType(raw))
 }
 
-fn skip_type_map(cursor: &mut Cursor<'_>) -> Result<(), WireError> {
+fn skip_type_map(cursor: &mut Cursor<'_>) -> Result<Option<usize>, WireError> {
     match cursor.read_u8()? {
-        0 => Ok(()),
+        0 => Ok(None),
         1 => {
             cursor.read_bool()?;
             let local_count = cursor.read_u32()? as usize;
@@ -96,7 +100,7 @@ fn skip_type_map(cursor: &mut Cursor<'_>) -> Result<(), WireError> {
                 read_value_type(cursor.read_u8()?)?;
                 read_value_type(cursor.read_u8()?)?;
             }
-            Ok(())
+            Ok(Some(local_count))
         }
         value => Err(WireError::InvalidTypeMapFlag(value)),
     }

pd-vm-nostd/tests/embedded_vmbc.rs

@@ -1,5 +1,9 @@
 use pd_vm_nostd::{Value as EmbeddedValue, WireError, decode_program};
-use vm::{HostImport, OpCode, Program, Value, ValueType, compile_source, encode_program};
+use vm::compiler::TypeSchema;
+use vm::{
+    HostImport, OpCode, Program, ReplLocalBinding, Value, ValueType, compile_source,
+    compile_source_for_repl_with_locals, encode_program,
+};
 
 fn encoded_scalar_program() -> Vec<u8> {
     let mut program = Program::new(
@@ -57,6 +61,31 @@ fn embedded_decoder_accepts_compiler_type_and_debug_metadata() {
     assert_eq!(program.local_count(), compiled.locals);
 }
 
+#[test]
+fn embedded_decoder_preserves_metadata_only_repl_locals() {
+    let compiled = compile_source_for_repl_with_locals(
+        "print(42);",
+        &[ReplLocalBinding {
+            name: "saved".to_string(),
+            mutable: false,
+            schema: Some(TypeSchema::Int),
+            optional: false,
+        }],
+    )
+    .expect("REPL source should compile");
+    let bytes = encode_program(
+        &compiled
+            .compiled
+            .program
+            .with_local_count(compiled.compiled.locals),
+    )
+    .expect("REPL output should encode");
+
+    let program = decode_program(&bytes).expect("embedded decoder should accept REPL VMBC");
+    assert_eq!(program.local_count(), compiled.compiled.locals);
+    assert_eq!(program.local_count(), 1);
+}
+
 #[test]
 fn embedded_decoder_rejects_trailing_bytes() {
     let mut bytes = encoded_scalar_program();