[mypyc] Refactor: extract code from mypyc.irbuild.function (#8508)

This extracts these things to new modules:

* Code related to generator functions (not all of it, since some was hard to extract)
* Code related to callable classes
* Code related to environment classes

Work on mypyc/mypyc#714.

Author: JukkaL

mypyc/irbuild/builder.py

@@ -64,6 +64,7 @@
 from mypyc.irbuild.context import FuncInfo, ImplicitClass
 from mypyc.irbuild.mapper import Mapper
 from mypyc.irbuild.ll_builder import LowLevelIRBuilder
+from mypyc.irbuild.util import is_constant
 
 GenFunc = Callable[[], None]
 
@@ -1187,3 +1188,36 @@ def warning(self, msg: str, line: int) -> None:
 
     def error(self, msg: str, line: int) -> None:
         self.errors.error(msg, self.module_path, line)
+
+
+def gen_arg_defaults(builder: IRBuilder) -> None:
+    """Generate blocks for arguments that have default values.
+
+    If the passed value is an error value, then assign the default
+    value to the argument.
+    """
+    fitem = builder.fn_info.fitem
+    for arg in fitem.arguments:
+        if arg.initializer:
+            target = builder.environment.lookup(arg.variable)
+
+            def get_default() -> Value:
+                assert arg.initializer is not None
+
+                # If it is constant, don't bother storing it
+                if is_constant(arg.initializer):
+                    return builder.accept(arg.initializer)
+
+                # Because gen_arg_defaults runs before calculate_arg_defaults, we
+                # add the static/attribute to final_names/the class here.
+                elif not builder.fn_info.is_nested:
+                    name = fitem.fullname + '.' + arg.variable.name
+                    builder.final_names.append((name, target.type))
+                    return builder.add(LoadStatic(target.type, name, builder.module_name))
+                else:
+                    name = arg.variable.name
+                    builder.fn_info.callable_class.ir.attributes[name] = target.type
+                    return builder.add(
+                        GetAttr(builder.fn_info.callable_class.self_reg, name, arg.line))
+            assert isinstance(target, AssignmentTargetRegister)
+            builder.assign_if_null(target, get_default, arg.initializer.line)

mypyc/irbuild/callable_class.py

@@ -0,0 +1,161 @@
+"""Generate a class that represents a nested function.
+
+The class defines __call__ for calling the function and allows access to variables
+defined in outer scopes.
+"""
+
+from typing import List
+
+from mypy.nodes import Var
+
+from mypyc.common import SELF_NAME, ENV_ATTR_NAME
+from mypyc.ir.ops import BasicBlock, Return, Call, SetAttr, Value, Environment
+from mypyc.ir.rtypes import RInstance, object_rprimitive
+from mypyc.ir.func_ir import FuncIR, FuncSignature, RuntimeArg, FuncDecl
+from mypyc.ir.class_ir import ClassIR
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.context import FuncInfo, ImplicitClass
+from mypyc.irbuild.util import add_self_to_env
+from mypyc.primitives.misc_ops import method_new_op
+
+
+def setup_callable_class(builder: IRBuilder) -> None:
+    """Generates a callable class representing a nested function or a function within a
+    non-extension class and sets up the 'self' variable for that class.
+
+    This takes the most recently visited function and returns a ClassIR to represent that
+    function. Each callable class contains an environment attribute with points to another
+    ClassIR representing the environment class where some of its variables can be accessed.
+    Note that its '__call__' method is not yet implemented, and is implemented in the
+    add_call_to_callable_class function.
+
+    Returns a newly constructed ClassIR representing the callable class for the nested
+    function.
+    """
+
+    # Check to see that the name has not already been taken. If so, rename the class. We allow
+    # multiple uses of the same function name because this is valid in if-else blocks. Example:
+    #     if True:
+    #         def foo():          ---->    foo_obj()
+    #             return True
+    #     else:
+    #         def foo():          ---->    foo_obj_0()
+    #             return False
+    name = base_name = '{}_obj'.format(builder.fn_info.namespaced_name())
+    count = 0
+    while name in builder.callable_class_names:
+        name = base_name + '_' + str(count)
+        count += 1
+    builder.callable_class_names.add(name)
+
+    # Define the actual callable class ClassIR, and set its environment to point at the
+    # previously defined environment class.
+    callable_class_ir = ClassIR(name, builder.module_name, is_generated=True)
+
+    # The functools @wraps decorator attempts to call setattr on nested functions, so
+    # we create a dict for these nested functions.
+    # https://github.com/python/cpython/blob/3.7/Lib/functools.py#L58
+    if builder.fn_info.is_nested:
+        callable_class_ir.has_dict = True
+
+    # If the enclosing class doesn't contain nested (which will happen if
+    # 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
+        )
+    callable_class_ir.mro = [callable_class_ir]
+    builder.fn_info.callable_class = ImplicitClass(callable_class_ir)
+    builder.classes.append(callable_class_ir)
+
+    # Add a 'self' variable to the callable class' environment, and store that variable in a
+    # register to be accessed later.
+    self_target = add_self_to_env(builder.environment, callable_class_ir)
+    builder.fn_info.callable_class.self_reg = builder.read(self_target, builder.fn_info.fitem.line)
+
+
+def add_call_to_callable_class(builder: IRBuilder,
+                               blocks: List[BasicBlock],
+                               sig: FuncSignature,
+                               env: Environment,
+                               fn_info: FuncInfo) -> FuncIR:
+    """Generates a '__call__' method for a callable class representing a nested function.
+
+    This takes the blocks, signature, and environment associated with a function definition and
+    uses those to build the '__call__' method of a given callable class, used to represent that
+    function. Note that a 'self' parameter is added to its list of arguments, as the nested
+    function becomes a class method.
+    """
+    sig = FuncSignature((RuntimeArg(SELF_NAME, object_rprimitive),) + sig.args, sig.ret_type)
+    call_fn_decl = FuncDecl('__call__', fn_info.callable_class.ir.name, builder.module_name, sig)
+    call_fn_ir = FuncIR(call_fn_decl, blocks, env,
+                        fn_info.fitem.line, traceback_name=fn_info.fitem.name)
+    fn_info.callable_class.ir.methods['__call__'] = call_fn_ir
+    return call_fn_ir
+
+
+def add_get_to_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
+    """Generates the '__get__' method for a callable class."""
+    line = fn_info.fitem.line
+    builder.enter(fn_info)
+
+    vself = builder.read(
+        builder.environment.add_local_reg(Var(SELF_NAME), object_rprimitive, True)
+    )
+    instance = builder.environment.add_local_reg(Var('instance'), object_rprimitive, True)
+    builder.environment.add_local_reg(Var('owner'), object_rprimitive, True)
+
+    # If accessed through the class, just return the callable
+    # object. If accessed through an object, create a new bound
+    # instance method object.
+    instance_block, class_block = BasicBlock(), BasicBlock()
+    comparison = builder.binary_op(
+        builder.read(instance), builder.none_object(), 'is', line
+    )
+    builder.add_bool_branch(comparison, class_block, instance_block)
+
+    builder.activate_block(class_block)
+    builder.add(Return(vself))
+
+    builder.activate_block(instance_block)
+    builder.add(Return(builder.primitive_op(method_new_op, [vself, builder.read(instance)], line)))
+
+    blocks, env, _, fn_info = builder.leave()
+
+    sig = FuncSignature((RuntimeArg(SELF_NAME, object_rprimitive),
+                         RuntimeArg('instance', object_rprimitive),
+                         RuntimeArg('owner', object_rprimitive)),
+                        object_rprimitive)
+    get_fn_decl = FuncDecl('__get__', fn_info.callable_class.ir.name, builder.module_name, sig)
+    get_fn_ir = FuncIR(get_fn_decl, blocks, env)
+    fn_info.callable_class.ir.methods['__get__'] = get_fn_ir
+    builder.functions.append(get_fn_ir)
+
+
+def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:
+    """
+    Assigns a callable class to a register named after the given function definition. Note
+    that fn_info refers to the function being assigned, whereas builder.fn_info refers to the
+    function encapsulating the function being turned into a callable class.
+    """
+    fitem = fn_info.fitem
+    func_reg = builder.add(Call(fn_info.callable_class.ir.ctor, [], fitem.line))
+
+    # Set the callable class' environment attribute to point at the environment class
+    # defined in the callable class' immediate outer scope. Note that there are three possible
+    # environment class registers we may use. If the encapsulating function is:
+    # - a generator function, then the callable class is instantiated from the generator class'
+    #   __next__' function, and hence the generator class' environment register is used.
+    # - a nested function, then the callable class is instantiated from the current callable
+    #   class' '__call__' function, and hence the callable class' environment register is used.
+    # - neither, then we use the environment register of the original function.
+    curr_env_reg = None
+    if builder.fn_info.is_generator:
+        curr_env_reg = builder.fn_info.generator_class.curr_env_reg
+    elif builder.fn_info.is_nested:
+        curr_env_reg = builder.fn_info.callable_class.curr_env_reg
+    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))
+    return func_reg

mypyc/irbuild/env_class.py

@@ -0,0 +1,182 @@
+"""Generate classes representing function environments (+ related operations)."""
+
+from typing import 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, Environment, AssignmentTargetAttr
+from mypyc.ir.rtypes import RInstance, object_rprimitive
+from mypyc.ir.class_ir import ClassIR
+from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.context import FuncInfo, ImplicitClass, GeneratorClass
+
+
+def setup_env_class(builder: IRBuilder) -> ClassIR:
+    """Generates 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.
+
+    Returns 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:
+    """Generates, instantiates, and sets 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:
+    """Assigns 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:
+    """Loads 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: Environment) -> Value:
+    """Loads 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.symtable.items():
+        env.type.class_ir.attributes[symbol.name] = target.type
+        symbol_target = AssignmentTargetAttr(env, symbol.name)
+        builder.environment.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.builders[index].environment
+        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.builders[index].environment
+        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.environment.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:
+    """
+    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.environment.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]
+    )