Streamline non-specialized FOR_ITER

Author: markshannon

Include/internal/pycore_opcode_metadata.h

@@ -906,7 +906,7 @@ def loop_test():
               LIST_EXTEND              1
               LOAD_SMALL_INT           3
               BINARY_OP                5 (*)
-              GET_ITER_LIST_OR_TUPLE
+              GET_ITER_INDEX
       L1:     FOR_ITER_LIST           14 (to L2)
               STORE_FAST               0 (i)
 

Include/internal/pycore_uop_ids.h

@@ -1891,7 +1891,7 @@ def for_iter_list():
                     self.assertIn(i, L)
 
         for_iter_list()
-        self.assert_specialized(for_iter_list, "GET_ITER_LIST_OR_TUPLE")
+        self.assert_specialized(for_iter_list, "GET_ITER_INDEX")
         self.assert_no_opcode(for_iter_list, "GET_ITER")
 
         t = tuple(range(2))
@@ -1903,9 +1903,21 @@ def for_iter_tuple():
                     self.assertIn(i, t)
 
         for_iter_tuple()
-        self.assert_specialized(for_iter_tuple, "GET_ITER_LIST_OR_TUPLE")
+        self.assert_specialized(for_iter_tuple, "GET_ITER_INDEX")
         self.assert_no_opcode(for_iter_tuple, "GET_ITER")
 
+        s = "hello"
+        def for_iter_str():
+            n = 0
+            while n < _testinternalcapi.SPECIALIZATION_THRESHOLD:
+                n += 1
+                for c in s:
+                    self.assertIn(c, s)
+
+        for_iter_str()
+        self.assert_specialized(for_iter_str, "GET_ITER_INDEX")
+        self.assert_no_opcode(for_iter_str, "GET_ITER")
+
         def for_iter_generator():
             n = 0
             while n < _testinternalcapi.SPECIALIZATION_THRESHOLD:

Include/internal/pycore_uop_metadata.h

@@ -3053,7 +3053,7 @@ dummy_func(
 
 
         family(GET_ITER, 1) = {
-            GET_ITER_LIST_OR_TUPLE,
+            GET_ITER_INDEX,
             GET_ITER_SELF,
             GET_ITER_RANGE,
         };
@@ -3108,11 +3108,9 @@ dummy_func(
             null = PyStackRef_NULL;
         }
 
-        inst(GET_ITER_LIST_OR_TUPLE, (unused/1, iter -- iter, index0)) {
+        inst(GET_ITER_INDEX, (unused/1, iter -- iter, index0)) {
             PyTypeObject *tp = PyStackRef_TYPE(iter);
-            if (tp != &PyList_Type) {
-                DEOPT_IF(tp != &PyTuple_Type);
-            }
+            DEOPT_IF(tp->tp_iterindex == NULL);
             index0 = PyStackRef_TagInt(0);
         }
 
@@ -3171,20 +3169,20 @@ dummy_func(
             FOR_ITER_INDEX,
         };
 
-        specializing op(_SPECIALIZE_FOR_ITER, (counter/1, iter, null_or_index -- iter, null_or_index)) {
+        specializing op(_SPECIALIZE_FOR_ITER, (counter/1, iter, maybe_index -- iter, maybe_index)) {
             #if ENABLE_SPECIALIZATION_FT
             if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                 next_instr = this_instr;
-                _Py_Specialize_ForIter(iter, null_or_index, next_instr, oparg);
+                _Py_Specialize_ForIter(iter, maybe_index, next_instr, oparg);
                 DISPATCH_SAME_OPARG();
             }
             OPCODE_DEFERRED_INC(FOR_ITER);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
             #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
-        replaced op(_FOR_ITER, (iter, null_or_index -- iter, null_or_index, next)) {
-            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
+        replaced op(_FOR_ITER, (iter, null_or_index[1] -- iter, null_or_index[1], next)) {
+            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, null_or_index);
             if (!PyStackRef_IsValid(item)) {
                 if (PyStackRef_IsError(item)) {
                     ERROR_NO_POP();
@@ -3230,8 +3228,8 @@ dummy_func(
             _ITER_CHECK_INDEX +
             _FOR_ITER_INDEX;
 
-        op(_FOR_ITER_TIER_TWO, (iter, null_or_index -- iter, null_or_index, next)) {
-            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
+        op(_FOR_ITER_TIER_TWO, (iter, null_or_index[1] -- iter, null_or_index[1], next)) {
+            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, null_or_index);
             if (!PyStackRef_IsValid(item)) {
                 if (PyStackRef_IsError(item)) {
                     ERROR_NO_POP();
@@ -3246,8 +3244,8 @@ dummy_func(
         macro(FOR_ITER) = _SPECIALIZE_FOR_ITER + _FOR_ITER;
 
 
-        inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
-            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, &null_or_index);
+        inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index[1] -- iter, null_or_index[1], next)) {
+            _PyStackRef item = _PyForIter_VirtualIteratorNext(tstate, frame, iter, null_or_index);
             if (!PyStackRef_IsValid(item)) {
                 if (PyStackRef_IsError(item)) {
                     ERROR_NO_POP();

Include/opcode_ids.h

@@ -3439,13 +3439,44 @@ _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *na
     return value;
 }
 
+static _PyStackRef
+not_iterator(PyObject *obj)
+{
+    PyErr_Format(PyExc_TypeError,
+        "'%.200s' object is not an async iterator",
+        Py_TYPE(obj)->tp_name);
+    return PyStackRef_ERROR;
+}
+
 _PyStackRef
 _PyForIter_VirtualIteratorNext(
     PyThreadState* tstate, _PyInterpreterFrame* frame,
     _PyStackRef iter, _PyStackRef* index_ptr
 ) {
     _PyStackRef index = *index_ptr;
-    if (PyStackRef_IsTaggedInt(index)) {
+    if (PyStackRef_IsNull(index)) {
+        PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
+        iternextfunc func = Py_TYPE(iter_o)->tp_iternext;
+        if (func == NULL) {
+            return not_iterator(iter_o);
+        }
+        PyObject *next_o = func(iter_o);
+        if (next_o == NULL) {
+            if (_PyErr_Occurred(tstate)) {
+                if (_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
+                    _PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
+                    _PyErr_Clear(tstate);
+                }
+                else {
+                    return PyStackRef_ERROR;
+                }
+            }
+            return PyStackRef_NULL;
+        }
+        return PyStackRef_FromPyObjectSteal(next_o);
+    }
+    else {
+        assert(PyStackRef_IsTaggedInt(index));
         if (PyStackRef_IsTaggedInt(iter)) {
             if (!PyStackRef_TaggedIntLessThan(index, iter)) {
                 return PyStackRef_NULL;
@@ -3465,21 +3496,6 @@ _PyForIter_VirtualIteratorNext(
             return PyStackRef_FromPyObjectSteal(item);
         }
     }
-    PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
-    PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
-    if (next_o == NULL) {
-        if (_PyErr_Occurred(tstate)) {
-            if (_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
-                _PyEval_MonitorRaise(tstate, frame, frame->instr_ptr);
-                _PyErr_Clear(tstate);
-            }
-            else {
-                return PyStackRef_ERROR;
-            }
-        }
-        return PyStackRef_NULL;
-    }
-    return PyStackRef_FromPyObjectSteal(next_o);
 }
 
 /* Check if a 'cls' provides the given special method. */