Fix handling of EXTENDED_ARG

Author: Fidget-Spinner

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);
             }
         }

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);

Python/executor_cases.c.h

@@ -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;