Preserve extended opcode arguments under sys.monitoring (#9203)

* Preserve extended opcode arguments under sys.monitoring

Assisted-by: Codex:GPT-6

* Restart monitored execution after line callback jumps

Assisted-by: Codex:GPT-6

* Keep monitoring argument state out of the dispatch loop

Assisted-by: Codex:GPT-6

crates/vm/src/frame.rs

@@ -8237,20 +8237,30 @@ impl ExecutingFrame<'_> {
             }
             // All INSTRUMENTED_* opcodes delegate to a cold function to keep
             // the hot instruction loop free of monitoring overhead.
-            _ => self.execute_instrumented(instruction, arg, vm),
+            _ => {
+                let (result, do_extend_arg) = self.execute_instrumented(instruction, arg, vm);
+                *extend_arg = do_extend_arg;
+                result
+            }
         }
     }
 
     /// Handle all INSTRUMENTED_* opcodes. This function is cold — it only
     /// runs when sys.monitoring has rewritten the bytecode.
     #[cold]
+    #[inline(never)]
     fn execute_instrumented(
         &mut self,
         instruction: Instruction,
         arg: bytecode::OpArg,
         vm: &VirtualMachine,
-    ) -> FrameResult {
-        self.execute_instrumented_from(instruction, arg, false, vm)
+    ) -> (FrameResult, bool) {
+        // Keep the mutable flag inside this cold function so its address does
+        // not escape the dispatch loop and force a stack access for every opcode.
+        let mut do_extend_arg = false;
+        let result =
+            self.execute_instrumented_from(instruction, arg, false, &mut do_extend_arg, vm);
+        (result, do_extend_arg)
     }
 
     #[cold]
@@ -8259,6 +8269,7 @@ impl ExecutingFrame<'_> {
         instruction: Instruction,
         arg: bytecode::OpArg,
         wrapper_dispatch: bool,
+        do_extend_arg: &mut bool,
         vm: &VirtualMachine,
     ) -> FrameResult {
         debug_assert!(
@@ -8625,6 +8636,13 @@ impl ExecutingFrame<'_> {
                     if line != self.prev_line.get() && line > 0 {
                         self.prev_line.set(line);
                         monitoring::fire_line(vm, self.code, offset, line)?;
+                        if self.lasti() != idx as u32 + 1 {
+                            self.apply_pending_lineno_jump(vm);
+                            // Restart at the callback's target without executing the
+                            // old opcode or carrying its EXTENDED_ARG prefix there.
+                            *do_extend_arg = false;
+                            return Ok(None);
+                        }
                     }
                 }
 
@@ -8654,10 +8672,15 @@ impl ExecutingFrame<'_> {
                     .expect("invalid opcode in side-table chain");
                 let lasti_before_dispatch = self.lasti();
                 let result = if original_op.to_base().is_some() {
-                    self.execute_instrumented_from(original_op, arg, wrapper_dispatch, vm)
+                    self.execute_instrumented_from(
+                        original_op,
+                        arg,
+                        wrapper_dispatch,
+                        do_extend_arg,
+                        vm,
+                    )
                 } else {
-                    let mut do_extend_arg = false;
-                    self.execute_instruction(original_op, arg, &mut do_extend_arg, vm)
+                    self.execute_instruction(original_op, arg, do_extend_arg, vm)
                 };
                 let orig_caches = original_op.to_base().unwrap_or(original_op).cache_entries();
                 if orig_caches > 0 && self.lasti() == lasti_before_dispatch {
@@ -8694,10 +8717,15 @@ impl ExecutingFrame<'_> {
                     .expect("invalid opcode in instruction side-table");
                 let lasti_before_dispatch = self.lasti();
                 let result = if original_op.to_base().is_some() {
-                    self.execute_instrumented_from(original_op, arg, wrapper_dispatch, vm)
+                    self.execute_instrumented_from(
+                        original_op,
+                        arg,
+                        wrapper_dispatch,
+                        do_extend_arg,
+                        vm,
+                    )
                 } else {
-                    let mut do_extend_arg = false;
-                    self.execute_instruction(original_op, arg, &mut do_extend_arg, vm)
+                    self.execute_instruction(original_op, arg, do_extend_arg, vm)
                 };
                 let orig_caches = original_op.to_base().unwrap_or(original_op).cache_entries();
                 if orig_caches > 0 && self.lasti() == lasti_before_dispatch {

extra_tests/snippets/stdlib_sys.py

@@ -155,6 +155,45 @@ finally:
     monitoring.free_tool_id(2)
 assert records == ["RAISE", "PY_UNWIND"], records
 
+# A loop body large enough to require EXTENDED_ARG on its jump targets.
+namespace = {}
+exec(
+    "def large_loop(skip=False):\n if skip:\n  return 'done'\n"
+    " total = 0\n for x in (1, 2):\n" + "  total += x\n" * 100 + " return total\n",
+    namespace,
+)
+large_loop = namespace["large_loop"]
+
+
+def jump_out_of_large_loop(code, line):
+    global loop_hits
+    if code is large_loop.__code__ and line == 5:
+        loop_hits += 1
+        if loop_hits == 2:
+            sys._getframe(1).f_lineno = 3
+
+
+monitoring.use_tool_id(2, "extended-arg")
+try:
+    for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION):
+        monitoring.register_callback(2, event, lambda *args: None)
+    for events in (
+        monitoring.events.LINE,
+        monitoring.events.LINE | monitoring.events.INSTRUCTION,
+    ):
+        monitoring.set_local_events(2, large_loop.__code__, events)
+        assert large_loop() == 300
+    loop_hits = 0
+    monitoring.register_callback(2, monitoring.events.LINE, jump_out_of_large_loop)
+    monitoring.set_local_events(2, large_loop.__code__, monitoring.events.LINE)
+    assert large_loop() == "done"
+    assert loop_hits == 2
+finally:
+    monitoring.set_local_events(2, large_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: