Preserve loop stack state for monitored line jumps (#9232)

Assisted-by: Codex:GPT-6

crates/vm/src/builtins/frame.rs

@@ -155,9 +155,24 @@ pub(crate) mod stack_analysis {
     }
 
     /// Analyze bytecode and compute the stack state at each instruction index.
-    pub(crate) fn mark_stacks<C: Constant>(code: &bytecode::CodeObject<C>) -> Vec<i64> {
+    pub(crate) fn mark_stacks(code: &PyCode) -> Vec<i64> {
         let instructions = &*code.instructions;
         let len = instructions.len();
+        let monitoring_data = code.monitoring_data.lock();
+        let base_opcode = |i| {
+            let mut op = code.instructions.read_op(i);
+            if let Some(data) = monitoring_data.as_ref() {
+                if matches!(op, Instruction::InstrumentedLine) {
+                    op = Instruction::try_from(data.line_opcodes[i])
+                        .expect("INSTRUMENTED_LINE without original opcode");
+                }
+                if matches!(op, Instruction::InstrumentedInstruction) {
+                    op = Instruction::try_from(data.per_instruction_opcodes[i])
+                        .expect("INSTRUMENTED_INSTRUCTION without original opcode");
+                }
+            }
+            op.to_base().unwrap_or(op).deoptimize()
+        };
 
         let mut stacks = vec![UNINITIALIZED; len + 1];
         stacks[0] = EMPTY_STACK;
@@ -169,7 +184,7 @@ pub(crate) mod stack_analysis {
             let mut i = 0;
             while i < len {
                 let mut next_stack = stacks[i];
-                let mut opcode = instructions[i].op;
+                let mut opcode = base_opcode(i);
                 let mut oparg: u32 = 0;
 
                 // Accumulate EXTENDED_ARG prefixes
@@ -180,16 +195,13 @@ pub(crate) mod stack_analysis {
                         break;
                     }
                     stacks[i] = next_stack;
-                    opcode = instructions[i].op;
+                    opcode = base_opcode(i);
                 }
                 if i >= len {
                     break;
                 }
                 oparg = (oparg << 8) | u32::from(u8::from(instructions[i].arg));
 
-                // De-instrument and de-specialize: get the underlying base instruction
-                let opcode = opcode.to_base().unwrap_or(opcode).deoptimize();
-
                 let caches = opcode.cache_entries();
                 let next_i = i + 1 + caches;
 
@@ -784,7 +796,7 @@ impl Py<FrameObject> {
             )));
         }
 
-        let stacks = mark_stacks(code);
+        let stacks = mark_stacks(py_code);
         let len = self.iframe().code().instructions.len();
 
         // lasti points past the current instruction (already incremented).

extra_tests/snippets/stdlib_sys.py

@@ -173,6 +173,11 @@ def jump_out_of_large_loop(code, line):
             sys._getframe(1).f_lineno = 3
 
 
+def jump_after_large_loop(code, line):
+    if code is large_loop.__code__ and line == 5:
+        sys._getframe(1).f_lineno = 106
+
+
 monitoring.use_tool_id(2, "extended-arg")
 try:
     for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION):
@@ -188,6 +193,8 @@ try:
     monitoring.set_local_events(2, large_loop.__code__, monitoring.events.LINE)
     assert large_loop() == "done"
     assert loop_hits == 2
+    monitoring.register_callback(2, monitoring.events.LINE, jump_after_large_loop)
+    assert large_loop() == 0
 finally:
     monitoring.set_local_events(2, large_loop.__code__, 0)
     for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION):