diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h index 8cd141545bbb..b98cd5966f0b 100644 --- a/mypyc/lib-rt/CPy.h +++ b/mypyc/lib-rt/CPy.h @@ -139,6 +139,7 @@ CPyTagged CPyTagged_Remainder_(CPyTagged left, CPyTagged right); CPyTagged CPyTagged_BitwiseLongOp_(CPyTagged a, CPyTagged b, char op); CPyTagged CPyTagged_Rshift_(CPyTagged left, CPyTagged right); CPyTagged CPyTagged_Lshift_(CPyTagged left, CPyTagged right); +CPyTagged CPyTagged_BitLength(CPyTagged self); PyObject *CPyTagged_Str(CPyTagged n); CPyTagged CPyTagged_FromFloat(double f); diff --git a/mypyc/lib-rt/int_ops.c b/mypyc/lib-rt/int_ops.c index e2c302eea576..dd76955320cb 100644 --- a/mypyc/lib-rt/int_ops.c +++ b/mypyc/lib-rt/int_ops.c @@ -5,6 +5,10 @@ #include #include "CPy.h" +#ifdef _MSC_VER +#include +#endif + #ifndef _WIN32 // On 64-bit Linux and macOS, ssize_t and long are both 64 bits, and // PyLong_FromLong is faster than PyLong_FromSsize_t, so use the faster one @@ -15,6 +19,17 @@ #define CPyLong_FromSsize_t PyLong_FromSsize_t #endif +#if defined(__GNUC__) || defined(__clang__) +# if defined(__x86_64__) || defined(_M_X64) || defined(__aarch64__) || (defined(__SIZEOF_POINTER__) && __SIZEOF_POINTER__ == 8) +# define CPY_CLZ(x) __builtin_clzll((unsigned long long)(x)) +# define CPY_BITS 64 +# else +# define CPY_CLZ(x) __builtin_clz((unsigned int)(x)) +# define CPY_BITS 32 +# endif +#endif + + CPyTagged CPyTagged_FromSsize_t(Py_ssize_t value) { // We use a Python object if the value shifted left by 1 is too // large for Py_ssize_t @@ -581,3 +596,52 @@ double CPyTagged_TrueDivide(CPyTagged x, CPyTagged y) { } return 1.0; } + +// int.bit_length() +CPyTagged CPyTagged_BitLength(CPyTagged self) { + // Handle zero + if (self == 0) { + return 0; + } + + // Fast path for small (tagged) ints + if (CPyTagged_CheckShort(self)) { + Py_ssize_t val = CPyTagged_ShortAsSsize_t(self); + Py_ssize_t absval = val < 0 ? -val : val; + int bits = 0; + if (absval) { +#if defined(_MSC_VER) + #if defined(_WIN64) + unsigned long idx; + if (_BitScanReverse64(&idx, (unsigned __int64)absval)) { + bits = (int)(idx + 1); + } + #else + unsigned long idx; + if (_BitScanReverse(&idx, (unsigned long)absval)) { + bits = (int)(idx + 1); + } + #endif +#elif defined(__GNUC__) || defined(__clang__) + bits = (int)(CPY_BITS - CPY_CLZ(absval)); +#else + // Fallback to loop if no builtin + while (absval) { + absval >>= 1; + bits++; + } +#endif + } + return bits << 1; + } + + // Slow path for big ints + PyObject *pyint = CPyTagged_StealAsObject(self); + int bits = _PyLong_NumBits(pyint); + Py_DECREF(pyint); + if (bits < 0) { + // _PyLong_NumBits sets an error on failure + return CPY_INT_TAG; + } + return bits << 1; +} diff --git a/mypyc/primitives/int_ops.py b/mypyc/primitives/int_ops.py index d723c9b63a86..8f43140dd255 100644 --- a/mypyc/primitives/int_ops.py +++ b/mypyc/primitives/int_ops.py @@ -31,7 +31,14 @@ str_rprimitive, void_rtype, ) -from mypyc.primitives.registry import binary_op, custom_op, function_op, load_address_op, unary_op +from mypyc.primitives.registry import ( + binary_op, + custom_op, + function_op, + load_address_op, + method_op, + unary_op, +) # Constructors for builtins.int and native int types have the same behavior. In # interpreted mode, native int types are just aliases to 'int'. @@ -305,3 +312,12 @@ def int_unary_op(name: str, c_function_name: str) -> PrimitiveDescription: c_function_name="PyLong_Check", error_kind=ERR_NEVER, ) + +# int.bit_length() +method_op( + name="bit_length", + arg_types=[int_rprimitive], + return_type=int_rprimitive, + c_function_name="CPyTagged_BitLength", + error_kind=ERR_MAGIC, +) diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py index c041c661741c..b4d131ca44a4 100644 --- a/mypyc/test-data/fixtures/ir.py +++ b/mypyc/test-data/fixtures/ir.py @@ -82,6 +82,7 @@ def __lt__(self, n: int) -> bool: pass def __gt__(self, n: int) -> bool: pass def __le__(self, n: int) -> bool: pass def __ge__(self, n: int) -> bool: pass + def bit_length(self) -> int: pass class str: @overload diff --git a/mypyc/test-data/irbuild-int.test b/mypyc/test-data/irbuild-int.test index bdf9127b722a..184c66fafb7c 100644 --- a/mypyc/test-data/irbuild-int.test +++ b/mypyc/test-data/irbuild-int.test @@ -210,3 +210,13 @@ L0: r0 = CPyTagged_Invert(n) x = r0 return x + +[case testIntBitLength] +def f(x: int) -> int: + return x.bit_length() +[out] +def f(x): + x, r0 :: int +L0: + r0 = CPyTagged_BitLength(x) + return r0 diff --git a/mypyc/test-data/run-integers.test b/mypyc/test-data/run-integers.test index 1163c9d942f7..c32e609ad8de 100644 --- a/mypyc/test-data/run-integers.test +++ b/mypyc/test-data/run-integers.test @@ -572,3 +572,21 @@ class subc(int): [file userdefinedint.py] class int: pass + +[case testBitLength] +def bit_length(n: int) -> int: + return n.bit_length() +def bit_length_python(n: int) -> int: + return getattr(n, "bit_length")() +def test_bit_length() -> None: + assert bit_length(0) == bit_length_python(0) + assert bit_length(1) == bit_length_python(1) + assert bit_length(255) == bit_length_python(255) + assert bit_length(256) == bit_length_python(256) + assert bit_length(-256) == bit_length_python(-256) + # Large positive int + assert bit_length(1 << 70) == bit_length_python(1 << 70) + # Large negative int + assert bit_length(-(1 << 70)) == bit_length_python(-(1 << 70)) + # Large int with all bits set + assert bit_length((1 << 100) - 1) == bit_length_python((1 << 100) - 1)