Implement FOR_END

sweeneyde · Feb 20, 2022 · Feb 20, 2022 · Feb 20, 2022 · Mar 1, 2022 · Mar 20, 2022
commit f7c220a20cfb6392ceb71cde024b9e023cda77dc
diff --git a/Include/opcode.h b/Include/opcode.h
@@ -403,7 +403,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3479).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3485).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'

diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -194,6 +194,8 @@ def jabs_op(name, op):
 def_op('KW_NAMES', 172)
 hasconst.append(172)
 
+jrel_op('FOR_END', 173)
+
 del def_op, name_op, jrel_op, jabs_op
 
 _nb_ops = [
@@ -264,6 +266,7 @@ def jabs_op(name, op):
     "CALL_NO_KW_TYPE_1",
     "CALL_NO_KW_METHOD_DESCRIPTOR_FAST",
     "JUMP_ABSOLUTE_QUICK",
+    "FOR_END_QUICK",
     "LOAD_ATTR_ADAPTIVE",
     "LOAD_ATTR_INSTANCE_VALUE",
     "LOAD_ATTR_WITH_HINT",

@@ -1378,7 +1378,6 @@ eval_frame_handle_pending(PyThreadState *tstate)
 /* The stack can grow at most MAXINT deep, as co_nlocals and
    co_stacksize are ints. */
 #define STACK_LEVEL()     ((int)(stack_pointer - _PyFrame_Stackbase(frame)))
-#define EMPTY()           (STACK_LEVEL() == 0)
 #define TOP()             (stack_pointer[-1])
 #define SECOND()          (stack_pointer[-2])
 #define THIRD()           (stack_pointer[-3])
@@ -1390,6 +1389,22 @@ eval_frame_handle_pending(PyThreadState *tstate)
 #define BASIC_PUSH(v)     (*stack_pointer++ = (v))
 #define BASIC_POP()       (*--stack_pointer)
 
+#ifdef Py_DEBUG
+#define ASSERT_EMPTY() do {                              \
+    if (STACK_LEVEL()) {                                 \
+        fprintf(stderr, "----\n");                       \
+        while (STACK_LEVEL()) {                          \
+            _PyObject_Dump(BASIC_POP());                 \
+            fprintf(stderr, "----\n");                   \
+        }                                                \
+        Py_FatalError("Unexpected items on the stack."); \
+    }                                                    \
+} while (0)
+#else
+#define ASSERT_EMPTY() ((void)0)
+#endif
+
+
 #ifdef LLTRACE
 #define PUSH(v)         { (void)(BASIC_PUSH(v), \
                           lltrace && prtrace(tstate, TOP(), "push")); \
@@ -1434,6 +1449,19 @@ eval_frame_handle_pending(PyThreadState *tstate)
 #define GET_CACHE() \
     _GetSpecializedCacheEntryForInstruction(first_instr, INSTR_OFFSET(), oparg)
 
+#define MAYBE_QUICKEN() \
+    do { \
+        int err = _Py_IncrementCountAndMaybeQuicken(frame->f_code); \
+        if (err) { \
+            if (err < 0) { \
+                goto error; \
+            } \
+            /* Update first_instr and next_instr to point to newly quickened code */ \
+            int nexti = INSTR_OFFSET(); \
+            first_instr = frame->f_code->co_firstinstr; \
+            next_instr = first_instr + nexti; \
+        } \
+    } while (0)
 
 #define DEOPT_IF(cond, instname) if (cond) { goto instname ## _miss; }
 
@@ -1734,16 +1762,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
         }
 
         TARGET(RESUME) {
-            int err = _Py_IncrementCountAndMaybeQuicken(frame->f_code);
-            if (err) {
-                if (err < 0) {
-                    goto error;
-                }
-                /* Update first_instr and next_instr to point to newly quickened code */
-                int nexti = INSTR_OFFSET();
-                first_instr = frame->f_code->co_firstinstr;
-                next_instr = first_instr + nexti;
-            }
+            MAYBE_QUICKEN();
             JUMP_TO_INSTRUCTION(RESUME_QUICK);
         }
 
@@ -2374,7 +2393,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
 
         TARGET(RETURN_VALUE) {
             PyObject *retval = POP();
-            assert(EMPTY());
+            ASSERT_EMPTY();
             frame->f_state = FRAME_RETURNED;
             _PyFrame_SetStackPointer(frame, stack_pointer);
             TRACE_FUNCTION_EXIT();
@@ -4046,16 +4065,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
 
         TARGET(JUMP_ABSOLUTE) {
             PREDICTED(JUMP_ABSOLUTE);
-            int err = _Py_IncrementCountAndMaybeQuicken(frame->f_code);
-            if (err) {
-                if (err < 0) {
-                    goto error;
-                }
-                /* Update first_instr and next_instr to point to newly quickened code */
-                int nexti = INSTR_OFFSET();
-                first_instr = frame->f_code->co_firstinstr;
-                next_instr = first_instr + nexti;
-            }
+            MAYBE_QUICKEN();
             JUMP_TO_INSTRUCTION(JUMP_ABSOLUTE_QUICK);
         }
 
@@ -4225,6 +4235,39 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
             DISPATCH();
         }
 
+        TARGET(FOR_END) {
+            /* before: [iter]; after: [iter, iter()] *or* [] */
+            MAYBE_QUICKEN();
+            JUMP_TO_INSTRUCTION(FOR_END_QUICK);
+        }
+
+        TARGET(FOR_END_QUICK) {
+            PREDICTED(FOR_END_QUICK);
+            PyObject *iter = TOP();
+            PyObject *next = (*Py_TYPE(iter)->tp_iternext)(iter);
+            if (next != NULL) {
+                PUSH(next);
+                JUMPBY(oparg);
+                CHECK_EVAL_BREAKER();
+                PREDICT(STORE_FAST);
+                PREDICT(UNPACK_SEQUENCE);
+                DISPATCH();
+            }
+            if (_PyErr_Occurred(tstate)) {
+                if (!_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
+                    goto error;
+                }
+                else if (tstate->c_tracefunc != NULL) {
+                    call_exc_trace(tstate->c_tracefunc, tstate->c_traceobj, tstate, frame);
+                }
+                _PyErr_Clear(tstate);
+            }
+            /* iterator ended normally */
+            STACK_SHRINK(1);
+            Py_DECREF(iter);
+            DISPATCH();
+        }
+
         TARGET(BEFORE_ASYNC_WITH) {
             PyObject *mgr = TOP();
             PyObject *res;
@@ -5152,7 +5195,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, InterpreterFrame *frame, int thr
             if (gen == NULL) {
                 goto error;
             }
-            assert(EMPTY());
+            ASSERT_EMPTY();
             _PyFrame_SetStackPointer(frame, stack_pointer);
             InterpreterFrame *gen_frame = (InterpreterFrame *)gen->gi_iframe;
             _PyFrame_Copy(frame, gen_frame);

@@ -915,6 +915,9 @@ stack_effect(int opcode, int oparg, int jump)
         case FOR_ITER:
             /* -1 at end of iterator, 1 if continue iterating. */
             return jump > 0 ? -1 : 1;
+        case FOR_END:
+            /* -1 at end of iterator, 1 if continue iterating. */
+            return jump == 0 ? -1 : 1;
         case SEND:
             return jump > 0 ? -1 : 0;
         case STORE_ATTR:
@@ -2970,11 +2973,9 @@ compiler_for(struct compiler *c, stmt_ty s)
     compiler_use_next_block(c, body);
     VISIT(c, expr, s->v.For.target);
     VISIT_SEQ(c, stmt, s->v.For.body);
-    /* Mark jump as artificial */
-    UNSET_LOC(c);
-    ADDOP_JUMP(c, JUMP_ABSOLUTE, start);
+    SET_LOC(c, s->v.For.iter);
+    ADDOP_JUMP(c, FOR_END, body);
     compiler_use_next_block(c, cleanup);
-
     compiler_pop_fblock(c, FOR_LOOP, start);
 
     VISIT_SEQ(c, stmt, s->v.For.orelse);