GH-96793: Specialize FOR_ITER for generators. (GH-98772)

4 files changed, 47 insertions, 44 deletions

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
index cf8573a..61f4cf4 100644
--- a/Include/internal/pycore_code.h
+++ b/Include/internal/pycore_code.h
@@ -230,7 +230,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);
+extern void _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr, int oparg);
 
 /* Finalizer function for static codeobjects used in deepfreeze.py */
 extern void _PyStaticCode_Fini(PyCodeObject *co);
diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
index 5bd0a7f..feee692 100644
--- a/Include/internal/pycore_frame.h
+++ b/Include/internal/pycore_frame.h
@@ -61,6 +61,7 @@ typedef struct _PyInterpreterFrame {
     // over, or (in the case of a newly-created frame) a totally invalid value:
     _Py_CODEUNIT *prev_instr;
     int stacktop;     /* Offset of TOS from localsplus  */
+    uint16_t yield_offset;
     bool is_entry;  // Whether this is the "root" frame for the current _PyCFrame.
     char owner;
     /* Locals and stack */
@@ -110,6 +111,7 @@ _PyFrame_InitializeSpecials(
     frame->frame_obj = NULL;
     frame->prev_instr = _PyCode_CODE(code) - 1;
     frame->is_entry = false;
+    frame->yield_offset = 0;
     frame->owner = FRAME_OWNED_BY_THREAD;
 }
 
diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
index 59276c8..2ae1fc8 100644
--- a/Include/internal/pycore_opcode.h
+++ b/Include/internal/pycore_opcode.h
@@ -128,6 +128,7 @@ const uint8_t _PyOpcode_Deopt[256] = {
     [FORMAT_VALUE] = FORMAT_VALUE,
     [FOR_ITER] = FOR_ITER,
     [FOR_ITER_ADAPTIVE] = FOR_ITER,
+    [FOR_ITER_GEN] = FOR_ITER,
     [FOR_ITER_LIST] = FOR_ITER,
     [FOR_ITER_RANGE] = FOR_ITER,
     [GET_AITER] = GET_AITER,
@@ -313,20 +314,20 @@ static const char *const _PyOpcode_OpName[263] = {
     [PRINT_EXPR] = "PRINT_EXPR",
     [LOAD_BUILD_CLASS] = "LOAD_BUILD_CLASS",
     [FOR_ITER_RANGE] = "FOR_ITER_RANGE",
-    [LOAD_ATTR_ADAPTIVE] = "LOAD_ATTR_ADAPTIVE",
+    [FOR_ITER_GEN] = "FOR_ITER_GEN",
     [LOAD_ASSERTION_ERROR] = "LOAD_ASSERTION_ERROR",
     [RETURN_GENERATOR] = "RETURN_GENERATOR",
+    [LOAD_ATTR_ADAPTIVE] = "LOAD_ATTR_ADAPTIVE",
     [LOAD_ATTR_CLASS] = "LOAD_ATTR_CLASS",
     [LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN] = "LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN",
     [LOAD_ATTR_INSTANCE_VALUE] = "LOAD_ATTR_INSTANCE_VALUE",
     [LOAD_ATTR_MODULE] = "LOAD_ATTR_MODULE",
     [LOAD_ATTR_PROPERTY] = "LOAD_ATTR_PROPERTY",
-    [LOAD_ATTR_SLOT] = "LOAD_ATTR_SLOT",
     [LIST_TO_TUPLE] = "LIST_TO_TUPLE",
     [RETURN_VALUE] = "RETURN_VALUE",
     [IMPORT_STAR] = "IMPORT_STAR",
     [SETUP_ANNOTATIONS] = "SETUP_ANNOTATIONS",
-    [LOAD_ATTR_WITH_HINT] = "LOAD_ATTR_WITH_HINT",
+    [LOAD_ATTR_SLOT] = "LOAD_ATTR_SLOT",
     [ASYNC_GEN_WRAP] = "ASYNC_GEN_WRAP",
     [PREP_RERAISE_STAR] = "PREP_RERAISE_STAR",
     [POP_EXCEPT] = "POP_EXCEPT",
@@ -353,7 +354,7 @@ static const char *const _PyOpcode_OpName[263] = {
     [JUMP_FORWARD] = "JUMP_FORWARD",
     [JUMP_IF_FALSE_OR_POP] = "JUMP_IF_FALSE_OR_POP",
     [JUMP_IF_TRUE_OR_POP] = "JUMP_IF_TRUE_OR_POP",
-    [LOAD_ATTR_METHOD_LAZY_DICT] = "LOAD_ATTR_METHOD_LAZY_DICT",
+    [LOAD_ATTR_WITH_HINT] = "LOAD_ATTR_WITH_HINT",
     [POP_JUMP_IF_FALSE] = "POP_JUMP_IF_FALSE",
     [POP_JUMP_IF_TRUE] = "POP_JUMP_IF_TRUE",
     [LOAD_GLOBAL] = "LOAD_GLOBAL",
@@ -361,7 +362,7 @@ static const char *const _PyOpcode_OpName[263] = {
     [CONTAINS_OP] = "CONTAINS_OP",
     [RERAISE] = "RERAISE",
     [COPY] = "COPY",
-    [LOAD_ATTR_METHOD_NO_DICT] = "LOAD_ATTR_METHOD_NO_DICT",
+    [LOAD_ATTR_METHOD_LAZY_DICT] = "LOAD_ATTR_METHOD_LAZY_DICT",
     [BINARY_OP] = "BINARY_OP",
     [SEND] = "SEND",
     [LOAD_FAST] = "LOAD_FAST",
@@ -381,9 +382,9 @@ static const char *const _PyOpcode_OpName[263] = {
     [STORE_DEREF] = "STORE_DEREF",
     [DELETE_DEREF] = "DELETE_DEREF",
     [JUMP_BACKWARD] = "JUMP_BACKWARD",
-    [LOAD_ATTR_METHOD_WITH_DICT] = "LOAD_ATTR_METHOD_WITH_DICT",
+    [LOAD_ATTR_METHOD_NO_DICT] = "LOAD_ATTR_METHOD_NO_DICT",
     [CALL_FUNCTION_EX] = "CALL_FUNCTION_EX",
-    [LOAD_ATTR_METHOD_WITH_VALUES] = "LOAD_ATTR_METHOD_WITH_VALUES",
+    [LOAD_ATTR_METHOD_WITH_DICT] = "LOAD_ATTR_METHOD_WITH_DICT",
     [EXTENDED_ARG] = "EXTENDED_ARG",
     [LIST_APPEND] = "LIST_APPEND",
     [SET_ADD] = "SET_ADD",
@@ -393,26 +394,27 @@ static const char *const _PyOpcode_OpName[263] = {
     [YIELD_VALUE] = "YIELD_VALUE",
     [RESUME] = "RESUME",
     [MATCH_CLASS] = "MATCH_CLASS",
+    [LOAD_ATTR_METHOD_WITH_VALUES] = "LOAD_ATTR_METHOD_WITH_VALUES",
     [LOAD_CONST__LOAD_FAST] = "LOAD_CONST__LOAD_FAST",
-    [LOAD_FAST__LOAD_CONST] = "LOAD_FAST__LOAD_CONST",
     [FORMAT_VALUE] = "FORMAT_VALUE",
     [BUILD_CONST_KEY_MAP] = "BUILD_CONST_KEY_MAP",
     [BUILD_STRING] = "BUILD_STRING",
+    [LOAD_FAST__LOAD_CONST] = "LOAD_FAST__LOAD_CONST",
     [LOAD_FAST__LOAD_FAST] = "LOAD_FAST__LOAD_FAST",
     [LOAD_GLOBAL_ADAPTIVE] = "LOAD_GLOBAL_ADAPTIVE",
     [LOAD_GLOBAL_BUILTIN] = "LOAD_GLOBAL_BUILTIN",
-    [LOAD_GLOBAL_MODULE] = "LOAD_GLOBAL_MODULE",
     [LIST_EXTEND] = "LIST_EXTEND",
     [SET_UPDATE] = "SET_UPDATE",
     [DICT_MERGE] = "DICT_MERGE",
     [DICT_UPDATE] = "DICT_UPDATE",
+    [LOAD_GLOBAL_MODULE] = "LOAD_GLOBAL_MODULE",
     [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
     [STORE_ATTR_INSTANCE_VALUE] = "STORE_ATTR_INSTANCE_VALUE",
     [STORE_ATTR_SLOT] = "STORE_ATTR_SLOT",
     [STORE_ATTR_WITH_HINT] = "STORE_ATTR_WITH_HINT",
-    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [CALL] = "CALL",
     [KW_NAMES] = "KW_NAMES",
+    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
     [STORE_SUBSCR_ADAPTIVE] = "STORE_SUBSCR_ADAPTIVE",
     [STORE_SUBSCR_DICT] = "STORE_SUBSCR_DICT",
@@ -421,7 +423,6 @@ static const char *const _PyOpcode_OpName[263] = {
     [UNPACK_SEQUENCE_LIST] = "UNPACK_SEQUENCE_LIST",
     [UNPACK_SEQUENCE_TUPLE] = "UNPACK_SEQUENCE_TUPLE",
     [UNPACK_SEQUENCE_TWO_TUPLE] = "UNPACK_SEQUENCE_TWO_TUPLE",
-    [181] = "<181>",
     [182] = "<182>",
     [183] = "<183>",
     [184] = "<184>",
@@ -507,7 +508,6 @@ static const char *const _PyOpcode_OpName[263] = {
 #endif
 
 #define EXTRA_CASES \
-    case 181: \
     case 182: \
     case 183: \
     case 184: \
diff --git a/Include/opcode.h b/Include/opcode.h
index 4efa357..e1978a5 100644
--- a/Include/opcode.h
+++ b/Include/opcode.h
@@ -168,37 +168,38 @@ extern "C" {
 #define FOR_ITER_ADAPTIVE                       66
 #define FOR_ITER_LIST                           67
 #define FOR_ITER_RANGE                          72
-#define LOAD_ATTR_ADAPTIVE                      73
-#define LOAD_ATTR_CLASS                         76
-#define LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN       77
-#define LOAD_ATTR_INSTANCE_VALUE                78
-#define LOAD_ATTR_MODULE                        79
-#define LOAD_ATTR_PROPERTY                      80
-#define LOAD_ATTR_SLOT                          81
-#define LOAD_ATTR_WITH_HINT                     86
-#define LOAD_ATTR_METHOD_LAZY_DICT             113
-#define LOAD_ATTR_METHOD_NO_DICT               121
-#define LOAD_ATTR_METHOD_WITH_DICT             141
-#define LOAD_ATTR_METHOD_WITH_VALUES           143
-#define LOAD_CONST__LOAD_FAST                  153
-#define LOAD_FAST__LOAD_CONST                  154
-#define LOAD_FAST__LOAD_FAST                   158
-#define LOAD_GLOBAL_ADAPTIVE                   159
-#define LOAD_GLOBAL_BUILTIN                    160
-#define LOAD_GLOBAL_MODULE                     161
-#define STORE_ATTR_ADAPTIVE                    166
-#define STORE_ATTR_INSTANCE_VALUE              167
-#define STORE_ATTR_SLOT                        168
-#define STORE_ATTR_WITH_HINT                   169
-#define STORE_FAST__LOAD_FAST                  170
-#define STORE_FAST__STORE_FAST                 173
-#define STORE_SUBSCR_ADAPTIVE                  174
-#define STORE_SUBSCR_DICT                      175
-#define STORE_SUBSCR_LIST_INT                  176
-#define UNPACK_SEQUENCE_ADAPTIVE               177
-#define UNPACK_SEQUENCE_LIST                   178
-#define UNPACK_SEQUENCE_TUPLE                  179
-#define UNPACK_SEQUENCE_TWO_TUPLE              180
+#define FOR_ITER_GEN                            73
+#define LOAD_ATTR_ADAPTIVE                      76
+#define LOAD_ATTR_CLASS                         77
+#define LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN       78
+#define LOAD_ATTR_INSTANCE_VALUE                79
+#define LOAD_ATTR_MODULE                        80
+#define LOAD_ATTR_PROPERTY                      81
+#define LOAD_ATTR_SLOT                          86
+#define LOAD_ATTR_WITH_HINT                    113
+#define LOAD_ATTR_METHOD_LAZY_DICT             121
+#define LOAD_ATTR_METHOD_NO_DICT               141
+#define LOAD_ATTR_METHOD_WITH_DICT             143
+#define LOAD_ATTR_METHOD_WITH_VALUES           153
+#define LOAD_CONST__LOAD_FAST                  154
+#define LOAD_FAST__LOAD_CONST                  158
+#define LOAD_FAST__LOAD_FAST                   159
+#define LOAD_GLOBAL_ADAPTIVE                   160
+#define LOAD_GLOBAL_BUILTIN                    161
+#define LOAD_GLOBAL_MODULE                     166
+#define STORE_ATTR_ADAPTIVE                    167
+#define STORE_ATTR_INSTANCE_VALUE              168
+#define STORE_ATTR_SLOT                        169
+#define STORE_ATTR_WITH_HINT                   170
+#define STORE_FAST__LOAD_FAST                  173
+#define STORE_FAST__STORE_FAST                 174
+#define STORE_SUBSCR_ADAPTIVE                  175
+#define STORE_SUBSCR_DICT                      176
+#define STORE_SUBSCR_LIST_INT                  177
+#define UNPACK_SEQUENCE_ADAPTIVE               178
+#define UNPACK_SEQUENCE_LIST                   179
+#define UNPACK_SEQUENCE_TUPLE                  180
+#define UNPACK_SEQUENCE_TWO_TUPLE              181
 #define DO_TRACING                             255
 
 #define HAS_ARG(op) ((((op) >= HAVE_ARGUMENT) && (!IS_PSEUDO_OPCODE(op)))\