[mypyc] Do all vtable setup dynamically at runtime (python#7629)

msullivan · web-flow · commit 38f199d25db9 · 2019-10-04T14:55:42.000-07:00
This lets us eliminate the CPy_FixupTraitVtable hack where we fixed up
vtables at runtime in a hokey manner and will allow us to populate
vtables with entries loaded dynamically from other modules once we
have separate compilation.
diff --git a/mypyc/emitclass.py b/mypyc/emitclass.py
@@ -198,7 +198,7 @@ def emit_line() -> None:
         generate_dealloc_for_class(cl, dealloc_name, clear_name, emitter)
         emit_line()
         generate_native_getters_and_setters(cl, emitter)
-        vtable_name = generate_vtables(cl, vtable_name, emitter)
+        vtable_name = generate_vtables(cl, vtable_setup_name, vtable_name, emitter)
         emit_line()
     if needs_getseters:
         generate_getseter_declarations(cl, emitter)
@@ -222,7 +222,6 @@ def emit_line() -> None:
         t=emitter.type_struct_name(cl)))
 
     emitter.emit_line()
-    generate_trait_vtable_setup(cl, vtable_setup_name, vtable_name, emitter)
     if generate_full:
         generate_setup_for_class(cl, setup_name, defaults_fn, vtable_name, emitter)
         emitter.emit_line()
@@ -307,39 +306,63 @@ def generate_native_getters_and_setters(cl: ClassIR,
 
 
 def generate_vtables(base: ClassIR,
+                     vtable_setup_name: str,
                      vtable_name: str,
                      emitter: Emitter) -> str:
-    """Emit the vtables for a class.
+    """Emit the vtables and vtable setup functions for a class.
 
     This includes both the primary vtable and any trait implementation vtables.
 
+    To account for both dynamic loading and dynamic class creation,
+    vtables are populated dynamically at class creation time, so we
+    emit empty array definitions to store the vtables and a function to
+    populate them.
+
     Returns the expression to use to refer to the vtable, which might be
-    different than the name, if there are trait vtables."""
+    different than the name, if there are trait vtables.
+    """
+
+    def trait_vtable_name(trait: ClassIR) -> str:
+        return '{}_{}_trait_vtable'.format(
+            base.name_prefix(emitter.names), trait.name_prefix(emitter.names))
+
+    # Emit array definitions with enough space for all the entries
+    emitter.emit_line('static CPyVTableItem {}[{}];'.format(
+        vtable_name,
+        max(1, len(base.vtable_entries) + 2 * len(base.trait_vtables))))
+    for trait, vtable in base.trait_vtables.items():
+        emitter.emit_line('static CPyVTableItem {}[{}];'.format(
+            trait_vtable_name(trait),
+            max(1, len(vtable))))
+
+    # Emit vtable setup function
+    emitter.emit_line('static bool')
+    emitter.emit_line('{}{}(void)'.format(NATIVE_PREFIX, vtable_setup_name))
+    emitter.emit_line('{')
 
     subtables = []
     for trait, vtable in base.trait_vtables.items():
-        name = '{}_{}_trait_vtable'.format(
-            base.name_prefix(emitter.names), trait.name_prefix(emitter.names))
+        name = trait_vtable_name(trait)
         generate_vtable(vtable, name, emitter, [])
         subtables.append((trait, name))
 
     generate_vtable(base.vtable_entries, vtable_name, emitter, subtables)
 
+    emitter.emit_line('return 1;')
+    emitter.emit_line('}')
+
     return vtable_name if not subtables else "{} + {}".format(vtable_name, len(subtables) * 2)
 
 
 def generate_vtable(entries: VTableEntries,
                     vtable_name: str,
                     emitter: Emitter,
                     subtables: List[Tuple[ClassIR, str]]) -> None:
-    emitter.emit_line('static CPyVTableItem {}[] = {{'.format(vtable_name))
+    emitter.emit_line('CPyVTableItem {}_scratch[] = {{'.format(vtable_name))
     if subtables:
         emitter.emit_line('/* Array of trait vtables */')
         for trait, table in subtables:
-            # N.B: C only lets us store constant values. We do a nasty hack of
-            # storing a pointer to the location, which we will then dynamically
-            # patch up on module load in CPy_FixupTraitVtable.
-            emitter.emit_line('(CPyVTableItem)&{}, (CPyVTableItem){},'.format(
+            emitter.emit_line('(CPyVTableItem){}, (CPyVTableItem){},'.format(
                 emitter.type_struct_name(trait), table))
         emitter.emit_line('/* Start of real vtable */')
 
@@ -355,24 +378,7 @@ def generate_vtable(entries: VTableEntries,
     if not entries:
         emitter.emit_line('NULL')
     emitter.emit_line('};')
-
-
-def generate_trait_vtable_setup(cl: ClassIR,
-                                vtable_setup_name: str,
-                                vtable_name: str,
-                                emitter: Emitter) -> None:
-    """Generate a native function that fixes up the trait vtables of a class.
-
-    This needs to be called before a class is used.
-    """
-    emitter.emit_line('static bool')
-    emitter.emit_line('{}{}(void)'.format(NATIVE_PREFIX, vtable_setup_name))
-    emitter.emit_line('{')
-    if cl.trait_vtables and not cl.is_trait:
-        emitter.emit_lines('CPy_FixupTraitVtable({}_vtable, {});'.format(
-            cl.name_prefix(emitter.names), len(cl.trait_vtables)))
-    emitter.emit_line('return 1;')
-    emitter.emit_line('}')
+    emitter.emit_line('memcpy({name}, {name}_scratch, sizeof({name}));'.format(name=vtable_name))
 
 
 def generate_setup_for_class(cl: ClassIR,
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -1483,10 +1483,12 @@ def allocate_class(self, cdef: ClassDef) -> None:
         tp = self.primitive_op(pytype_from_template_op,
                                [template, tp_bases, modname], cdef.line)
         # Immediately fix up the trait vtables, before doing anything with the class.
-        self.add(Call(
-            FuncDecl(cdef.name + '_trait_vtable_setup',
-                     None, self.module_name,
-                     FuncSignature([], bool_rprimitive)), [], -1))
+        ir = self.mapper.type_to_ir[cdef.info]
+        if not ir.is_trait and not ir.builtin_base:
+            self.add(Call(
+                FuncDecl(cdef.name + '_trait_vtable_setup',
+                         None, self.module_name,
+                         FuncSignature([], bool_rprimitive)), [], -1))
         # Populate a '__mypyc_attrs__' field containing the list of attrs
         self.primitive_op(py_setattr_op, [
             tp, self.load_static_unicode('__mypyc_attrs__'),
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -57,15 +57,6 @@ static inline CPyVTableItem *CPy_FindTraitVtable(PyTypeObject *trait, CPyVTableI
     }
 }
 
-// At load time, we need to patch up trait vtables to contain actual pointers
-// to the type objects of the trait, rather than an indirection.
-static inline void CPy_FixupTraitVtable(CPyVTableItem *vtable, int count) {
-    int i;
-    for (i = 0; i < count; i++) {
-        vtable[i*2] = *(CPyVTableItem *)vtable[i*2];
-    }
-}
-
 static bool _CPy_IsSafeMetaClass(PyTypeObject *metaclass) {
     // mypyc classes can't work with metaclasses in
     // general. Through some various nasty hacks we *do*
diff --git a/mypyc/test-data/genops-classes.test b/mypyc/test-data/genops-classes.test
@@ -369,31 +369,30 @@ def __top_level__():
     r50 :: object
     r51 :: str
     r52, r53 :: object
-    r54 :: bool
-    r55 :: str
-    r56 :: tuple
-    r57 :: bool
-    r58 :: dict
-    r59 :: str
-    r60 :: bool
-    r61, r62 :: object
-    r63 :: dict
-    r64 :: str
-    r65 :: object
-    r66 :: dict
-    r67 :: str
-    r68, r69 :: object
-    r70 :: tuple
-    r71 :: str
-    r72, r73 :: object
-    r74 :: bool
-    r75, r76 :: str
-    r77 :: tuple
-    r78 :: bool
-    r79 :: dict
-    r80 :: str
-    r81 :: bool
-    r82 :: None
+    r54 :: str
+    r55 :: tuple
+    r56 :: bool
+    r57 :: dict
+    r58 :: str
+    r59 :: bool
+    r60, r61 :: object
+    r62 :: dict
+    r63 :: str
+    r64 :: object
+    r65 :: dict
+    r66 :: str
+    r67, r68 :: object
+    r69 :: tuple
+    r70 :: str
+    r71, r72 :: object
+    r73 :: bool
+    r74, r75 :: str
+    r76 :: tuple
+    r77 :: bool
+    r78 :: dict
+    r79 :: str
+    r80 :: bool
+    r81 :: None
 L0:
     r0 = builtins :: module
     r1 = builtins.None :: object
@@ -462,38 +461,37 @@ L6:
     r51 = unicode_7 :: static  ('__main__')
     r52 = __main__.S_template :: type
     r53 = pytype_from_template(r52, r50, r51)
-    r54 = S_trait_vtable_setup()
-    r55 = unicode_8 :: static  ('__mypyc_attrs__')
-    r56 = () :: tuple
-    r57 = setattr r53, r55, r56
+    r54 = unicode_8 :: static  ('__mypyc_attrs__')
+    r55 = () :: tuple
+    r56 = setattr r53, r54, r55
     __main__.S = r53 :: type
-    r58 = __main__.globals :: static
-    r59 = unicode_10 :: static  ('S')
-    r60 = r58.__setitem__(r59, r53) :: dict
-    r61 = __main__.C :: type
-    r62 = __main__.S :: type
-    r63 = __main__.globals :: static
-    r64 = unicode_3 :: static  ('Generic')
-    r65 = r63[r64] :: dict
-    r66 = __main__.globals :: static
-    r67 = unicode_6 :: static  ('T')
-    r68 = r66[r67] :: dict
-    r69 = r65[r68] :: object
-    r70 = (r61, r62, r69) :: tuple
-    r71 = unicode_7 :: static  ('__main__')
-    r72 = __main__.D_template :: type
-    r73 = pytype_from_template(r72, r70, r71)
-    r74 = D_trait_vtable_setup()
-    r75 = unicode_8 :: static  ('__mypyc_attrs__')
-    r76 = unicode_11 :: static  ('__dict__')
-    r77 = (r76) :: tuple
-    r78 = setattr r73, r75, r77
-    __main__.D = r73 :: type
-    r79 = __main__.globals :: static
-    r80 = unicode_12 :: static  ('D')
-    r81 = r79.__setitem__(r80, r73) :: dict
-    r82 = None
-    return r82
+    r57 = __main__.globals :: static
+    r58 = unicode_10 :: static  ('S')
+    r59 = r57.__setitem__(r58, r53) :: dict
+    r60 = __main__.C :: type
+    r61 = __main__.S :: type
+    r62 = __main__.globals :: static
+    r63 = unicode_3 :: static  ('Generic')
+    r64 = r62[r63] :: dict
+    r65 = __main__.globals :: static
+    r66 = unicode_6 :: static  ('T')
+    r67 = r65[r66] :: dict
+    r68 = r64[r67] :: object
+    r69 = (r60, r61, r68) :: tuple
+    r70 = unicode_7 :: static  ('__main__')
+    r71 = __main__.D_template :: type
+    r72 = pytype_from_template(r71, r69, r70)
+    r73 = D_trait_vtable_setup()
+    r74 = unicode_8 :: static  ('__mypyc_attrs__')
+    r75 = unicode_11 :: static  ('__dict__')
+    r76 = (r75) :: tuple
+    r77 = setattr r72, r74, r76
+    __main__.D = r72 :: type
+    r78 = __main__.globals :: static
+    r79 = unicode_12 :: static  ('D')
+    r80 = r78.__setitem__(r79, r72) :: dict
+    r81 = None
+    return r81
 
 [case testIsInstance]
 class A: pass
