Optimize calls to final classes

Author: jairov4

mypyc/codegen/emitclass.py

@@ -380,7 +380,9 @@ def setter_name(cl: ClassIR, attribute: str, names: NameGenerator) -> str:
 def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
     seen_attrs: set[tuple[str, RType]] = set()
     lines: list[str] = []
-    lines += ["typedef struct {", "PyObject_HEAD", "CPyVTableItem *vtable;"]
+    lines += ["typedef struct {", "PyObject_HEAD"]
+    if not cl.is_final_class:
+        lines.append("CPyVTableItem *vtable;")
     if cl.has_method("__call__") and emitter.use_vectorcall():
         lines.append("vectorcallfunc vectorcall;")
     bitmap_attrs = []
@@ -563,14 +565,16 @@ def generate_setup_for_class(
     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};")
+    if not cl.is_final_class:
+        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};")
+
     for i in range(0, len(cl.bitmap_attrs), BITMAP_BITS):
         field = emitter.bitmap_field(i)
         emitter.emit_line(f"self->{field} = 0;")

mypyc/codegen/emitfunc.py

@@ -361,21 +361,24 @@ def visit_get_attr(self, op: GetAttr) -> None:
         attr_rtype, decl_cl = cl.attr_details(op.attr)
         prefer_method = cl.is_trait and attr_rtype.error_overlap
         if cl.get_method(op.attr, prefer_method=prefer_method):
-            # Properties are essentially methods, so use vtable access for them.
-            version = "_TRAIT" if cl.is_trait else ""
-            self.emit_line(
-                "%s = CPY_GET_ATTR%s(%s, %s, %d, %s, %s); /* %s */"
-                % (
-                    dest,
-                    version,
-                    obj,
-                    self.emitter.type_struct_name(rtype.class_ir),
-                    rtype.getter_index(op.attr),
-                    rtype.struct_name(self.names),
-                    self.ctype(rtype.attr_type(op.attr)),
-                    op.attr,
+            if cl.is_method_final(op.attr):
+                self.emit_method_call(f"{dest} = ", op.obj, op.attr, [])
+            else:
+                # Properties are essentially methods, so use vtable access for them.
+                version = "_TRAIT" if cl.is_trait else ""
+                self.emit_line(
+                    "%s = CPY_GET_ATTR%s(%s, %s, %d, %s, %s); /* %s */"
+                    % (
+                        dest,
+                        version,
+                        obj,
+                        self.emitter.type_struct_name(rtype.class_ir),
+                        rtype.getter_index(op.attr),
+                        rtype.struct_name(self.names),
+                        self.ctype(rtype.attr_type(op.attr)),
+                        op.attr,
+                    )
                 )
-            )
         else:
             # Otherwise, use direct or offset struct access.
             attr_expr = self.get_attr_expr(obj, op, decl_cl)
@@ -529,11 +532,12 @@ def visit_call(self, op: Call) -> None:
     def visit_method_call(self, op: MethodCall) -> None:
         """Call native method."""
         dest = self.get_dest_assign(op)
-        obj = self.reg(op.obj)
+        self.emit_method_call(dest, op.obj, op.method, op.args)
 
-        rtype = op.receiver_type
+    def emit_method_call(self, dest: str, op_obj: Value, name: str, op_args: list[Value]) -> None:
+        obj = self.reg(op_obj)
+        rtype = op_obj.type
         class_ir = rtype.class_ir
-        name = op.method
         method = rtype.class_ir.get_method(name)
         assert method is not None
 
@@ -547,7 +551,7 @@ def visit_method_call(self, op: MethodCall) -> None:
             if method.decl.kind == FUNC_STATICMETHOD
             else [f"(PyObject *)Py_TYPE({obj})"] if method.decl.kind == FUNC_CLASSMETHOD else [obj]
         )
-        args = ", ".join(obj_args + [self.reg(arg) for arg in op.args])
+        args = ", ".join(obj_args + [self.reg(arg) for arg in op_args])
         mtype = native_function_type(method, self.emitter)
         version = "_TRAIT" if rtype.class_ir.is_trait else ""
         if is_direct:
@@ -567,7 +571,7 @@ def visit_method_call(self, op: MethodCall) -> None:
                     rtype.struct_name(self.names),
                     mtype,
                     args,
-                    op.method,
+                    name,
                 )
             )
 

mypyc/ir/class_ir.py

@@ -93,13 +93,15 @@ def __init__(
         is_generated: bool = False,
         is_abstract: bool = False,
         is_ext_class: bool = True,
+        is_final_class: bool = False,
     ) -> None:
         self.name = name
         self.module_name = module_name
         self.is_trait = is_trait
         self.is_generated = is_generated
         self.is_abstract = is_abstract
         self.is_ext_class = is_ext_class
+        self.is_final_class = is_final_class
         # An augmented class has additional methods separate from what mypyc generates.
         # Right now the only one is dataclasses.
         self.is_augmented = False
@@ -248,8 +250,7 @@ def has_method(self, name: str) -> bool:
     def is_method_final(self, name: str) -> bool:
         subs = self.subclasses()
         if subs is None:
-            # TODO: Look at the final attribute!
-            return False
+            return self.is_final_class
 
         if self.has_method(name):
             method_decl = self.method_decl(name)

mypyc/irbuild/classdef.py

@@ -63,7 +63,7 @@
     handle_non_ext_method,
     load_type,
 )
-from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator
+from mypyc.irbuild.util import dataclass_type, get_func_def, is_constant, is_dataclass_decorator, is_final_class
 from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
 from mypyc.primitives.generic_ops import (
     iter_op,
@@ -294,6 +294,7 @@ def add_attr(self, lvalue: NameExpr, stmt: AssignmentStmt) -> None:
             self.builder.init_final_static(lvalue, value, self.cdef.name)
 
     def finalize(self, ir: ClassIR) -> None:
+        ir.is_final_class = is_final_class(self.cdef)
         attrs_with_defaults, default_assignments = find_attr_initializers(
             self.builder, self.cdef, self.skip_attr_default
         )

mypyc/irbuild/util.py

@@ -31,6 +31,14 @@
 DATACLASS_DECORATORS = {"dataclasses.dataclass", "attr.s", "attr.attrs"}
 
 
+def is_final_decorator(d: Expression) -> bool:
+    return isinstance(d, RefExpr) and d.fullname == "typing.final"
+
+
+def is_final_class(cdef: ClassDef) -> bool:
+    return any(is_final_decorator(d) for d in cdef.decorators)
+
+
 def is_trait_decorator(d: Expression) -> bool:
     return isinstance(d, RefExpr) and d.fullname == "mypy_extensions.trait"
 
@@ -119,7 +127,7 @@ def get_mypyc_attrs(stmt: ClassDef | Decorator) -> dict[str, Any]:
 
 def is_extension_class(cdef: ClassDef) -> bool:
     if any(
-        not is_trait_decorator(d) and not is_dataclass_decorator(d) and not get_mypyc_attr_call(d)
+        not is_trait_decorator(d) and not is_dataclass_decorator(d) and not get_mypyc_attr_call(d) and not is_final_decorator(d)
         for d in cdef.decorators
     ):
         return False