Implement initial support for single inheritance (mypyc/mypyc#139)

msullivan · JukkaL · commit b2ca9848cbd4 · 2018-06-12T16:13:39.000+01:00
This doesn't handle subclasses changing method signatures (which 
will require generating compatibility methods, and using the 
compatibility method in the base class part of the vtable and the 
real one in the subclass part).

* Make ir.attributes be a dictionary
* Centralize vtable computation into ClassIR
* Support basic class inheritance
diff --git a/mypyc/emitclass.py b/mypyc/emitclass.py
@@ -26,12 +26,14 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     getseters_name = '{}_getseters'.format(name)
     methods_name = '{}_methods'.format(name)
     vtable_name = '{}_vtable'.format(name)
+    base_arg = "&{}".format(type_struct_name(cl.base.name)) if cl.base else "0"
 
     def emit_line() -> None:
         emitter.emit_line()
 
     emit_line()
     generate_object_struct(cl, emitter)
+    emit_line()
 
     # If there is a __init__ method, generate a function for tp_init and
     # extract the args (which we'll use for the native constructor)
@@ -100,7 +102,7 @@ def emit_line() -> None:
             {methods_name},            /* tp_methods */
             0,                         /* tp_members */
             {getseters_name},          /* tp_getset */
-            0,                         /* tp_base */
+            {base_arg},                /* tp_base */
             0,                         /* tp_dict */
             0,                         /* tp_descr_get */
             0,                         /* tp_descr_set */
@@ -119,6 +121,7 @@ def emit_line() -> None:
                     methods_name=methods_name,
                     getseters_name=getseters_name,
                     init_name=init_name,
+                    base_arg=base_arg,
                     ))
     emitter.emit_line()
     generate_setup_for_class(cl, setup_name, vtable_name, emitter)
@@ -148,14 +151,15 @@ def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
     emitter.emit_lines('typedef struct {',
                        'PyObject_HEAD',
                        'CPyVTableItem *vtable;')
-    for attr, rtype in cl.attributes:
-        emitter.emit_line('{}{};'.format(rtype.ctype_spaced(), attr))
+    for base in reversed(cl.mro):
+        for attr, rtype in base.attributes.items():
+            emitter.emit_line('{}{};'.format(rtype.ctype_spaced(), attr))
     emitter.emit_line('}} {};'.format(cl.struct_name()))
 
 
 def generate_native_getters_and_setters(cl: ClassIR,
                                         emitter: Emitter) -> None:
-    for attr, rtype in cl.attributes:
+    for attr, rtype in cl.attributes.items():
         # Native getter
         emitter.emit_line('{}{}({} *self)'.format(rtype.ctype_spaced(),
                                                native_getter_name(cl.name, attr),
@@ -189,15 +193,23 @@ def generate_native_getters_and_setters(cl: ClassIR,
         emitter.emit_line()
 
 
-def generate_vtable(cl: ClassIR,
+def generate_vtable(base: ClassIR,
                     vtable_name: str,
                     emitter: Emitter) -> None:
     emitter.emit_line('static CPyVTableItem {}[] = {{'.format(vtable_name))
-    for attr, rtype in cl.attributes:
-        emitter.emit_line('(CPyVTableItem){},'.format(native_getter_name(cl.name, attr)))
-        emitter.emit_line('(CPyVTableItem){},'.format(native_setter_name(cl.name, attr)))
-    for fn in cl.methods:
-        emitter.emit_line('(CPyVTableItem){}{},'.format(NATIVE_PREFIX, fn.cname))
+    for cl in reversed(base.mro):
+        for attr in cl.attributes:
+            emitter.emit_line('(CPyVTableItem){},'.format(native_getter_name(cl.name, attr)))
+            emitter.emit_line('(CPyVTableItem){},'.format(native_setter_name(cl.name, attr)))
+        for fn in cl.methods:
+            # TODO: This is gross, and inefficient, and wrong if the type changes.
+            # This logic should all live on the genops side, I think
+            search = base.mro if fn.name != '__init__' else [cl]
+            for cl2 in search:
+                m = cl2.get_method(fn.name)
+                if m:
+                    emitter.emit_line('(CPyVTableItem){}{},'.format(NATIVE_PREFIX, m.cname))
+                    break
     emitter.emit_line('};')
 
 
@@ -216,7 +228,7 @@ def generate_setup_for_class(cl: ClassIR,
     emitter.emit_line('if (self == NULL)')
     emitter.emit_line('    return NULL;')
     emitter.emit_line('self->vtable = {};'.format(vtable_name))
-    for attr, rtype in cl.attributes:
+    for attr, rtype in cl.attributes.items():
         emitter.emit_line('self->{} = {};'.format(attr, rtype.c_undefined_value()))
     emitter.emit_line('return (PyObject *)self;')
     emitter.emit_line('}')
@@ -282,7 +294,7 @@ def generate_traverse_for_class(cl: ClassIR,
     emitter.emit_line('{}({} *self, visitproc visit, void *arg)'.format(func_name,
                                                                         cl.struct_name()))
     emitter.emit_line('{')
-    for attr, rtype in cl.attributes:
+    for attr, rtype in cl.attributes.items():
         emitter.emit_gc_visit('self->{}'.format(attr), rtype)
     emitter.emit_line('return 0;')
     emitter.emit_line('}')
@@ -294,7 +306,7 @@ def generate_clear_for_class(cl: ClassIR,
     emitter.emit_line('static int')
     emitter.emit_line('{}({} *self)'.format(func_name, cl.struct_name()))
     emitter.emit_line('{')
-    for attr, rtype in cl.attributes:
+    for attr, rtype in cl.attributes.items():
         emitter.emit_gc_clear('self->{}'.format(attr), rtype)
     emitter.emit_line('return 0;')
     emitter.emit_line('}')
@@ -327,7 +339,7 @@ def generate_methods_table(cl: ClassIR,
 
 
 def generate_getseter_declarations(cl: ClassIR, emitter: Emitter) -> None:
-    for attr, rtype in cl.attributes:
+    for attr in cl.attributes:
         emitter.emit_line('static PyObject *')
         emitter.emit_line('{}({} *self, void *closure);'.format(getter_name(cl.name, attr),
                                                             cl.struct_name()))
@@ -342,7 +354,7 @@ def generate_getseters_table(cl: ClassIR,
                              emitter: Emitter) -> None:
 
     emitter.emit_line('static PyGetSetDef {}[] = {{'.format(name))
-    for attr, rtype in cl.attributes:
+    for attr in cl.attributes:
         emitter.emit_line('{{"{}",'.format(attr))
         emitter.emit_line(' (getter){}, (setter){},'.format(getter_name(cl.name, attr),
                                                             setter_name(cl.name, attr)))
@@ -352,7 +364,7 @@ def generate_getseters_table(cl: ClassIR,
 
 
 def generate_getseters(cl: ClassIR, emitter: Emitter) -> None:
-    for i, (attr, rtype) in enumerate(cl.attributes):
+    for i, (attr, rtype) in enumerate(cl.attributes.items()):
         generate_getter(cl, attr, rtype, emitter)
         emitter.emit_line('')
         generate_setter(cl, attr, rtype, emitter)
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -183,29 +183,52 @@ def visit_mypy_file(self, mypyfile: MypyFile) -> Value:
             # built-in primitives.
             return INVALID_VALUE
 
-        # First pass: Build ClassIRs and TypeInfo-to-ClassIR mapping.
-        for node in mypyfile.defs:
-            if isinstance(node, ClassDef):
-                self.prepare_class_def(node)
+        classes = [node for node in mypyfile.defs if isinstance(node, ClassDef)]
+
+        # Build ClassIRs and TypeInfo-to-ClassIR mapping.
+        for cls in classes:
+            self.create_class_def(cls)
 
-        # Second pass: Generate ops.
+        # Do class def setup
+        for cls in classes:
+            self.prepare_class_def(cls)
+
+        # Generate ops.
         self.current_module_name = mypyfile.fullname()
         for node in mypyfile.defs:
             node.accept(self)
 
+        # Compute vtables.
+        for cls in classes:
+            self.mapper.type_to_ir[cls.info].compute_vtable()
+
         return INVALID_VALUE
 
-    def prepare_class_def(self, cdef: ClassDef) -> None:
+    def create_class_def(self, cdef: ClassDef) -> None:
         # We want to collect the attributes first so they are available
         # while generating the methods
         ir = ClassIR(cdef.name)
         self.classes.append(ir)
         self.mapper.type_to_ir[cdef.info] = ir
 
-        for name, node in cdef.info.names.items():
+    def prepare_class_def(self, cdef: ClassDef) -> None:
+        ir = self.mapper.type_to_ir[cdef.info]
+        info = cdef.info
+        for name, node in info.names.items():
             if isinstance(node.node, Var):
                 assert node.node.type, "Class member missing type"
-                ir.attributes.append((name, self.type_to_rtype(node.node.type)))
+                ir.attributes[name] = self.type_to_rtype(node.node.type)
+
+        # Set up the parent class
+        assert len(info.bases) == 1, "Only single inheritance is supported"
+        mro = []
+        for cls in info.mro:
+            if cls.fullname() == 'builtins.object': continue
+            assert cls in self.mapper.type_to_ir, "Can't subclass cpython types yet"
+            mro.append(self.mapper.type_to_ir[cls])
+        if len(mro) > 1:
+            ir.base = mro[1]
+        ir.mro = mro
 
     def visit_class_def(self, cdef: ClassDef) -> Value:
         ir = self.mapper.type_to_ir[cdef.info]
diff --git a/mypyc/ops.py b/mypyc/ops.py
@@ -273,26 +273,16 @@ def struct_name(self) -> str:
         return self.class_ir.struct_name()
 
     def getter_index(self, name: str) -> int:
-        for i, (attr, _) in enumerate(self.class_ir.attributes):
-            if attr == name:
-                return i * 2
-        assert False, '%r has no attribute %r' % (self.name, name)
+        return self.class_ir.vtable_entry(name)
 
     def setter_index(self, name: str) -> int:
         return self.getter_index(name) + 1
 
     def method_index(self, name: str) -> int:
-        base = len(self.class_ir.attributes) * 2
-        for i, fn in enumerate(self.class_ir.methods):
-            if fn.name == name:
-                return base + i
-        assert False, '%r has no attribute %r' % (self.name, name)
+        return self.class_ir.vtable_entry(name)
 
     def attr_type(self, name: str) -> RType:
-        for i, (attr, rtype) in enumerate(self.class_ir.attributes):
-            if attr == name:
-                return rtype
-        assert False, '%r has no attribute %r' % (self.name, name)
+        return self.class_ir.attr_type(name)
 
     def __repr__(self) -> str:
         return '<RInstance %s>' % self.name
@@ -1232,8 +1222,39 @@ class ClassIR:
 
     def __init__(self, name: str) -> None:
         self.name = name
-        self.attributes = []  # type: List[Tuple[str, RType]]
+        self.attributes = OrderedDict()  # type: OrderedDict[str, RType]
         self.methods = []  # type: List[FuncIR]
+        self.vtable = None  # type: Optional[Dict[str, int]]
+        self.vtable_size = 0
+        self.base = None  # type: Optional[ClassIR]
+        self.mro = []  # type: List[ClassIR]
+
+    def compute_vtable(self) -> None:
+        if self.vtable is not None: return
+        self.vtable = {}
+        base = 0
+        if self.base:
+            self.base.compute_vtable()
+            assert self.base.vtable is not None
+            self.vtable.update(self.base.vtable)
+            base = self.base.vtable_size
+
+        for i, attr in enumerate(self.attributes):
+            self.vtable[attr] = base + i * 2
+        base += len(self.attributes) * 2
+        for i, fn in enumerate(self.methods):
+            self.vtable[fn.name] = base + i
+        self.vtable_size = base + len(self.methods)
+
+    def vtable_entry(self, name: str) -> int:
+        assert self.vtable is not None, "vtable not computed yet"
+        assert name in self.vtable, '%r has no attribute %r' % (self.name, name)
+        return self.vtable[name]
+
+    def attr_type(self, name: str) -> RType:
+        for ir in self.mro:
+            if name in ir.attributes: return ir.attributes[name]
+        assert False, '%r has no attribute %r' % (self.name, name)
 
     def struct_name(self) -> str:
         return '{}Object'.format(self.name)
diff --git a/mypyc/subtype.py b/mypyc/subtype.py
@@ -27,8 +27,7 @@ def __init__(self, right: RType) -> None:
         self.right = right
 
     def visit_rinstance(self, left: RInstance) -> bool:
-        # TODO: Inheritance
-        return isinstance(self.right, RInstance) and self.right.name == left.name
+        return isinstance(self.right, RInstance) and self.right.class_ir in left.class_ir.mro
 
     def visit_roptional(self, left: ROptional) -> bool:
         return isinstance(self.right, ROptional) and is_subtype(left.value_type,
diff --git a/mypyc/test/test_emitfunc.py b/mypyc/test/test_emitfunc.py
@@ -1,5 +1,7 @@
 import unittest
 
+from collections import OrderedDict
+
 from mypy.nodes import Var
 from mypy.test.helpers import assert_string_arrays_equal
 
@@ -38,7 +40,9 @@ def setUp(self) -> None:
         self.tt = self.env.add_local(
             Var('tt'), RTuple([RTuple([int_rprimitive, bool_rprimitive]), bool_rprimitive]))
         ir = ClassIR('A')
-        ir.attributes = [('x', bool_rprimitive), ('y', int_rprimitive)]
+        ir.attributes = OrderedDict([('x', bool_rprimitive), ('y', int_rprimitive)])
+        ir.compute_vtable()
+        ir.mro = [ir]
         self.r = self.env.add_local(Var('r'), RInstance(ir))
 
         self.context = EmitterContext()
diff --git a/test-data/genops-classes.test b/test-data/genops-classes.test
@@ -141,3 +141,38 @@ L1:
 L2:
     r6 = 1
     return r6
+
+[case testSubclass]
+class A:
+    def __init__(self) -> None:
+        self.x = 10
+class B(A):
+    def __init__(self) -> None:
+        self.x = 20
+        self.y = 30
+[out]
+-- This is totally useless because we don't print any class info!!
+def __init__(self):
+    self :: A
+    r0 :: int
+    r1 :: bool
+    r2 :: None
+L0:
+    r0 = 10
+    self.x = r0; r1 = is_error
+    r2 = None
+    return r2
+def __init__(self):
+    self :: B
+    r0 :: int
+    r1 :: bool
+    r2 :: int
+    r3 :: bool
+    r4 :: None
+L0:
+    r0 = 20
+    self.x = r0; r1 = is_error
+    r2 = 30
+    self.y = r2; r3 = is_error
+    r4 = None
+    return r4
diff --git a/test-data/run-bench.test b/test-data/run-bench.test
diff --git a/test-data/run-classes.test b/test-data/run-classes.test
