diff --git a/Include/internal/pycore_bitutils.h b/Include/internal/pycore_bitutils.h index 50f69377523818..352112eb98344d 100644 --- a/Include/internal/pycore_bitutils.h +++ b/Include/internal/pycore_bitutils.h @@ -180,6 +180,43 @@ _Py_bit_length(unsigned long x) } +static inline int +_Py_bit_length64(uint64_t x) +{ +#if (defined(__clang__) || defined(__GNUC__)) + if (x != 0) { + // __builtin_clzll() is available since GCC 3.4. + // Undefined behavior for x == 0. + return (int)sizeof(uint64_t) * 8 - __builtin_clzll(x); + } + else { + return 0; + } +#elif defined(_MSC_VER) && defined(_WIN64) + // _BitScanReverse64() is documented to search 64 bits. + unsigned long msb; + if (_BitScanReverse64(&msb, x)) { + return (int)msb + 1; + } + else { + return 0; + } +#else + const int BIT_LENGTH_TABLE[32] = { + 0, 1, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, + 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5 + }; + int msb = 0; + while (x >= 32) { + msb += 6; + x >>= 6; + } + msb += BIT_LENGTH_TABLE[x]; + return msb; +#endif +} + + #ifdef __cplusplus } #endif diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index 80fe9cff98509d..3d7434ddaa5626 100644 --- a/Modules/_heapqmodule.c +++ b/Modules/_heapqmodule.c @@ -12,6 +12,7 @@ annotated by François Pinard, and converted to C by Raymond Hettinger. #include "Python.h" #include "pycore_list.h" // _PyList_ITEMS() +#include "pycore_bitutils.h" // _Py_bit_length(), _Py_bit_length64() #include "clinic/_heapqmodule.c.h" @@ -279,16 +280,15 @@ _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item) return returnitem; } -static Py_ssize_t +static inline Py_ssize_t keep_top_bit(Py_ssize_t n) { - int i = 0; - - while (n > 1) { - n >>= 1; - i++; - } - return n << i; + Py_BUILD_ASSERT(sizeof(Py_ssize_t) <= sizeof(uint64_t)); +#if defined(SIZEOF_SIZE_T) && SIZEOF_SIZE_T <= 4 + return (Py_ssize_t)1 << (_Py_bit_length((unsigned long)n) - 1); +#else + return (Py_ssize_t)1 << (_Py_bit_length64((uint64_t)n) - 1); +#endif } /* Cache friendly version of heapify() diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c index 0451688a46c75f..235f4755f66c41 100644 --- a/Modules/_testinternalcapi.c +++ b/Modules/_testinternalcapi.c @@ -227,6 +227,47 @@ test_bit_length(PyObject *self, PyObject *Py_UNUSED(args)) } +static int +check_bit_length64(uint64_t x, int expected) +{ + // Use volatile to prevent the compiler to optimize out the whole test + volatile uint64_t u = x; + int len = _Py_bit_length64(u); + if (len != expected) { + PyErr_Format(PyExc_AssertionError, + "_Py_bit_length(%lu) returns %i, expected %i", + x, len, expected); + return -1; + } + return 0; +} + + +static PyObject* +test_bit_length64(PyObject *self, PyObject *Py_UNUSED(args)) +{ +#define CHECK(X, RESULT) \ + do { \ + if (check_bit_length64(X, RESULT) < 0) { \ + return NULL; \ + } \ + } while (0) + + CHECK(0, 0); + CHECK(1, 1); + CHECK(0x1000, 13); + CHECK(0x1234, 13); + CHECK(0x54321, 19); + CHECK(0x7FFFFFFF, 31); + CHECK(0xFFFFFFFF, 32); + CHECK(0x7FFFFFFFFFFFFFFFULL, 63); + CHECK(0xFFFFFFFFFFFFFFFFULL, 64); + Py_RETURN_NONE; + +#undef CHECK +} + + #define TO_PTR(ch) ((void*)(uintptr_t)ch) #define FROM_PTR(ptr) ((uintptr_t)ptr) #define VALUE(key) (1 + ((int)(key) - 'a')) @@ -2056,6 +2097,7 @@ static PyMethodDef module_functions[] = { {"test_bswap", test_bswap, METH_NOARGS}, {"test_popcount", test_popcount, METH_NOARGS}, {"test_bit_length", test_bit_length, METH_NOARGS}, + {"test_bit_length64", test_bit_length64, METH_NOARGS}, {"test_hashtable", test_hashtable, METH_NOARGS}, {"get_config", test_get_config, METH_NOARGS}, {"set_config", test_set_config, METH_O},