Fix vtable pointer with inherited dunder new

Author: p-sawicki

mypyc/codegen/emitclass.py

@@ -229,6 +229,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     name = cl.name
     name_prefix = cl.name_prefix(emitter.names)
 
+    allocate_name = emitter.native_function_name(cl.allocate)
     setup_name = emitter.native_function_name(cl.setup)
     new_name = f"{name_prefix}_new"
     finalize_name = f"{name_prefix}_finalize"
@@ -322,14 +323,15 @@ def emit_line() -> None:
         fields["tp_basicsize"] = base_size
 
     if generate_full:
-        assert cl.setup is not None
-        emitter.emit_line(native_function_header(cl.setup, emitter) + ";")
-        assert cl.ctor is not None
-        emitter.emit_line(native_function_header(cl.ctor, emitter) + ";")
+        for decl in [cl.allocate, cl.setup, cl.ctor]:
+            assert decl is not None
+            emitter.emit_line(native_function_header(decl, emitter) + ";")
 
         emit_line()
         init_fn = cl.get_method("__init__")
-        generate_new_for_class(cl, new_name, vtable_name, setup_name, init_fn, emitter)
+        generate_new_for_class(
+            cl, new_name, vtable_name, allocate_name, setup_name, init_fn, emitter
+        )
         emit_line()
         generate_traverse_for_class(cl, traverse_name, emitter)
         emit_line()
@@ -393,9 +395,13 @@ def emit_line() -> None:
 
     emitter.emit_line()
     if generate_full:
-        generate_setup_for_class(cl, defaults_fn, vtable_name, shadow_vtable_name, emitter)
+        generate_allocate_for_class(cl, defaults_fn, emitter)
+        emitter.emit_line()
+        generate_setup_for_class(cl, vtable_name, shadow_vtable_name, emitter)
         emitter.emit_line()
-        generate_constructor_for_class(cl, cl.ctor, init_fn, setup_name, vtable_name, emitter)
+        generate_constructor_for_class(
+            cl, cl.ctor, init_fn, allocate_name, setup_name, vtable_name, emitter
+        )
         emitter.emit_line()
     if needs_getseters:
         generate_getseters(cl, emitter)
@@ -580,17 +586,11 @@ def generate_vtable(
     emitter.emit_line("memcpy({name}, {name}_scratch, sizeof({name}));".format(name=vtable_name))
 
 
-def generate_setup_for_class(
-    cl: ClassIR,
-    defaults_fn: FuncIR | None,
-    vtable_name: str,
-    shadow_vtable_name: str | None,
-    emitter: Emitter,
-) -> None:
+def generate_allocate_for_class(cl: ClassIR, defaults_fn: FuncIR | None, emitter: Emitter) -> None:
     """Generate a native function that allocates an instance of a class."""
-    emitter.emit_line(native_function_header(cl.setup, emitter))
+    emitter.emit_line(native_function_header(cl.allocate, emitter))
     emitter.emit_line("{")
-    type_arg_name = REG_PREFIX + cl.setup.sig.args[0].name
+    type_arg_name = REG_PREFIX + cl.allocate.sig.args[0].name
     emitter.emit_line(f"PyTypeObject *type = (PyTypeObject*){type_arg_name};")
     struct_name = cl.struct_name(emitter.names)
     emitter.emit_line(f"{struct_name} *self;")
@@ -609,24 +609,8 @@ def generate_setup_for_class(
         emitter.emit_line("}")
 
     emitter.emit_line(f"self = ({cl.struct_name(emitter.names)} *)type->tp_alloc(type, 0);")
-    emitter.emit_line("if (self == NULL)")
-    emitter.emit_line("    return NULL;")
-
-    if shadow_vtable_name:
-        emitter.emit_line(f"if (type != {emitter.type_struct_name(cl)}) {{")
-        emitter.emit_line(f"self->vtable = {shadow_vtable_name};")
-        emitter.emit_line("} else {")
-        emitter.emit_line(f"self->vtable = {vtable_name};")
-        emitter.emit_line("}")
-    else:
-        emitter.emit_line(f"self->vtable = {vtable_name};")
-
     emit_clear_bitmaps(cl, emitter)
 
-    if cl.has_method("__call__"):
-        name = cl.method_decl("__call__").cname(emitter.names)
-        emitter.emit_line(f"self->vectorcall = {PREFIX}{name};")
-
     for base in reversed(cl.base_mro):
         for attr, rtype in base.attributes.items():
             value = emitter.c_undefined_value(rtype)
@@ -644,6 +628,34 @@ def generate_setup_for_class(
     emitter.emit_line("}")
 
 
+def generate_setup_for_class(
+    cl: ClassIR, vtable_name: str, shadow_vtable_name: str | None, emitter: Emitter
+) -> None:
+    """Generate a native function that performs setup needed after allocation."""
+    emitter.emit_line(native_function_header(cl.setup, emitter))
+    emitter.emit_line("{")
+    obj_arg_name = REG_PREFIX + cl.setup.sig.args[0].name
+    emitter.emit_line(f"PyTypeObject *type = Py_TYPE({obj_arg_name});")
+    struct_name = cl.struct_name(emitter.names)
+    emitter.emit_line(f"{struct_name} *self = ({struct_name} *){obj_arg_name};")
+
+    if shadow_vtable_name:
+        emitter.emit_line(f"if (type != {emitter.type_struct_name(cl)}) {{")
+        emitter.emit_line(f"self->vtable = {shadow_vtable_name};")
+        emitter.emit_line("} else {")
+        emitter.emit_line(f"self->vtable = {vtable_name};")
+        emitter.emit_line("}")
+    else:
+        emitter.emit_line(f"self->vtable = {vtable_name};")
+
+    if cl.has_method("__call__"):
+        name = cl.method_decl("__call__").cname(emitter.names)
+        emitter.emit_line(f"self->vectorcall = {PREFIX}{name};")
+
+    emitter.emit_line("return (PyObject *)self;")
+    emitter.emit_line("}")
+
+
 def emit_clear_bitmaps(cl: ClassIR, emitter: Emitter) -> None:
     """Emit C code to clear bitmaps that track if attributes have an assigned value."""
     for i in range(0, len(cl.bitmap_attrs), BITMAP_BITS):
@@ -666,8 +678,9 @@ def emit_attr_defaults_func_call(defaults_fn: FuncIR, self_name: str, emitter: E
     )
 
 
-def emit_setup_or_dunder_new_call(
+def emit_allocate_or_dunder_new_call(
     cl: ClassIR,
+    allocate_name: str,
     setup_name: str,
     type_arg: str,
     native_prefix: bool,
@@ -680,8 +693,9 @@ def emit_null_check() -> None:
 
     new_fn = cl.get_method("__new__")
     if not new_fn:
-        emitter.emit_line(f"PyObject *self = {setup_name}({type_arg});")
+        emitter.emit_line(f"PyObject *self = {allocate_name}({type_arg});")
         emit_null_check()
+        emitter.emit_line(f"self = {setup_name}(self);")
         return
     prefix = emitter.get_group_prefix(new_fn.decl) + NATIVE_PREFIX if native_prefix else PREFIX
     all_args = type_arg
@@ -694,11 +708,14 @@ def emit_null_check() -> None:
     emitter.emit_line(f"if (Py_TYPE(self) != {emitter.type_struct_name(cl)})")
     emitter.emit_line("    return self;")
 
+    emitter.emit_line(f"self = {setup_name}(self);")
+
 
 def generate_constructor_for_class(
     cl: ClassIR,
     fn: FuncDecl,
     init_fn: FuncIR | None,
+    allocate_name: str,
     setup_name: str,
     vtable_name: str,
     emitter: Emitter,
@@ -717,7 +734,9 @@ def generate_constructor_for_class(
         and fn.sig.args[0].kind == ARG_STAR
         and fn.sig.args[1].kind == ARG_STAR2
     )
-    emit_setup_or_dunder_new_call(cl, setup_name, type_arg, not use_wrapper, new_args, emitter)
+    emit_allocate_or_dunder_new_call(
+        cl, allocate_name, setup_name, type_arg, not use_wrapper, new_args, emitter
+    )
 
     args = ", ".join(["self"] + fn_args)
     if init_fn is not None:
@@ -779,6 +798,7 @@ def generate_new_for_class(
     cl: ClassIR,
     func_name: str,
     vtable_name: str,
+    allocate_name: str,
     setup_name: str,
     init_fn: FuncIR | None,
     emitter: Emitter,
@@ -797,7 +817,9 @@ def generate_new_for_class(
 
     type_arg = "(PyObject*)type"
     new_args = "args, kwds"
-    emit_setup_or_dunder_new_call(cl, setup_name, type_arg, False, new_args, emitter)
+    emit_allocate_or_dunder_new_call(
+        cl, allocate_name, setup_name, type_arg, False, new_args, emitter
+    )
     if (
         not init_fn
         or cl.allow_interpreted_subclasses

mypyc/ir/class_ir.py

@@ -133,15 +133,25 @@ def __init__(
         self.builtin_base: str | None = None
         # Default empty constructor
         self.ctor = FuncDecl(name, None, module_name, FuncSignature([], RInstance(self)))
-        # Declare setup method that allocates and initializes an object. type is the
+        # Declare allocate method that allocates and initializes an object. type is the
         # type of the class being initialized, which could be another class if there
         # is an interpreted subclass.
         # TODO: Make it a regular method and generate its body in IR
+        self.allocate = FuncDecl(
+            "__mypyc__" + name + "_allocate",
+            None,
+            module_name,
+            FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)),
+        )
+        # Declare setup method that performs initialization that needs to happen after
+        # allocation in case of inherited dunder new methods, eg. setting the vtable pointer.
+        # The argument is an already allocated object.
+        # TODO: Make it a regular method and generate its body in IR
         self.setup = FuncDecl(
             "__mypyc__" + name + "_setup",
             None,
             module_name,
-            FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)),
+            FuncSignature([RuntimeArg("obj", object_rprimitive)], RInstance(self)),
         )
         # Attributes defined in the class (not inherited)
         self.attributes: dict[str, RType] = {}

mypyc/irbuild/specialize.py

@@ -1103,7 +1103,7 @@ def translate_object_new(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
     method_args = fn.fitem.arg_names
     if isinstance(typ_arg, NameExpr) and len(method_args) > 0 and method_args[0] == typ_arg.name:
         subtype = builder.accept(expr.args[0])
-        return builder.add(Call(ir.setup, [subtype], expr.line))
+        return builder.add(Call(ir.allocate, [subtype], expr.line))
 
     return None
 

mypyc/test-data/irbuild-classes.test

@@ -1715,7 +1715,7 @@ def Test.__new__(cls, val):
     r0, obj :: __main__.Test
     r1 :: bool
 L0:
-    r0 = __mypyc__Test_setup(cls)
+    r0 = __mypyc__Test_allocate(cls)
     obj = r0
     obj.val = val; r1 = is_error
     return obj
@@ -1732,7 +1732,7 @@ def NewClassMethod.__new__(cls, val):
     r0, obj :: __main__.NewClassMethod
     r1 :: bool
 L0:
-    r0 = __mypyc__NewClassMethod_setup(cls)
+    r0 = __mypyc__NewClassMethod_allocate(cls)
     obj = r0
     obj.val = val; r1 = is_error
     return obj
@@ -1870,7 +1870,7 @@ def Test.__new__(cls, val):
     r0, obj :: __main__.Test
     r1 :: bool
 L0:
-    r0 = __mypyc__Test_setup(cls)
+    r0 = __mypyc__Test_allocate(cls)
     obj = r0
     obj.val = val; r1 = is_error
     return obj
@@ -1887,7 +1887,7 @@ def NewClassMethod.__new__(cls, val):
     r0, obj :: __main__.NewClassMethod
     r1 :: bool
 L0:
-    r0 = __mypyc__NewClassMethod_setup(cls)
+    r0 = __mypyc__NewClassMethod_allocate(cls)
     obj = r0
     obj.val = val; r1 = is_error
     return obj