[mypyc] Ensure a generator can't be entered while it's being executed

mypyc/codegen/emitclass.py

@@ -34,6 +34,7 @@
     NATIVE_PREFIX,
     PREFIX,
     REG_PREFIX,
+    RUNNING_FIELD,
     short_id_from_name,
 )
 from mypyc.ir.class_ir import ClassIR, VTableEntries
@@ -264,7 +265,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     fields: dict[str, str] = {"tp_name": f'"{name}"'}
 
     generate_full = not cl.is_trait and not cl.builtin_base
-    needs_getseters = cl.needs_getseters or not cl.is_generated or cl.has_dict
+    needs_getseters = cl.needs_getseters_table
 
     if not cl.builtin_base:
         fields["tp_new"] = new_name
@@ -484,6 +485,13 @@ def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
     lines += ["typedef struct {", "PyObject_HEAD", "CPyVTableItem *vtable;"]
     if cl.has_method("__call__"):
         lines.append("vectorcallfunc vectorcall;")
+    # The flag affects attribute offsets, so concrete bases must agree on its presence.
+    assert all(
+        base.has_running_flag == cl.has_running_flag for base in cl.base_mro
+    ), f"{cl.name} disagrees with a base class about the running flag"
+    if cl.has_running_flag:
+        # This implementation field is intentionally absent from the IR attributes.
+        lines.append(f"uint32_t {RUNNING_FIELD};")
     bitmap_attrs = []
     for base in reversed(cl.base_mro):
         if not base.is_trait:

mypyc/codegen/emitfunc.py

@@ -18,6 +18,7 @@
     IS_FREE_THREADED,
     NATIVE_PREFIX,
     REG_PREFIX,
+    RUNNING_FIELD,
 )
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR, all_values
@@ -129,12 +130,18 @@ def native_function_header(fn: FuncDecl, emitter: Emitter) -> str:
 
 
 def generate_native_function(
-    fn: FuncIR, emitter: Emitter, source_path: str, module_name: str
+    fn: FuncIR,
+    emitter: Emitter,
+    source_path: str,
+    module_name: str,
+    running_flag_class: ClassIR | None = None,
 ) -> None:
     declarations = Emitter(emitter.context)
     names = generate_names_for_ir(fn.arg_regs, fn.blocks)
     body = Emitter(emitter.context, names)
-    visitor = FunctionEmitterVisitor(body, declarations, source_path, module_name)
+    visitor = FunctionEmitterVisitor(
+        body, declarations, source_path, module_name, running_flag_class
+    )
 
     declarations.emit_line(f"{native_function_header(fn.decl, emitter)} {{")
     body.indent()
@@ -183,6 +190,10 @@ def generate_native_function(
             if not is_next_block or is_problematic_op:
                 fn.blocks[target.label].referenced = True
 
+    if running_flag_class is not None:
+        # Place this before all labels so it runs on resume, but not on internal jumps.
+        visitor.emit_claim_running_flag(fn)
+
     common = frequently_executed_blocks(fn.blocks[0])
 
     for i in range(len(blocks)):
@@ -209,13 +220,21 @@ def generate_native_function(
 
 class FunctionEmitterVisitor(OpVisitor[None]):
     def __init__(
-        self, emitter: Emitter, declarations: Emitter, source_path: str, module_name: str
+        self,
+        emitter: Emitter,
+        declarations: Emitter,
+        source_path: str,
+        module_name: str,
+        running_flag_class: ClassIR | None = None,
     ) -> None:
         self.emitter = emitter
         self.names = emitter.names
         self.declarations = declarations
         self.source_path = source_path
         self.module_name = module_name
+        # Set while emitting a generator helper protected by its running flag.
+        self.running_flag_class = running_flag_class
+        self.running_flag_ptr: str | None = None
         self.literals = emitter.context.literals
         self.rare = False
         # Next basic block to be processed after the current one (if any), set by caller
@@ -291,8 +310,28 @@ def visit_branch(self, op: Branch) -> None:
 
             self.emit_lines("} else", "    goto %s;" % self.label(false))
 
+    def emit_claim_running_flag(self, fn: FuncIR) -> None:
+        """Claim the generator's running flag or raise ValueError."""
+        cl = self.running_flag_class
+        assert cl is not None
+        struct = cl.struct_name(self.names)
+        self_str = self.reg(fn.arg_regs[0])
+        self.running_flag_ptr = f"&(({struct} *){self_str})->{RUNNING_FIELD}"
+        flag = self.running_flag_ptr
+        is_coroutine = 1 if cl.has_method("__await__") else 0
+        self.emit_line(f"if (unlikely(!CPyGen_TryEnter({flag}))) {{")
+        self.emit_line(f"return CPyGen_AlreadyExecutingError({is_coroutine});")
+        self.emit_line("}")
+
+    def emit_release_running_flag(self) -> None:
+        """Release the flag; every helper exit is represented by Return."""
+        assert self.running_flag_ptr is not None
+        self.emit_line(f"CPyGen_Exit({self.running_flag_ptr});")
+
     def visit_return(self, op: Return) -> None:
         value_str = self.reg(op.value)
+        if self.running_flag_class is not None:
+            self.emit_release_running_flag()
         self.emit_line("return %s;" % value_str)
 
     def visit_tuple_set(self, op: TupleSet) -> None:
@@ -418,9 +457,15 @@ def emit_load_attr_take_ref(
         for attributes safe to borrow on free-threaded builds (Final and vec attrs -- see
         transform_member_expr in irbuild), whose values live as long as their container.
         The default (GIL) build always takes the plain-load path and increfs separately.
+
+        Thread-confined attributes also use plain loads; see
+        ClassIR.attrs_are_thread_confined.
         """
         use_get_attr_ref = (
-            IS_FREE_THREADED and is_simple_refcounted_pointer(attr_rtype) and not op.is_borrowed
+            IS_FREE_THREADED
+            and is_simple_refcounted_pointer(attr_rtype)
+            and not op.is_borrowed
+            and not cl.attrs_are_thread_confined()
         )
         if use_get_attr_ref and cl.is_final_attr(op.attr):
             self.emitter.emit_line(f"{dest} = CPy_GetAttrRefFinal((PyObject **)&{attr_expr});")
@@ -578,7 +623,11 @@ def visit_set_attr(self, op: SetAttr) -> None:
             )
             self.emit_line(f"{dest} = 1;")
             self.emitter.emit_error_check(tmp, ret_type, f"{dest} = 0;")
-        elif IS_FREE_THREADED and is_simple_refcounted_pointer(attr_rtype):
+        elif (
+            IS_FREE_THREADED
+            and is_simple_refcounted_pointer(attr_rtype)
+            and not cl.attrs_are_thread_confined()
+        ):
             # In free-threaded builds, publishing a single reference-counted
             # 'PyObject *' field must be atomic so a concurrent reader (see
             # CPy_GetAttrRef) never observes a torn pointer or a freed value.

mypyc/codegen/emitmodule.py

@@ -75,7 +75,7 @@
 from mypyc.ir.rtypes import RType
 from mypyc.irbuild.main import build_ir
 from mypyc.irbuild.mapper import Mapper
-from mypyc.irbuild.prepare import load_type_map
+from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME, load_type_map
 from mypyc.namegen import NameGenerator, exported_name
 from mypyc.options import CompilerOptions
 from mypyc.transform.copy_propagation import do_copy_propagation
@@ -708,12 +708,19 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
                 if cl.is_ext_class:
                     generate_class(cl, module_name, emitter)
 
+            running_flag_classes = {cl.name: cl for cl in module.classes if cl.has_running_flag}
+
             # Generate Python extension module definitions and module initialization functions.
             self.generate_module_def(emitter, module_name, module)
 
             for fn in module.functions:
                 emitter.emit_line()
-                generate_native_function(fn, emitter, self.source_paths[module_name], module_name)
+                running_flag_class = None
+                if fn.decl.name == GENERATOR_HELPER_NAME and fn.class_name is not None:
+                    running_flag_class = running_flag_classes.get(fn.class_name)
+                generate_native_function(
+                    fn, emitter, self.source_paths[module_name], module_name, running_flag_class
+                )
                 if fn.name != TOP_LEVEL_NAME and not fn.internal:
                     emitter.emit_line()
                     if is_fastcall_supported(fn, emitter.capi_version):

mypyc/common.py

@@ -28,6 +28,9 @@
 GENERATOR_ATTRIBUTE_PREFIX: Final = "__mypyc_generator_attribute__"
 CPYFUNCTION_NAME = "__cpyfunction__"
 
+# Omits the prefix added to user attribute fields, so it cannot collide with one.
+RUNNING_FIELD: Final = "mypyc_running"
+
 # Max short int we accept as a literal is based on 32-bit platforms,
 # so that we can just always emit the same code.
 

mypyc/ir/class_ir.py

@@ -253,6 +253,12 @@ def __init__(
         # Name of the function if this a callable class representing a coroutine.
         self.coroutine_name: str | None = None
 
+        # Does this generator or coroutine helper serialize execution using an instance flag?
+        self.has_running_flag = False
+
+        # Does this generator object contain its merged environment?
+        self.has_merged_generator_env = False
+
     def __repr__(self) -> str:
         return (
             "ClassIR("
@@ -305,6 +311,25 @@ def is_final_attr(self, name: str) -> bool:
                 return False
         return False
 
+    @property
+    def needs_getseters_table(self) -> bool:
+        """Do we generate a tp_getset table exposing the attributes to Python?"""
+        return self.needs_getseters or not self.is_generated or self.has_dict
+
+    def attrs_are_thread_confined(self) -> bool:
+        """Can these attributes safely use plain access in free-threaded builds?
+
+        This requires locals to live directly in the generator object, execution to be
+        serialized by its running flag, and no Python getseters exposing the attributes.
+        A separate environment does not qualify because captured locals may be accessed
+        by nested functions.
+        """
+        return (
+            self.has_merged_generator_env
+            and self.has_running_flag
+            and not self.needs_getseters_table
+        )
+
     def class_final_attr_details(self, name: str) -> tuple[RType, ClassIR] | None:
         """Look up a (possibly inherited) class-body Final attribute.
 
@@ -495,6 +520,8 @@ def serialize(self) -> JsonDict:
             "init_self_leak": self.init_self_leak,
             "env_user_function": self.env_user_function.id if self.env_user_function else None,
             "reuse_freed_instance": self.reuse_freed_instance,
+            "has_running_flag": self.has_running_flag,
+            "has_merged_generator_env": self.has_merged_generator_env,
             "is_acyclic": self.is_acyclic,
             "is_enum": self.is_enum,
             "is_coroutine": self.coroutine_name,
@@ -561,6 +588,8 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
             ctx.functions[data["env_user_function"]] if data["env_user_function"] else None
         )
         ir.reuse_freed_instance = data["reuse_freed_instance"]
+        ir.has_running_flag = data["has_running_flag"]
+        ir.has_merged_generator_env = data["has_merged_generator_env"]
         ir.is_acyclic = data.get("is_acyclic", False)
         ir.is_enum = data["is_enum"]
         ir.coroutine_name = data["is_coroutine"]

mypyc/irbuild/generator.py

@@ -170,8 +170,11 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
     mapper = builder.mapper
     assert isinstance(builder.fn_info.fitem, FuncDef), builder.fn_info.fitem
     generator_class_ir = mapper.fdef_to_generator[builder.fn_info.fitem]
+    generator_class_ir.has_running_flag = True
     if builder.fn_info.can_merge_generator_and_env_classes():
         builder.fn_info.env_class = generator_class_ir
+        # The merged environment can be thread-confined; see attrs_are_thread_confined.
+        generator_class_ir.has_merged_generator_env = True
     else:
         generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)
         if not builder.fn_info.fitem.is_coroutine:

mypyc/lib-rt/CPy.h

@@ -990,6 +990,7 @@ static inline PyObject *CPy_TYPE(PyObject *obj) {
 
 PyObject *CPy_CalculateMetaclass(PyObject *type, PyObject *o);
 PyObject *CPy_GetCoro(PyObject *obj);
+PyObject *CPyGen_AlreadyExecutingError(int is_coroutine);
 PyObject *CPyIter_Send(PyObject *iter, PyObject *val);
 int CPy_YieldFromErrorHandle(PyObject *iter, PyObject **outp);
 PyObject *CPy_FetchStopIterationValue(void);

mypyc/lib-rt/misc_ops.c

@@ -22,6 +22,15 @@ PyObject *CPy_GetCoro(PyObject *obj)
     }
 }
 
+// Raise CPython-compatible errors after a failed CPyGen_TryEnter.
+PyObject *CPyGen_AlreadyExecutingError(int is_coroutine)
+{
+    PyErr_SetString(PyExc_ValueError,
+                    is_coroutine ? "coroutine already executing"
+                                 : "generator already executing");
+    return NULL;
+}
+
 PyObject *CPyIter_Send(PyObject *iter, PyObject *val)
 {
     // Do a send, or a next if second arg is None.

mypyc/lib-rt/pythonsupport.h

@@ -166,6 +166,37 @@ static inline void CPy_SetAttrRef(PyObject **field, PyObject *value) {
 static inline void CPy_InitAttrRef(PyObject **field, PyObject *value) {
     _Py_atomic_store_ptr_relaxed(field, value);
 }
+
+#endif
+
+// Generated generator and coroutine helpers claim this flag while executing,
+// rejecting reentrant or concurrent resumes.
+//
+// On free-threaded builds, the atomic exchange provides mutual exclusion and acquire
+// ordering; the release store publishes body writes to the next resume. This also
+// permits plain access to private generator attributes. Claiming must be a single
+// atomic operation, or two threads could both observe a clear flag and enter. Under
+// the GIL, plain accesses suffice.
+#ifdef Py_GIL_DISABLED
+static inline int CPyGen_TryEnter(uint32_t *running) {
+    return _Py_atomic_exchange_uint32(running, 1) == 0;
+}
+
+static inline void CPyGen_Exit(uint32_t *running) {
+    _Py_atomic_store_uint32_release(running, 0);
+}
+#else
+static inline int CPyGen_TryEnter(uint32_t *running) {
+    if (*running) {
+        return 0;
+    }
+    *running = 1;
+    return 1;
+}
+
+static inline void CPyGen_Exit(uint32_t *running) {
+    *running = 0;
+}
 #endif
 
 PyObject* update_bases(PyObject *bases);

mypyc/test-data/run-generators.test

@@ -1002,3 +1002,75 @@ def test_borrow_across_yield_from() -> None:
             assert False
 
 [typing fixtures/typing-full.pyi]
+
+[case testReentrantResume]
+from typing import Any, Iterator, Optional
+from testutil import assertRaises
+
+box: list[Optional[Iterator[int]]] = [None]
+
+def self_resuming() -> Iterator[int]:
+    g = box[0]
+    assert g is not None
+    yield next(g)
+
+def test_reentrant_next() -> None:
+    g = self_resuming()
+    box[0] = g
+    with assertRaises(ValueError, "generator already executing"):
+        next(g)
+    # The generator was left in a completed state by the propagating exception.
+    with assertRaises(StopIteration):
+        next(g)
+
+cbox: list[Any] = [None]
+
+async def self_sending() -> int:
+    cbox[0].send(None)
+    return 1
+
+def test_reentrant_send_to_coroutine() -> None:
+    c: Any = self_sending()
+    cbox[0] = c
+    with assertRaises(ValueError, "coroutine already executing"):
+        c.send(None)
+
+[case testReentrantResumeWithSeparateEnvironment]
+from typing import Any, Iterator, Optional
+from testutil import assertRaises
+
+def outer() -> int:
+    box: list[Any] = []
+
+    def nested() -> Iterator[int]:
+        # Captures 'box', so the environment is a separate class and the locals
+        # are not private to the generator object. The running flag still applies.
+        yield next(box[0])
+
+    g = nested()
+    box.append(g)
+    return next(g)
+
+def test_reentrant_next_in_nested_generator() -> None:
+    with assertRaises(ValueError, "generator already executing"):
+        outer()
+
+class Base:
+    def gen(self) -> Iterator[int]:
+        yield 1
+
+class Derived(Base):
+    def gen(self) -> Iterator[int]:
+        g = box[0]
+        assert g is not None
+        yield next(g)
+
+box: list[Optional[Iterator[int]]] = [None]
+
+def test_generator_method_override_still_works() -> None:
+    assert list(Base().gen()) == [1]
+    b: Base = Derived()
+    g = b.gen()
+    box[0] = g
+    with assertRaises(ValueError, "generator already executing"):
+        next(g)