"""Generate classes representing function environments (+ related operations). If we have a nested function that has non-local (free) variables, access to the non-locals is via an instance of an environment class. Example: def f() -> int: x = 0 # Make 'x' an attribute of an environment class instance def g() -> int: # We have access to the environment class instance to # allow accessing 'x' return x + 2 x = x + 1 # Modify the attribute return g() """ from __future__ import annotations from mypy.nodes import Argument, FuncDef, SymbolNode, Var from mypyc.common import ( BITMAP_BITS, ENV_ATTR_NAME, GENERATOR_ATTRIBUTE_PREFIX, SELF_NAME, bitmap_name, ) from mypyc.ir.class_ir import ClassIR from mypyc.ir.ops import Call, GetAttr, SetAttr, Value from mypyc.ir.rtypes import RInstance, bitmap_rprimitive, object_rprimitive from mypyc.irbuild.builder import IRBuilder, SymbolTarget from mypyc.irbuild.context import FuncInfo, GeneratorClass, ImplicitClass from mypyc.irbuild.targets import AssignmentTargetAttr from mypyc.irbuild.util import get_func_target_symbol def setup_env_class(builder: IRBuilder) -> ClassIR: """Generate a class representing a function environment. Note that the variables in the function environment are not actually populated here. This is because when the environment class is generated, the function environment has not yet been visited. This behavior is allowed so that when the compiler visits nested functions, it can use the returned ClassIR instance to figure out free variables it needs to access. The remaining attributes of the environment class are populated when the environment registers are loaded. Return a ClassIR representing an environment for a function containing a nested function. """ env_class = ClassIR( f"{builder.fn_info.namespaced_name()}_env", builder.module_name, is_generated=True, is_final_class=True, ) env_class.reuse_freed_instance = True env_class.attributes[SELF_NAME] = RInstance(env_class) if builder.fn_info.is_nested and builder.fn_infos[-2]._env_class is not None: # If the function is nested, its environment class must contain an environment # attribute pointing to its encapsulating functions' environment class. env_class.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_infos[-2].env_class) if builder.fn_info.contains_nested: env_class.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME) env_class.mro = [env_class] builder.fn_info.env_class = env_class builder.classes.append(env_class) return env_class def finalize_env_class(builder: IRBuilder, prefix: str = "", *, add_args: bool = True) -> Value: """Generate, instantiate, and set up the environment of an environment class.""" if not builder.fn_info.can_merge_generator_and_env_classes(): env_reg = instantiate_env_class(builder) else: env_reg = builder.fn_info.curr_env_reg # Iterate through the function arguments and replace local definitions (using registers) # that were previously added to the environment with references to the function's # environment class. Comprehension scopes have no arguments to add. if add_args and not builder.fn_info.is_comprehension_scope: if builder.fn_info.is_nested: add_args_to_env( builder, local=False, base=builder.fn_info.callable_class, prefix=prefix ) else: add_args_to_env(builder, local=False, base=builder.fn_info, prefix=prefix) return env_reg def instantiate_env_class(builder: IRBuilder) -> Value: """Assign an environment class to a register named after the given function definition.""" curr_env_reg = builder.add( Call(builder.fn_info.env_class.ctor, [], builder.fn_info.fitem.line) ) if builder.fn_info.is_nested and not builder.fn_info.is_comprehension_scope: builder.fn_info.callable_class._curr_env_reg = curr_env_reg builder.add( SetAttr( curr_env_reg, ENV_ATTR_NAME, builder.fn_info.callable_class.prev_env_reg, builder.fn_info.fitem.line, ) ) else: # Top-level functions and comprehension scopes store env reg directly. builder.fn_info._curr_env_reg = curr_env_reg # Comprehension scopes link to parent env if it exists. parent = builder.fn_infos[-2] if builder.fn_info.is_nested else None if parent is not None and parent._env_class is not None: if parent.is_generator: parent_env_reg = parent.generator_class.curr_env_reg else: parent_env_reg = parent._curr_env_reg else: parent_env_reg = None if parent_env_reg is not None: builder.add( SetAttr(curr_env_reg, ENV_ATTR_NAME, parent_env_reg, builder.fn_info.fitem.line) ) return curr_env_reg def load_env_registers(builder: IRBuilder, prefix: str = "") -> None: """Load the registers for the current FuncItem being visited. Adds the arguments of the FuncItem to the environment. If the FuncItem is nested inside of another function, then this also loads all of the outer environments of the FuncItem into registers so that they can be used when accessing free variables. """ add_args_to_env(builder, local=True, prefix=prefix) fn_info = builder.fn_info fitem = fn_info.fitem if fn_info.is_nested and builder.fn_infos[-2]._env_class is not None: load_outer_envs(builder, fn_info.callable_class) # If this is a FuncDef, then make sure to load the FuncDef into its own environment # class so that the function can be called recursively. if isinstance(fitem, FuncDef) and fn_info.add_nested_funcs_to_env: setup_func_for_recursive_call(builder, fitem, fn_info.callable_class) def load_outer_env( builder: IRBuilder, base: Value, outer_env: dict[SymbolNode, SymbolTarget], *, borrow: bool = False, ) -> Value: """Load the environment referenced by "base" and bind its symbols. "outer_env" may contain registers and attributes from multiple environment levels. Bind only attributes stored directly on the loaded environment, since the other targets cannot be accessed through it. If "borrow" is true, return a borrowed reference. In that case the environment link on "base" must be final so that the reference remains valid. """ 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(): # Only targets actually stored in this environment are visible through it. if ( not isinstance(target, AssignmentTargetAttr) or not isinstance(target.obj.type, RInstance) or target.obj.type.class_ir is not env.type.class_ir ): continue attr_name = target.attr env.type.class_ir.attributes[attr_name] = target.type symbol_target = AssignmentTargetAttr(env, attr_name) builder.add_target(symbol, symbol_target) return env def load_outer_envs(builder: IRBuilder, base: ImplicitClass) -> None: index = len(builder.builders) - 2 # Load the first outer environment. This one is special because it gets saved in the # FuncInfo instance's prev_env_reg field. has_outer = index > 1 or (index == 1 and builder.fn_infos[1].contains_nested) if has_outer and builder.fn_infos[index]._env_class is not None: # outer_env = builder.fn_infos[index].environment outer_env = builder.symtables[index] if isinstance(base, GeneratorClass): base.prev_env_reg = load_outer_env(builder, base.curr_env_reg, outer_env) else: # 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 # Load the remaining outer environments into registers. while index > 1: if builder.fn_infos[index]._env_class is None: break # outer_env = builder.fn_infos[index].environment outer_env = builder.symtables[index] env_reg = load_outer_env(builder, env_reg, outer_env) index -= 1 def num_bitmap_args(builder: IRBuilder, args: list[Argument]) -> int: n = 0 for arg in args: t = builder.type_to_rtype(arg.variable.type) if t.error_overlap and arg.kind.is_optional(): n += 1 return (n + (BITMAP_BITS - 1)) // BITMAP_BITS def add_args_to_env( builder: IRBuilder, local: bool = True, base: FuncInfo | ImplicitClass | None = None, reassign: bool = True, prefix: str = "", ) -> None: fn_info = builder.fn_info args = fn_info.fitem.arguments if local: nb = num_bitmap_args(builder, args) for arg in args: rtype = builder.type_to_rtype(arg.variable.type) builder.add_local_reg(arg.variable, rtype, is_arg=True) for i in reversed(range(nb)): builder.add_local_reg(Var(bitmap_name(i)), bitmap_rprimitive, is_arg=True) else: for arg in args: if is_free_variable(builder, arg.variable): rtype = builder.type_to_rtype(arg.variable.type) assert base is not None, "base cannot be None for adding nonlocal args" builder.add_var_to_env_class( arg.variable, rtype, base, reassign=reassign, keep_alive_on_completion=builder.is_captured_by_nested_func(arg.variable), prefix=prefix, ) def add_generator_args( builder: IRBuilder, frame_reg: Value, env_reg: Value, reassign: bool ) -> None: """Put generator arguments on either the closure environment or private frame.""" for arg in builder.fn_info.fitem.arguments: rtype = builder.type_to_rtype(arg.variable.type) if builder.is_captured_by_nested_func(arg.variable): builder.add_var_to_class( arg.variable, rtype, builder.fn_info.env_class, env_reg, reassign=reassign, keep_alive_on_completion=True, prefix=GENERATOR_ATTRIBUTE_PREFIX, ) else: builder.add_var_to_generator_frame(arg.variable, rtype, frame_reg, reassign=reassign) def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None: """Add relevant local variables and nested functions to persistent storage. Captured bindings go in the environment class. Generator-owned bindings that aren't captured go on the private generator frame. This is done before visiting the function body so nested functions can reference declarations that occur later in the body. """ fn_info = builder.fn_info env_for_func: FuncInfo | ImplicitClass = fn_info if fn_info.is_generator: env_for_func = fn_info.generator_class elif (fn_info.is_nested or fn_info.in_non_ext) and not fn_info.is_comprehension_scope: env_for_func = fn_info.callable_class if fn_info.fitem in builder.free_variables: # Sort the variables to keep things deterministic for var in sorted(builder.free_variables[fn_info.fitem], key=lambda x: x.name): if isinstance(var, Var): builder.add_var_to_env_class( var, builder.type_to_rtype(var.type), env_for_func, reassign=False, keep_alive_on_completion=True, prefix=prefix, ) if fn_info.fitem in builder.encapsulating_funcs: added: set[SymbolNode] = set() for nested_fn in builder.encapsulating_funcs[fn_info.fitem]: if isinstance(nested_fn, FuncDef): target_symbol = get_func_target_symbol(nested_fn) if target_symbol in added: continue added.add(target_symbol) # The return type is 'object' instead of an RInstance of the # callable class because differently defined functions with # the same name and signature across conditional blocks # will generate different callable classes, so the callable # class that gets instantiated must be generic. nested_prefix = prefix if nested_fn.is_generator or nested_fn.is_coroutine: nested_prefix = GENERATOR_ATTRIBUTE_PREFIX keep_alive = is_free_variable(builder, target_symbol) if fn_info.is_generator and not builder.is_captured_by_nested_func(target_symbol): builder.add_var_to_generator_frame( target_symbol, object_rprimitive, fn_info.generator_class.self_reg, reassign=False, keep_alive_on_completion=keep_alive, ) else: builder.add_var_to_env_class( target_symbol, object_rprimitive, env_for_func, reassign=False, keep_alive_on_completion=keep_alive, prefix=nested_prefix, ) def setup_func_for_recursive_call(builder: IRBuilder, fdef: FuncDef, base: ImplicitClass) -> None: """Enable calling a nested function (with a callable class) recursively. Adds the instance of the callable class representing the given FuncDef to a register in the environment so that the function can be called recursively. Note that this needs to be done only for nested functions. """ prev_env = builder.fn_infos[-2].env_class target_symbol = get_func_target_symbol(fdef) target = builder.symtables[-2].get(target_symbol) if ( not isinstance(target_symbol, FuncDef) or not isinstance(target, AssignmentTargetAttr) or not isinstance(target.obj.type, RInstance) or target.obj.type.class_ir is not prev_env ): return attr_name = target.attr line = fdef.line if isinstance(base, GeneratorClass): # If we are dealing with a generator class, then we need to first get the register # holding the current environment class, and load the previous environment class from # there. prev_env_reg = builder.add(GetAttr(base.curr_env_reg, ENV_ATTR_NAME, line)) else: prev_env_reg = base.prev_env_reg # Obtain the instance of the callable class representing the FuncDef, and add it to the # current environment. val = builder.add(GetAttr(prev_env_reg, attr_name, line)) target = builder.add_local_reg(target_symbol, object_rprimitive) builder.assign(target, val, line) def is_free_variable(builder: IRBuilder, symbol: SymbolNode) -> bool: fitem = builder.fn_info.fitem return fitem in builder.free_variables and symbol in builder.free_variables[fitem]