diff --git a/mypyc/codegen/emitclass.py b/mypyc/codegen/emitclass.py index f8d3d2696b16..9a1b2efca815 100644 --- a/mypyc/codegen/emitclass.py +++ b/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: diff --git a/mypyc/codegen/emitfunc.py b/mypyc/codegen/emitfunc.py index 78746e143433..0cef3fd78691 100644 --- a/mypyc/codegen/emitfunc.py +++ b/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. diff --git a/mypyc/codegen/emitmodule.py b/mypyc/codegen/emitmodule.py index 9130007e3f6e..01ba7e0554c5 100644 --- a/mypyc/codegen/emitmodule.py +++ b/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): diff --git a/mypyc/common.py b/mypyc/common.py index fa34647c5c72..f76a3f189dc4 100644 --- a/mypyc/common.py +++ b/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. diff --git a/mypyc/ir/class_ir.py b/mypyc/ir/class_ir.py index dd2d0806295b..1124a0ee613d 100644 --- a/mypyc/ir/class_ir.py +++ b/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"] diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py index 555baaada0e8..cd030f82d421 100644 --- a/mypyc/irbuild/generator.py +++ b/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: diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h index e2b4ff0e8c75..e3a8be30952a 100644 --- a/mypyc/lib-rt/CPy.h +++ b/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); diff --git a/mypyc/lib-rt/misc_ops.c b/mypyc/lib-rt/misc_ops.c index 392dba0deca4..621237063ae0 100644 --- a/mypyc/lib-rt/misc_ops.c +++ b/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. diff --git a/mypyc/lib-rt/pythonsupport.h b/mypyc/lib-rt/pythonsupport.h index 33c5a596d174..5fc759b7f8b6 100644 --- a/mypyc/lib-rt/pythonsupport.h +++ b/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); diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test index f4cef93ea8f3..56f342edb354 100644 --- a/mypyc/test-data/run-generators.test +++ b/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)