GH-113853: Guarantee forward progress in executors (GH-113854)

1 files changed, 13 insertions, 5 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index f53ddae..2011d96 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -2325,12 +2325,18 @@ dummy_func(
             // Double-check that the opcode isn't instrumented or something:
             if (ucounter > threshold && this_instr->op.code == JUMP_BACKWARD) {
                 OPT_STAT_INC(attempts);
-                int optimized = _PyOptimizer_BackEdge(frame, this_instr, next_instr, stack_pointer);
+                _Py_CODEUNIT *start = this_instr;
+                /* Back up over EXTENDED_ARGs so optimizer sees the whole instruction */
+                while (oparg > 255) {
+                    oparg >>= 8;
+                    start--;
+                }
+                int optimized = _PyOptimizer_Optimize(frame, start, stack_pointer);
                 ERROR_IF(optimized < 0, error);
                 if (optimized) {
                     // Rewind and enter the executor:
-                    assert(this_instr->op.code == ENTER_EXECUTOR);
-                    next_instr = this_instr;
+                    assert(start->op.code == ENTER_EXECUTOR);
+                    next_instr = start;
                     this_instr[1].cache &= ((1 << OPTIMIZER_BITS_IN_COUNTER) - 1);
                 }
                 else {
@@ -2370,10 +2376,12 @@ dummy_func(
                 GOTO_TIER_TWO();
             }
             else {
-                code->co_executors->executors[oparg & 255] = NULL;
+                /* ENTER_EXECUTOR will be the first code unit of the instruction */
+                assert(oparg < 256);
+                code->co_executors->executors[oparg] = NULL;
                 opcode = this_instr->op.code = executor->vm_data.opcode;
                 this_instr->op.arg = executor->vm_data.oparg;
-                oparg = (oparg & (~255)) | executor->vm_data.oparg;
+                oparg = executor->vm_data.oparg;
                 Py_DECREF(executor);
                 next_instr = this_instr;
                 DISPATCH_GOTO();