code object linetable (#6150)

youknowone · web-flow · commit b666c52df96d · 2025-09-16T21:49:54.000+09:00
* Code.replace

* implement linetable
diff --git a/Lib/test/test_code.py b/Lib/test/test_code.py
@@ -347,8 +347,6 @@ def func(arg):
         newcode = code.replace(co_name="func")  # Should not raise SystemError
         self.assertEqual(code, newcode)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_empty_linetable(self):
         def func():
             pass
@@ -468,8 +466,6 @@ def f():
 
     # co_positions behavior when info is missing.
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     # @requires_debug_ranges()
     def test_co_positions_empty_linetable(self):
         def func():
@@ -480,8 +476,6 @@ def func():
             self.assertIsNone(line)
             self.assertEqual(end_line, new_code.co_firstlineno + 1)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_code_equality(self):
         def f():
             try:
@@ -522,8 +516,6 @@ def test_code_hash_uses_order(self):
         self.assertNotEqual(c, swapped)
         self.assertNotEqual(hash(c), hash(swapped))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_code_hash_uses_bytecode(self):
         c = (lambda x, y: x + y).__code__
         d = (lambda x, y: x * y).__code__
@@ -735,8 +727,6 @@ def check_positions(self, func):
             self.assertEqual(l1, l2)
         self.assertEqual(len(pos1), len(pos2))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_positions(self):
         self.check_positions(parse_location_table)
         self.check_positions(misshappen)
@@ -751,8 +741,6 @@ def check_lines(self, func):
             self.assertEqual(l1, l2)
         self.assertEqual(len(lines1), len(lines2))
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_lines(self):
         self.check_lines(parse_location_table)
         self.check_lines(misshappen)
diff --git a/Lib/test/test_compile.py b/Lib/test/test_compile.py
@@ -885,8 +885,6 @@ def foo(x):
                 self.assertIn('LOAD_ATTR', instructions)
                 self.assertIn('PRECALL', instructions)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_lineno_procedure_call(self):
         def call():
             (
diff --git a/compiler/codegen/src/ir.rs b/compiler/codegen/src/ir.rs
@@ -5,7 +5,7 @@ use rustpython_compiler_core::{
     OneIndexed, SourceLocation,
     bytecode::{
         CodeFlags, CodeObject, CodeUnit, ConstantData, InstrDisplayContext, Instruction, Label,
-        OpArg,
+        OpArg, PyCodeLocationInfoKind,
     },
 };
 
@@ -72,6 +72,7 @@ pub struct InstructionInfo {
     pub target: BlockIdx,
     // pub range: TextRange,
     pub location: SourceLocation,
+    // TODO: end_location for debug ranges
 }
 
 // spell-checker:ignore petgraph
@@ -199,6 +200,9 @@ impl CodeInfo {
             locations.clear()
         }
 
+        // Generate linetable from locations
+        let linetable = generate_linetable(&locations, first_line_number.get() as i32);
+
         Ok(CodeObject {
             flags,
             posonlyarg_count,
@@ -218,6 +222,8 @@ impl CodeInfo {
             cellvars: cellvar_cache.into_iter().collect(),
             freevars: freevar_cache.into_iter().collect(),
             cell2arg,
+            linetable,
+            exceptiontable: Box::new([]), // TODO: Generate actual exception table
         })
     }
 
@@ -388,3 +394,134 @@ fn iter_blocks(blocks: &[Block]) -> impl Iterator<Item = (BlockIdx, &Block)> + '
         Some((idx, b))
     })
 }
+
+/// Generate CPython 3.11+ format linetable from source locations
+fn generate_linetable(locations: &[SourceLocation], first_line: i32) -> Box<[u8]> {
+    if locations.is_empty() {
+        return Box::new([]);
+    }
+
+    let mut linetable = Vec::new();
+    // Initialize prev_line to first_line
+    // The first entry's delta is relative to co_firstlineno
+    let mut prev_line = first_line;
+    let mut i = 0;
+
+    while i < locations.len() {
+        let loc = &locations[i];
+
+        // Count consecutive instructions with the same location
+        let mut length = 1;
+        while i + length < locations.len() && locations[i + length] == locations[i] {
+            length += 1;
+        }
+
+        // Process in chunks of up to 8 instructions
+        while length > 0 {
+            let entry_length = length.min(8);
+
+            // Get line and column information
+            // SourceLocation always has row and column (both are OneIndexed)
+            let line = loc.row.get() as i32;
+            let col = (loc.column.get() as i32) - 1; // Convert 1-based to 0-based
+
+            let line_delta = line - prev_line;
+
+            // Choose the appropriate encoding based on line delta and column info
+            // Note: SourceLocation always has valid column, so we never get NO_COLUMNS case
+            if line_delta == 0 {
+                let end_col = col; // Use same column for end (no range info available)
+
+                if col < 80 && end_col - col < 16 && end_col >= col {
+                    // Short form (codes 0-9) for common cases
+                    let code = (col / 8).min(9) as u8; // Short0 to Short9
+                    linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
+                    let col_byte = (((col % 8) as u8) << 4) | ((end_col - col) as u8 & 0xf);
+                    linetable.push(col_byte);
+                } else if col < 128 && end_col < 128 {
+                    // One-line form (code 10) for same line
+                    linetable.push(
+                        0x80 | ((PyCodeLocationInfoKind::OneLine0 as u8) << 3)
+                            | ((entry_length - 1) as u8),
+                    );
+                    linetable.push(col as u8);
+                    linetable.push(end_col as u8);
+                } else {
+                    // Long form for columns >= 128
+                    linetable.push(
+                        0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
+                            | ((entry_length - 1) as u8),
+                    );
+                    write_signed_varint(&mut linetable, 0); // line_delta = 0
+                    write_varint(&mut linetable, 0); // end_line delta = 0
+                    write_varint(&mut linetable, (col as u32) + 1); // column + 1 for encoding
+                    write_varint(&mut linetable, (end_col as u32) + 1); // end_col + 1
+                }
+            } else if line_delta > 0 && line_delta < 3
+            /* && column.is_some() */
+            {
+                // One-line form (codes 11-12) for line deltas 1-2
+                let end_col = col; // Use same column for end
+
+                if col < 128 && end_col < 128 {
+                    let code = (PyCodeLocationInfoKind::OneLine0 as u8) + (line_delta as u8); // 11 for delta=1, 12 for delta=2
+                    linetable.push(0x80 | (code << 3) | ((entry_length - 1) as u8));
+                    linetable.push(col as u8);
+                    linetable.push(end_col as u8);
+                } else {
+                    // Long form for columns >= 128 or negative line delta
+                    linetable.push(
+                        0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3)
+                            | ((entry_length - 1) as u8),
+                    );
+                    write_signed_varint(&mut linetable, line_delta);
+                    write_varint(&mut linetable, 0); // end_line delta = 0
+                    write_varint(&mut linetable, (col as u32) + 1); // column + 1 for encoding
+                    write_varint(&mut linetable, (end_col as u32) + 1); // end_col + 1
+                }
+            } else {
+                // Long form (code 14) for all other cases
+                // This handles: line_delta < 0, line_delta >= 3, or columns >= 128
+                let end_col = col; // Use same column for end
+                linetable.push(
+                    0x80 | ((PyCodeLocationInfoKind::Long as u8) << 3) | ((entry_length - 1) as u8),
+                );
+                write_signed_varint(&mut linetable, line_delta);
+                write_varint(&mut linetable, 0); // end_line delta = 0
+                write_varint(&mut linetable, (col as u32) + 1); // column + 1 for encoding
+                write_varint(&mut linetable, (end_col as u32) + 1); // end_col + 1
+            }
+
+            prev_line = line;
+            length -= entry_length;
+            i += entry_length;
+        }
+    }
+
+    linetable.into_boxed_slice()
+}
+
+/// Write a variable-length unsigned integer (6-bit chunks)
+/// Returns the number of bytes written
+fn write_varint(buf: &mut Vec<u8>, mut val: u32) -> usize {
+    let start_len = buf.len();
+    while val >= 64 {
+        buf.push(0x40 | (val & 0x3f) as u8);
+        val >>= 6;
+    }
+    buf.push(val as u8);
+    buf.len() - start_len
+}
+
+/// Write a variable-length signed integer
+/// Returns the number of bytes written
+fn write_signed_varint(buf: &mut Vec<u8>, val: i32) -> usize {
+    let uval = if val < 0 {
+        // (unsigned int)(-val) has an undefined behavior for INT_MIN
+        // So we use (0 - val as u32) to handle it correctly
+        ((0u32.wrapping_sub(val as u32)) << 1) | 1
+    } else {
+        (val as u32) << 1
+    };
+    write_varint(buf, uval)
+}
diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
@@ -33,6 +33,75 @@ pub enum ResumeType {
     AfterAwait = 3,
 }
 
+/// CPython 3.11+ linetable location info codes
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[repr(u8)]
+pub enum PyCodeLocationInfoKind {
+    // Short forms are 0 to 9
+    Short0 = 0,
+    Short1 = 1,
+    Short2 = 2,
+    Short3 = 3,
+    Short4 = 4,
+    Short5 = 5,
+    Short6 = 6,
+    Short7 = 7,
+    Short8 = 8,
+    Short9 = 9,
+    // One line forms are 10 to 12
+    OneLine0 = 10,
+    OneLine1 = 11,
+    OneLine2 = 12,
+    NoColumns = 13,
+    Long = 14,
+    None = 15,
+}
+
+impl PyCodeLocationInfoKind {
+    pub fn from_code(code: u8) -> Option<Self> {
+        match code {
+            0 => Some(Self::Short0),
+            1 => Some(Self::Short1),
+            2 => Some(Self::Short2),
+            3 => Some(Self::Short3),
+            4 => Some(Self::Short4),
+            5 => Some(Self::Short5),
+            6 => Some(Self::Short6),
+            7 => Some(Self::Short7),
+            8 => Some(Self::Short8),
+            9 => Some(Self::Short9),
+            10 => Some(Self::OneLine0),
+            11 => Some(Self::OneLine1),
+            12 => Some(Self::OneLine2),
+            13 => Some(Self::NoColumns),
+            14 => Some(Self::Long),
+            15 => Some(Self::None),
+            _ => Option::None,
+        }
+    }
+
+    pub fn is_short(&self) -> bool {
+        (*self as u8) <= 9
+    }
+
+    pub fn short_column_group(&self) -> Option<u8> {
+        if self.is_short() {
+            Some(*self as u8)
+        } else {
+            Option::None
+        }
+    }
+
+    pub fn one_line_delta(&self) -> Option<i32> {
+        match self {
+            Self::OneLine0 => Some(0),
+            Self::OneLine1 => Some(1),
+            Self::OneLine2 => Some(2),
+            _ => Option::None,
+        }
+    }
+}
+
 pub trait Constant: Sized {
     type Name: AsRef<str>;
 
@@ -146,6 +215,10 @@ pub struct CodeObject<C: Constant = ConstantData> {
     pub varnames: Box<[C::Name]>,
     pub cellvars: Box<[C::Name]>,
     pub freevars: Box<[C::Name]>,
+    pub linetable: Box<[u8]>,
+    // Line number table (CPython 3.11+ format)
+    pub exceptiontable: Box<[u8]>,
+    // Exception handling table
 }
 
 bitflags! {
@@ -1202,6 +1275,8 @@ impl<C: Constant> CodeObject<C> {
             first_line_number: self.first_line_number,
             max_stackdepth: self.max_stackdepth,
             cell2arg: self.cell2arg,
+            linetable: self.linetable,
+            exceptiontable: self.exceptiontable,
         }
     }
 
@@ -1232,6 +1307,8 @@ impl<C: Constant> CodeObject<C> {
             first_line_number: self.first_line_number,
             max_stackdepth: self.max_stackdepth,
             cell2arg: self.cell2arg.clone(),
+            linetable: self.linetable.clone(),
+            exceptiontable: self.exceptiontable.clone(),
         }
     }
 }
diff --git a/compiler/core/src/marshal.rs b/compiler/core/src/marshal.rs
@@ -251,6 +251,16 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
     let cellvars = read_names()?;
     let freevars = read_names()?;
 
+    // Read linetable and exceptiontable
+    let linetable_len = rdr.read_u32()?;
+    let linetable = rdr.read_slice(linetable_len)?.to_vec().into_boxed_slice();
+
+    let exceptiontable_len = rdr.read_u32()?;
+    let exceptiontable = rdr
+        .read_slice(exceptiontable_len)?
+        .to_vec()
+        .into_boxed_slice();
+
     Ok(CodeObject {
         instructions,
         locations,
@@ -269,6 +279,8 @@ pub fn deserialize_code<R: Read, Bag: ConstantBag>(
         varnames,
         cellvars,
         freevars,
+        linetable,
+        exceptiontable,
     })
 }
 
@@ -684,4 +696,8 @@ pub fn serialize_code<W: Write, C: Constant>(buf: &mut W, code: &CodeObject<C>)
     write_names(&code.varnames);
     write_names(&code.cellvars);
     write_names(&code.freevars);
+
+    // Serialize linetable and exceptiontable
+    write_vec(buf, &code.linetable);
+    write_vec(buf, &code.exceptiontable);
 }
diff --git a/vm/src/builtins/code.rs b/vm/src/builtins/code.rs

Original file line number	Diff line number	Diff line change
`@@ -885,8 +885,6 @@ def foo(x):`
`885`	`885`	`self.assertIn('LOAD_ATTR', instructions)`
`886`	`886`	`self.assertIn('PRECALL', instructions)`
`887`	`887`
`888`		`- # TODO: RUSTPYTHON`
`889`		`- @unittest.expectedFailure`
`890`	`888`	`def test_lineno_procedure_call(self):`
`891`	`889`	`def call():`
`892`	`890`	`(`