[mypyc] feat: new primitive for int.bit_length (#19673)

BobTheBuidler · web-flow · commit 8e57622c45b1 · 2025-10-14T13:12:06.000+01:00
This PR adds a new primitive for `int.bit_length`.
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -148,6 +148,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
@@ -5,6 +5,10 @@
 #include <Python.h>
 #include "CPy.h"
 
+#ifdef _MSC_VER
+#include <intrin.h>
+#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_AsObject(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
@@ -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
@@ -86,6 +86,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
@@ -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
@@ -572,3 +572,17 @@ 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:
+    for n in range(256):
+        i = 1 << n
+        assert bit_length(i) == bit_length_python(i)
+        assert bit_length(-(i)) == bit_length_python(-(i))
+        i -= 1
+        assert bit_length(i) == bit_length_python(i)
+        assert bit_length(-(i)) == bit_length_python(-(i))
