Merge remote-tracking branch 'upstream/master' into default-native-parser

Author: ilevkivskyi

mypy-requirements.txt

@@ -6,4 +6,4 @@ mypy_extensions>=1.0.0
 pathspec>=1.0.0
 tomli>=1.1.0; python_version<'3.11'
 librt>=0.15.0; platform_python_implementation != 'PyPy'
-ast-serialize>=0.10.0,<1.0.0
+ast-serialize>=0.11.0,<1.0.0

mypyc/codegen/emitmodule.py

@@ -279,11 +279,11 @@ def compile_scc_to_ir(
     if errors.num_errors > 0:
         return modules
 
-    env_user_functions = {}
+    generator_spill_owners = {}
     for module in modules.values():
         for cls in module.classes:
             if cls.env_user_function:
-                env_user_functions[cls.env_user_function] = cls
+                generator_spill_owners[cls.env_user_function] = cls
 
     for module in modules.values():
         module_path = result.graph[module.fullname].xpath
@@ -296,8 +296,8 @@ def compile_scc_to_ir(
                 # Insert reference count handling.
                 insert_ref_count_opcodes(fn)
 
-                if fn in env_user_functions:
-                    insert_spills(fn, env_user_functions[fn])
+                if fn in generator_spill_owners:
+                    insert_spills(fn, generator_spill_owners[fn])
 
                 if compiler_options.log_trace:
                     insert_event_trace_logging(fn, compiler_options)

mypyc/ir/class_ir.py

@@ -233,7 +233,8 @@ def __init__(
         # value of an attribute is the same as the error value.
         self.bitmap_attrs: list[str] = []
 
-        # If this is a generator environment class, what is the actual method for it
+        # If this class owns a generator helper's compiler-generated spill slots, what is the
+        # actual helper method for it.
         self.env_user_function: FuncIR | None = None
 
         # If True, keep one freed, cleared instance available for immediate reuse to
@@ -256,8 +257,10 @@ def __init__(
         # 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
+        # Are this generator object's implementation attributes only accessed while its
+        # running flag is held (or before publication/during destruction)? This also applies
+        # when captured source variables live in a separate environment object.
+        self.has_private_generator_frame = False
 
     def __repr__(self) -> str:
         return (
@@ -319,13 +322,13 @@ def needs_getseters_table(self) -> bool:
     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.
+        This requires the attributes to belong to a private generator frame, execution to
+        be serialized by its running flag, and no Python getseters exposing the attributes.
+        Captured locals in a separate environment don't qualify, since nested functions may
+        access them independently.
         """
         return (
-            self.has_merged_generator_env
+            self.has_private_generator_frame
             and self.has_running_flag
             and not self.needs_getseters_table
         )
@@ -521,7 +524,7 @@ def serialize(self) -> JsonDict:
             "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,
+            "has_private_generator_frame": self.has_private_generator_frame,
             "is_acyclic": self.is_acyclic,
             "is_enum": self.is_enum,
             "is_coroutine": self.coroutine_name,
@@ -589,7 +592,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         )
         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.has_private_generator_frame = data["has_private_generator_frame"]
         ir.is_acyclic = data.get("is_acyclic", False)
         ir.is_enum = data["is_enum"]
         ir.coroutine_name = data["is_coroutine"]

mypyc/irbuild/builder.py

@@ -303,7 +303,7 @@ def __init__(
         # Whether the current top-level expression contains a suspension point
         # (await, yield or yield from). A whole-expression borrow can't span such a
         # point, since the borrowed value (and its root) live in registers that are
-        # not spilled into the generator environment across the suspend.
+        # not spilled into the generator frame across the suspend.
         self.expr_has_suspend = False
         # Saved expression state for enclosing functions (see enter()/leave()).
         self.expression_depth_stack: list[int] = []
@@ -1046,22 +1046,26 @@ def pop_loop_stack(self) -> None:
         self.nonlocal_control.pop()
 
     def make_spill_target(self, type: RType) -> AssignmentTarget:
-        """Moves a given Value instance into the generator class' environment class."""
-        name = f"{TEMP_ATTR_NAME}{self.temp_counter}"
+        """Moves a given Value instance into the private generator frame."""
+        frame = self.fn_info.generator_class
+        # Generator classes for overriding methods can inherit from one another. Include the
+        # module-qualified owning class name so unrelated helper spills don't alias an inherited
+        # struct field.
+        name = f"{TEMP_ATTR_NAME}1_{exported_name(frame.ir.fullname)}_{self.temp_counter}"
         self.temp_counter += 1
-        target = self.add_var_to_env_class(Var(name), type, self.fn_info.generator_class)
+        target = self.add_var_to_class(Var(name), type, frame.ir, frame.self_reg)
         return target
 
     def spill(self, value: Value) -> AssignmentTarget:
-        """Moves a given Value instance into the generator class' environment class."""
+        """Moves a given Value instance into the private generator frame."""
         target = self.make_spill_target(value.type)
         # Shouldn't be able to fail
         self.assign(target, value, NO_TRACEBACK_LINE_NO)
         return target
 
     def maybe_spill(self, value: Value) -> Value | AssignmentTarget:
         """
-        Moves a given Value instance into the environment class for generator functions. For
+        Moves a given Value instance into the private frame for generator functions. For
         non-generator functions, leaves the Value instance as it is.
 
         Returns an AssignmentTarget associated with the Value for generator functions and the
@@ -1073,7 +1077,7 @@ def maybe_spill(self, value: Value) -> Value | AssignmentTarget:
 
     def maybe_spill_assignable(self, value: Value) -> Register | AssignmentTarget:
         """
-        Moves a given Value instance into the environment class for generator functions. For
+        Moves a given Value instance into the private frame for generator functions. For
         non-generator functions, allocate a temporary Register.
 
         Returns an AssignmentTarget associated with the Value for generator functions and an
@@ -1633,24 +1637,45 @@ def add_var_to_env_class(
         keep_alive_on_completion: bool = False,
         prefix: str = "",
     ) -> AssignmentTarget:
-        # First, define the variable name as an attribute of the environment class, and then
-        # construct a target for that attribute.
+        return self.add_var_to_class(
+            var,
+            rtype,
+            self.fn_info.env_class,
+            base.curr_env_reg,
+            reassign=reassign,
+            always_defined=always_defined,
+            keep_alive_on_completion=keep_alive_on_completion,
+            prefix=prefix,
+        )
+
+    def add_var_to_class(
+        self,
+        var: SymbolNode,
+        rtype: RType,
+        cls: ClassIR,
+        base: Value,
+        reassign: bool = False,
+        always_defined: bool = False,
+        keep_alive_on_completion: bool = False,
+        prefix: str = "",
+    ) -> AssignmentTarget:
+        """Declare an attribute on a class and construct a target using an explicit base."""
         name = prefix + remangle_redefinition_name(var.name)
-        self.fn_info.env_class.attributes[name] = rtype
+        cls.attributes[name] = rtype
         if keep_alive_on_completion:
-            self.fn_info.env_class.attrs_to_keep_alive_on_completion.add(name)
+            cls.attrs_to_keep_alive_on_completion.add(name)
         if always_defined:
-            self.fn_info.env_class.attrs_with_defaults.add(name)
-        attr_target = AssignmentTargetAttr(base.curr_env_reg, name)
+            cls.attrs_with_defaults.add(name)
+        attr_target = AssignmentTargetAttr(base, name)
 
         if reassign:
             # Read the local definition of the variable, and set the corresponding attribute of
-            # the environment class' variable to be that value.
+            # the class' variable to be that value.
             reg = self.read(self.lookup(var), self.fn_info.fitem.line)
-            self.add(SetAttr(base.curr_env_reg, name, reg, self.fn_info.fitem.line))
+            self.add(SetAttr(base, name, reg, self.fn_info.fitem.line))
 
         # Override the local definition of the variable to instead point at the variable in
-        # the environment class.
+        # the class.
         return self.add_target(var, attr_target)
 
     def is_builtin_ref_expr(self, expr: RefExpr) -> bool:

mypyc/irbuild/callable_class.py

@@ -78,6 +78,9 @@ class for the nested function.
     # this is a toplevel lambda), don't set up an environment.
     if builder.fn_infos[-2].contains_nested:
         callable_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class)
+        # The link is initialized before the callable is published and is never rebound.
+        # Treating it as Final permits plain loads on free-threaded builds.
+        callable_class_ir.final_attributes.add(ENV_ATTR_NAME)
     callable_class_ir.mro = [callable_class_ir]
     builder.fn_info.callable_class = ImplicitClass(callable_class_ir)
     builder.classes.append(callable_class_ir)
@@ -235,7 +238,11 @@ def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:
     elif builder.fn_info.contains_nested:
         curr_env_reg = builder.fn_info.curr_env_reg
     if curr_env_reg:
-        builder.add(SetAttr(func_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
+        set_env = SetAttr(func_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line)
+        # A new or freelist-reused callable has had all of its fields cleared, and this store
+        # happens before the callable can escape.
+        set_env.mark_as_initializer()
+        builder.add(set_env)
     # Initialize function wrapper for callable classes. As opposed to regular functions,
     # each instance of a callable class needs its own wrapper because they might be instantiated
     # inside other functions.

mypyc/irbuild/env_class.py

@@ -143,7 +143,11 @@ def load_env_registers(builder: IRBuilder, prefix: str = "") -> None:
 
 
 def load_outer_env(
-    builder: IRBuilder, base: Value, outer_env: dict[SymbolNode, SymbolTarget]
+    builder: IRBuilder,
+    base: Value,
+    outer_env: dict[SymbolNode, SymbolTarget],
+    *,
+    borrow: bool = False,
 ) -> Value:
     """Load the environment class for a given base into a register.
 
@@ -156,7 +160,10 @@ def load_outer_env(
 
     Returns the register where the environment class was loaded.
     """
-    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line))
+    if borrow:
+        assert isinstance(base.type, RInstance)
+        assert base.type.class_ir.is_final_attr(ENV_ATTR_NAME)
+    env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line, borrow=borrow))
     assert isinstance(env.type, RInstance), f"{env} must be of type RInstance"
 
     for symbol, target in outer_env.items():
@@ -182,7 +189,9 @@ def load_outer_envs(builder: IRBuilder, base: ImplicitClass) -> None:
         if isinstance(base, GeneratorClass):
             base.prev_env_reg = load_outer_env(builder, base.curr_env_reg, outer_env)
         else:
-            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env)
+            # The callable stays alive throughout __call__, and its environment link is Final,
+            # so the environment can be borrowed for the duration of the call.
+            base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env, borrow=True)
         env_reg = base.prev_env_reg
         index -= 1
 

mypyc/irbuild/for_helpers.py

@@ -705,7 +705,7 @@ def need_cleanup(self) -> bool:
     def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define targets to contain the expression, along with the iterator that will be used
         # for the for-loop. If we are inside of a generator function, spill these into the
-        # environment class.
+        # private generator frame.
         builder = self.builder
         iter_reg = builder.primitive_op(iter_op, [expr_reg], self.line)
         builder.maybe_spill(expr_reg)
@@ -753,7 +753,7 @@ def need_cleanup(self) -> bool:
 
     def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define target to contains the generator expression. It's also the iterator.
-        # If we are inside a generator function, spill these into the environment class.
+        # If we are inside a generator function, spill these into the private generator frame.
         builder = self.builder
         self.iter_target = builder.maybe_spill(expr_reg)
         self.target_type = target_type
@@ -811,7 +811,7 @@ def init(self, expr_reg: Value, target_type: RType) -> None:
         # Define targets to contain the expression, along with the
         # iterator that will be used for the for-loop. We are inside
         # of a generator function, so we will spill these into
-        # environment class.
+        # the private generator frame.
         builder = self.builder
         iter_reg = builder.call_c(aiter_op, [expr_reg], self.line)
         builder.maybe_spill(expr_reg)
@@ -910,7 +910,7 @@ def init(
         self.reverse = reverse
         # Define target to contain the expression, along with the index that will be used
         # for the for-loop. If we are inside of a generator function, spill these into the
-        # environment class.
+        # private generator frame.
         self.expr_target = builder.maybe_spill(expr_reg)
         if is_immutable_rprimitive(expr_reg.type):
             # If the expression is an immutable type, we can load the length just once.
@@ -1011,7 +1011,7 @@ def init(self, expr_reg: Value, target_type: RType) -> None:
         builder = self.builder
         self.target_type = target_type
 
-        # We add some variables to environment class, so they can be read across yield.
+        # Spill some values so they can be read across yield.
         self.expr_target = builder.maybe_spill(expr_reg)
         offset = Integer(0)
         self.offset_target = builder.maybe_spill_assignable(offset)

mypyc/irbuild/generator.py

@@ -52,6 +52,7 @@
 )
 from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup
 from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
+from mypyc.irbuild.targets import AssignmentTargetAttr
 from mypyc.primitives.exc_ops import (
     error_catch_op,
     exc_matches_op,
@@ -142,11 +143,7 @@ def instantiate_generator_class(builder: IRBuilder) -> Value:
     fitem = builder.fn_info.fitem
     generator_reg = builder.add(Call(builder.fn_info.generator_class.ir.ctor, [], fitem.line))
 
-    if builder.fn_info.can_merge_generator_and_env_classes():
-        # Set the generator instance to the initial state (zero).
-        zero = Integer(0)
-        builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
-    else:
+    if not builder.fn_info.can_merge_generator_and_env_classes():
         # Get the current environment register. If the current function is nested, then the
         # generator class gets instantiated from the callable class' '__call__' method, and hence
         # we use the callable class' environment register. Otherwise, we use the original
@@ -160,9 +157,10 @@ def instantiate_generator_class(builder: IRBuilder) -> Value:
         # defined in the current scope.
         builder.add(SetAttr(generator_reg, ENV_ATTR_NAME, curr_env_reg, fitem.line))
 
-        # Set the generator instance's environment to the initial state (zero).
-        zero = Integer(0)
-        builder.add(SetAttr(curr_env_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
+    # The continuation label is private generator state even when captured source variables
+    # require a separate environment.
+    zero = Integer(0)
+    builder.add(SetAttr(generator_reg, NEXT_LABEL_ATTR_NAME, zero, fitem.line))
     return generator_reg
 
 
@@ -171,15 +169,14 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
     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
+    generator_class_ir.has_private_generator_frame = 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:
-            # After completion generators still need generator.__mypyc_env__ for subsequent
-            # __next__() calls to observe the terminal next-label and raise StopIteration.
+            # The helper currently loads generator.__mypyc_env__ before terminal dispatch, so
+            # an exhausted generator still needs the link on subsequent __next__() calls.
             # Coroutines can't be resumed after completion, so keeping the environment alive
             # there would just extend local lifetimes unnecessarily.
             generator_class_ir.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME)
@@ -260,7 +257,9 @@ def add_helper_to_generator_class(
     )
     fn_info.generator_class.ir.methods[GENERATOR_HELPER_NAME] = helper_fn_ir
     builder.functions.append(helper_fn_ir)
-    fn_info.env_class.env_user_function = helper_fn_ir
+    # Compiler-generated values live on the private generator frame even if source-level
+    # captured variables require a separate environment.
+    fn_info.generator_class.ir.env_user_function = helper_fn_ir
 
     return helper_fn_decl
 
@@ -438,12 +437,13 @@ def setup_env_for_generator_class(builder: IRBuilder) -> None:
     else:
         cls.curr_env_reg = load_outer_env(builder, cls.self_reg, builder.symtables[-1])
 
-    # Define a variable representing the label to go to the next time
-    # the '__next__' function of the generator is called, and add it
-    # as an attribute to the environment class.
-    cls.next_label_target = builder.add_var_to_env_class(
-        Var(NEXT_LABEL_ATTR_NAME), int32_rprimitive, cls, reassign=False, always_defined=True
-    )
+    # The continuation label identifies where execution resumes when the generator is next
+    # advanced. Only the serialized generator helper accesses it, so keep it on the private
+    # generator frame instead of a potentially shared closure environment.
+    cls.ir.attributes[NEXT_LABEL_ATTR_NAME] = int32_rprimitive
+    cls.ir.attrs_with_defaults.add(NEXT_LABEL_ATTR_NAME)
+    next_label_target = AssignmentTargetAttr(cls.self_reg, NEXT_LABEL_ATTR_NAME)
+    cls.next_label_target = builder.add_target(Var(NEXT_LABEL_ATTR_NAME), next_label_target)
 
     # Add arguments from the original generator function to the
     # environment of the generator class.