[mypyc] Allow specifying primitives as pure

mypyc/ir/ops.py

@@ -600,6 +600,7 @@ def __init__(
         ordering: list[int] | None,
         extra_int_constants: list[tuple[int, RType]],
         priority: int,
+        is_pure: bool,
     ) -> None:
         # Each primitive much have a distinct name, but otherwise they are arbitrary.
         self.name: Final = name
@@ -617,6 +618,11 @@ def __init__(
         self.ordering: Final = ordering
         self.extra_int_constants: Final = extra_int_constants
         self.priority: Final = priority
+        # Pure primitives have no side effects, take immutable arguments, and
+        # never fail. They support additional optimizations.
+        self.is_pure: Final = is_pure
+        if is_pure:
+            assert error_kind == ERR_NEVER
 
     def __repr__(self) -> str:
         return f"<PrimitiveDescription {self.name}>"
@@ -1036,6 +1042,8 @@ def __init__(
         error_kind: int,
         line: int,
         var_arg_idx: int = -1,
+        *,
+        is_pure: bool = False,
     ) -> None:
         self.error_kind = error_kind
         super().__init__(line)
@@ -1046,6 +1054,12 @@ def __init__(
         self.is_borrowed = is_borrowed
         # The position of the first variable argument in args (if >= 0)
         self.var_arg_idx = var_arg_idx
+        # Is the function pure? Pure functions have no side effects
+        # and all the arguments are immutable. Pure functions support
+        # additional optimizations. Pure functions never fail.
+        self.is_pure = is_pure
+        if is_pure:
+            assert error_kind == ERR_NEVER
 
     def sources(self) -> list[Value]:
         return self.args

mypyc/irbuild/ll_builder.py

@@ -1821,6 +1821,7 @@ def call_c(
                 error_kind,
                 line,
                 var_arg_idx,
+                is_pure=desc.is_pure,
             )
         )
         if desc.is_borrowed:
@@ -1903,6 +1904,7 @@ def primitive_op(
                 desc.ordering,
                 desc.extra_int_constants,
                 desc.priority,
+                is_pure=desc.is_pure,
             )
             return self.call_c(c_desc, args, line, result_type)
 

mypyc/primitives/int_ops.py

@@ -199,6 +199,7 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
     return_type=bit_rprimitive,
     c_function_name="CPyTagged_IsEq_",
     error_kind=ERR_NEVER,
+    is_pure=True,
 )
 
 # Less than operation on two boxed tagged integers
@@ -207,6 +208,7 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
     return_type=bit_rprimitive,
     c_function_name="CPyTagged_IsLt_",
     error_kind=ERR_NEVER,
+    is_pure=True,
 )
 
 int64_divide_op = custom_op(

mypyc/primitives/registry.py

@@ -60,6 +60,7 @@ class CFunctionDescription(NamedTuple):
     ordering: list[int] | None
     extra_int_constants: list[tuple[int, RType]]
     priority: int
+    is_pure: bool
 
 
 # A description for C load operations including LoadGlobal and LoadAddress
@@ -97,6 +98,7 @@ def method_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     priority: int = 1,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Define a c function call op that replaces a method call.
 
@@ -121,6 +123,8 @@ def method_op(
         steals: description of arguments that this steals (ref count wise)
         is_borrowed: if True, returned value is borrowed (no need to decrease refcount)
         priority: if multiple ops match, the one with the highest priority is picked
+        is_pure: if True, declare that the C function has no side effects, takes immutable
+                 arguments, and never raises an exception
     """
     if extra_int_constants is None:
         extra_int_constants = []
@@ -138,6 +142,7 @@ def method_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=is_pure,
     )
     ops.append(desc)
     return desc
@@ -183,6 +188,7 @@ def function_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=False,
     )
     ops.append(desc)
     return desc
@@ -228,6 +234,7 @@ def binary_op(
         ordering=ordering,
         extra_int_constants=extra_int_constants,
         priority=priority,
+        is_pure=False,
     )
     ops.append(desc)
     return desc
@@ -244,6 +251,8 @@ def custom_op(
     extra_int_constants: list[tuple[int, RType]] | None = None,
     steals: StealsDescription = False,
     is_borrowed: bool = False,
+    *,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Create a one-off CallC op that can't be automatically generated from the AST.
 
@@ -264,6 +273,7 @@ def custom_op(
         ordering,
         extra_int_constants,
         0,
+        is_pure=is_pure,
     )
 
 
@@ -279,6 +289,7 @@ def custom_primitive_op(
     extra_int_constants: list[tuple[int, RType]] | None = None,
     steals: StealsDescription = False,
     is_borrowed: bool = False,
+    is_pure: bool = False,
 ) -> PrimitiveDescription:
     """Define a primitive op that can't be automatically generated based on the AST.
 
@@ -299,6 +310,7 @@ def custom_primitive_op(
         ordering=ordering,
         extra_int_constants=extra_int_constants,
         priority=0,
+        is_pure=is_pure,
     )
 
 
@@ -314,6 +326,7 @@ def unary_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     priority: int = 1,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Define a c function call op for an unary operation.
 
@@ -338,6 +351,7 @@ def unary_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=is_pure,
     )
     ops.append(desc)
     return desc