Nested functions working without free variables (mypyc/mypyc#149)

hughhan1 · msullivan · commit b913793370c8 · 2018-06-21T11:15:40.000-07:00
Nested functions that do not reference free variables now work.
Functions can be multiply nested, and same-named functions in
different scopes won't conflict.

The environment symbol table can now take FuncDef node
types in addition to Var node types.
diff --git a/mypyc/emitclass.py b/mypyc/emitclass.py
@@ -48,6 +48,9 @@ def emit_line() -> None:
         init_name = '0'
         init_args = []
 
+    call_fn = cl.get_method('__call__')
+    call_name = '{}{}'.format(PREFIX, call_fn.cname(emitter.names)) if call_fn else '0'
+
     emitter.emit_line('static PyObject *{}(void);'.format(setup_name))
     # TODO: Use RInstance
     ctor = FuncIR(cl.name, None, module, init_args, object_rprimitive, [], Environment())
@@ -88,7 +91,7 @@ def emit_line() -> None:
             0,                         /* tp_as_sequence */
             0,                         /* tp_as_mapping */
             0,                         /* tp_hash  */
-            0,                         /* tp_call */
+            {tp_call},                 /* tp_call */
             0,                         /* tp_str */
             0,                         /* tp_getattro */
             0,                         /* tp_setattro */
@@ -119,6 +122,7 @@ def emit_line() -> None:
                     traverse_name=traverse_name,
                     clear_name=clear_name,
                     dealloc_name=dealloc_name,
+                    tp_call=call_name,
                     new_name=new_name,
                     methods_name=methods_name,
                     getseters_name=getseters_name,
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -17,8 +17,8 @@ def f(x: int) -> int:
 from typing import Dict, List, Tuple, Optional, Union
 
 from mypy.nodes import (
-    Node, MypyFile, FuncDef, ReturnStmt, AssignmentStmt, OpExpr, IntExpr, NameExpr, LDEF, Var,
-    IfStmt, UnaryExpr, ComparisonExpr, WhileStmt, Argument, CallExpr, IndexExpr, Block,
+    Node, MypyFile, SymbolNode, FuncDef, ReturnStmt, AssignmentStmt, OpExpr, IntExpr, NameExpr,
+    LDEF, Var, IfStmt, UnaryExpr, ComparisonExpr, WhileStmt, Argument, CallExpr, IndexExpr, Block,
     Expression, ListExpr, ExpressionStmt, MemberExpr, ForStmt, RefExpr, Lvalue, BreakStmt,
     ContinueStmt, ConditionalExpr, OperatorAssignmentStmt, TupleExpr, ClassDef, TypeInfo,
     Import, ImportFrom, ImportAll, DictExpr, StrExpr, CastExpr, TempNode, ARG_POS, MODULE_REF,
@@ -235,6 +235,7 @@ def __init__(self,
         self.types = types
         self.environment = Environment()
         self.environments = [self.environment]
+        self.ret_types = []  # type: List[RType]
         self.blocks = []  # type: List[List[BasicBlock]]
         self.functions = []  # type: List[FuncIR]
         self.classes = []  # type: List[ClassIR]
@@ -324,23 +325,45 @@ def visit_import_all(self, node: ImportAll) -> Value:
         return INVALID_VALUE
 
     def gen_func_def(self, fdef: FuncDef, class_name: Optional[str] = None) -> FuncIR:
-        self.enter()
+        # If there is more than one environment in the environment stack, then we are visiting a
+        # non-global function.
+        is_nested = len(self.environments) > 1
 
+        self.enter(fdef.name())
+
+        if is_nested:
+            # If this is a nested function, then add a 'self' field to the environment, since we
+            # will be instantiating the function as a method of a new class representing that
+            # original function.
+            self.environment.add_local(Var('self'), object_rprimitive, is_arg=True)
         for arg in fdef.arguments:
             assert arg.variable.type, "Function argument missing type"
             self.environment.add_local(arg.variable, self.type_to_rtype(arg.variable.type),
                                        is_arg=True)
-        self.ret_type = self.convert_return_type(fdef)
+        self.ret_types[-1] = self.convert_return_type(fdef)
+
         fdef.body.accept(self)
 
-        if is_none_rprimitive(self.ret_type) or is_object_rprimitive(self.ret_type):
+        if (is_none_rprimitive(self.ret_types[-1]) or
+                is_object_rprimitive(self.ret_types[-1])):
             self.add_implicit_return()
         else:
             self.add_implicit_unreachable()
 
-        blocks, env = self.leave()
+        blocks, env, ret_type = self.leave()
         args = self.convert_args(fdef)
-        return FuncIR(fdef.name(), class_name, self.module_name, args, self.ret_type, blocks, env)
+
+        if is_nested:
+            namespace = self.generate_function_namespace()
+            func_ir = self.generate_function_class(fdef, namespace, blocks, env, ret_type)
+
+            # Instantiate the callable class and load it into a register in the current environment
+            # immediately so that it does not have to be loaded every time the function is called.
+            self.instantiate_function_class(fdef, namespace)
+        else:
+            func_ir = FuncIR(fdef.name(), class_name, self.module_name, args, ret_type, blocks,
+                             env)
+        return func_ir
 
     def visit_func_def(self, fdef: FuncDef) -> Value:
         self.functions.append(self.gen_func_def(fdef))
@@ -379,7 +402,7 @@ def visit_expression_stmt(self, stmt: ExpressionStmt) -> Value:
     def visit_return_stmt(self, stmt: ReturnStmt) -> Value:
         if stmt.expr:
             retval = self.accept(stmt.expr)
-            retval = self.coerce(retval, self.ret_type, stmt.line)
+            retval = self.coerce(retval, self.ret_types[-1], stmt.line)
         else:
             retval = self.add(PrimitiveOp([], none_op, line=-1))
         self.add(Return(retval))
@@ -813,15 +836,15 @@ def visit_name_expr(self, expr: NameExpr) -> Value:
         if not self.is_native_name_expr(expr):
             return self.load_static_module_attr(expr)
 
-        # TODO: We assume that this is a Var or FuncDef node, which is very limited
-        if isinstance(expr.node, Var):
-            return self.environment.lookup(expr.node)
-        if isinstance(expr.node, FuncDef):
-            # If we have a function, then we can look it up in the global variables dictionary.
+        # TODO: Behavior currently only defined for Var and FuncDef node types.
+        if expr.kind == LDEF:
+            try:
+                return self.environment.lookup(expr.node)
+            except KeyError:
+                assert False, 'expression %s not defined in current scope'.format(expr.name)
+        else:
             return self.load_global(expr)
 
-        assert False, 'node must be of either Var or FuncDef type'
-
     def is_global_name(self, name: str) -> bool:
         # TODO: this is pretty hokey
         for _, names in self.from_imports.items():
@@ -1264,9 +1287,10 @@ def visit_yield_expr(self, o: YieldExpr) -> Value:
 
     # Helpers
 
-    def enter(self) -> None:
-        self.environment = Environment()
+    def enter(self, name: Optional[str] = None) -> None:
+        self.environment = Environment(name)
         self.environments.append(self.environment)
+        self.ret_types.append(none_rprimitive)
         self.blocks.append([])
         self.new_block()
 
@@ -1282,18 +1306,22 @@ def goto_new_block(self) -> BasicBlock:
         goto.label = block.label
         return block
 
-    def leave(self) -> Tuple[List[BasicBlock], Environment]:
+    def leave(self) -> Tuple[List[BasicBlock], Environment, RType]:
         blocks = self.blocks.pop()
         env = self.environments.pop()
+        ret_type = self.ret_types.pop()
         self.environment = self.environments[-1]
-        return blocks, env
+        return blocks, env, ret_type
 
     def add(self, op: Op) -> Value:
         self.blocks[-1][-1].ops.append(op)
         if isinstance(op, RegisterOp):
             self.environment.add_op(op)
         return op
 
+    def generate_function_namespace(self) -> str:
+        return '_'.join(env.name for env in self.environments if env.name)
+
     def primitive_op(self, desc: OpDescription, args: List[Value], line: int) -> Value:
         assert desc.result_type is not None
         coerced = []
@@ -1345,6 +1373,44 @@ def unbox_or_cast(self, src: Value, target_type: RType, line: int) -> Value:
     def box_expr(self, expr: Expression) -> Value:
         return self.box(self.accept(expr))
 
+    def generate_function_class(self,
+                                fdef: FuncDef,
+                                namespace: str,
+                                blocks: List[BasicBlock],
+                                env: Environment,
+                                ret_type: RType) -> FuncIR:
+        """Generates a callable class representing a nested function.
+
+        This takes a FuncDef and its associated namespace, blocks, environment, and return type and
+        builds a ClassIR with its '__call__' method implemented to represent the function. Note
+        that the name of the function is changed to be '__call__', and a 'self' parameter is added
+        to its list of arguments, as it becomes a class method. The name of the newly constructed
+        class is generated using the names of the functions that enclose the given nested function.
+
+        Returns a newly constructed FuncIR associated with the given FuncDef.
+        """
+        class_name = '{}_{}_obj'.format(fdef.name(), namespace)
+        args = self.convert_args(fdef)
+        args.insert(0, RuntimeArg('self', object_rprimitive))
+        func_ir = FuncIR('__call__', class_name, self.module_name, args, ret_type, blocks, env)
+        class_ir = ClassIR(class_name, self.module_name)
+        class_ir.methods.append(func_ir)
+        self.classes.append(class_ir)
+        return func_ir
+
+    def instantiate_function_class(self, fdef: FuncDef, namespace: str) -> Value:
+        """Assigns a callable class to a register named after the given function definition."""
+        temp_reg = self.load_function_class(fdef, namespace)
+        func_reg = self.environment.add_local(fdef, object_rprimitive)
+        return self.add(Assign(func_reg, temp_reg))
+
+    def load_function_class(self, fdef: FuncDef, namespace: str) -> Value:
+        """Loads a callable class representing a nested function into a register."""
+        return self.add(Call(self.convert_return_type(fdef),
+                             '{}.{}_{}_obj'.format(self.module_name, fdef.name(), namespace),
+                             [],
+                             fdef.line))
+
     def load_global(self, expr: NameExpr) -> Value:
         """Loads a Python-level global.
 
diff --git a/mypyc/ops.py b/mypyc/ops.py
@@ -14,11 +14,11 @@
 import re
 from typing import (
     List, Dict, Generic, TypeVar, Optional, Any, NamedTuple, Tuple, NewType, Callable, Union,
-    Iterable, Type,
+    Iterable, Type, Sequence,
 )
 from collections import OrderedDict
 
-from mypy.nodes import Var
+from mypy.nodes import SymbolNode, Var, FuncDef
 
 from mypyc.namegen import NameGenerator
 
@@ -322,9 +322,10 @@ def __hash__(self) -> int:
 class Environment:
     """Maintain the register symbol table and manage temp generation"""
 
-    def __init__(self) -> None:
+    def __init__(self, name: Optional[str] = None) -> None:
+        self.name = name
         self.indexes = OrderedDict()  # type: Dict[Value, int]
-        self.symtable = {}  # type: Dict[Var, Register]
+        self.symtable = {}  # type: Dict[SymbolNode, Register]
         self.temp_index = 0
 
     def regs(self) -> Iterable['Value']:
@@ -334,16 +335,16 @@ def add(self, reg: 'Value', name: str) -> None:
         reg.name = name
         self.indexes[reg] = len(self.indexes)
 
-    def add_local(self, var: Var, typ: RType, is_arg: bool = False) -> 'Register':
-        assert isinstance(var, Var)
-        reg = Register(typ, var.line, is_arg = is_arg)
+    def add_local(self, symbol: SymbolNode, typ: RType, is_arg: bool = False) -> 'Register':
+        assert isinstance(symbol, SymbolNode)
+        reg = Register(typ, symbol.line, is_arg = is_arg)
 
-        self.symtable[var] = reg
-        self.add(reg, var.name())
+        self.symtable[symbol] = reg
+        self.add(reg, symbol.name())
         return reg
 
-    def lookup(self, var: Var) -> 'Register':
-        return self.symtable[var]
+    def lookup(self, symbol: SymbolNode) -> 'Register':
+        return self.symtable[symbol]
 
     def add_temp(self, typ: RType) -> 'Register':
         assert isinstance(typ, RType)
@@ -663,7 +664,7 @@ class Call(RegisterOp):
     error_kind = ERR_MAGIC
 
     # TODO: take a FuncIR and extract the ret type
-    def __init__(self, ret_type: RType, fn: str, args: List[Value], line: int) -> None:
+    def __init__(self, ret_type: RType, fn: str, args: Sequence[Value], line: int) -> None:
         super().__init__(line)
         self.fn = fn
         self.args = args
@@ -679,7 +680,7 @@ def to_str(self, env: Environment) -> str:
         return s
 
     def sources(self) -> List[Value]:
-        return self.args[:]
+        return list(self.args[:])
 
     def accept(self, visitor: 'OpVisitor[T]') -> T:
         return visitor.visit_call(self)
diff --git a/mypyc/test/test_emitfunc.py b/mypyc/test/test_emitfunc.py
@@ -268,7 +268,8 @@ def setUp(self) -> None:
 
     def test_simple(self) -> None:
         self.block.ops.append(Return(self.reg))
-        fn = FuncIR('myfunc', None, 'mod', [self.arg], int_rprimitive, [self.block], self.env)
+        fn = FuncIR('myfunc', None, 'mod', [self.arg], int_rprimitive, [self.block],
+                    self.env)
         emitter = Emitter(EmitterContext(['mod']))
         generate_native_function(fn, emitter, 'prog.py')
         result = emitter.fragments
@@ -286,7 +287,8 @@ def test_register(self) -> None:
         op = LoadInt(5)
         self.block.ops.append(op)
         self.env.add_op(op)
-        fn = FuncIR('myfunc', None, 'mod', [self.arg], list_rprimitive, [self.block], self.env)
+        fn = FuncIR('myfunc', None, 'mod', [self.arg], list_rprimitive, [self.block],
+                    self.env)
         emitter = Emitter(EmitterContext(['mod']))
         generate_native_function(fn, emitter, 'prog.py')
         result = emitter.fragments
diff --git a/test-data/genops-basic.test b/test-data/genops-basic.test
diff --git a/test-data/run.test b/test-data/run.test
