Fix jump tracing

Author: Fidget-Spinner

Include/internal/pycore_optimizer.h

@@ -366,7 +366,8 @@ _PyJIT_translate_single_bytecode_to_trace(
     PyCodeObject *code,
     PyFunctionObject *func,
     int opcode,
-    int oparg);
+    int oparg,
+    int jump_taken);
 
 void
 _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int curr_stackdepth, int chain_depth);

Python/bytecodes.c

@@ -3230,6 +3230,7 @@ dummy_func(
                 }
                 // Jump forward by oparg and skip the following END_FOR
                 JUMPBY(oparg + 1);
+                RECORD_JUMP_TAKEN();
                 DISPATCH();
             }
             next = item;
@@ -3291,6 +3292,7 @@ dummy_func(
                 null_or_index = PyStackRef_TagInt(-1);
                 /* Jump forward oparg, then skip following END_FOR instruction */
                 JUMPBY(oparg + 1);
+                RECORD_JUMP_TAKEN();
                 DISPATCH();
             }
 #endif
@@ -3368,6 +3370,7 @@ dummy_func(
                 null_or_index = PyStackRef_TagInt(-1);
                 /* Jump forward oparg, then skip following END_FOR instruction */
                 JUMPBY(oparg + 1);
+                RECORD_JUMP_TAKEN();
                 DISPATCH();
             }
         }
@@ -3412,6 +3415,7 @@ dummy_func(
             if (r->len <= 0) {
                 // Jump over END_FOR instruction.
                 JUMPBY(oparg + 1);
+                RECORD_JUMP_TAKEN();
                 DISPATCH();
             }
         }

Python/ceval.c

@@ -975,12 +975,12 @@ _PyObjectArray_Free(PyObject **array, PyObject **scratch)
 
 // 1 for trace full, 0 for successful write.
 static int
-add_to_code_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, PyCodeObject *old_code, _Py_CODEUNIT *this_instr, _Py_CODEUNIT *next_instr, int opcode, int oparg)
+add_to_code_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, PyCodeObject *old_code, _Py_CODEUNIT *this_instr, _Py_CODEUNIT *next_instr, int opcode, int oparg, int jump_taken)
 {
     assert(frame != NULL);
     assert(tstate->interp->jit_tracer_code_curr_size < UOP_MAX_TRACE_LENGTH);
     PyFunctionObject *func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
-    return !_PyJIT_translate_single_bytecode_to_trace(tstate, frame, this_instr, next_instr, old_code, func, opcode, oparg);
+    return !_PyJIT_translate_single_bytecode_to_trace(tstate, frame, this_instr, next_instr, old_code, func, opcode, oparg, jump_taken);
 }
 
 /* _PyEval_EvalFrameDefault is too large to optimize for speed with PGO on MSVC.

Python/ceval_macros.h

@@ -144,7 +144,7 @@
     DISPATCH();
 #  define RECORD_TRACE() do { \
         frame->instr_ptr = next_instr; \
-        if (add_to_code_trace(tstate, frame, old_code, this_instr, next_instr, opcode, oparg)) { \
+        if (add_to_code_trace(tstate, frame, old_code, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING(); \
         } \
     } while (0);
@@ -339,6 +339,8 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 #define RECORD_BRANCH_TAKEN(bitset, flag)
 #endif
 
+#define RECORD_JUMP_TAKEN() _jump_taken = 1;
+
 #define UNBOUNDLOCAL_ERROR_MSG \
     "cannot access local variable '%s' where it is not associated with a value"
 #define UNBOUNDFREE_ERROR_MSG \