[mypyc] Faster bool and integer conversions (python#14422)

Speed up various conversions between bool and integer types. These cases
are covered:
* `bool(x)` for various types
* `int(x)` for `bool` and integer arguments
* Implicit coercion from `bool` to an integer type

Support both regular `int` values and native int values.

(Various small optimizations, including these, together netted a 6%
performance improvement in self check.)

mypyc/irbuild/ll_builder.py

@@ -326,6 +326,17 @@ def coerce(
             ):
                 # Equivalent types
                 return src
+            elif (
+                is_bool_rprimitive(src_type) or is_bit_rprimitive(src_type)
+            ) and is_int_rprimitive(target_type):
+                shifted = self.int_op(
+                    bool_rprimitive, src, Integer(1, bool_rprimitive), IntOp.LEFT_SHIFT
+                )
+                return self.add(Extend(shifted, int_rprimitive, signed=False))
+            elif (
+                is_bool_rprimitive(src_type) or is_bit_rprimitive(src_type)
+            ) and is_fixed_width_rtype(target_type):
+                return self.add(Extend(src, target_type, signed=False))
             else:
                 # To go from one unboxed type to another, we go through a boxed
                 # in-between value, for simplicity.
@@ -1642,35 +1653,38 @@ def shortcircuit_helper(
         self.activate_block(next_block)
         return target
 
-    def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) -> None:
-        if is_runtime_subtype(value.type, int_rprimitive):
+    def bool_value(self, value: Value) -> Value:
+        """Return bool(value).
+
+        The result type can be bit_rprimitive or bool_rprimitive.
+        """
+        if is_bool_rprimitive(value.type) or is_bit_rprimitive(value.type):
+            result = value
+        elif is_runtime_subtype(value.type, int_rprimitive):
             zero = Integer(0, short_int_rprimitive)
-            self.compare_tagged_condition(value, zero, "!=", true, false, value.line)
-            return
+            result = self.comparison_op(value, zero, ComparisonOp.NEQ, value.line)
         elif is_fixed_width_rtype(value.type):
             zero = Integer(0, value.type)
-            value = self.add(ComparisonOp(value, zero, ComparisonOp.NEQ))
+            result = self.add(ComparisonOp(value, zero, ComparisonOp.NEQ))
         elif is_same_type(value.type, str_rprimitive):
-            value = self.call_c(str_check_if_true, [value], value.line)
+            result = self.call_c(str_check_if_true, [value], value.line)
         elif is_same_type(value.type, list_rprimitive) or is_same_type(
             value.type, dict_rprimitive
         ):
             length = self.builtin_len(value, value.line)
             zero = Integer(0)
-            value = self.binary_op(length, zero, "!=", value.line)
+            result = self.binary_op(length, zero, "!=", value.line)
         elif (
             isinstance(value.type, RInstance)
             and value.type.class_ir.is_ext_class
             and value.type.class_ir.has_method("__bool__")
         ):
             # Directly call the __bool__ method on classes that have it.
-            value = self.gen_method_call(value, "__bool__", [], bool_rprimitive, value.line)
+            result = self.gen_method_call(value, "__bool__", [], bool_rprimitive, value.line)
         else:
             value_type = optional_value_type(value.type)
             if value_type is not None:
-                is_none = self.translate_is_op(value, self.none_object(), "is not", value.line)
-                branch = Branch(is_none, true, false, Branch.BOOL)
-                self.add(branch)
+                not_none = self.translate_is_op(value, self.none_object(), "is not", value.line)
                 always_truthy = False
                 if isinstance(value_type, RInstance):
                     # check whether X.__bool__ is always just the default (object.__bool__)
@@ -1679,18 +1693,55 @@ def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) ->
                     ) and value_type.class_ir.is_method_final("__bool__"):
                         always_truthy = True
 
-                if not always_truthy:
-                    # Optional[X] where X may be falsey and requires a check
-                    branch.true = BasicBlock()
-                    self.activate_block(branch.true)
+                if always_truthy:
+                    result = not_none
+                else:
+                    # "X | None" where X may be falsey and requires a check
+                    result = Register(bit_rprimitive)
+                    true, false, end = BasicBlock(), BasicBlock(), BasicBlock()
+                    branch = Branch(not_none, true, false, Branch.BOOL)
+                    self.add(branch)
+                    self.activate_block(true)
                     # unbox_or_cast instead of coerce because we want the
                     # type to change even if it is a subtype.
                     remaining = self.unbox_or_cast(value, value_type, value.line)
-                    self.add_bool_branch(remaining, true, false)
-                return
-            elif not is_bool_rprimitive(value.type) and not is_bit_rprimitive(value.type):
-                value = self.call_c(bool_op, [value], value.line)
-        self.add(Branch(value, true, false, Branch.BOOL))
+                    as_bool = self.bool_value(remaining)
+                    self.add(Assign(result, as_bool))
+                    self.goto(end)
+                    self.activate_block(false)
+                    self.add(Assign(result, Integer(0, bit_rprimitive)))
+                    self.goto(end)
+                    self.activate_block(end)
+            else:
+                result = self.call_c(bool_op, [value], value.line)
+        return result
+
+    def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) -> None:
+        opt_value_type = optional_value_type(value.type)
+        if opt_value_type is None:
+            bool_value = self.bool_value(value)
+            self.add(Branch(bool_value, true, false, Branch.BOOL))
+        else:
+            # Special-case optional types
+            is_none = self.translate_is_op(value, self.none_object(), "is not", value.line)
+            branch = Branch(is_none, true, false, Branch.BOOL)
+            self.add(branch)
+            always_truthy = False
+            if isinstance(opt_value_type, RInstance):
+                # check whether X.__bool__ is always just the default (object.__bool__)
+                if not opt_value_type.class_ir.has_method(
+                    "__bool__"
+                ) and opt_value_type.class_ir.is_method_final("__bool__"):
+                    always_truthy = True
+
+            if not always_truthy:
+                # Optional[X] where X may be falsey and requires a check
+                branch.true = BasicBlock()
+                self.activate_block(branch.true)
+                # unbox_or_cast instead of coerce because we want the
+                # type to change even if it is a subtype.
+                remaining = self.unbox_or_cast(value, opt_value_type, value.line)
+                self.add_bool_branch(remaining, true, false)
 
     def call_c(
         self,
@@ -1795,7 +1846,7 @@ def matching_call_c(
             return target
         return None
 
-    def int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int) -> Value:
+    def int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int = -1) -> Value:
         """Generate a native integer binary op.
 
         Use native/C semantics, which sometimes differ from Python

mypyc/irbuild/specialize.py

@@ -51,7 +51,10 @@
     dict_rprimitive,
     int32_rprimitive,
     int64_rprimitive,
+    int_rprimitive,
+    is_bool_rprimitive,
     is_dict_rprimitive,
+    is_fixed_width_rtype,
     is_int32_rprimitive,
     is_int64_rprimitive,
     is_int_rprimitive,
@@ -688,3 +691,28 @@ def translate_i32(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
         val = builder.accept(arg)
         return builder.coerce(val, int32_rprimitive, expr.line)
     return None
+
+
+@specialize_function("builtins.int")
+def translate_int(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    arg_type = builder.node_type(arg)
+    if (
+        is_bool_rprimitive(arg_type)
+        or is_int_rprimitive(arg_type)
+        or is_fixed_width_rtype(arg_type)
+    ):
+        src = builder.accept(arg)
+        return builder.coerce(src, int_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("builtins.bool")
+def translate_bool(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) != 1 or expr.arg_kinds[0] != ARG_POS:
+        return None
+    arg = expr.args[0]
+    src = builder.accept(arg)
+    return builder.builder.bool_value(src)

mypyc/test-data/irbuild-basic.test

@@ -1108,15 +1108,17 @@ L0:
     return 1
 
 [case testCallableTypes]
-from typing import Callable
+from typing import Callable, Any
+from m import f
+
 def absolute_value(x: int) -> int:
     return x if x > 0 else -x
 
 def call_native_function(x: int) -> int:
     return absolute_value(x)
 
 def call_python_function(x: int) -> int:
-    return int(x)
+    return f(x)
 
 def return_float() -> float:
     return 5.0
@@ -1127,6 +1129,9 @@ def return_callable_type() -> Callable[[], float]:
 def call_callable_type() -> float:
     f = return_callable_type()
     return f()
+[file m.py]
+def f(x: int) -> int:
+    return x
 [out]
 def absolute_value(x):
     x :: int
@@ -1158,14 +1163,18 @@ L0:
     return r0
 def call_python_function(x):
     x :: int
-    r0, r1, r2 :: object
-    r3 :: int
+    r0 :: dict
+    r1 :: str
+    r2, r3, r4 :: object
+    r5 :: int
 L0:
-    r0 = load_address PyLong_Type
-    r1 = box(int, x)
-    r2 = PyObject_CallFunctionObjArgs(r0, r1, 0)
-    r3 = unbox(int, r2)
-    return r3
+    r0 = __main__.globals :: static
+    r1 = 'f'
+    r2 = CPyDict_GetItem(r0, r1)
+    r3 = box(int, x)
+    r4 = PyObject_CallFunctionObjArgs(r2, r3, 0)
+    r5 = unbox(int, r4)
+    return r5
 def return_float():
     r0 :: float
 L0:
@@ -3068,8 +3077,7 @@ def call_sum(l, comparison):
     r1, r2 :: object
     r3, x :: int
     r4, r5 :: object
-    r6 :: bool
-    r7 :: object
+    r6, r7 :: bool
     r8, r9 :: int
     r10 :: bit
 L0:
@@ -3084,8 +3092,8 @@ L2:
     r4 = box(int, x)
     r5 = PyObject_CallFunctionObjArgs(comparison, r4, 0)
     r6 = unbox(bool, r5)
-    r7 = box(bool, r6)
-    r8 = unbox(int, r7)
+    r7 = r6 << 1
+    r8 = extend r7: builtins.bool to builtins.int
     r9 = CPyTagged_Add(r0, r8)
     r0 = r9
 L3: