gh-106529: Support FOR_ITER specializations as uops

Lib/test/test_capi/test_misc.py

@@ -2347,11 +2347,12 @@ def func():
 
 @contextlib.contextmanager
 def temporary_optimizer(opt):
+    old_opt = _testinternalcapi.get_optimizer()
     _testinternalcapi.set_optimizer(opt)
     try:
         yield
     finally:
-        _testinternalcapi.set_optimizer(None)
+        _testinternalcapi.set_optimizer(old_opt)
 
 
 @contextlib.contextmanager
@@ -2420,8 +2421,8 @@ def long_loop():
             self.assertEqual(opt.get_count(), 10)
 
 
-
-def get_first_executor(code):
+def get_first_executor(func):
+    code = func.__code__
     co_code = code.co_code
     JUMP_BACKWARD = opcode.opmap["JUMP_BACKWARD"]
     for i in range(0, len(co_code), 2):
@@ -2446,13 +2447,7 @@ def testfunc(x):
         with temporary_optimizer(opt):
             testfunc(1000)
 
-        ex = None
-        for offset in range(0, len(testfunc.__code__.co_code), 2):
-            try:
-                ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
-                break
-            except ValueError:
-                pass
+        ex = get_first_executor(testfunc)
         self.assertIsNotNone(ex)
         uops = {opname for opname, _ in ex}
         self.assertIn("SAVE_IP", uops)
@@ -2493,11 +2488,13 @@ def many_vars():
 
         opt = _testinternalcapi.get_uop_optimizer()
         with temporary_optimizer(opt):
-            ex = get_first_executor(many_vars.__code__)
+            ex = get_first_executor(many_vars)
             self.assertIsNone(ex)
             many_vars()
-            ex = get_first_executor(many_vars.__code__)
-            self.assertIn(("LOAD_FAST", 259), list(ex))
+
+        ex = get_first_executor(many_vars)
+        self.assertIsNotNone(ex)
+        self.assertIn(("LOAD_FAST", 259), list(ex))
 
     def test_unspecialized_unpack(self):
         # An example of an unspecialized opcode
@@ -2516,17 +2513,26 @@ def testfunc(x):
         with temporary_optimizer(opt):
             testfunc(10)
 
-        ex = None
-        for offset in range(0, len(testfunc.__code__.co_code), 2):
-            try:
-                ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
-                break
-            except ValueError:
-                pass
+        ex = get_first_executor(testfunc)
         self.assertIsNotNone(ex)
         uops = {opname for opname, _ in ex}
         self.assertIn("UNPACK_SEQUENCE", uops)
 
+    def test_for_iter(self):
+        def testfunc(x):
+            for i in range(x):
+                i += 1
+
+        opt = _testinternalcapi.get_uop_optimizer()
+
+        with temporary_optimizer(opt):
+            testfunc(10)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _ in ex}
+        self.assertIn("FOR_ITER_RANGE", uops)
+
 
 if __name__ == "__main__":
     unittest.main()

Python/bytecodes.c

@@ -2468,11 +2468,8 @@ dummy_func(
                 Py_DECREF(seq);
             }
             Py_DECREF(iter);
-            STACK_SHRINK(1);
-            SKIP_OVER(INLINE_CACHE_ENTRIES_FOR_ITER);
             /* Jump forward oparg, then skip following END_FOR instruction */
-            JUMPBY(oparg + 1);
-            DISPATCH();
+            JUMP_POP_DISPATCH(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
         end_for_iter_list:
             // Common case: no jump, leave it to the code generator
         }
@@ -2491,11 +2488,8 @@ dummy_func(
                 Py_DECREF(seq);
             }
             Py_DECREF(iter);
-            STACK_SHRINK(1);
-            SKIP_OVER(INLINE_CACHE_ENTRIES_FOR_ITER);
             /* Jump forward oparg, then skip following END_FOR instruction */
-            JUMPBY(oparg + 1);
-            DISPATCH();
+            JUMP_POP_DISPATCH(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
         end_for_iter_tuple:
             // Common case: no jump, leave it to the code generator
         }
@@ -2505,12 +2499,8 @@ dummy_func(
             DEOPT_IF(Py_TYPE(r) != &PyRangeIter_Type, FOR_ITER);
             STAT_INC(FOR_ITER, hit);
             if (r->len <= 0) {
-                STACK_SHRINK(1);
                 Py_DECREF(r);
-                SKIP_OVER(INLINE_CACHE_ENTRIES_FOR_ITER);
-                // Jump over END_FOR instruction.
-                JUMPBY(oparg + 1);
-                DISPATCH();
+                JUMP_POP_DISPATCH(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
             }
             long value = r->start;
             r->start = value + r->step;

Python/ceval.c

@@ -2710,6 +2710,14 @@ void Py_LeaveRecursiveCall(void)
 
 ///////////////////// Experimental UOp Interpreter /////////////////////
 
+#undef JUMP_POP_DISPATCH
+#define JUMP_POP_DISPATCH(x) \
+    do { \
+        frame->prev_instr += (x); \
+        stack_pointer--; \
+        goto exit; \
+    } while (0)
+
 #undef DEOPT_IF
 #define DEOPT_IF(COND, INSTNAME) \
     if ((COND)) {                \
@@ -2790,6 +2798,12 @@ _PyUopExecute(_PyExecutorObject *executor, _PyInterpreterFrame *frame, PyObject
         }
     }
 
+exit:
+    DPRINTF(2, "Jumping out of trace\n");
+    _PyFrame_SetStackPointer(frame, stack_pointer);
+    Py_DECREF(self);
+    return frame;
+
 unbound_local_error:
     format_exc_check_arg(tstate, PyExc_UnboundLocalError,
         UNBOUNDLOCAL_ERROR_MSG,

Python/ceval_macros.h

@@ -155,6 +155,15 @@ GETITEM(PyObject *v, Py_ssize_t i) {
 #define JUMPBY(x)       (next_instr += (x))
 #define SKIP_OVER(x)    (next_instr += (x))
 
+// Helper for FOR_ITER and specializations.
+// This macro is defined differently in the Tier 2 (uops) interpreter.
+#define JUMP_POP_DISPATCH(x) \
+    do { \
+        JUMPBY(x); \
+        stack_pointer--; \
+        DISPATCH(); \
+    } while (0)
+
 /* OpCode prediction macros
     Some opcodes tend to come in pairs thus making it possible to
     predict the second code when the first is run.  For example,