Put limited RESUME at start of genexpr for free-threading. Fix up exception handling in genexpr

Author: markshannon

Include/internal/pycore_magic_number.h

@@ -294,7 +294,7 @@ Known values:
     Python 3.15a4 3661 (Lazy imports IMPORT_NAME opcode changes)
     Python 3.15a8 3662 (Add counter to RESUME)
     Python 3.15a8 3663 (Merge GET_ITER and GET_YIELD_FROM_ITER. Modify SEND to make it a bit more like FOR_ITER)
-    Python 3.15a8 3664 (Add FOR_ITER_VIRTUAL and GET_ITER specializations)
+    Python 3.15a8 3665 (Add FOR_ITER_VIRTUAL and GET_ITER specializations)
 
 
     Python 3.16 will start with 3700
@@ -308,7 +308,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3664
+#define PYC_MAGIC_NUMBER 3665
 /* This is equivalent to converting PYC_MAGIC_NUMBER to 2 bytes
    (little-endian) and then appending b'\r\n'. */
 #define PYC_MAGIC_NUMBER_TOKEN \

Include/internal/pycore_opcode_utils.h

@@ -82,9 +82,10 @@ extern "C" {
 #define RESUME_AFTER_YIELD 1
 #define RESUME_AFTER_YIELD_FROM 2
 #define RESUME_AFTER_AWAIT 3
+#define RESUME_AT_GEN_EXPR_START 4
 
-#define RESUME_OPARG_LOCATION_MASK 0x3
-#define RESUME_OPARG_DEPTH1_MASK 0x4
+#define RESUME_OPARG_LOCATION_MASK 0x7
+#define RESUME_OPARG_DEPTH1_MASK 0x8
 
 #define GET_ITER_YIELD_FROM 1
 #define GET_ITER_YIELD_FROM_NO_CHECK 2

Objects/codeobject.c

@@ -574,8 +574,12 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     co->co_tlbc->entries[0] = co->co_code_adaptive;
 #endif
     int entry_point = 0;
-    while (entry_point < Py_SIZE(co) &&
-        _PyCode_CODE(co)[entry_point].op.code != RESUME) {
+    while (entry_point < Py_SIZE(co)) {
+        if (_PyCode_CODE(co)[entry_point].op.code == RESUME &&
+           (_PyCode_CODE(co)[entry_point].op.arg & RESUME_OPARG_LOCATION_MASK) != RESUME_AT_GEN_EXPR_START
+        ) {
+            break;
+        }
         entry_point++;
     }
     co->_co_firsttraceable = entry_point;

Python/bytecodes.c

@@ -199,7 +199,7 @@ dummy_func(
             }
         }
 
-        op(_LOAD_BYTECODE, (--)) {
+        replaced op(_LOAD_BYTECODE, (--)) {
             #ifdef Py_GIL_DISABLED
             if (frame->tlbc_index !=
                 ((_PyThreadStateImpl *)tstate)->tlbc_index) {

Python/codegen.c

@@ -1224,12 +1224,17 @@ codegen_wrap_in_stopiteration_handler(compiler *c)
 {
     NEW_JUMP_TARGET_LABEL(c, handler);
 
-    /* Insert SETUP_CLEANUP just before RESUME */
+    /* Insert SETUP_CLEANUP just after the initial RETURN_GENERATOR; POP_TOP */
     instr_sequence *seq = INSTR_SEQUENCE(c);
     int resume = 0;
-    while (_PyInstructionSequence_GetInstruction(seq, resume).i_opcode != RESUME) {
+    while (_PyInstructionSequence_GetInstruction(seq, resume).i_opcode != RETURN_GENERATOR) {
         resume++;
+        assert(resume < seq->s_used);
     }
+    resume++;
+    assert(_PyInstructionSequence_GetInstruction(seq, resume).i_opcode == POP_TOP);
+    resume++;
+    assert(resume < seq->s_used);
     RETURN_IF_ERROR(
         _PyInstructionSequence_InsertInstruction(
             seq, resume,
@@ -4977,10 +4982,14 @@ codegen_comprehension(compiler *c, expr_ty e, int type,
             RETURN_IF_ERROR(
                 _PyInstructionSequence_InsertInstruction(
                     INSTR_SEQUENCE(c), 0,
-                    LOAD_FAST, 0, LOC(outermost->iter)));
+                    RESUME, RESUME_AT_GEN_EXPR_START, NO_LOCATION));
             RETURN_IF_ERROR(
                 _PyInstructionSequence_InsertInstruction(
                     INSTR_SEQUENCE(c), 1,
+                    LOAD_FAST, 0, LOC(outermost->iter)));
+            RETURN_IF_ERROR(
+                _PyInstructionSequence_InsertInstruction(
+                    INSTR_SEQUENCE(c), 2,
                     outermost->is_async ? GET_AITER : GET_ITER,
                     0, LOC(outermost->iter)));
             iter_state = ITERATOR_ON_STACK;

Python/executor_cases.c.h

@@ -963,7 +963,7 @@ label_exception_targets(basicblock *entryblock) {
             }
             else if (instr->i_opcode == RESUME) {
                 instr->i_except = handler;
-                if (instr->i_oparg != RESUME_AT_FUNC_START) {
+                if (instr->i_oparg != RESUME_AT_FUNC_START && instr->i_oparg != RESUME_AT_GEN_EXPR_START) {
                     assert(last_yield_except_depth >= 0);
                     if (last_yield_except_depth == 1) {
                         instr->i_oparg |= RESUME_OPARG_DEPTH1_MASK;

Python/flowgraph.c

@@ -499,6 +499,7 @@ _PyUOp_Replacements[MAX_UOP_ID + 1] = {
     [_FOR_ITER_VIRTUAL] = _FOR_ITER_VIRTUAL_TIER_TWO,
     [_ITER_NEXT_LIST] = _ITER_NEXT_LIST_TIER_TWO,
     [_CHECK_PERIODIC_AT_END] = _TIER2_RESUME_CHECK,
+    [_LOAD_BYTECODE] = _NOP,
 };
 
 static const uint8_t