bpo-29882: Add an efficient popcount method for integers (#771)

niklasf · mdickinson · web-flow · commit 8bd216dfede9 · 2020-05-29T17:28:02.000+01:00
* bpo-29882: Add an efficient popcount method for integers * Update 'sign bit' and versionadded in docs * Add entry to whatsnew document * Doc: use positive example, mention population count * Minor cleanups of the core code * Move popcount_digit closer to where it's used * Use z instead of self after conversion * Add 'absolute value' and 'population count' to docstring * Fix clinic error about missing summary line * Ensure popcount_digit is portable with 64-bit ints Co-authored-by: Mark Dickinson <dickinsm@gmail.com>
diff --git a/Doc/library/stdtypes.rst b/Doc/library/stdtypes.rst
@@ -478,6 +478,27 @@ class`. In addition, it provides a few more methods:
 
     .. versionadded:: 3.1
 
+.. method:: int.bit_count()
+
+    Return the number of ones in the binary representation of the absolute
+    value of the integer. This is also known as the population count.
+    Example::
+
+        >>> n = 19
+        >>> bin(n)
+        '0b10011'
+        >>> n.bit_count()
+        3
+        >>> (-n).bit_count()
+        3
+
+    Equivalent to::
+
+        def bit_count(self):
+            return bin(self).count("1")
+
+    .. versionadded:: 3.10
+
 .. method:: int.to_bytes(length, byteorder, \*, signed=False)
 
     Return an array of bytes representing an integer.
diff --git a/Doc/whatsnew/3.10.rst b/Doc/whatsnew/3.10.rst
@@ -70,6 +70,9 @@ Summary -- Release highlights
 New Features
 ============
 
+* The :class:`int` type has a new method :meth:`int.bit_count`, returning the
+  number of ones in the binary expansion of a given integer, also known
+  as the population count. (Contributed by Niklas Fiekas in :issue:`29882`.)
 
 
 Other Language Changes
diff --git a/Lib/test/test_doctest.py b/Lib/test/test_doctest.py
@@ -669,7 +669,7 @@ def non_Python_modules(): r"""
     True
     >>> real_tests = [t for t in tests if len(t.examples) > 0]
     >>> len(real_tests) # objects that actually have doctests
-    13
+    14
     >>> for t in real_tests:
     ...     print('{}  {}'.format(len(t.examples), t.name))
     ...
@@ -682,6 +682,7 @@ def non_Python_modules(): r"""
     1  builtins.hex
     1  builtins.int
     3  builtins.int.as_integer_ratio
+    2  builtins.int.bit_count
     2  builtins.int.bit_length
     5  builtins.memoryview.hex
     1  builtins.oct
diff --git a/Lib/test/test_long.py b/Lib/test/test_long.py
@@ -1016,6 +1016,17 @@ def test_bit_length(self):
             self.assertEqual((a+1).bit_length(), i+1)
             self.assertEqual((-a-1).bit_length(), i+1)
 
+    def test_bit_count(self):
+        for a in range(-1000, 1000):
+            self.assertEqual(a.bit_count(), bin(a).count("1"))
+
+        for exp in [10, 17, 63, 64, 65, 1009, 70234, 1234567]:
+            a = 2**exp
+            self.assertEqual(a.bit_count(), 1)
+            self.assertEqual((a - 1).bit_count(), exp)
+            self.assertEqual((a ^ 63).bit_count(), 7)
+            self.assertEqual(((a - 1) ^ 510).bit_count(), exp - 8)
+
     def test_round(self):
         # check round-half-even algorithm. For round to nearest ten;
         # rounding map is invariant under adding multiples of 20
diff --git a/Misc/NEWS.d/next/Core and Builtins/2019-06-02-11-29-15.bpo-29882.AkRzjb.rst b/Misc/NEWS.d/next/Core and Builtins/2019-06-02-11-29-15.bpo-29882.AkRzjb.rst
@@ -0,0 +1,2 @@
+Add :meth:`int.bit_count()`, counting the number of ones in the binary
+representation of an integer. Patch by Niklas Fiekas.
diff --git a/Objects/clinic/longobject.c.h b/Objects/clinic/longobject.c.h
diff --git a/Objects/longobject.c b/Objects/longobject.c
@@ -5304,6 +5304,75 @@ int_bit_length_impl(PyObject *self)
     return NULL;
 }
 
+static int
+popcount_digit(digit d)
+{
+    /* 32bit SWAR popcount. */
+    uint32_t u = d;
+    u -= (u >> 1) & 0x55555555U;
+    u = (u & 0x33333333U) + ((u >> 2) & 0x33333333U);
+    u = (u + (u >> 4)) & 0x0f0f0f0fU;
+    return (uint32_t)(u * 0x01010101U) >> 24;
+}
+
+/*[clinic input]
+int.bit_count
+
+Number of ones in the binary representation of the absolute value of self.
+
+Also known as the population count.
+
+>>> bin(13)
+'0b1101'
+>>> (13).bit_count()
+3
+[clinic start generated code]*/
+
+static PyObject *
+int_bit_count_impl(PyObject *self)
+/*[clinic end generated code: output=2e571970daf1e5c3 input=7e0adef8e8ccdf2e]*/
+{
+    assert(self != NULL);
+    assert(PyLong_Check(self));
+
+    PyLongObject *z = (PyLongObject *)self;
+    Py_ssize_t ndigits = Py_ABS(Py_SIZE(z));
+    Py_ssize_t bit_count = 0;
+
+    /* Each digit has up to PyLong_SHIFT ones, so the accumulated bit count
+       from the first PY_SSIZE_T_MAX/PyLong_SHIFT digits can't overflow a
+       Py_ssize_t. */
+    Py_ssize_t ndigits_fast = Py_MIN(ndigits, PY_SSIZE_T_MAX/PyLong_SHIFT);
+    for (Py_ssize_t i = 0; i < ndigits_fast; i++) {
+        bit_count += popcount_digit(z->ob_digit[i]);
+    }
+
+    PyObject *result = PyLong_FromSsize_t(bit_count);
+    if (result == NULL) {
+        return NULL;
+    }
+
+    /* Use Python integers if bit_count would overflow. */
+    for (Py_ssize_t i = ndigits_fast; i < ndigits; i++) {
+        PyObject *x = PyLong_FromLong(popcount_digit(z->ob_digit[i]));
+        if (x == NULL) {
+            goto error;
+        }
+        PyObject *y = long_add((PyLongObject *)result, (PyLongObject *)x);
+        Py_DECREF(x);
+        if (y == NULL) {
+            goto error;
+        }
+        Py_DECREF(result);
+        result = y;
+    }
+
+    return result;
+
+  error:
+    Py_DECREF(result);
+    return NULL;
+}
 
 /*[clinic input]
 int.as_integer_ratio
@@ -5460,6 +5529,7 @@ static PyMethodDef long_methods[] = {
     {"conjugate",       long_long_meth, METH_NOARGS,
      "Returns self, the complex conjugate of any int."},
     INT_BIT_LENGTH_METHODDEF
+    INT_BIT_COUNT_METHODDEF
     INT_TO_BYTES_METHODDEF
     INT_FROM_BYTES_METHODDEF
     INT_AS_INTEGER_RATIO_METHODDEF

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Add :meth:`int.bit_count()`, counting the number of ones in the binary
	`2`	`+representation of an integer. Patch by Niklas Fiekas.`