Merge branch 'master' into for-map

Author: BobTheBuidler

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_"  # Optimized methods in non-extension classes
 
 ENV_ATTR_NAME: Final = "__mypyc_env__"
 NEXT_LABEL_ATTR_NAME: Final = "__mypyc_next_label__"

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
 

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,9 @@ 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:
+            # 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
             func_decl = FuncDecl(
@@ -453,6 +457,15 @@ 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:
+        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."""

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,
@@ -788,6 +789,18 @@ def _construct_varargs(
         for value, kind, name in args:
             if kind == ARG_STAR:
                 if star_result is None:
+                    # fast path if star expr is a tuple:
+                    # we can pass the immutable tuple straight into the function call.
+                    if is_tuple_rprimitive(value.type):
+                        if len(args) == 1:
+                            # fn(*args)
+                            return value, self._create_dict([], [], line)
+                        elif len(args) == 2 and args[1][1] == ARG_STAR2:
+                            # fn(*args, **kwargs)
+                            star_result = value
+                            continue
+                        # elif ...: TODO extend this to optimize fn(*args, k=1, **kwargs) case
+                    # TODO optimize this case using the length utils - currently in review
                     star_result = self.new_list_op(star_values, line)
                 self.primitive_op(list_extend_op, [star_result, value], line)
             elif kind == ARG_STAR2:
@@ -885,9 +898,11 @@ def _construct_varargs(
             # tuple. Otherwise create the tuple from the list.
             if star_result is None:
                 star_result = self.new_tuple(star_values, line)
-            else:
+            elif not is_tuple_rprimitive(star_result.type):
+                # if star_result is a tuple we took the fast path
                 star_result = self.primitive_op(list_tuple_op, [star_result], line)
         if has_star2 and star2_result is None:
+            # TODO: use dict_copy_op for simple cases of **kwargs
             star2_result = self._create_dict(star2_keys, star2_values, line)
 
         return star_result, star2_result
@@ -1171,11 +1186,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:

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,36 @@ 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:
+    """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)
+        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 +610,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

mypyc/test-data/irbuild-classes.test

@@ -180,7 +180,7 @@ L0:
     o.x = r1; r2 = is_error
     return o
 
-[case testSubclass_toplevel]
+[case testSubclass_withgil_toplevel]
 from typing import TypeVar, Generic
 from mypy_extensions import trait
 T = TypeVar('T')
@@ -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

mypyc/test-data/irbuild-generics.test

@@ -166,25 +166,18 @@ execute(f, 1)
 def execute(func, args, kwargs):
     func :: object
     args :: tuple
-    kwargs :: dict
-    r0 :: list
-    r1 :: object
-    r2 :: dict
-    r3 :: i32
-    r4 :: bit
-    r5 :: tuple
-    r6 :: object
-    r7 :: int
+    kwargs, r0 :: dict
+    r1 :: i32
+    r2 :: bit
+    r3 :: object
+    r4 :: int
 L0:
-    r0 = PyList_New(0)
-    r1 = CPyList_Extend(r0, args)
-    r2 = PyDict_New()
-    r3 = CPyDict_UpdateInDisplay(r2, kwargs)
-    r4 = r3 >= 0 :: signed
-    r5 = PyList_AsTuple(r0)
-    r6 = PyObject_Call(func, r5, r2)
-    r7 = unbox(int, r6)
-    return r7
+    r0 = PyDict_New()
+    r1 = CPyDict_UpdateInDisplay(r0, kwargs)
+    r2 = r1 >= 0 :: signed
+    r3 = PyObject_Call(func, args, r0)
+    r4 = unbox(int, r3)
+    return r4
 def f(x):
     x :: int
 L0:
@@ -709,14 +702,11 @@ def inner_deco_obj.__call__(__mypyc_self__, args, kwargs):
     can_dictcomp :: dict
     r22, can_iter, r23, can_use_keys, r24, can_use_values :: list
     r25 :: object
-    r26 :: list
-    r27 :: object
-    r28 :: dict
-    r29 :: i32
-    r30 :: bit
-    r31 :: tuple
-    r32 :: object
-    r33 :: int
+    r26 :: dict
+    r27 :: i32
+    r28 :: bit
+    r29 :: object
+    r30 :: int
 L0:
     r0 = __mypyc_self__.__mypyc_env__
     r1 = var_object_size args
@@ -768,15 +758,12 @@ L9:
     r24 = CPyDict_Values(kwargs)
     can_use_values = r24
     r25 = r0.func
-    r26 = PyList_New(0)
-    r27 = CPyList_Extend(r26, args)
-    r28 = PyDict_New()
-    r29 = CPyDict_UpdateInDisplay(r28, kwargs)
-    r30 = r29 >= 0 :: signed
-    r31 = PyList_AsTuple(r26)
-    r32 = PyObject_Call(r25, r31, r28)
-    r33 = unbox(int, r32)
-    return r33
+    r26 = PyDict_New()
+    r27 = CPyDict_UpdateInDisplay(r26, kwargs)
+    r28 = r27 >= 0 :: signed
+    r29 = PyObject_Call(r25, args, r26)
+    r30 = unbox(int, r29)
+    return r30
 def deco(func):
     func :: object
     r0 :: __main__.deco_env