Fix PyCode Constructor

Author: youknowone

Lib/test/test_code.py

@@ -222,8 +222,6 @@ class List(list):
         obj = List([1, 2, 3])
         self.assertEqual(obj[0], "Foreign getitem: 1")
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_constructor(self):
         def func(): pass
         co = func.__code__
@@ -255,8 +253,6 @@ def test_qualname(self):
             CodeTest.test_qualname.__qualname__
         )
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_replace(self):
         def func():
             x = 1
@@ -297,8 +293,6 @@ def func2():
         self.assertEqual(new_code.co_varnames, code2.co_varnames)
         self.assertEqual(new_code.co_nlocals, code2.co_nlocals)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_nlocals_mismatch(self):
         def func():
             x = 1
@@ -362,8 +356,6 @@ def func():
         with self.assertWarns(DeprecationWarning):
             func.__code__.co_lnotab
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_invalid_bytecode(self):
         def foo():
             pass

compiler/codegen/src/ir.rs

@@ -4,7 +4,7 @@ use crate::{IndexMap, IndexSet, error::InternalError};
 use rustpython_compiler_core::{
     OneIndexed, SourceLocation,
     bytecode::{
-        CodeFlags, CodeObject, CodeUnit, ConstantData, InstrDisplayContext, Instruction, Label,
+        CodeFlags, CodeObject, CodeUnit, ConstantData, InstrDisplayContext, Instruction, Instructions, Label,
         OpArg, PyCodeLocationInfoKind,
     },
 };
@@ -214,7 +214,7 @@ impl CodeInfo {
             qualname: qualname.unwrap_or(obj_name),
 
             max_stackdepth,
-            instructions: instructions.into_boxed_slice(),
+            instructions: Instructions::Parsed(instructions.into_boxed_slice()),
             locations: locations.into_boxed_slice(),
             constants: constants.into_iter().collect(),
             names: name_cache.into_iter().collect(),

compiler/core/src/bytecode.rs

@@ -191,11 +191,55 @@ impl ConstantBag for BasicBag {
     }
 }
 
+/// Container for bytecode instructions - either parsed or raw
+#[derive(Clone)]
+pub enum Instructions {
+    /// Parsed bytecode instructions
+    Parsed(Box<[CodeUnit]>),
+    /// Raw bytecode bytes (may contain invalid opcodes)
+    Raw(Box<[u8]>),
+}
+
+impl Instructions {
+    /// Get as parsed instructions, panics if raw
+    pub fn as_parsed(&self) -> &[CodeUnit] {
+        match self {
+            Instructions::Parsed(units) => units,
+            Instructions::Raw(_) => panic!("Cannot get parsed instructions from raw bytecode"),
+        }
+    }
+
+    /// Get as raw bytes (borrowed)
+    pub fn as_bytes(&self) -> &[u8] {
+        match self {
+            Instructions::Parsed(units) => {
+                // SAFETY: it's ok to transmute CodeUnit to [u8; 2]
+                let (_, instructions_bytes, _) = unsafe { units.align_to() };
+                instructions_bytes
+            }
+            Instructions::Raw(bytes) => bytes,
+        }
+    }
+
+    /// Check if this contains raw bytecode
+    pub fn is_raw(&self) -> bool {
+        matches!(self, Instructions::Raw(_))
+    }
+
+    /// Get the number of instructions (or byte pairs for raw)
+    pub fn len(&self) -> usize {
+        match self {
+            Instructions::Parsed(units) => units.len(),
+            Instructions::Raw(bytes) => bytes.len() / 2,
+        }
+    }
+}
+
 /// Primary container of a single code object. Each python function has
 /// a code object. Also a module has a code object.
 #[derive(Clone)]
 pub struct CodeObject<C: Constant = ConstantData> {
-    pub instructions: Box<[CodeUnit]>,
+    pub instructions: Instructions,
     pub locations: Box<[SourceLocation]>,
     pub flags: CodeFlags,
     /// Number of positional-only arguments
@@ -1179,11 +1223,18 @@ impl<C: Constant> CodeObject<C> {
     /// Return the labels targeted by the instructions of this CodeObject
     pub fn label_targets(&self) -> BTreeSet<Label> {
         let mut label_targets = BTreeSet::new();
-        let mut arg_state = OpArgState::default();
-        for instruction in &*self.instructions {
-            let (instruction, arg) = arg_state.get(*instruction);
-            if let Some(l) = instruction.label_arg() {
-                label_targets.insert(l.get(arg));
+        match &self.instructions {
+            Instructions::Parsed(units) => {
+                let mut arg_state = OpArgState::default();
+                for instruction in &**units {
+                    let (instruction, arg) = arg_state.get(*instruction);
+                    if let Some(l) = instruction.label_arg() {
+                        label_targets.insert(l.get(arg));
+                    }
+                }
+            }
+            Instructions::Raw(_) => {
+                // Raw bytecode doesn't have parsed labels
             }
         }
         label_targets
@@ -1195,46 +1246,54 @@ impl<C: Constant> CodeObject<C> {
         expand_code_objects: bool,
         level: usize,
     ) -> fmt::Result {
-        let label_targets = self.label_targets();
-        let line_digits = (3).max(self.locations.last().unwrap().row.to_string().len());
-        let offset_digits = (4).max(self.instructions.len().to_string().len());
-        let mut last_line = OneIndexed::MAX;
-        let mut arg_state = OpArgState::default();
-        for (offset, &instruction) in self.instructions.iter().enumerate() {
-            let (instruction, arg) = arg_state.get(instruction);
-            // optional line number
-            let line = self.locations[offset].row;
-            if line != last_line {
-                if last_line != OneIndexed::MAX {
-                    writeln!(f)?;
-                }
-                last_line = line;
-                write!(f, "{line:line_digits$}")?;
-            } else {
-                for _ in 0..line_digits {
-                    write!(f, " ")?;
-                }
+        match &self.instructions {
+            Instructions::Raw(bytes) => {
+                writeln!(f, "Raw bytecode ({} bytes): {:?}", bytes.len(), bytes)?;
+                return Ok(());
             }
-            write!(f, " ")?;
+            Instructions::Parsed(units) => {
+                let label_targets = self.label_targets();
+                let line_digits = (3).max(self.locations.last().unwrap().row.to_string().len());
+                let offset_digits = (4).max(units.len().to_string().len());
+                let mut last_line = OneIndexed::MAX;
+                let mut arg_state = OpArgState::default();
+                for (offset, &instruction) in units.iter().enumerate() {
+                    let (instruction, arg) = arg_state.get(instruction);
+                    // optional line number
+                    let line = self.locations[offset].row;
+                    if line != last_line {
+                        if last_line != OneIndexed::MAX {
+                            writeln!(f)?;
+                        }
+                        last_line = line;
+                        write!(f, "{line:line_digits$}")?;
+                    } else {
+                        for _ in 0..line_digits {
+                            write!(f, " ")?;
+                        }
+                    }
+                    write!(f, " ")?;
 
-            // level indent
-            for _ in 0..level {
-                write!(f, "    ")?;
-            }
+                    // level indent
+                    for _ in 0..level {
+                        write!(f, "    ")?;
+                    }
 
-            // arrow and offset
-            let arrow = if label_targets.contains(&Label(offset as u32)) {
-                ">>"
-            } else {
-                "  "
-            };
-            write!(f, "{arrow} {offset:offset_digits$} ")?;
+                    // arrow and offset
+                    let arrow = if label_targets.contains(&Label(offset as u32)) {
+                        ">>"
+                    } else {
+                        "  "
+                    };
+                    write!(f, "{arrow} {offset:offset_digits$} ")?;
 
-            // instruction
-            instruction.fmt_dis(arg, f, self, expand_code_objects, 21, level)?;
-            writeln!(f)?;
+                    // instruction
+                    instruction.fmt_dis(arg, f, self, expand_code_objects, 21, level)?;
+                    writeln!(f)?;
+                }
+                Ok(())
+            }
         }
-        Ok(())
     }
 
     /// Recursively display this CodeObject

compiler/core/src/marshal.rs

@@ -184,15 +184,24 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
     bag: Bag,
 ) -> Result<CodeObject<Bag::Constant>> {
     let len = rdr.read_u32()?;
-    let instructions = rdr.read_slice(len * 2)?;
-    let instructions = instructions
+    let instructions_bytes = rdr.read_slice(len * 2)?;
+
+    // Try to parse as valid instructions
+    let instructions = match instructions_bytes
         .chunks_exact(2)
         .map(|cu| {
             let op = Instruction::try_from(cu[0])?;
             let arg = OpArgByte(cu[1]);
             Ok(CodeUnit { op, arg })
         })
-        .collect::<Result<Box<[CodeUnit]>>>()?;
+        .collect::<Result<Box<[CodeUnit]>>>()
+    {
+        Ok(parsed) => Instructions::Parsed(parsed),
+        Err(_) => {
+            // Contains invalid opcodes, store as raw
+            Instructions::Raw(instructions_bytes.to_vec().into_boxed_slice())
+        }
+    };
 
     let len = rdr.read_u32()?;
     let locations = (0..len)
@@ -649,9 +658,9 @@ pub fn serialize_value<W: Write, D: Dumpable>(
 }
 
 pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>) {
-    write_len(buf, code.instructions.len());
-    // SAFETY: it's ok to transmute CodeUnit to [u8; 2]
-    let (_, instructions_bytes, _) = unsafe { code.instructions.align_to() };
+    // Write instructions as bytes
+    let instructions_bytes = code.instructions.as_bytes();
+    write_len(buf, instructions_bytes.len() / 2);
     buf.write_slice(instructions_bytes);
 
     write_len(buf, code.locations.len());

stdlib/src/dis.rs

@@ -38,7 +38,7 @@ mod decl {
 
     #[pyfunction]
     fn disassemble(co: PyRef<PyCode>) -> PyResult<()> {
-        print!("{}", &co.code);
+        print!("{}", &**co);
         Ok(())
     }