init

.venv/lib/python3.8/site-packages/mypyc/irbuild/env_class.py

@@ -0,0 +1,207 @@
+"""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]
+    )