gh-131798: Support generator frames in the JIT optimizer (GH-143340)

diff --git a/Include/internal/pycore_tstate.h b/Include/internal/pycore_tstate.h
index 27299d1ebcb57a9fbaa2d5d04d8a240e55cc1195..d8f4bfef98af7e72594233de947d0a8b6c9f492b 100644 (file)
--- a/Include/internal/pycore_tstate.h
+++ b/Include/internal/pycore_tstate.h
@@ -47,10 +47,15 @@ typedef struct _PyJitTracerPreviousState {
     _PyBloomFilter dependencies;
 } _PyJitTracerPreviousState;
 
+typedef struct _PyJitTracerTranslatorState {
+    int jump_backward_seen;
+} _PyJitTracerTranslatorState;
+
 typedef struct _PyJitTracerState {
     _PyUOpInstruction *code_buffer;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
+    _PyJitTracerTranslatorState translator_state;
 } _PyJitTracerState;
 
 #endif

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 92b3b3454e4bf75ee3e965d1d08a9578815de8bd..4f092beeea179b4e96d77724a2f107958d2f557f 100644 (file)
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -1168,22 +1168,6 @@ class TestUopsOptimization(unittest.TestCase):
         self.assertIsNotNone(ex)
         self.assertIn("_FOR_ITER_TIER_TWO", get_opnames(ex))
 
-    @unittest.skip("Tracing into generators currently isn't supported.")
-    def test_for_iter_gen(self):
-        def gen(n):
-            for i in range(n):
-                yield i
-        def testfunc(n):
-            g = gen(n)
-            s = 0
-            for i in g:
-                s += i
-            return s
-        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
-        self.assertEqual(res, sum(range(TIER2_THRESHOLD)))
-        self.assertIsNotNone(ex)
-        self.assertIn("_FOR_ITER_GEN_FRAME", get_opnames(ex))
-
     def test_modified_local_is_seen_by_optimized_code(self):
         l = sys._getframe().f_locals
         a = 1
@@ -3404,6 +3388,47 @@ class TestUopsOptimization(unittest.TestCase):
         self.assertIn("_POP_TOP_NOP", uops)
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
+    def test_for_iter_gen_frame(self):
+        def f(n):
+            for i in range(n):
+                # Should be optimized to POP_TOP_NOP
+                yield i + i
+        def testfunc(n):
+            for _ in f(n):
+                pass
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD*2)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_FOR_ITER_GEN_FRAME", uops)
+        # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom.
+        self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
+
+    def test_send_gen_frame(self):
+
+        def gen(n):
+            for i in range(n):
+                yield i + i
+        def send_gen(n):
+            yield from gen(n)
+        def testfunc(n):
+            for _ in send_gen(n):
+                pass
+
+        for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+            # Ensure SEND is specialized to SEND_GEN
+            send_gen(10)
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD*2)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_FOR_ITER_GEN_FRAME", uops)
+        self.assertIn("_SEND_GEN_FRAME", uops)
+        # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom.
+        self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
+
     def test_143026(self):
         # https://github.com/python/cpython/issues/143026
 

diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst b/Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst
new file mode 100644 (file)
index 0000000..528eb70
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst
@@ -0,0 +1 @@
+The JIT optimizer now understands more generator instructions.

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 7f8d2dbd8e83b8af50414d132b7d7723eea189c8..cbcae5075a279d5421efb99b71d4e6d3526202cd 100644 (file)
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -939,15 +939,35 @@ dummy_func(void) {
     }
 
     op(_FOR_ITER_GEN_FRAME, (unused, unused -- unused, unused, gen_frame)) {
-        gen_frame = PyJitRef_NULL;
-        /* We are about to hit the end of the trace */
-        ctx->done = true;
+        assert((this_instr + 1)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging((this_instr + 1));
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+        if (new_frame == NULL) {
+            ctx->done = true;
+            break;
+        }
+        new_frame->stack[0] = sym_new_const(ctx, Py_None);
+        gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
     }
 
-    op(_SEND_GEN_FRAME, (unused, unused -- unused, gen_frame)) {
-        gen_frame = PyJitRef_NULL;
-        // We are about to hit the end of the trace:
-        ctx->done = true;
+    op(_SEND_GEN_FRAME, (unused, v -- unused, gen_frame)) {
+        assert((this_instr + 1)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging((this_instr + 1));
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+        if (new_frame == NULL) {
+            ctx->done = true;
+            break;
+        }
+        new_frame->stack[0] = PyJitRef_StripReferenceInfo(v);
+        gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
     }
 
     op(_CHECK_STACK_SPACE, (unused, unused, unused[oparg] -- unused, unused, unused[oparg])) {

diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 37e3ba29fa20f04cbf04d9a494ad41e5f7586ae5..96aee7830b9e249a5cf7c53167549505c44f6ee3 100644 (file)
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -1081,9 +1081,22 @@
         /* _SEND is not a viable micro-op for tier 2 */
 
         case _SEND_GEN_FRAME: {
+            JitOptRef v;
             JitOptRef gen_frame;
-            gen_frame = PyJitRef_NULL;
-            ctx->done = true;
+            v = stack_pointer[-1];
+            assert((this_instr + 1)->opcode == _PUSH_FRAME);
+            PyCodeObject *co = get_code_with_logging((this_instr + 1));
+            if (co == NULL) {
+                ctx->done = true;
+                break;
+            }
+            _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+            if (new_frame == NULL) {
+                ctx->done = true;
+                break;
+            }
+            new_frame->stack[0] = PyJitRef_StripReferenceInfo(v);
+            gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
             stack_pointer[-1] = gen_frame;
             break;
         }
@@ -2267,8 +2280,19 @@
 
         case _FOR_ITER_GEN_FRAME: {
             JitOptRef gen_frame;
-            gen_frame = PyJitRef_NULL;
-            ctx->done = true;
+            assert((this_instr + 1)->opcode == _PUSH_FRAME);
+            PyCodeObject *co = get_code_with_logging((this_instr + 1));
+            if (co == NULL) {
+                ctx->done = true;
+                break;
+            }
+            _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+            if (new_frame == NULL) {
+                ctx->done = true;
+                break;
+            }
+            new_frame->stack[0] = sym_new_const(ctx, Py_None);
+            gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
             CHECK_STACK_BOUNDS(1);
             stack_pointer[0] = gen_frame;
             stack_pointer += 1;