From 95eee89121e0db3bf230f58c67c079cfaf2ada76 Mon Sep 17 00:00:00 2001
From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
Date: Thu, 9 Oct 2025 20:54:43 +0100
Subject: [PATCH] Fix handling of EXTENDED_ARG

---
 Python/bytecodes.c        |  9 +++------
 Python/ceval_macros.h     |  4 +++-
 Python/executor_cases.c.h |  9 +++------
 Python/optimizer.c        | 14 +++-----------
 4 files changed, 12 insertions(+), 24 deletions(-)

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 15dbab0bef70..812b323c137b 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -5466,7 +5466,6 @@ dummy_func(
             _Py_CODEUNIT *target = frame->instr_ptr;
             _PyExitData *exit = (_PyExitData *)exit_p;
             _Py_BackoffCounter temperature = exit->temperature;
-            tstate->jit_exit = exit;
 #if defined(Py_DEBUG) && !defined(_Py_JIT)
             OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
             if (frame->lltrace >= 2) {
@@ -5481,10 +5480,8 @@ dummy_func(
             if (target->op.code == ENTER_EXECUTOR) {
                 PyCodeObject *code = _PyFrame_GetCode(frame);
                 _PyExecutorObject *executor = code->co_executors->executors[target->op.arg];
-                Py_INCREF(executor);
-                assert(tstate->jit_exit == exit);
-                exit->executor = executor;
-                TIER2_TO_TIER2(exit->executor);
+                tstate->jit_exit = NULL;
+                TIER2_TO_TIER2(executor);
             }
             else {
                 if (frame->owner >= FRAME_OWNED_BY_INTERPRETER) {
@@ -5497,7 +5494,7 @@ dummy_func(
                 exit->temperature = initial_temperature_backoff_counter();
                 _PyExecutorObject *previous_executor = _PyExecutor_FromExit(exit);
                 assert(tstate->current_executor == (PyObject *)previous_executor);
-                _PyJIT_InitializeTracing(tstate, frame, target, STACK_LEVEL(), 0, exit);
+                _PyJIT_InitializeTracing(tstate, frame, target, STACK_LEVEL(), 0, NULL);
                 GOTO_TIER_ONE(target, 1);
             }
         }
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
index edacf4418b97..8085fb3f1928 100644
--- a/Python/ceval_macros.h
+++ b/Python/ceval_macros.h
@@ -133,6 +133,8 @@
 #define TRACING_JUMP_TO_LABEL(label) \
     RECORD_JUMP_TAKEN() \
     RECORD_TRACE_NO_DISPATCH() \
+    assert(!IS_JIT_TRACING()); \
+    RELOAD_TRACING(); \
     JUMP_TO_LABEL(label);
 
 #if _Py_TAIL_CALL_INTERP || USE_COMPUTED_GOTOS
@@ -161,7 +163,7 @@
         JUMP_TO_LABEL(error); \
     }
 #  define RECORD_TRACE_NO_DISPATCH() do { \
-        if (DISPATCH_TABLE_VAR == TRACING_DISPATCH_TABLE && add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
+        if (IS_JIT_TRACING() && add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING_NO_DISPATCH(); \
         } \
     } while (0);
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 86c47b560104..58c57ea6b96f 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -7528,7 +7528,6 @@
             _Py_CODEUNIT *target = frame->instr_ptr;
             _PyExitData *exit = (_PyExitData *)exit_p;
             _Py_BackoffCounter temperature = exit->temperature;
-            tstate->jit_exit = exit;
             #if defined(Py_DEBUG) && !defined(_Py_JIT)
             OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
             if (frame->lltrace >= 2) {
@@ -7545,10 +7544,8 @@
             if (target->op.code == ENTER_EXECUTOR) {
                 PyCodeObject *code = _PyFrame_GetCode(frame);
                 _PyExecutorObject *executor = code->co_executors->executors[target->op.arg];
-                Py_INCREF(executor);
-                assert(tstate->jit_exit == exit);
-                exit->executor = executor;
-                TIER2_TO_TIER2(exit->executor);
+                tstate->jit_exit = NULL;
+                TIER2_TO_TIER2(executor);
             }
             else {
                 if (frame->owner >= FRAME_OWNED_BY_INTERPRETER) {
@@ -7561,7 +7558,7 @@
                 exit->temperature = initial_temperature_backoff_counter();
                 _PyExecutorObject *previous_executor = _PyExecutor_FromExit(exit);
                 assert(tstate->current_executor == (PyObject *)previous_executor);
-                _PyJIT_InitializeTracing(tstate, frame, target, STACK_LEVEL(), 0, exit);
+                _PyJIT_InitializeTracing(tstate, frame, target, STACK_LEVEL(), 0, NULL);
                 GOTO_TIER_ONE(target, 1);
             }
             break;
diff --git a/Python/optimizer.c b/Python/optimizer.c
index 03cd070c163e..a8df28cd2a86 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -468,7 +468,6 @@ is_for_iter_test[MAX_UOP_ID + 1] = {
     [_GUARD_NOT_EXHAUSTED_RANGE] = 1,
     [_GUARD_NOT_EXHAUSTED_LIST] = 1,
     [_GUARD_NOT_EXHAUSTED_TUPLE] = 1,
-    [_FOR_ITER_TIER_TWO] = 1,
 };
 
 static const uint16_t
@@ -586,18 +585,11 @@ _PyJIT_translate_single_bytecode_to_trace(
 
     DPRINTF(2, "%d: %s(%d)\n", target, _PyOpcode_OpName[opcode], oparg);
 
-    // One for possible _DEOPT, one because _CHECK_VALIDITY itself might _DEOPT
-    max_length -= 2;
     if ((uint16_t)oparg != (uint64_t)oparg) {
-        // Back up over EXTENDED_ARGs
-        _PyUOpInstruction *curr = &trace[trace_length-1];
-        while (oparg > 0) {
-            oparg >>= 8;
-            trace_length--;
-        }
-        goto full;
+        goto unsupported;
     }
-
+    // One for possible _DEOPT, one because _CHECK_VALIDITY itself might _DEOPT
+    max_length -= 2;
 
     if (opcode == EXTENDED_ARG) {
         return 1;
-- 
2.47.3