[mypyc] Make explicit conversions i64(x) and i32(x) faster (#14504)

These behave the same as `int(x)` and we want them to be no slower than
the corresponding `int` conversions. Optimize them for bool, float, str
and RInstance arguments.

Work on mypyc/mypyc#837.

mypyc/irbuild/specialize.py

@@ -160,14 +160,20 @@ def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va
 @specialize_function("builtins.int")
 @specialize_function("builtins.float")
 @specialize_function("builtins.complex")
+@specialize_function("mypy_extensions.i64")
+@specialize_function("mypy_extensions.i32")
 def translate_builtins_with_unary_dunder(
     builder: IRBuilder, expr: CallExpr, callee: RefExpr
 ) -> Value | None:
     """Specialize calls on native classes that implement the associated dunder."""
     if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(callee, NameExpr):
         arg = expr.args[0]
         arg_typ = builder.node_type(arg)
-        method = f"__{callee.name}__"
+        shortname = callee.fullname.split(".")[1]
+        if shortname in ("i64", "i32"):
+            method = "__int__"
+        else:
+            method = f"__{shortname}__"
         if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method):
             obj = builder.accept(arg)
             return builder.gen_method_call(obj, method, [], None, expr.line)
@@ -676,7 +682,7 @@ def translate_i64(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
     elif is_int32_rprimitive(arg_type):
         val = builder.accept(arg)
         return builder.add(Extend(val, int64_rprimitive, signed=True, line=expr.line))
-    elif is_int_rprimitive(arg_type):
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
         val = builder.accept(arg)
         return builder.coerce(val, int64_rprimitive, expr.line)
     return None
@@ -693,7 +699,7 @@ def translate_i32(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value
     elif is_int64_rprimitive(arg_type):
         val = builder.accept(arg)
         return builder.add(Truncate(val, int32_rprimitive, line=expr.line))
-    elif is_int_rprimitive(arg_type):
+    elif is_int_rprimitive(arg_type) or is_bool_rprimitive(arg_type):
         val = builder.accept(arg)
         return builder.coerce(val, int32_rprimitive, expr.line)
     return None

mypyc/primitives/int_ops.py

@@ -35,39 +35,43 @@
     unary_op,
 )
 
-# These int constructors produce object_rprimitives that then need to be unboxed
-# I guess unboxing ourselves would save a check and branch though?
-
-# Get the type object for 'builtins.int'.
-# For ordinary calls to int() we use a load_address to the type
-load_address_op(name="builtins.int", type=object_rprimitive, src="PyLong_Type")
-
-# int(float). We could do a bit better directly.
-function_op(
-    name="builtins.int",
-    arg_types=[float_rprimitive],
-    return_type=object_rprimitive,
-    c_function_name="CPyLong_FromFloat",
-    error_kind=ERR_MAGIC,
-)
+# Constructors for builtins.int and native int types have the same behavior. In
+# interpreted mode, native int types are just aliases to 'int'.
+for int_name in ("builtins.int", "mypy_extensions.i64", "mypy_extensions.i32"):
+    # These int constructors produce object_rprimitives that then need to be unboxed
+    # I guess unboxing ourselves would save a check and branch though?
+
+    # Get the type object for 'builtins.int' or a native int type.
+    # For ordinary calls to int() we use a load_address to the type.
+    # Native ints don't have a separate type object -- we just use 'builtins.int'.
+    load_address_op(name=int_name, type=object_rprimitive, src="PyLong_Type")
+
+    # int(float). We could do a bit better directly.
+    function_op(
+        name=int_name,
+        arg_types=[float_rprimitive],
+        return_type=object_rprimitive,
+        c_function_name="CPyLong_FromFloat",
+        error_kind=ERR_MAGIC,
+    )
 
-# int(string)
-function_op(
-    name="builtins.int",
-    arg_types=[str_rprimitive],
-    return_type=object_rprimitive,
-    c_function_name="CPyLong_FromStr",
-    error_kind=ERR_MAGIC,
-)
+    # int(string)
+    function_op(
+        name=int_name,
+        arg_types=[str_rprimitive],
+        return_type=object_rprimitive,
+        c_function_name="CPyLong_FromStr",
+        error_kind=ERR_MAGIC,
+    )
 
-# int(string, base)
-function_op(
-    name="builtins.int",
-    arg_types=[str_rprimitive, int_rprimitive],
-    return_type=object_rprimitive,
-    c_function_name="CPyLong_FromStrWithBase",
-    error_kind=ERR_MAGIC,
-)
+    # int(string, base)
+    function_op(
+        name=int_name,
+        arg_types=[str_rprimitive, int_rprimitive],
+        return_type=object_rprimitive,
+        c_function_name="CPyLong_FromStrWithBase",
+        error_kind=ERR_MAGIC,
+    )
 
 # str(int)
 int_to_str_op = function_op(

mypyc/test-data/irbuild-i32.test

@@ -480,3 +480,55 @@ L0:
     y = 11
     z = -3
     return 1
+
+[case testI32ExplicitConversionFromVariousTypes]
+from mypy_extensions import i32
+
+def bool_to_i32(b: bool) -> i32:
+    return i32(b)
+
+def str_to_i32(s: str) -> i32:
+    return i32(s)
+
+class C:
+    def __int__(self) -> i32:
+        return 5
+
+def instance_to_i32(c: C) -> i32:
+    return i32(c)
+
+def float_to_i32(x: float) -> i32:
+    return i32(x)
+[out]
+def bool_to_i32(b):
+    b :: bool
+    r0 :: int32
+L0:
+    r0 = extend b: builtins.bool to int32
+    return r0
+def str_to_i32(s):
+    s :: str
+    r0 :: object
+    r1 :: int32
+L0:
+    r0 = CPyLong_FromStr(s)
+    r1 = unbox(int32, r0)
+    return r1
+def C.__int__(self):
+    self :: __main__.C
+L0:
+    return 5
+def instance_to_i32(c):
+    c :: __main__.C
+    r0 :: int32
+L0:
+    r0 = c.__int__()
+    return r0
+def float_to_i32(x):
+    x :: float
+    r0 :: object
+    r1 :: int32
+L0:
+    r0 = CPyLong_FromFloat(x)
+    r1 = unbox(int32, r0)
+    return r1

mypyc/test-data/irbuild-i64.test

@@ -1771,3 +1771,66 @@ L2:
     keep_alive x
 L3:
     return r3
+
+[case testI64ExplicitConversionFromVariousTypes]
+from mypy_extensions import i64
+
+def bool_to_i64(b: bool) -> i64:
+    return i64(b)
+
+def str_to_i64(s: str) -> i64:
+    return i64(s)
+
+def str_to_i64_with_base(s: str) -> i64:
+    return i64(s, 2)
+
+class C:
+    def __int__(self) -> i64:
+        return 5
+
+def instance_to_i64(c: C) -> i64:
+    return i64(c)
+
+def float_to_i64(x: float) -> i64:
+    return i64(x)
+[out]
+def bool_to_i64(b):
+    b :: bool
+    r0 :: int64
+L0:
+    r0 = extend b: builtins.bool to int64
+    return r0
+def str_to_i64(s):
+    s :: str
+    r0 :: object
+    r1 :: int64
+L0:
+    r0 = CPyLong_FromStr(s)
+    r1 = unbox(int64, r0)
+    return r1
+def str_to_i64_with_base(s):
+    s :: str
+    r0 :: object
+    r1 :: int64
+L0:
+    r0 = CPyLong_FromStrWithBase(s, 4)
+    r1 = unbox(int64, r0)
+    return r1
+def C.__int__(self):
+    self :: __main__.C
+L0:
+    return 5
+def instance_to_i64(c):
+    c :: __main__.C
+    r0 :: int64
+L0:
+    r0 = c.__int__()
+    return r0
+def float_to_i64(x):
+    x :: float
+    r0 :: object
+    r1 :: int64
+L0:
+    r0 = CPyLong_FromFloat(x)
+    r1 = unbox(int64, r0)
+    return r1

mypyc/test-data/run-i32.test

@@ -306,6 +306,22 @@ def test_i32_truncate_from_i64() -> None:
     x = i32(small2)
     assert x == 2**31 - 1
 
+def from_float(x: float) -> i32:
+    return i32(x)
+
+def test_explicit_conversion_from_float() -> None:
+    assert from_float(0.0) == 0
+    assert from_float(1.456) == 1
+    assert from_float(-1234.567) == -1234
+    assert from_float(2**31 - 1) == 2**31 - 1
+    assert from_float(-2**31) == -2**31
+    # The error message could be better, but this is acceptable
+    with assertRaises(OverflowError, "int too large to convert to i32"):
+        assert from_float(float(2**31))
+    with assertRaises(OverflowError, "int too large to convert to i32"):
+        # One ulp below the lowest valid i64 value
+        from_float(float(-2**31 - 2048))
+
 def test_tuple_i32() -> None:
     a: i32 = 1
     b: i32 = 2

mypyc/test-data/run-i64.test

@@ -310,6 +310,68 @@ def test_i64_from_large_small_literal() -> None:
     x = i64(-2**63)
     assert x == -2**63
 
+def from_float(x: float) -> i64:
+    return i64(x)
+
+def test_explicit_conversion_from_float() -> None:
+    assert from_float(0.0) == 0
+    assert from_float(1.456) == 1
+    assert from_float(-1234.567) == -1234
+    assert from_float(2**63 - 1) == 2**63 - 1
+    assert from_float(-2**63) == -2**63
+    # The error message could be better, but this is acceptable
+    with assertRaises(OverflowError, "int too large to convert to i64"):
+        assert from_float(float(2**63))
+    with assertRaises(OverflowError, "int too large to convert to i64"):
+        # One ulp below the lowest valid i64 value
+        from_float(float(-2**63 - 2048))
+
+def from_str(s: str) -> i64:
+    return i64(s)
+
+def test_explicit_conversion_from_str() -> None:
+    assert from_str("0") == 0
+    assert from_str("1") == 1
+    assert from_str("-1234") == -1234
+    with assertRaises(ValueError):
+        from_str("1.2")
+
+def from_str_with_base(s: str, base: int) -> i64:
+    return i64(s, base)
+
+def test_explicit_conversion_from_str_with_base() -> None:
+    assert from_str_with_base("101", 2) == 5
+    assert from_str_with_base("109", 10) == 109
+    assert from_str_with_base("-f0A", 16) == -3850
+    assert from_str_with_base("0x1a", 16) == 26
+    assert from_str_with_base("0X1A", 16) == 26
+    with assertRaises(ValueError):
+        from_str_with_base("1.2", 16)
+
+def from_bool(b: bool) -> i64:
+    return i64(b)
+
+def test_explicit_conversion_from_bool() -> None:
+    assert from_bool(True) == 1
+    assert from_bool(False) == 0
+
+class IntConv:
+    def __init__(self, x: i64) -> None:
+        self.x = x
+
+    def __int__(self) -> i64:
+        return self.x + 1
+
+def test_explicit_conversion_from_instance() -> None:
+    assert i64(IntConv(0)) == 1
+    assert i64(IntConv(12345)) == 12346
+    assert i64(IntConv(-23)) == -22
+
+def test_explicit_conversion_from_any() -> None:
+    # This can't be specialized
+    a: Any = "101"
+    assert i64(a, base=2) == 5
+
 def test_tuple_i64() -> None:
     a: i64 = 1
     b: i64 = 2

test-data/unit/lib-stub/mypy_extensions.pyi

@@ -1,7 +1,7 @@
 # NOTE: Requires fixtures/dict.pyi
 from typing import (
     Any, Dict, Type, TypeVar, Optional, Any, Generic, Mapping, NoReturn as NoReturn, Iterator,
-    Union
+    Union, Protocol
 )
 import sys
 
@@ -51,10 +51,13 @@ mypyc_attr: Any
 class FlexibleAlias(Generic[_T, _U]): ...
 
 if sys.version_info >= (3, 0):
+    class __SupportsInt(Protocol[T_co]):
+        def __int__(self) -> int: pass
+
     _Int = Union[int, i32, i64]
 
     class i32:
-        def __init__(self, x: _Int) -> None: ...
+        def __init__(self, x: Union[_Int, str, bytes, SupportsInt], base: int = 10) -> None: ...
         def __add__(self, x: i32) -> i32: ...
         def __radd__(self, x: i32) -> i32: ...
         def __sub__(self, x: i32) -> i32: ...
@@ -84,7 +87,7 @@ if sys.version_info >= (3, 0):
         def __gt__(self, x: i32) -> bool: ...
 
     class i64:
-        def __init__(self, x: _Int) -> None: ...
+        def __init__(self, x: Union[_Int, str, bytes, SupportsInt], base: int = 10) -> None: ...
         def __add__(self, x: i64) -> i64: ...
         def __radd__(self, x: i64) -> i64: ...
         def __sub__(self, x: i64) -> i64: ...