[mypyc] Optimize calls of final methods in classes allowing interpreted subclasses

mypyc/ir/class_ir.py

@@ -369,12 +369,21 @@ def has_method(self, name: str) -> bool:
         return True
 
     def is_method_final(self, name: str) -> bool:
+        method_decl: FuncDecl | None = None
+        try:
+            method_decl = self.method_decl(name)
+        except KeyError:
+            pass
+        # A declared @final method cannot be overridden in checked code. Trust
+        # this even when interpreted subclasses cannot be enumerated.
+        if method_decl is not None and method_decl.is_final:
+            return True
+
         subs = self.subclasses()
         if subs is None:
             return self.is_final_class
 
-        if self.has_method(name):
-            method_decl = self.method_decl(name)
+        if method_decl is not None:
             for subc in subs:
                 if subc.method_decl(name) != method_decl:
                     return False

mypyc/ir/func_ir.py

@@ -154,6 +154,7 @@ def __init__(
         is_prop_getter: bool = False,
         is_generator: bool = False,
         is_coroutine: bool = False,
+        is_final: bool = False,
         implicit: bool = False,
         internal: bool = False,
     ) -> None:
@@ -166,6 +167,8 @@ def __init__(
         self.is_prop_getter = is_prop_getter
         self.is_generator = is_generator
         self.is_coroutine = is_coroutine
+        # A declared @final method may be called without vtable dispatch.
+        self.is_final = is_final
         if class_name is None:
             self.bound_sig: FuncSignature | None = None
         else:
@@ -226,6 +229,7 @@ def serialize(self) -> JsonDict:
             "is_prop_getter": self.is_prop_getter,
             "is_generator": self.is_generator,
             "is_coroutine": self.is_coroutine,
+            "is_final": self.is_final,
             "implicit": self.implicit,
             "internal": self.internal,
         }
@@ -251,6 +255,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> FuncDecl:
             is_prop_getter=data["is_prop_getter"],
             is_generator=data["is_generator"],
             is_coroutine=data["is_coroutine"],
+            is_final=data["is_final"],
             implicit=data["implicit"],
             internal=data["internal"],
         )

mypyc/irbuild/function.py

@@ -369,6 +369,7 @@ def gen_func_ir(
                 is_prop_setter=func_decl.is_prop_setter,
                 is_generator=func_decl.is_generator,
                 is_coroutine=func_decl.is_coroutine,
+                is_final=func_decl.is_final,
                 implicit=func_decl.implicit,
                 internal=func_decl.internal,
             )

mypyc/irbuild/prepare.py

@@ -256,6 +256,7 @@ def prepare_func_def(
         kind,
         is_generator=fdef.is_generator,
         is_coroutine=fdef.is_coroutine,
+        is_final=fdef.is_final,
     )
     mapper.func_to_decl[fdef] = decl
     return decl

mypyc/test-data/run-multimodule.test

@@ -2106,6 +2106,115 @@ class CompiledBase:
     def value(self) -> int:
         raise NotImplementedError
 
+[case testFinalMethodCallsWithinAndAcrossModules]
+from other import FinalMethods
+
+def cross_module_calls(obj: FinalMethods) -> tuple[int, int, int, int, int, int, int]:
+    return (
+        obj.regular(2),
+        obj.class_method(2),
+        obj.static_method(2),
+        sum(obj.generator(2)),
+        FinalMethods.class_method(3),
+        FinalMethods.static_method(3),
+        obj.decorated_final(2),
+    )
+
+async def cross_module_async(obj: FinalMethods) -> tuple[int, int, int, int, int]:
+    return (
+        await obj.async_method(2),
+        await obj.async_class_method(2),
+        await obj.async_static_method(2),
+        await FinalMethods.async_class_method(3),
+        await FinalMethods.async_static_method(3),
+    )
+
+[file other.py]
+from collections.abc import Iterator
+from typing import TypeVar, final
+from mypy_extensions import mypyc_attr
+
+T = TypeVar("T")
+
+def identity(value: T) -> T:
+    return value
+
+@mypyc_attr(allow_interpreted_subclasses=True)
+class FinalMethods:
+    def __init__(self, base: int) -> None:
+        self.base = base
+
+    @final
+    def regular(self, value: int) -> int:
+        return self.base + value
+
+    @classmethod
+    @final
+    def class_method(cls, value: int) -> int:
+        return value + 20
+
+    @staticmethod
+    @final
+    def static_method(value: int) -> int:
+        return value + 30
+
+    @final
+    async def async_method(self, value: int) -> int:
+        return self.base + value + 40
+
+    @classmethod
+    @final
+    async def async_class_method(cls, value: int) -> int:
+        return value + 70
+
+    @staticmethod
+    @final
+    async def async_static_method(value: int) -> int:
+        return value + 80
+
+    @final
+    def generator(self, value: int) -> Iterator[int]:
+        yield self.base + value + 50
+        yield self.base + value + 60
+
+    @final
+    @identity
+    def decorated_final(self, value: int) -> int:
+        return self.base + value + 90
+
+def same_module_calls(obj: FinalMethods) -> tuple[int, int, int, int, int, int, int]:
+    return (
+        obj.regular(2),
+        obj.class_method(2),
+        obj.static_method(2),
+        sum(obj.generator(2)),
+        FinalMethods.class_method(3),
+        FinalMethods.static_method(3),
+        obj.decorated_final(2),
+    )
+
+async def same_module_async(obj: FinalMethods) -> tuple[int, int, int, int, int]:
+    return (
+        await obj.async_method(2),
+        await obj.async_class_method(2),
+        await obj.async_static_method(2),
+        await FinalMethods.async_class_method(3),
+        await FinalMethods.async_static_method(3),
+    )
+
+[file driver.py]
+import asyncio
+from native import cross_module_calls, cross_module_async
+from other import FinalMethods, same_module_calls, same_module_async
+
+obj = FinalMethods(10)
+expected = (12, 22, 32, 134, 23, 33, 102)
+assert same_module_calls(obj) == expected
+assert cross_module_calls(obj) == expected
+expected_async = (52, 72, 82, 73, 83)
+assert asyncio.run(same_module_async(obj)) == expected_async
+assert asyncio.run(cross_module_async(obj)) == expected_async
+
 [case testCallFunctionInLazilyImportedModule]
 import my_lib
 

mypyc/test/test_emitfunc.py

@@ -7,7 +7,14 @@
 from mypyc.codegen.emitfunc import FunctionEmitterVisitor, generate_native_function
 from mypyc.common import HAVE_IMMORTAL, IS_FREE_THREADED, PLATFORM_SIZE
 from mypyc.ir.class_ir import ClassIR
-from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
+from mypyc.ir.func_ir import (
+    FUNC_CLASSMETHOD,
+    FUNC_STATICMETHOD,
+    FuncDecl,
+    FuncIR,
+    FuncSignature,
+    RuntimeArg,
+)
 from mypyc.ir.ops import (
     ERR_NEVER,
     Assign,
@@ -21,6 +28,7 @@
     ComparisonOp,
     CString,
     DecRef,
+    DeserMaps,
     Extend,
     GetAttr,
     GetElementPtr,
@@ -31,6 +39,7 @@
     LoadAddress,
     LoadLiteral,
     LoadMem,
+    MethodCall,
     Op,
     Register,
     Return,
@@ -300,6 +309,48 @@ def test_call_two_args(self) -> None:
             Call(decl, [self.m, self.k], 55), "cpy_r_r0 = CPyDef_myfn(cpy_r_m, cpy_r_k);"
         )
 
+    def test_final_method_calls_bypass_vtable(self) -> None:
+        assert isinstance(self.r.type, RInstance)
+        cl = self.r.type.class_ir
+        cl.allow_interpreted_subclasses = True
+        sig = FuncSignature(
+            [RuntimeArg("self", self.r.type), RuntimeArg("n", int_rprimitive)], int_rprimitive
+        )
+        final_decl = FuncDecl("final_method", "A", "mod", sig, is_final=True)
+        cl.method_decls["final_method"] = final_decl
+        assert FuncDecl.deserialize(final_decl.serialize(), DeserMaps({"mod.A": cl}, {})).is_final
+        assert cl.is_method_final("final_method")
+        self.assert_emit(
+            MethodCall(self.r, "final_method", [self.n]),
+            "cpy_r_r0 = CPyDef_A___final_method(cpy_r_r, cpy_r_n);",
+        )
+
+        cl.method_decls["final_class_method"] = FuncDecl(
+            "final_class_method", "A", "mod", sig, FUNC_CLASSMETHOD, is_final=True
+        )
+        self.assert_emit(
+            MethodCall(self.r, "final_class_method", [self.n]),
+            "cpy_r_r0 = CPyDef_A___final_class_method((PyObject *)Py_TYPE(cpy_r_r), cpy_r_n);",
+        )
+
+        static_sig = FuncSignature([RuntimeArg("n", int_rprimitive)], int_rprimitive)
+        cl.method_decls["final_static_method"] = FuncDecl(
+            "final_static_method", "A", "mod", static_sig, FUNC_STATICMETHOD, is_final=True
+        )
+        self.assert_emit(
+            MethodCall(self.r, "final_static_method", [self.n]),
+            "cpy_r_r0 = CPyDef_A___final_static_method(cpy_r_n);",
+        )
+
+        cl.method_decls["virtual_method"] = FuncDecl("virtual_method", "A", "mod", sig)
+        cl.vtable = {"virtual_method": 0}
+        assert not cl.is_method_final("virtual_method")
+        self.assert_emit(
+            MethodCall(self.r, "virtual_method", [self.n]),
+            "cpy_r_r0 = CPY_GET_METHOD(cpy_r_r, CPyType_A, 0, mod___AObject, "
+            "CPyTagged (*)(PyObject *, CPyTagged))(cpy_r_r, cpy_r_n); /* virtual_method */",
+        )
+
     def test_inc_ref(self) -> None:
         self.assert_emit(IncRef(self.o), "CPy_INCREF(cpy_r_o);")
         self.assert_emit(IncRef(self.o), "CPy_INCREF(cpy_r_o);", rare=True)