Support FOR_ITER specializations as uops

Author: gvanrossum

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,