gh-91869: Fix tracing of specialized instructions with extended args (alternate)

Include/internal/pycore_opcode.h

@@ -104,7 +104,8 @@ const uint8_t _PyOpcode_Deopt[256] = {
     [DICT_MERGE] = DICT_MERGE,
     [DICT_UPDATE] = DICT_UPDATE,
     [END_ASYNC_FOR] = END_ASYNC_FOR,
-    [EXTENDED_ARG] = EXTENDED_ARG,
+    [EXTENDED_ARG] = EXTENDED_ARG_TRACE,
+    [EXTENDED_ARG_TRACE] = EXTENDED_ARG_TRACE,
     [FORMAT_VALUE] = FORMAT_VALUE,
     [FOR_ITER] = FOR_ITER,
     [GET_AITER] = GET_AITER,
@@ -381,44 +382,44 @@ static const char *const _PyOpcode_OpName[256] = {
     [JUMP_BACKWARD] = "JUMP_BACKWARD",
     [PRECALL_PYFUNC] = "PRECALL_PYFUNC",
     [CALL_FUNCTION_EX] = "CALL_FUNCTION_EX",
-    [RESUME_QUICK] = "RESUME_QUICK",
+    [EXTENDED_ARG_TRACE] = "EXTENDED_ARG_TRACE",
     [EXTENDED_ARG] = "EXTENDED_ARG",
     [LIST_APPEND] = "LIST_APPEND",
     [SET_ADD] = "SET_ADD",
     [MAP_ADD] = "MAP_ADD",
     [LOAD_CLASSDEREF] = "LOAD_CLASSDEREF",
     [COPY_FREE_VARS] = "COPY_FREE_VARS",
-    [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
+    [RESUME_QUICK] = "RESUME_QUICK",
     [RESUME] = "RESUME",
     [MATCH_CLASS] = "MATCH_CLASS",
+    [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
     [STORE_ATTR_INSTANCE_VALUE] = "STORE_ATTR_INSTANCE_VALUE",
-    [STORE_ATTR_SLOT] = "STORE_ATTR_SLOT",
     [FORMAT_VALUE] = "FORMAT_VALUE",
     [BUILD_CONST_KEY_MAP] = "BUILD_CONST_KEY_MAP",
     [BUILD_STRING] = "BUILD_STRING",
+    [STORE_ATTR_SLOT] = "STORE_ATTR_SLOT",
     [STORE_ATTR_WITH_HINT] = "STORE_ATTR_WITH_HINT",
-    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [LOAD_METHOD] = "LOAD_METHOD",
-    [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
+    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [LIST_EXTEND] = "LIST_EXTEND",
     [SET_UPDATE] = "SET_UPDATE",
     [DICT_MERGE] = "DICT_MERGE",
     [DICT_UPDATE] = "DICT_UPDATE",
     [PRECALL] = "PRECALL",
+    [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
     [STORE_SUBSCR_ADAPTIVE] = "STORE_SUBSCR_ADAPTIVE",
     [STORE_SUBSCR_DICT] = "STORE_SUBSCR_DICT",
     [STORE_SUBSCR_LIST_INT] = "STORE_SUBSCR_LIST_INT",
-    [UNPACK_SEQUENCE_ADAPTIVE] = "UNPACK_SEQUENCE_ADAPTIVE",
     [CALL] = "CALL",
     [KW_NAMES] = "KW_NAMES",
     [POP_JUMP_BACKWARD_IF_NOT_NONE] = "POP_JUMP_BACKWARD_IF_NOT_NONE",
     [POP_JUMP_BACKWARD_IF_NONE] = "POP_JUMP_BACKWARD_IF_NONE",
     [POP_JUMP_BACKWARD_IF_FALSE] = "POP_JUMP_BACKWARD_IF_FALSE",
     [POP_JUMP_BACKWARD_IF_TRUE] = "POP_JUMP_BACKWARD_IF_TRUE",
+    [UNPACK_SEQUENCE_ADAPTIVE] = "UNPACK_SEQUENCE_ADAPTIVE",
     [UNPACK_SEQUENCE_LIST] = "UNPACK_SEQUENCE_LIST",
     [UNPACK_SEQUENCE_TUPLE] = "UNPACK_SEQUENCE_TUPLE",
     [UNPACK_SEQUENCE_TWO_TUPLE] = "UNPACK_SEQUENCE_TWO_TUPLE",
-    [180] = "<180>",
     [181] = "<181>",
     [182] = "<182>",
     [183] = "<183>",
@@ -498,7 +499,6 @@ static const char *const _PyOpcode_OpName[256] = {
 #endif
 
 #define EXTRA_CASES \
-    case 180: \
     case 181: \
     case 182: \
     case 183: \

Lib/importlib/_bootstrap_external.py

@@ -416,7 +416,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3493).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3494).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Lib/opcode.py

@@ -169,8 +169,10 @@ def jabs_op(name, op, entries=0):
 
 def_op('CALL_FUNCTION_EX', 142)  # Flags
 
+def_op('EXTENDED_ARG_TRACE', 143)
 def_op('EXTENDED_ARG', 144)
 EXTENDED_ARG = 144
+
 def_op('LIST_APPEND', 145)
 def_op('SET_ADD', 146)
 def_op('MAP_ADD', 147)

Lib/test/test__opcode.py

@@ -18,7 +18,10 @@ def test_stack_effect(self):
         self.assertRaises(ValueError, stack_effect, dis.opmap['BUILD_SLICE'])
         self.assertRaises(ValueError, stack_effect, dis.opmap['POP_TOP'], 0)
         # All defined opcodes
-        for name, code in filter(lambda item: item[0] not in dis.deoptmap, dis.opmap.items()):
+        opnames = {name: op for name, op in dis.opmap.items()
+                   if name not in dis.deoptmap
+                   and name != 'EXTENDED_ARG_TRACE'}
+        for name, code in opnames.items():
             with self.subTest(opname=name):
                 if code < dis.HAVE_ARGUMENT:
                     stack_effect(code)
@@ -31,6 +34,8 @@ def test_stack_effect(self):
             with self.subTest(opcode=code):
                 self.assertRaises(ValueError, stack_effect, code)
                 self.assertRaises(ValueError, stack_effect, code, 0)
+        self.assertRaises(ValueError, stack_effect,
+                          dis.opmap["EXTENDED_ARG_TRACE"])
 
     def test_stack_effect_jump(self):
         JUMP_IF_TRUE_OR_POP = dis.opmap['JUMP_IF_TRUE_OR_POP']
@@ -47,7 +52,10 @@ def test_stack_effect_jump(self):
         self.assertEqual(stack_effect(JUMP_FORWARD, 0, jump=False), 0)
         # All defined opcodes
         has_jump = dis.hasjabs + dis.hasjrel
-        for name, code in filter(lambda item: item[0] not in dis.deoptmap, dis.opmap.items()):
+        opnames = {name: op for name, op in dis.opmap.items()
+                   if name not in dis.deoptmap
+                   and name != 'EXTENDED_ARG_TRACE'}
+        for name, code in opnames.items():
             with self.subTest(opname=name):
                 if code < dis.HAVE_ARGUMENT:
                     common = stack_effect(code)

Lib/test/test_sys_settrace.py

@@ -2432,5 +2432,47 @@ def gen():
         output.append(5)
 
 
+class TestExtendedArgs(unittest.TestCase):
+
+    def setUp(self):
+        self.addCleanup(sys.settrace, sys.gettrace())
+        sys.settrace(None)
+
+    def count_traces(self, func):
+        # warmup
+        for _ in range(20):
+            func()
+
+        counts = {"call": 0, "line": 0, "return": 0}
+        def trace(frame, event, arg):
+            counts[event] += 1
+            return trace
+
+        sys.settrace(trace)
+        func()
+        sys.settrace(None)
+
+        return counts
+
+    def test_trace_unpack_long_sequence(self):
+        ns = {}
+        code = "def f():\n  (" + "y,\n   "*300 + ") = range(300)"
+        exec(code, ns)
+        counts = self.count_traces(ns["f"])
+        self.assertEqual(counts, {'call': 1, 'line': 301, 'return': 1})
+
+    def test_trace_lots_of_globals(self):
+        code = """if 1:
+            def f():
+                return (
+                    {}
+                )
+        """.format("\n+\n".join(f"var{i}\n" for i in range(1000)))
+        ns = {f"var{i}": i for i in range(1000)}
+        exec(code, ns)
+        counts = self.count_traces(ns["f"])
+        self.assertEqual(counts, {'call': 1, 'line': 2000, 'return': 1})
+
+
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/Core and Builtins/2022-04-26-05-05-32.gh-issue-91869.ELbTXl.rst

@@ -0,0 +1 @@
+Fix an issue where specialized opcodes with extended arguments could produce incorrect tracing output or lead to assertion failures.

Python/ceval.c

@@ -5630,6 +5630,15 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             NOTRACE_DISPATCH_SAME_OPARG();
         }
 
+        TARGET(EXTENDED_ARG_TRACE) {
+            assert(oparg);
+            oparg <<= 8;
+            oparg |= _Py_OPARG(*next_instr);
+            opcode = _PyOpcode_Deopt[_Py_OPCODE(*next_instr)];
+            PRE_DISPATCH_GOTO();
+            DISPATCH_GOTO();
+        }
+
         TARGET(CACHE) {
             Py_UNREACHABLE();
         }

Tools/scripts/generate_opcode_h.py

@@ -133,6 +133,7 @@ def main(opcode_py, outfile='Include/opcode.h', internaloutfile='Include/interna
         for basic, family in opcode["_specializations"].items():
             for specialized in family:
                 deoptcodes[specialized] = basic
+        deoptcodes["EXTENDED_ARG"] = "EXTENDED_ARG_TRACE"
         iobj.write("\nconst uint8_t _PyOpcode_Deopt[256] = {\n")
         for opt, deopt in sorted(deoptcodes.items()):
             iobj.write(f"    [{opt}] = {deopt},\n")