| """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 typing import Dict, Optional, Union |
| |
| from mypy.nodes import FuncDef, SymbolNode |
| |
| from mypyc.common import SELF_NAME, ENV_ATTR_NAME |
| from mypyc.ir.ops import Call, GetAttr, SetAttr, Value |
| from mypyc.ir.rtypes import RInstance, object_rprimitive |
| from mypyc.ir.class_ir import ClassIR |
| from mypyc.irbuild.builder import IRBuilder, SymbolTarget |
| from mypyc.irbuild.targets import AssignmentTargetAttr |
| from mypyc.irbuild.context import FuncInfo, ImplicitClass, GeneratorClass |
| |
| |
| 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('{}_env'.format(builder.fn_info.namespaced_name()), |
| builder.module_name, is_generated=True) |
| env_class.attributes[SELF_NAME] = RInstance(env_class) |
| if builder.fn_info.is_nested: |
| # 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) |
| 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) -> None: |
| """Generate, instantiate, and set up the environment of an environment class.""" |
| instantiate_env_class(builder) |
| |
| # 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. |
| if builder.fn_info.is_nested: |
| add_args_to_env(builder, local=False, base=builder.fn_info.callable_class) |
| else: |
| add_args_to_env(builder, local=False, base=builder.fn_info) |
| |
| |
| 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: |
| 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: |
| builder.fn_info._curr_env_reg = curr_env_reg |
| |
| return curr_env_reg |
| |
| |
| def load_env_registers(builder: IRBuilder) -> 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) |
| |
| fn_info = builder.fn_info |
| fitem = fn_info.fitem |
| if fn_info.is_nested: |
| 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): |
| setup_func_for_recursive_call(builder, fitem, fn_info.callable_class) |
| |
| |
| def load_outer_env(builder: IRBuilder, |
| base: Value, |
| outer_env: Dict[SymbolNode, SymbolTarget]) -> Value: |
| """Load the environment class for a given base into a register. |
| |
| Additionally, iterates through all of the SymbolNode and |
| AssignmentTarget instances of the environment at the given index's |
| symtable, and adds those instances to the environment of the |
| current environment. This is done so that the current environment |
| can access outer environment variables without having to reload |
| all of the environment registers. |
| |
| Returns the register where the environment class was loaded. |
| """ |
| env = builder.add(GetAttr(base, ENV_ATTR_NAME, builder.fn_info.fitem.line)) |
| assert isinstance(env.type, RInstance), '{} must be of type RInstance'.format(env) |
| |
| for symbol, target in outer_env.items(): |
| env.type.class_ir.attributes[symbol.name] = target.type |
| symbol_target = AssignmentTargetAttr(env, symbol.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. |
| if index > 1: |
| # 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: |
| base.prev_env_reg = load_outer_env(builder, base.self_reg, outer_env) |
| env_reg = base.prev_env_reg |
| index -= 1 |
| |
| # Load the remaining outer environments into registers. |
| while index > 1: |
| # 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 add_args_to_env(builder: IRBuilder, |
| local: bool = True, |
| base: Optional[Union[FuncInfo, ImplicitClass]] = None, |
| reassign: bool = True) -> None: |
| fn_info = builder.fn_info |
| if local: |
| for arg in fn_info.fitem.arguments: |
| rtype = builder.type_to_rtype(arg.variable.type) |
| builder.add_local_reg(arg.variable, rtype, is_arg=True) |
| else: |
| for arg in fn_info.fitem.arguments: |
| if is_free_variable(builder, arg.variable) or fn_info.is_generator: |
| 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) |
| |
| |
| 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. |
| """ |
| # First, set the attribute of the environment class so that GetAttr can be called on it. |
| prev_env = builder.fn_infos[-2].env_class |
| prev_env.attributes[fdef.name] = builder.type_to_rtype(fdef.type) |
| |
| 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, -1)) |
| 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, fdef.name, -1)) |
| target = builder.add_local_reg(fdef, object_rprimitive) |
| builder.assign(target, val, -1) |
| |
| |
| 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] |
| ) |