[mypyc] Generate efficient code for (some) conversions i64(x) and i32(x)

mypyc/irbuild/specialize.py

@@ -32,15 +32,29 @@
     TupleExpr,
 )
 from mypy.types import AnyType, TypeOfAny
-from mypyc.ir.ops import BasicBlock, Integer, RaiseStandardError, Register, Unreachable, Value
+from mypyc.ir.ops import (
+    BasicBlock,
+    Extend,
+    Integer,
+    RaiseStandardError,
+    Register,
+    Truncate,
+    Unreachable,
+    Value,
+)
 from mypyc.ir.rtypes import (
     RInstance,
     RTuple,
     RType,
     bool_rprimitive,
     c_int_rprimitive,
     dict_rprimitive,
+    int32_rprimitive,
+    int64_rprimitive,
     is_dict_rprimitive,
+    is_int32_rprimitive,
+    is_int64_rprimitive,
+    is_int_rprimitive,
     is_list_rprimitive,
     list_rprimitive,
     set_rprimitive,
@@ -640,3 +654,37 @@ def translate_fstring(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va
 
         return join_formatted_strings(builder, None, substitutions, expr.line)
     return None
+
+
+@specialize_function("mypy_extensions.i64")
+def translate_i64(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_int64_rprimitive(arg_type):
+        return builder.accept(arg)
+    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):
+        val = builder.accept(arg)
+        return builder.coerce(val, int64_rprimitive, expr.line)
+    return None
+
+
+@specialize_function("mypy_extensions.i32")
+def translate_i32(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_int32_rprimitive(arg_type):
+        return builder.accept(arg)
+    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):
+        val = builder.accept(arg)
+        return builder.coerce(val, int32_rprimitive, expr.line)
+    return None

mypyc/test-data/irbuild-i32.test

@@ -411,3 +411,72 @@ L0:
     y = -127
     z = 12
     return 1
+
+[case testI32ExplicitConversionFromNativeInt]
+from mypy_extensions import i64, i32
+
+def from_i32(x: i32) -> i32:
+    return i32(x)
+
+def from_i64(x: i64) -> i32:
+    return i32(x)
+[out]
+def from_i32(x):
+    x :: int32
+L0:
+    return x
+def from_i64(x):
+    x :: int64
+    r0 :: int32
+L0:
+    r0 = truncate x: int64 to int32
+    return r0
+
+[case testI32ExplicitConversionFromInt_64bit]
+from mypy_extensions import i32
+
+def f(x: int) -> i32:
+    return i32(x)
+[out]
+def f(x):
+    x :: int
+    r0 :: native_int
+    r1, r2, r3 :: bit
+    r4 :: native_int
+    r5, r6 :: int32
+L0:
+    r0 = x & 1
+    r1 = r0 == 0
+    if r1 goto L1 else goto L4 :: bool
+L1:
+    r2 = x < 4294967296 :: signed
+    if r2 goto L2 else goto L4 :: bool
+L2:
+    r3 = x >= -4294967296 :: signed
+    if r3 goto L3 else goto L4 :: bool
+L3:
+    r4 = x >> 1
+    r5 = truncate r4: native_int to int32
+    r6 = r5
+    goto L5
+L4:
+    CPyInt32_Overflow()
+    unreachable
+L5:
+    return r6
+
+[case testI32ExplicitConversionFromLiteral]
+from mypy_extensions import i32
+
+def f() -> None:
+    x = i32(0)
+    y = i32(11)
+    z = i32(-3)
+[out]
+def f():
+    x, y, z :: int32
+L0:
+    x = 0
+    y = 11
+    z = -3
+    return 1

mypyc/test-data/irbuild-i64.test

@@ -1507,3 +1507,70 @@ L0:
     r0 = c.m(0, 0)
     r1 = c.m(6, 1)
     return 1
+
+[case testI64ExplicitConversionFromNativeInt]
+from mypy_extensions import i64, i32
+
+def from_i32(x: i32) -> i64:
+    return i64(x)
+
+def from_i64(x: i64) -> i64:
+    return i64(x)
+[out]
+def from_i32(x):
+    x :: int32
+    r0 :: int64
+L0:
+    r0 = extend signed x: int32 to int64
+    return r0
+def from_i64(x):
+    x :: int64
+L0:
+    return x
+
+[case testI64ExplicitConversionFromInt_64bit]
+from mypy_extensions import i64
+
+def f(x: int) -> i64:
+    return i64(x)
+[out]
+def f(x):
+    x :: int
+    r0 :: native_int
+    r1 :: bit
+    r2, r3 :: int64
+    r4 :: ptr
+    r5 :: c_ptr
+    r6 :: int64
+L0:
+    r0 = x & 1
+    r1 = r0 == 0
+    if r1 goto L1 else goto L2 :: bool
+L1:
+    r2 = x >> 1
+    r3 = r2
+    goto L3
+L2:
+    r4 = x ^ 1
+    r5 = r4
+    r6 = CPyLong_AsInt64(r5)
+    r3 = r6
+    keep_alive x
+L3:
+    return r3
+
+[case testI64ExplicitConversionFromLiteral]
+from mypy_extensions import i64
+
+def f() -> None:
+    x = i64(0)
+    y = i64(11)
+    z = i64(-3)
+[out]
+def f():
+    x, y, z :: int64
+L0:
+    x = 0
+    y = 11
+    z = -3
+    return 1

mypyc/test-data/run-i32.test

@@ -3,7 +3,7 @@ from typing import Any, Tuple
 
 MYPY = False
 if MYPY:
-    from mypy_extensions import i32
+    from mypy_extensions import i32, i64
 
 from testutil import assertRaises
 
@@ -252,6 +252,60 @@ def test_coerce_to_and_from_int() -> None:
                     m: int = x
                     assert m == n
 
+def test_explicit_conversion_to_i32() -> None:
+    x = i32(5)
+    assert x == 5
+    y = int() - 113
+    x = i32(y)
+    assert x == -113
+    n64: i64 = 1733
+    x = i32(n64)
+    assert x == 1733
+    n32 = -1733
+    x = i32(n32)
+    assert x == -1733
+    z = i32(x)
+    assert z == -1733
+
+def test_explicit_conversion_overflow() -> None:
+    max_i32 = int() + 2**31 - 1
+    x = i32(max_i32)
+    assert x == 2**31 - 1
+    assert int(x) == max_i32
+
+    min_i32 = int() - 2**31
+    y = i32(min_i32)
+    assert y == -2**31
+    assert int(y) == min_i32
+
+    too_big = int() + 2**31
+    with assertRaises(OverflowError):
+        x = i32(too_big)
+
+    too_small = int() - 2**31 - 1
+    with assertRaises(OverflowError):
+        x = i32(too_small)
+
+def test_i32_from_large_small_literal() -> None:
+    x = i32(2**31 - 1)
+    assert x == 2**31 - 1
+    x = i32(-2**31)
+    assert x == -2**31
+
+def test_i32_truncate_from_i64() -> None:
+    large = i64(2**32 + 157 + int())
+    x = i32(large)
+    assert x == 157
+    small = i64(-2**32 - 157 + int())
+    x = i32(small)
+    assert x == -157
+    large2 = i64(2**31 + int())
+    x = i32(large2)
+    assert x == -2**31
+    small2 = i64(-2**31 - 1 - int())
+    x = i32(small2)
+    assert x == 2**31 - 1
+
 def test_tuple_i32() -> None:
     a: i32 = 1
     b: i32 = 2

mypyc/test-data/run-i64.test

@@ -3,7 +3,7 @@ from typing import List, Any, Tuple
 
 MYPY = False
 if MYPY:
-    from mypy_extensions import i64
+    from mypy_extensions import i64, i32
 
 from testutil import assertRaises
 
@@ -229,6 +229,46 @@ def test_coerce_to_and_from_int() -> None:
                     m: int = x
                     assert m == n
 
+def test_explicit_conversion_to_i64() -> None:
+    x = i64(5)
+    assert x == 5
+    y = int() - 113
+    x = i64(y)
+    assert x == -113
+    n32: i32 = 1733
+    x = i64(n32)
+    assert x == 1733
+    n32 = -1733
+    x = i64(n32)
+    assert x == -1733
+    z = i64(x)
+    assert z == -1733
+
+def test_explicit_conversion_overflow() -> None:
+    max_i64 = int() + 2**63 - 1
+    x = i64(max_i64)
+    assert x == 2**63 - 1
+    assert int(x) == max_i64
+
+    min_i64 = int() - 2**63
+    y = i64(min_i64)
+    assert y == -2**63
+    assert int(y) == min_i64
+
+    too_big = int() + 2**63
+    with assertRaises(OverflowError):
+        x = i64(too_big)
+
+    too_small = int() - 2**63 - 1
+    with assertRaises(OverflowError):
+        x = i64(too_small)
+
+def test_i64_from_large_small_literal() -> None:
+    x = i64(2**63 - 1)
+    assert x == 2**63 - 1
+    x = i64(-2**63)
+    assert x == -2**63
+
 def test_tuple_i64() -> None:
     a: i64 = 1
     b: i64 = 2