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

3 files changed, 30 insertions, 20 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 6e7a248..1d2c643 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -5627,6 +5627,15 @@ handle_eval_breaker:
             assert(oparg);
             oparg <<= 8;
             oparg |= _Py_OPARG(*next_instr);
+            opcode = _PyOpcode_Deopt[_Py_OPCODE(*next_instr)];
+            PRE_DISPATCH_GOTO();
+            DISPATCH_GOTO();
+        }
+
+        TARGET(EXTENDED_ARG_QUICK) {
+            assert(oparg);
+            oparg <<= 8;
+            oparg |= _Py_OPARG(*next_instr);
             NOTRACE_DISPATCH_SAME_OPARG();
         }
 
diff --git a/Python/compile.c b/Python/compile.c
index 42b011c..10d6307 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -211,13 +211,13 @@ write_instr(_Py_CODEUNIT *codestr, struct instr *instruction, int ilen)
     int caches = _PyOpcode_Caches[opcode];
     switch (ilen - caches) {
         case 4:
-            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG, (oparg >> 24) & 0xFF);
+            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG_QUICK, (oparg >> 24) & 0xFF);
             /* fall through */
         case 3:
-            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG, (oparg >> 16) & 0xFF);
+            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG_QUICK, (oparg >> 16) & 0xFF);
             /* fall through */
         case 2:
-            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG, (oparg >> 8) & 0xFF);
+            *codestr++ = _Py_MAKECODEUNIT(EXTENDED_ARG_QUICK, (oparg >> 8) & 0xFF);
             /* fall through */
         case 1:
             *codestr++ = _Py_MAKECODEUNIT(opcode, oparg & 0xFF);
@@ -8254,6 +8254,7 @@ fix_cell_offsets(struct compiler *c, basicblock *entryblock, int *fixedmap)
             struct instr *inst = &b->b_instr[i];
             // This is called before extended args are generated.
             assert(inst->i_opcode != EXTENDED_ARG);
+            assert(inst->i_opcode != EXTENDED_ARG_QUICK);
             int oldoffset = inst->i_oparg;
             switch(inst->i_opcode) {
                 case MAKE_CELL:
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 5268c4f..d37c132 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -33,10 +33,11 @@ static void *opcode_targets[256] = {
     &&TARGET_MATCH_MAPPING,
     &&TARGET_MATCH_SEQUENCE,
     &&TARGET_MATCH_KEYS,
-    &&TARGET_JUMP_BACKWARD_QUICK,
+    &&TARGET_EXTENDED_ARG_QUICK,
     &&TARGET_PUSH_EXC_INFO,
     &&TARGET_CHECK_EXC_MATCH,
     &&TARGET_CHECK_EG_MATCH,
+    &&TARGET_JUMP_BACKWARD_QUICK,
     &&TARGET_LOAD_ATTR_ADAPTIVE,
     &&TARGET_LOAD_ATTR_INSTANCE_VALUE,
     &&TARGET_LOAD_ATTR_MODULE,
@@ -47,40 +48,39 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_LOAD_GLOBAL_ADAPTIVE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
-    &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_WITH_EXCEPT_START,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
     &&TARGET_BEFORE_ASYNC_WITH,
     &&TARGET_BEFORE_WITH,
     &&TARGET_END_ASYNC_FOR,
+    &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_LOAD_METHOD_ADAPTIVE,
     &&TARGET_LOAD_METHOD_CLASS,
     &&TARGET_LOAD_METHOD_MODULE,
     &&TARGET_LOAD_METHOD_NO_DICT,
-    &&TARGET_LOAD_METHOD_WITH_DICT,
     &&TARGET_STORE_SUBSCR,
     &&TARGET_DELETE_SUBSCR,
+    &&TARGET_LOAD_METHOD_WITH_DICT,
     &&TARGET_LOAD_METHOD_WITH_VALUES,
     &&TARGET_PRECALL_ADAPTIVE,
     &&TARGET_PRECALL_BOUND_METHOD,
     &&TARGET_PRECALL_BUILTIN_CLASS,
     &&TARGET_PRECALL_BUILTIN_FAST_WITH_KEYWORDS,
-    &&TARGET_PRECALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS,
     &&TARGET_GET_ITER,
     &&TARGET_GET_YIELD_FROM_ITER,
     &&TARGET_PRINT_EXPR,
     &&TARGET_LOAD_BUILD_CLASS,
+    &&TARGET_PRECALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS,
     &&TARGET_PRECALL_NO_KW_BUILTIN_FAST,
-    &&TARGET_PRECALL_NO_KW_BUILTIN_O,
     &&TARGET_LOAD_ASSERTION_ERROR,
     &&TARGET_RETURN_GENERATOR,
+    &&TARGET_PRECALL_NO_KW_BUILTIN_O,
     &&TARGET_PRECALL_NO_KW_ISINSTANCE,
     &&TARGET_PRECALL_NO_KW_LEN,
     &&TARGET_PRECALL_NO_KW_LIST_APPEND,
     &&TARGET_PRECALL_NO_KW_METHOD_DESCRIPTOR_FAST,
     &&TARGET_PRECALL_NO_KW_METHOD_DESCRIPTOR_NOARGS,
-    &&TARGET_PRECALL_NO_KW_METHOD_DESCRIPTOR_O,
     &&TARGET_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
@@ -112,7 +112,7 @@ static void *opcode_targets[256] = {
     &&TARGET_JUMP_FORWARD,
     &&TARGET_JUMP_IF_FALSE_OR_POP,
     &&TARGET_JUMP_IF_TRUE_OR_POP,
-    &&TARGET_PRECALL_NO_KW_STR_1,
+    &&TARGET_PRECALL_NO_KW_METHOD_DESCRIPTOR_O,
     &&TARGET_POP_JUMP_FORWARD_IF_FALSE,
     &&TARGET_POP_JUMP_FORWARD_IF_TRUE,
     &&TARGET_LOAD_GLOBAL,
@@ -120,13 +120,13 @@ static void *opcode_targets[256] = {
     &&TARGET_CONTAINS_OP,
     &&TARGET_RERAISE,
     &&TARGET_COPY,
-    &&TARGET_PRECALL_NO_KW_TUPLE_1,
+    &&TARGET_PRECALL_NO_KW_STR_1,
     &&TARGET_BINARY_OP,
     &&TARGET_SEND,
     &&TARGET_LOAD_FAST,
     &&TARGET_STORE_FAST,
     &&TARGET_DELETE_FAST,
-    &&TARGET_PRECALL_NO_KW_TYPE_1,
+    &&TARGET_PRECALL_NO_KW_TUPLE_1,
     &&TARGET_POP_JUMP_FORWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_FORWARD_IF_NONE,
     &&TARGET_RAISE_VARARGS,
@@ -140,42 +140,43 @@ static void *opcode_targets[256] = {
     &&TARGET_STORE_DEREF,
     &&TARGET_DELETE_DEREF,
     &&TARGET_JUMP_BACKWARD,
-    &&TARGET_PRECALL_PYFUNC,
+    &&TARGET_PRECALL_NO_KW_TYPE_1,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_RESUME_QUICK,
+    &&TARGET_PRECALL_PYFUNC,
     &&TARGET_EXTENDED_ARG,
     &&TARGET_LIST_APPEND,
     &&TARGET_SET_ADD,
     &&TARGET_MAP_ADD,
     &&TARGET_LOAD_CLASSDEREF,
     &&TARGET_COPY_FREE_VARS,
-    &&TARGET_STORE_ATTR_ADAPTIVE,
+    &&TARGET_RESUME_QUICK,
     &&TARGET_RESUME,
     &&TARGET_MATCH_CLASS,
+    &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
-    &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_STORE_ATTR_WITH_HINT,
-    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_LOAD_METHOD,
-    &&TARGET_STORE_FAST__STORE_FAST,
+    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
     &&TARGET_PRECALL,
+    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_STORE_SUBSCR_ADAPTIVE,
     &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_STORE_SUBSCR_LIST_INT,
-    &&TARGET_UNPACK_SEQUENCE_ADAPTIVE,
     &&TARGET_CALL,
     &&TARGET_KW_NAMES,
     &&TARGET_POP_JUMP_BACKWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_FALSE,
     &&TARGET_POP_JUMP_BACKWARD_IF_TRUE,
+    &&TARGET_UNPACK_SEQUENCE_ADAPTIVE,
     &&TARGET_UNPACK_SEQUENCE_LIST,
     &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_UNPACK_SEQUENCE_TWO_TUPLE,
@@ -253,6 +254,5 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_DO_TRACING
 };