[mypyc] Refactor method IR generation (#9846)

Add new helpers for generating methods that reduces the required
boilerplate. Switch several existing use cases to use the new helpers.

Summary of the new helpers:

1. Call `builder.enter_method(...)` to begin generating IR for a method.
2. Call `builder.add_argument(..)` for each non-self argument.
3. Generate method body normally.
4. Call `builder.leave_method()`.

Previously `enter()` and `leave()` had to be used, along with several
additional steps, such as defining `self` explicitly. The new approach 
is much easier to use and less error-prone. 

Since not all uses of the old interface have been removed, there is 
still some duplication. It would be great to eventually always use the 
higher-level interface to generate methods and functions.

Work on mypyc/mypyc#781.

mypyc/irbuild/builder.py

@@ -41,7 +41,7 @@
     none_rprimitive, is_none_rprimitive, object_rprimitive, is_object_rprimitive,
     str_rprimitive, is_tagged, is_list_rprimitive, is_tuple_rprimitive, c_pyssize_t_rprimitive
 )
-from mypyc.ir.func_ir import FuncIR, INVALID_FUNC_DEF
+from mypyc.ir.func_ir import FuncIR, INVALID_FUNC_DEF, RuntimeArg, FuncSignature, FuncDecl
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.primitives.registry import CFunctionDescription, function_ops
 from mypyc.primitives.list_ops import to_list, list_pop_last, list_get_item_unsafe_op
@@ -81,7 +81,9 @@ def __init__(self,
         self.builder = LowLevelIRBuilder(current_module, mapper)
         self.builders = [self.builder]
         self.symtables = [OrderedDict()]  # type: List[OrderedDict[SymbolNode, AssignmentTarget]]
-        self.args = [[]]  # type: List[List[Register]]
+        self.runtime_args = [[]]  # type: List[List[RuntimeArg]]
+        self.function_name_stack = []  # type: List[str]
+        self.class_ir_stack = []  # type: List[ClassIR]
 
         self.current_module = current_module
         self.mapper = mapper
@@ -178,6 +180,9 @@ def activate_block(self, block: BasicBlock) -> None:
     def goto_and_activate(self, block: BasicBlock) -> None:
         self.builder.goto_and_activate(block)
 
+    def self(self) -> Register:
+        return self.builder.self()
+
     def py_get_attr(self, obj: Value, attr: str, line: int) -> Value:
         return self.builder.py_get_attr(obj, attr, line)
 
@@ -289,13 +294,13 @@ def gen_import(self, id: str, line: int) -> None:
         self.add(InitStatic(value, id, namespace=NAMESPACE_MODULE))
         self.goto_and_activate(out)
 
-    def assign_if_null(self, target: AssignmentTargetRegister,
+    def assign_if_null(self, target: Register,
                        get_val: Callable[[], Value], line: int) -> None:
-        """Generate blocks for registers that NULL values."""
+        """If target is NULL, assign value produced by get_val to it."""
         error_block, body_block = BasicBlock(), BasicBlock()
-        self.add(Branch(target.register, error_block, body_block, Branch.IS_ERROR))
+        self.add(Branch(target, error_block, body_block, Branch.IS_ERROR))
         self.activate_block(error_block)
-        self.add(Assign(target.register, self.coerce(get_val(), target.register.type, line)))
+        self.add(Assign(target, self.coerce(get_val(), target.type, line)))
         self.goto(body_block)
         self.activate_block(body_block)
 
@@ -897,7 +902,7 @@ def enter(self, fn_info: Union[FuncInfo, str] = '') -> None:
         self.builder = LowLevelIRBuilder(self.current_module, self.mapper)
         self.builders.append(self.builder)
         self.symtables.append(OrderedDict())
-        self.args.append([])
+        self.runtime_args.append([])
         self.fn_info = fn_info
         self.fn_infos.append(self.fn_info)
         self.ret_types.append(none_rprimitive)
@@ -907,16 +912,71 @@ def enter(self, fn_info: Union[FuncInfo, str] = '') -> None:
             self.nonlocal_control.append(BaseNonlocalControl())
         self.activate_block(BasicBlock())
 
-    def leave(self) -> Tuple[List[Register], List[BasicBlock], RType, FuncInfo]:
+    def leave(self) -> Tuple[List[Register], List[RuntimeArg], List[BasicBlock], RType, FuncInfo]:
         builder = self.builders.pop()
         self.symtables.pop()
-        args = self.args.pop()
+        runtime_args = self.runtime_args.pop()
         ret_type = self.ret_types.pop()
         fn_info = self.fn_infos.pop()
         self.nonlocal_control.pop()
         self.builder = self.builders[-1]
         self.fn_info = self.fn_infos[-1]
-        return args, builder.blocks, ret_type, fn_info
+        return builder.args, runtime_args, builder.blocks, ret_type, fn_info
+
+    def enter_method(self,
+                     class_ir: ClassIR,
+                     name: str,
+                     ret_type: RType,
+                     fn_info: Union[FuncInfo, str] = '',
+                     self_type: Optional[RType] = None) -> None:
+        """Begin generating IR for a method.
+
+        If the method takes arguments, you should immediately afterwards call
+        add_argument() for each non-self argument (self is created implicitly).
+
+        Call leave_method() to finish the generation of the method.
+
+        You can enter multiple methods at a time. They are maintained in a
+        stack, and leave_method() leaves the topmost one.
+
+        Args:
+            class_ir: Add method to this class
+            name: Short name of the method
+            ret_type: Return type of the method
+            fn_info: Optionally, additional information about the method
+            self_type: If not None, override default type of the implicit 'self'
+                argument (by default, derive type from class_ir)
+        """
+        self.enter(fn_info)
+        self.function_name_stack.append(name)
+        self.class_ir_stack.append(class_ir)
+        self.ret_types[-1] = ret_type
+        if self_type is None:
+            self_type = RInstance(class_ir)
+        self.add_argument(SELF_NAME, self_type)
+
+    def add_argument(self, var: Union[str, Var], typ: RType, kind: int = ARG_POS) -> Register:
+        """Declare an argument in the current function.
+
+        You should use this instead of directly calling add_local() in new code.
+        """
+        if isinstance(var, str):
+            var = Var(var)
+        reg = self.add_local(var, typ, is_arg=True)
+        self.runtime_args[-1].append(RuntimeArg(var.name, typ, kind))
+        return reg
+
+    def leave_method(self) -> None:
+        """Finish the generation of IR for a method."""
+        arg_regs, args, blocks, ret_type, fn_info = self.leave()
+        sig = FuncSignature(args, ret_type)
+        name = self.function_name_stack.pop()
+        class_ir = self.class_ir_stack.pop()
+        decl = FuncDecl(name, class_ir.name, self.module_name, sig)
+        ir = FuncIR(decl, arg_regs, blocks)
+        class_ir.methods[name] = ir
+        class_ir.method_decls[name] = ir.decl
+        self.functions.append(ir)
 
     def lookup(self, symbol: SymbolNode) -> AssignmentTarget:
         return self.symtables[-1][symbol]
@@ -931,7 +991,7 @@ def add_local(self, symbol: SymbolNode, typ: RType, is_arg: bool = False) -> 'Re
         reg = Register(typ, symbol.name, is_arg=is_arg, line=symbol.line)
         self.symtables[-1][symbol] = AssignmentTargetRegister(reg)
         if is_arg:
-            self.args[-1].append(reg)
+            self.builder.args.append(reg)
         return reg
 
     def add_local_reg(self,
@@ -945,6 +1005,10 @@ def add_local_reg(self,
         return target
 
     def add_self_to_env(self, cls: ClassIR) -> AssignmentTargetRegister:
+        """Low-level function that adds a 'self' argument.
+
+        This is only useful if using enter() instead of enter_method().
+        """
         return self.add_local_reg(Var(SELF_NAME), RInstance(cls), is_arg=True)
 
     def add_target(self, symbol: SymbolNode, target: AssignmentTarget) -> AssignmentTarget:
@@ -1057,4 +1121,4 @@ def get_default() -> Value:
                     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)
+            builder.assign_if_null(target.register, get_default, arg.initializer.line)

mypyc/irbuild/callable_class.py

@@ -6,8 +6,6 @@
 
 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, Register
 from mypyc.ir.rtypes import RInstance, object_rprimitive
@@ -105,13 +103,13 @@ def add_call_to_callable_class(builder: IRBuilder,
 def add_get_to_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
     """Generate the '__get__' method for a callable class."""
     line = fn_info.fitem.line
-    builder.enter(fn_info)
 
-    vself = builder.read(
-        builder.add_local_reg(Var(SELF_NAME), object_rprimitive, True)
+    builder.enter_method(
+        fn_info.callable_class.ir, '__get__', object_rprimitive, fn_info,
+        self_type=object_rprimitive
     )
-    instance = builder.add_local_reg(Var('instance'), object_rprimitive, True)
-    builder.add_local_reg(Var('owner'), object_rprimitive, True)
+    instance = builder.add_argument('instance', object_rprimitive)
+    builder.add_argument('owner', object_rprimitive)
 
     # If accessed through the class, just return the callable
     # object. If accessed through an object, create a new bound
@@ -123,21 +121,13 @@ def add_get_to_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> None:
     builder.add_bool_branch(comparison, class_block, instance_block)
 
     builder.activate_block(class_block)
-    builder.add(Return(vself))
+    builder.add(Return(builder.self()))
 
     builder.activate_block(instance_block)
-    builder.add(Return(builder.call_c(method_new_op, [vself, builder.read(instance)], line)))
-
-    args, blocks, _, 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, args, blocks)
-    fn_info.callable_class.ir.methods['__get__'] = get_fn_ir
-    builder.functions.append(get_fn_ir)
+    builder.add(Return(builder.call_c(method_new_op,
+                                      [builder.self(), builder.read(instance)], line)))
+
+    builder.leave_method()
 
 
 def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:

mypyc/irbuild/classdef.py

@@ -4,25 +4,24 @@
 
 from mypy.nodes import (
     ClassDef, FuncDef, OverloadedFuncDef, PassStmt, AssignmentStmt, NameExpr, StrExpr,
-    ExpressionStmt, TempNode, Decorator, Lvalue, RefExpr, Var, is_class_var
+    ExpressionStmt, TempNode, Decorator, Lvalue, RefExpr, is_class_var
 )
 from mypyc.ir.ops import (
     Value, Register, Call, LoadErrorValue, LoadStatic, InitStatic, TupleSet, SetAttr, Return,
     BasicBlock, Branch, MethodCall, NAMESPACE_TYPE, LoadAddress
 )
 from mypyc.ir.rtypes import (
-    RInstance, object_rprimitive, bool_rprimitive, dict_rprimitive, is_optional_type,
+    object_rprimitive, bool_rprimitive, dict_rprimitive, is_optional_type,
     is_object_rprimitive, is_none_rprimitive
 )
-from mypyc.ir.func_ir import FuncIR, FuncDecl, FuncSignature, RuntimeArg
+from mypyc.ir.func_ir import FuncDecl, FuncSignature
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.primitives.generic_ops import py_setattr_op, py_hasattr_op
 from mypyc.primitives.misc_ops import (
     dataclass_sleight_of_hand, pytype_from_template_op, py_calc_meta_op, type_object_op,
     not_implemented_op
 )
 from mypyc.primitives.dict_ops import dict_set_item_op, dict_new_op
-from mypyc.common import SELF_NAME
 from mypyc.irbuild.util import (
     is_dataclass_decorator, get_func_def, is_dataclass, is_constant
 )
@@ -327,12 +326,9 @@ def generate_attr_defaults(builder: IRBuilder, cdef: ClassDef) -> None:
     if not default_assignments:
         return
 
-    builder.enter()
-    builder.ret_types[-1] = bool_rprimitive
-
-    rt_args = (RuntimeArg(SELF_NAME, RInstance(cls)),)
-    self_var = builder.read(builder.add_self_to_env(cls), -1)
+    builder.enter_method(cls, '__mypyc_defaults_setup', bool_rprimitive)
 
+    self_var = builder.self()
     for stmt in default_assignments:
         lvalue = stmt.lvalues[0]
         assert isinstance(lvalue, NameExpr)
@@ -351,43 +347,24 @@ def generate_attr_defaults(builder: IRBuilder, cdef: ClassDef) -> None:
 
     builder.add(Return(builder.true()))
 
-    args, blocks, ret_type, _ = builder.leave()
-    ir = FuncIR(
-        FuncDecl('__mypyc_defaults_setup',
-                 cls.name, builder.module_name,
-                 FuncSignature(rt_args, ret_type)),
-        args, blocks)
-    builder.functions.append(ir)
-    cls.methods[ir.name] = ir
+    builder.leave_method()
 
 
 def create_ne_from_eq(builder: IRBuilder, cdef: ClassDef) -> None:
     """Create a "__ne__" method from a "__eq__" method (if only latter exists)."""
     cls = builder.mapper.type_to_ir[cdef.info]
     if cls.has_method('__eq__') and not cls.has_method('__ne__'):
-        f = gen_glue_ne_method(builder, cls, cdef.line)
-        cls.method_decls['__ne__'] = f.decl
-        cls.methods['__ne__'] = f
-        builder.functions.append(f)
+        gen_glue_ne_method(builder, cls, cdef.line)
 
 
-def gen_glue_ne_method(builder: IRBuilder, cls: ClassIR, line: int) -> FuncIR:
+def gen_glue_ne_method(builder: IRBuilder, cls: ClassIR, line: int) -> None:
     """Generate a "__ne__" method from a "__eq__" method. """
-    builder.enter()
-
-    rt_args = (RuntimeArg("self", RInstance(cls)), RuntimeArg("rhs", object_rprimitive))
-
-    # The environment operates on Vars, so we make some up
-    fake_vars = [(Var(arg.name), arg.type) for arg in rt_args]
-    args = [
-        builder.read(builder.add_local_reg(var, type, is_arg=True), line)
-        for var, type in fake_vars
-    ]  # type: List[Value]
-    builder.ret_types[-1] = object_rprimitive
+    builder.enter_method(cls, '__ne__', object_rprimitive)
+    rhs_arg = builder.add_argument('rhs', object_rprimitive)
 
     # If __eq__ returns NotImplemented, then __ne__ should also
     not_implemented_block, regular_block = BasicBlock(), BasicBlock()
-    eqval = builder.add(MethodCall(args[0], '__eq__', [args[1]], line))
+    eqval = builder.add(MethodCall(builder.self(), '__eq__', [rhs_arg], line))
     not_implemented = builder.add(LoadAddress(not_implemented_op.type,
                                               not_implemented_op.src, line))
     builder.add(Branch(
@@ -405,11 +382,7 @@ def gen_glue_ne_method(builder: IRBuilder, cls: ClassIR, line: int) -> FuncIR:
     builder.activate_block(not_implemented_block)
     builder.add(Return(not_implemented))
 
-    arg_regs, blocks, ret_type, _ = builder.leave()
-    return FuncIR(
-        FuncDecl('__ne__', cls.name, builder.module_name,
-                 FuncSignature(rt_args, ret_type)),
-        arg_regs, blocks)
+    builder.leave_method()
 
 
 def load_non_ext_class(builder: IRBuilder,

mypyc/irbuild/expression.py

@@ -138,7 +138,7 @@ def transform_super_expr(builder: IRBuilder, o: SuperExpr) -> Value:
         assert o.info is not None
         typ = builder.load_native_type_object(o.info.fullname)
         ir = builder.mapper.type_to_ir[o.info]
-        iter_env = iter(builder.args[-1])
+        iter_env = iter(builder.builder.args)
         # Grab first argument
         vself = next(iter_env)  # type: Value
         if builder.fn_info.is_generator:
@@ -302,7 +302,7 @@ def translate_super_method_call(builder: IRBuilder, expr: CallExpr, callee: Supe
 
     if decl.kind != FUNC_STATICMETHOD:
         # Grab first argument
-        vself = builder.args[-1][0]  # type: Value
+        vself = builder.self()  # type: Value
         if decl.kind == FUNC_CLASSMETHOD:
             vself = builder.call_c(type_op, [vself], expr.line)
         elif builder.fn_info.is_generator:

mypyc/irbuild/function.py

@@ -225,7 +225,7 @@ def c() -> None:
     if builder.fn_info.is_generator:
         # Do a first-pass and generate a function that just returns a generator object.
         gen_generator_func(builder)
-        args, blocks, ret_type, fn_info = builder.leave()
+        args, _, blocks, ret_type, fn_info = builder.leave()
         func_ir, func_reg = gen_func_ir(builder, args, blocks, sig, fn_info, cdef)
 
         # Re-enter the FuncItem and visit the body of the function this time.
@@ -287,7 +287,7 @@ def c() -> None:
     # to calculate argument defaults below.
     symtable = builder.symtables[-1]
 
-    args, blocks, ret_type, fn_info = builder.leave()
+    args, _, blocks, ret_type, fn_info = builder.leave()
 
     if fn_info.is_generator:
         add_methods_to_generator_class(
@@ -675,7 +675,7 @@ def f(builder: IRBuilder, x: object) -> int: ...
     retval = builder.coerce(retval, sig.ret_type, line)
     builder.add(Return(retval))
 
-    arg_regs, blocks, ret_type, _ = builder.leave()
+    arg_regs, _, blocks, ret_type, _ = builder.leave()
     return FuncIR(
         FuncDecl(target.name + '__' + base.name + '_glue',
                  cls.name, builder.module_name,
@@ -713,7 +713,7 @@ def gen_glue_property(builder: IRBuilder,
     retbox = builder.coerce(retval, sig.ret_type, line)
     builder.add(Return(retbox))
 
-    args, blocks, return_type, _ = builder.leave()
+    args, _, blocks, return_type, _ = builder.leave()
     return FuncIR(
         FuncDecl(target.name + '__' + base.name + '_glue',
                  cls.name, builder.module_name, FuncSignature([rt_arg], return_type)),