GH-91432: Specialize FOR_ITER (GH-91713)

* Adds FOR_ITER_LIST and FOR_ITER_RANGE specializations.

* Adds _PyLong_AssignValue() internal function to avoid temporary boxing of ints.

Author: sweeneyde

Include/internal/pycore_code.h

@@ -85,6 +85,12 @@ typedef struct {
 
 #define INLINE_CACHE_ENTRIES_STORE_SUBSCR CACHE_ENTRIES(_PyStoreSubscrCache)
 
+typedef struct {
+    _Py_CODEUNIT counter;
+} _PyForIterCache;
+
+#define INLINE_CACHE_ENTRIES_FOR_ITER CACHE_ENTRIES(_PyForIterCache)
+
 #define QUICKENING_WARMUP_DELAY 8
 
 /* We want to compare to zero for efficiency, so we offset values accordingly */
@@ -243,6 +249,7 @@ extern void _Py_Specialize_CompareOp(PyObject *lhs, PyObject *rhs,
                                      _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_UnpackSequence(PyObject *seq, _Py_CODEUNIT *instr,
                                           int oparg);
+extern void _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr);
 
 /* Deallocator function for static codeobjects used in deepfreeze.py */
 extern void _PyStaticCode_Dealloc(PyCodeObject *co);

Include/internal/pycore_list.h

@@ -56,6 +56,12 @@ _PyList_AppendTakeRef(PyListObject *self, PyObject *newitem)
     return _PyList_AppendTakeRefListResize(self, newitem);
 }
 
+typedef struct {
+    PyObject_HEAD
+    Py_ssize_t it_index;
+    PyListObject *it_seq; /* Set to NULL when iterator is exhausted */
+} _PyListIterObject;
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_long.h

@@ -47,6 +47,8 @@ PyObject *_PyLong_Add(PyLongObject *left, PyLongObject *right);
 PyObject *_PyLong_Multiply(PyLongObject *left, PyLongObject *right);
 PyObject *_PyLong_Subtract(PyLongObject *left, PyLongObject *right);
 
+int _PyLong_AssignValue(PyObject **target, Py_ssize_t value);
+
 /* Used by Python/mystrtoul.c, _PyBytes_FromHex(),
    _PyBytes_DecodeEscape(), etc. */
 PyAPI_DATA(unsigned char) _PyLong_DigitValue[256];

Include/internal/pycore_opcode.h

@@ -0,0 +1,22 @@
+#ifndef Py_INTERNAL_RANGE_H
+#define Py_INTERNAL_RANGE_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+typedef struct {
+    PyObject_HEAD
+    long index;
+    long start;
+    long step;
+    long len;
+} _PyRangeIterObject;
+
+#ifdef __cplusplus
+}
+#endif
+#endif   /* !Py_INTERNAL_RANGE_H */

Include/internal/pycore_range.h

@@ -37,6 +37,7 @@
 LOAD_GLOBAL = opmap['LOAD_GLOBAL']
 BINARY_OP = opmap['BINARY_OP']
 JUMP_BACKWARD = opmap['JUMP_BACKWARD']
+FOR_ITER = opmap['FOR_ITER']
 LOAD_ATTR = opmap['LOAD_ATTR']
 
 CACHE = opmap["CACHE"]
@@ -476,6 +477,8 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
             elif deop in hasjrel:
                 signed_arg = -arg if _is_backward_jump(deop) else arg
                 argval = offset + 2 + signed_arg*2
+                if deop == FOR_ITER:
+                    argval += 2
                 argrepr = "to " + repr(argval)
             elif deop in haslocal or deop in hasfree:
                 argval, argrepr = _get_name_info(arg, varname_from_oparg)
@@ -629,11 +632,14 @@ def findlabels(code):
     labels = []
     for offset, op, arg in _unpack_opargs(code):
         if arg is not None:
-            if op in hasjrel:
-                if _is_backward_jump(op):
+            deop = _deoptop(op)
+            if deop in hasjrel:
+                if _is_backward_jump(deop):
                     arg = -arg
                 label = offset + 2 + arg*2
-            elif op in hasjabs:
+                if deop == FOR_ITER:
+                    label += 2
+            elif deop in hasjabs:
                 label = arg*2
             else:
                 continue

Include/opcode.h

@@ -403,12 +403,12 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.11a7 3492 (make POP_JUMP_IF_NONE/NOT_NONE/TRUE/FALSE relative)
 #     Python 3.11a7 3493 (Make JUMP_IF_TRUE_OR_POP/JUMP_IF_FALSE_OR_POP relative)
 #     Python 3.11a7 3494 (New location info table)
-
 #     Python 3.12a1 3500 (Remove PRECALL opcode)
 #     Python 3.12a1 3501 (YIELD_VALUE oparg == stack_depth)
 #     Python 3.12a1 3502 (LOAD_FAST_CHECK, no NULL-check in LOAD_FAST)
 #     Python 3.12a1 3503 (Shrink LOAD_METHOD cache)
 #     Python 3.12a1 3504 (Merge LOAD_METHOD back into LOAD_ATTR)
+#     Python 3.12a1 3505 (Specialization/Cache for FOR_ITER)
 
 #     Python 3.13 will start with 3550
 
@@ -422,7 +422,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 = (3504).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3505).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Lib/dis.py

@@ -278,6 +278,11 @@ def jabs_op(name, op):
     "EXTENDED_ARG": [
         "EXTENDED_ARG_QUICK",
     ],
+    "FOR_ITER": [
+        "FOR_ITER_ADAPTIVE",
+        "FOR_ITER_LIST",
+        "FOR_ITER_RANGE",
+    ],
     "JUMP_BACKWARD": [
         "JUMP_BACKWARD_QUICK",
     ],
@@ -367,6 +372,9 @@ def jabs_op(name, op):
         "type_version": 2,
         "func_version": 1,
     },
+    "FOR_ITER": {
+        "counter": 1,
+    },
     "LOAD_ATTR": {
         "counter": 1,
         "version": 2,