From 603d626db53c0363f2372a98c8f1f3dac8a3245a Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sat, 9 Aug 2025 21:02:55 +0100 Subject: [PATCH 1/5] Special-case enum method calls --- mypyc/common.py | 1 + mypyc/ir/class_ir.py | 5 +++ mypyc/irbuild/function.py | 14 ++++++- mypyc/irbuild/ll_builder.py | 5 ++- mypyc/irbuild/prepare.py | 27 ++++++++++++- mypyc/test-data/irbuild-classes.test | 60 ++++++++++++++++++++++++++++ mypyc/test-data/run-classes.test | 45 +++++++++++++++++++++ 7 files changed, 153 insertions(+), 4 deletions(-) diff --git a/mypyc/common.py b/mypyc/common.py index b5506eed89c2..982201566e0a 100644 --- a/mypyc/common.py +++ b/mypyc/common.py @@ -15,6 +15,7 @@ MODULE_PREFIX: Final = "CPyModule_" # Cached modules TYPE_VAR_PREFIX: Final = "CPyTypeVar_" # Type variables when using new-style Python 3.12 syntax ATTR_PREFIX: Final = "_" # Attributes +FAST_PREFIX: Final = "__mypyc_fast_" ENV_ATTR_NAME: Final = "__mypyc_env__" NEXT_LABEL_ATTR_NAME: Final = "__mypyc_next_label__" diff --git a/mypyc/ir/class_ir.py b/mypyc/ir/class_ir.py index 561dc9d438c4..f6015b64dcdd 100644 --- a/mypyc/ir/class_ir.py +++ b/mypyc/ir/class_ir.py @@ -210,6 +210,9 @@ def __init__( # per-type free "list" of up to length 1. self.reuse_freed_instance = False + # Is this a class inheriting from enum.Enum? Such classes can be special-cased. + self.is_enum = False + def __repr__(self) -> str: return ( "ClassIR(" @@ -410,6 +413,7 @@ def serialize(self) -> JsonDict: "init_self_leak": self.init_self_leak, "env_user_function": self.env_user_function.id if self.env_user_function else None, "reuse_freed_instance": self.reuse_freed_instance, + "is_enum": self.is_enum, } @classmethod @@ -466,6 +470,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR: ctx.functions[data["env_user_function"]] if data["env_user_function"] else None ) ir.reuse_freed_instance = data["reuse_freed_instance"] + ir.is_enum = data["is_enum"] return ir diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index 90506adde672..26d707b0b2f9 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -29,7 +29,7 @@ Var, ) from mypy.types import CallableType, Type, UnboundType, get_proper_type -from mypyc.common import LAMBDA_NAME, PROPSET_PREFIX, SELF_NAME +from mypyc.common import FAST_PREFIX, LAMBDA_NAME, PROPSET_PREFIX, SELF_NAME from mypyc.ir.class_ir import ClassIR, NonExtClassInfo from mypyc.ir.func_ir import ( FUNC_CLASSMETHOD, @@ -166,6 +166,7 @@ def gen_func_item( name: str, sig: FuncSignature, cdef: ClassDef | None = None, + make_ext_method: bool = False, ) -> tuple[FuncIR, Value | None]: """Generate and return the FuncIR for a given FuncDef. @@ -217,7 +218,7 @@ def c() -> None: class_name = None if cdef: ir = builder.mapper.type_to_ir[cdef.info] - in_non_ext = not ir.is_ext_class + in_non_ext = not ir.is_ext_class and not make_ext_method class_name = cdef.name if is_singledispatch: @@ -339,6 +340,8 @@ def gen_func_ir( fitem = fn_info.fitem assert isinstance(fitem, FuncDef), fitem func_decl = builder.mapper.func_to_decl[fitem] + if cdef and fn_info.name == FAST_PREFIX + func_decl.name: + func_decl = builder.mapper.type_to_ir[cdef.info].method_decls[fn_info.name] if fn_info.is_decorated or is_singledispatch_main_func: class_name = None if cdef is None else cdef.name func_decl = FuncDecl( @@ -453,6 +456,13 @@ def handle_non_ext_method( builder.add_to_non_ext_dict(non_ext, name, func_reg, fdef.line) + class_ir = builder.mapper.type_to_ir[cdef.info] + name = FAST_PREFIX + fdef.name + if name in class_ir.method_decls: + func_ir, func_reg = gen_func_item(builder, fdef, name, sig, cdef, make_ext_method=True) + class_ir.methods[name] = func_ir + builder.functions.append(func_ir) + def gen_func_ns(builder: IRBuilder) -> str: """Generate a namespace for a nested function using its outer function names.""" diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py index a5e28268efed..05d558e0822a 100644 --- a/mypyc/irbuild/ll_builder.py +++ b/mypyc/irbuild/ll_builder.py @@ -17,6 +17,7 @@ from mypyc.common import ( BITMAP_BITS, FAST_ISINSTANCE_MAX_SUBCLASSES, + FAST_PREFIX, IS_FREE_THREADED, MAX_LITERAL_SHORT_INT, MAX_SHORT_INT, @@ -1171,11 +1172,13 @@ def gen_method_call( return self.py_method_call(base, name, arg_values, line, arg_kinds, arg_names) # If the base type is one of ours, do a MethodCall + fast_name = FAST_PREFIX + name if ( isinstance(base.type, RInstance) - and base.type.class_ir.is_ext_class + and (base.type.class_ir.is_ext_class or base.type.class_ir.has_method(fast_name)) and not base.type.class_ir.builtin_base ): + name = name if base.type.class_ir.is_ext_class else fast_name if base.type.class_ir.has_method(name): decl = base.type.class_ir.method_decl(name) if arg_kinds is None: diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index 1d6117ab7b1e..ddd376617b32 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -38,7 +38,7 @@ from mypy.semanal import refers_to_fullname from mypy.traverser import TraverserVisitor from mypy.types import Instance, Type, get_proper_type -from mypyc.common import PROPSET_PREFIX, SELF_NAME, get_id_from_name +from mypyc.common import FAST_PREFIX, PROPSET_PREFIX, SELF_NAME, get_id_from_name from mypyc.crash import catch_errors from mypyc.errors import Errors from mypyc.ir.class_ir import ClassIR @@ -106,6 +106,7 @@ def build_type_map( class_ir.children = None mapper.type_to_ir[cdef.info] = class_ir mapper.symbol_fullnames.add(class_ir.fullname) + class_ir.is_enum = cdef.info.is_enum and len(cdef.info.enum_members) > 0 # Populate structural information in class IR for extension classes. for module, cdef in classes: @@ -270,6 +271,28 @@ def prepare_method_def( ir.property_types[node.name] = decl.sig.ret_type +def prepare_fast_path( + ir: ClassIR, + module_name: str, + cdef: ClassDef, + mapper: Mapper, + node: SymbolNode | None, + options: CompilerOptions, +) -> None: + if ir.is_enum: + if isinstance(node, OverloadedFuncDef): + if node.is_property: + return + node = node.impl + if not isinstance(node, FuncDef): + return + name = FAST_PREFIX + node.name + sig = mapper.fdef_to_sig(node, options.strict_dunders_typing) + decl = FuncDecl(name, cdef.name, module_name, sig, FUNC_NORMAL) + ir.method_decls[name] = decl + return + + def is_valid_multipart_property_def(prop: OverloadedFuncDef) -> bool: # Checks to ensure supported property decorator semantics if len(prop.items) != 2: @@ -579,6 +602,8 @@ def prepare_non_ext_class_def( else: prepare_method_def(ir, module_name, cdef, mapper, get_func_def(node.node), options) + prepare_fast_path(ir, module_name, cdef, mapper, node.node, options) + if any(cls in mapper.type_to_ir and mapper.type_to_ir[cls].is_ext_class for cls in info.mro): errors.error( "Non-extension classes may not inherit from extension classes", path, cdef.line diff --git a/mypyc/test-data/irbuild-classes.test b/mypyc/test-data/irbuild-classes.test index 1a2c237cc3c9..f8ea26cd41e8 100644 --- a/mypyc/test-data/irbuild-classes.test +++ b/mypyc/test-data/irbuild-classes.test @@ -1408,3 +1408,63 @@ class TestOverload: def __mypyc_generator_helper__(self, x: Any) -> Any: return x + +[case testEnumFastPath] +from enum import Enum + +def test(e: E) -> bool: + return e.is_one() + +class E(Enum): + ONE = 1 + TWO = 2 + + def is_one(self) -> bool: + return self == E.ONE +[out] +def test(e): + e :: __main__.E + r0 :: bool +L0: + r0 = e.__mypyc_fast_is_one() + return r0 +def is_one_E_obj.__get__(__mypyc_self__, instance, owner): + __mypyc_self__, instance, owner, r0 :: object + r1 :: bit + r2 :: object +L0: + r0 = load_address _Py_NoneStruct + r1 = instance == r0 + if r1 goto L1 else goto L2 :: bool +L1: + return __mypyc_self__ +L2: + r2 = PyMethod_New(__mypyc_self__, instance) + return r2 +def is_one_E_obj.__call__(__mypyc_self__, self): + __mypyc_self__ :: __main__.is_one_E_obj + self, r0 :: __main__.E + r1 :: bool + r2 :: bit +L0: + r0 = __main__.E.ONE :: static + if is_error(r0) goto L1 else goto L2 +L1: + r1 = raise NameError('value for final name "ONE" was not set') + unreachable +L2: + r2 = self == r0 + return r2 +def E.__mypyc_fast_is_one(self): + self, r0 :: __main__.E + r1 :: bool + r2 :: bit +L0: + r0 = __main__.E.ONE :: static + if is_error(r0) goto L1 else goto L2 +L1: + r1 = raise NameError('value for final name "ONE" was not set') + unreachable +L2: + r2 = self == r0 + return r2 diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index 54f5343bc7bb..488ae9136c0c 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -2710,6 +2710,51 @@ from native import Player [out] Player.MIN = +[case testEnumMethodCalls] +from enum import Enum +from typing import overload, Self + +class C: + def foo(self, x: Test) -> bool: + return x.is_one(inverse=True) + +class Test(Enum): + ONE = 1 + TWO = 2 + THREE = 3 + + def is_one(self, *, inverse: bool = False) -> bool: + if inverse: + return self != Test.ONE + return self == Test.ONE + + @classmethod + def next(cls, val: int) -> Self: + return cls(val + 1) + + @staticmethod + def prev(val: int) -> Test: + return Test(val - 1) + + @overload + def enigma(self, val: int) -> bool: ... + @overload + def enigma(self, val: str | None = None) -> int: ... + def enigma(self, val: int | str | None = None) -> int | bool: + if isinstance(val, int): + return self.is_one() + return 22 +[file driver.py] +from native import Test, C + +assert Test.ONE.is_one() +assert Test.TWO.is_one(inverse=True) +assert not C().foo(Test.ONE) +assert Test.next(2) == Test.THREE +assert Test.prev(2) == Test.ONE +assert Test.ONE.enigma(22) +assert Test.ONE.enigma("22") == 22 + [case testStaticCallsWithUnpackingArgs] from typing import Tuple From 1abbc3f6d0cc5cdd94b0c36bfa4701f5495dddb2 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sun, 10 Aug 2025 01:11:24 +0100 Subject: [PATCH 2/5] Add some comments --- mypyc/common.py | 2 +- mypyc/irbuild/function.py | 3 +++ mypyc/irbuild/prepare.py | 8 ++++++++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/mypyc/common.py b/mypyc/common.py index 982201566e0a..98edeaa68b61 100644 --- a/mypyc/common.py +++ b/mypyc/common.py @@ -15,7 +15,7 @@ MODULE_PREFIX: Final = "CPyModule_" # Cached modules TYPE_VAR_PREFIX: Final = "CPyTypeVar_" # Type variables when using new-style Python 3.12 syntax ATTR_PREFIX: Final = "_" # Attributes -FAST_PREFIX: Final = "__mypyc_fast_" +FAST_PREFIX: Final = "__mypyc_fast_" # Optimized methods in non-extension classes ENV_ATTR_NAME: Final = "__mypyc_env__" NEXT_LABEL_ATTR_NAME: Final = "__mypyc_next_label__" diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index 26d707b0b2f9..d70b16475503 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py @@ -341,6 +341,7 @@ def gen_func_ir( assert isinstance(fitem, FuncDef), fitem func_decl = builder.mapper.func_to_decl[fitem] if cdef and fn_info.name == FAST_PREFIX + func_decl.name: + # Special-cased version of a method has a separate FuncDecl, use that one. func_decl = builder.mapper.type_to_ir[cdef.info].method_decls[fn_info.name] if fn_info.is_decorated or is_singledispatch_main_func: class_name = None if cdef is None else cdef.name @@ -456,6 +457,8 @@ def handle_non_ext_method( builder.add_to_non_ext_dict(non_ext, name, func_reg, fdef.line) + # If we identified that this non-extension class method can be special-cased for + # direct access during prepare phase, generate a "static" version of it. class_ir = builder.mapper.type_to_ir[cdef.info] name = FAST_PREFIX + fdef.name if name in class_ir.method_decls: diff --git a/mypyc/irbuild/prepare.py b/mypyc/irbuild/prepare.py index ddd376617b32..83ec3f7c1d38 100644 --- a/mypyc/irbuild/prepare.py +++ b/mypyc/irbuild/prepare.py @@ -279,13 +279,21 @@ def prepare_fast_path( node: SymbolNode | None, options: CompilerOptions, ) -> None: + """Add fast (direct) variants of methods in non-extension classes.""" if ir.is_enum: + # We check that non-empty enums are implicitly final in mypy, so we + # can generate direct calls to enum methods. if isinstance(node, OverloadedFuncDef): if node.is_property: return node = node.impl if not isinstance(node, FuncDef): + # TODO: support decorated methods (at least @classmethod and @staticmethod). return + # The simplest case is a regular or overloaded method without decorators. In this + # case we can generate practically identical IR method body, but with a signature + # suitable for direct calls (usual non-extension class methods are converted to + # callable classes, and thus have an extra __mypyc_self__ argument). name = FAST_PREFIX + node.name sig = mapper.fdef_to_sig(node, options.strict_dunders_typing) decl = FuncDecl(name, cdef.name, module_name, sig, FUNC_NORMAL) From c5c6ef303aab25d73e27566d61a03bad137a8b5b Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sun, 10 Aug 2025 01:51:05 +0100 Subject: [PATCH 3/5] Argh 3.9 --- mypyc/test-data/run-classes.test | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index 488ae9136c0c..caabfa4272e9 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -2712,7 +2712,7 @@ Player.MIN = [case testEnumMethodCalls] from enum import Enum -from typing import overload, Self +from typing import overload, Optional, Self, Union class C: def foo(self, x: Test) -> bool: @@ -2739,8 +2739,8 @@ class Test(Enum): @overload def enigma(self, val: int) -> bool: ... @overload - def enigma(self, val: str | None = None) -> int: ... - def enigma(self, val: int | str | None = None) -> int | bool: + def enigma(self, val: Optional[str] = None) -> int: ... + def enigma(self, val: Union[int, str, None] = None) -> Union[int, bool]: if isinstance(val, int): return self.is_one() return 22 From 97b68746c858678bb3384bb1417e8689e240180e Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sun, 10 Aug 2025 13:18:39 +0100 Subject: [PATCH 4/5] More 3.9 --- mypyc/test-data/run-classes.test | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index caabfa4272e9..8e5895a5decc 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -2712,7 +2712,7 @@ Player.MIN = [case testEnumMethodCalls] from enum import Enum -from typing import overload, Optional, Self, Union +from typing import overload, Optional, Union class C: def foo(self, x: Test) -> bool: @@ -2729,7 +2729,7 @@ class Test(Enum): return self == Test.ONE @classmethod - def next(cls, val: int) -> Self: + def next(cls, val: int) -> Test: return cls(val + 1) @staticmethod From 6546051da4f6230b0695a8e117b31fe61a37e1a6 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Mon, 11 Aug 2025 16:29:14 +0100 Subject: [PATCH 5/5] Call all kinds of methods in the compiled code --- mypyc/test-data/run-classes.test | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test index 8e5895a5decc..1481f3e06871 100644 --- a/mypyc/test-data/run-classes.test +++ b/mypyc/test-data/run-classes.test @@ -2716,6 +2716,11 @@ from typing import overload, Optional, Union class C: def foo(self, x: Test) -> bool: + assert Test.ONE.is_one() + assert x.next(2) == Test.THREE + assert x.prev(2) == Test.ONE + assert x.enigma(22) + assert x.enigma("22") == 22 return x.is_one(inverse=True) class Test(Enum):