Preserve loop LINE events with INSTRUCTION monitoring (#9231)

Assisted-by: Codex:GPT-6

crates/vm/src/stdlib/sys/monitoring.rs

@@ -446,6 +446,14 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
     // Mirrors CPython's initialize_lines: first determine which positions
     // are line starts, then mark branch/jump targets, then place opcodes.
     if events & MonitoringEvent::Line.mask() != 0 {
+        let base_opcode = |i| {
+            let mut op = code.code.instructions.read_op(i);
+            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)
+        };
         // is_line_start[i] = true if position i should have INSTRUMENTED_LINE
         let mut is_line_start = vec![false; len];
         let line_locations = rustpython_compiler_core::marshal::linetable_to_locations(
@@ -459,16 +467,8 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
 
         // First pass: mark positions where the source line changes
         let mut prev_line: Option<u32> = None;
-        for (i, unit) in code
-            .code
-            .instructions
-            .iter()
-            .enumerate()
-            .take(len)
-            .skip(first_traceable)
-        {
-            let op = unit.op;
-            let base = op.to_base().unwrap_or(op);
+        for (i, is_start) in is_line_start.iter_mut().enumerate().skip(first_traceable) {
+            let base = base_opcode(i);
             if matches!(base, Instruction::ExtendedArg) {
                 continue;
             }
@@ -493,7 +493,7 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
                 let is_new = prev_line != Some(line);
                 prev_line = Some(line);
                 if is_new && line > 0 {
-                    is_line_start[i] = true;
+                    *is_start = true;
                 }
             }
         }
@@ -505,8 +505,8 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
         let mut arg_state = bytecode::OpArgState::default();
         let mut instr_idx = first_traceable;
         for unit in code.code.instructions[first_traceable..len].iter().copied() {
-            let (op, arg) = arg_state.get(unit);
-            let base = op.to_base().unwrap_or(op);
+            let (_, arg) = arg_state.get(unit);
+            let base = base_opcode(instr_idx);
 
             if matches!(base, Instruction::ExtendedArg) || matches!(base, Instruction::Cache) {
                 instr_idx += 1;
@@ -540,8 +540,7 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
                 && !is_line_start[target_idx]
                 && !no_loc_mask.get(target_idx).copied().unwrap_or(false)
             {
-                let target_op = code.code.instructions[target_idx].op;
-                let target_base = target_op.to_base().unwrap_or(target_op);
+                let target_base = base_opcode(target_idx);
                 // Skip synthetic cleanup targets.
                 if matches!(target_base, Instruction::PopIter) {
                     instr_idx += 1;
@@ -563,8 +562,7 @@ pub(crate) fn instrument_code(code: &Py<PyCode>, events: u32) {
                 && !is_line_start[target_idx]
                 && !no_loc_mask.get(target_idx).copied().unwrap_or(false)
             {
-                let target_op = code.code.instructions.read_op(target_idx);
-                let target_base = target_op.to_base().unwrap_or(target_op);
+                let target_base = base_opcode(target_idx);
                 if !matches!(target_base, Instruction::PopIter)
                     && let Some((loc, _)) = line_locations.get(target_idx)
                     && loc.line.get() > 0

extra_tests/snippets/stdlib_sys.py

@@ -194,6 +194,37 @@ finally:
         monitoring.register_callback(2, event, None)
     monitoring.free_tool_id(2)
 
+
+# Combined instruction events must preserve loop-header line events.
+def monitored_loop():
+    total = 0
+    for value in (1, 2):
+        total += value
+    return total
+
+
+lines = []
+monitoring.use_tool_id(2, "combined-lines")
+try:
+    monitoring.register_callback(
+        2,
+        monitoring.events.LINE,
+        lambda code, line: lines.append(line - code.co_firstlineno),
+    )
+    monitoring.register_callback(2, monitoring.events.INSTRUCTION, lambda *args: None)
+    monitoring.set_local_events(
+        2,
+        monitored_loop.__code__,
+        monitoring.events.LINE | monitoring.events.INSTRUCTION,
+    )
+    assert monitored_loop() == 3
+    assert lines == [1, 2, 3, 2, 3, 2, 4]
+finally:
+    monitoring.set_local_events(2, monitored_loop.__code__, 0)
+    for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION):
+        monitoring.register_callback(2, event, None)
+    monitoring.free_tool_id(2)
+
 assert sys.exc_info() == (None, None, None)
 
 try: