Move bit-twiddling helpers to their own file (#299)

This also deduplicates a repeated bit of code that iterates over bits

clifford/_bit_helpers.py

@@ -0,0 +1,48 @@
+""" Helpers for efficient bit operations
+
+Bitmaps are used in clifford as representations of basis blades, hence the need
+for these operations. These are only for internal use.
+"""
+
+from typing import Iterator
+
+import numba
+import numba.extending
+import numba.types
+import numba.config
+
+
+def set_bit_indices(x: int) -> Iterator[int]:
+    """ Iterate over the indices of bits set to 1 in `x`, in ascending order """
+    n = 0
+    while x > 0:
+        if x & 1:
+            yield n
+        x = x >> 1
+        n = n + 1
+
+
+@numba.extending.intrinsic
+def __builtin_popcnt(tyctx, x):
+    """ Emulate clang and GCC's `__builtin_popcnt` """
+    if isinstance(x, numba.types.Integer):
+        def impl(cgctx, builder, sig, args):
+            x, = args
+            return builder.ctpop(x)
+        sig = x(x)
+        return sig, impl
+
+
+if numba.config.DISABLE_JIT:
+    def count_set_bits(bitmap: int) -> int:
+        """ Counts the number of bits set to 1 in bitmap """
+        count = 0
+        for i in set_bit_indices(bitmap):
+            count += 1
+        return count
+
+else:
+    @numba.njit
+    def count_set_bits(x: int) -> int:
+        """ Counts the number of bits set to 1 in bitmap """
+        return __builtin_popcnt(x)

clifford/_layout_helpers.py

@@ -14,41 +14,7 @@
 import operator
 
 from . import _numba_utils
-import numba
-import numba.extending
-import numba.types
-import numba.config
-
-
-@numba.extending.intrinsic
-def __builtin_popcnt(tyctx, x):
-    """ Emulate clang and GCC's `__builtin_popcnt` """
-    if isinstance(x, numba.types.Integer):
-        def impl(cgctx, builder, sig, args):
-            x, = args
-            return builder.ctpop(x)
-        sig = x(x)
-        return sig, impl
-
-
-if numba.config.DISABLE_JIT:
-    def count_set_bits(bitmap):
-        """ Counts the number of bits set to 1 in bitmap """
-        bmp = bitmap
-        count = 0
-        n = 1
-        while bmp > 0:
-            if bmp & 1:
-                count += 1
-            bmp = bmp >> 1
-            n = n + 1
-        return count
-
-else:
-    @numba.njit
-    def count_set_bits(x):
-        """ Counts the number of bits set to 1 in bitmap """
-        return __builtin_popcnt(x)
+from ._bit_helpers import count_set_bits, set_bit_indices
 
 
 @_numba_utils.njit
@@ -192,15 +158,7 @@ def __init__(self, blade_ids: Sequence[IdT]):
 
     def bitmap_as_tuple(self, bitmap: int) -> Tuple[IdT]:
         """ Convert a bitmap representation into a tuple of ids. """
-        bmp = bitmap
-        blade = []
-        n = 0
-        while bmp > 0:
-            if bmp & 1:
-                blade.append(self.values[n])
-            bmp = bmp >> 1
-            n = n + 1
-        return tuple(blade)
+        return tuple(self.values[n] for n in set_bit_indices(bitmap))
 
     def id_as_bitmap(self, id: IdT) -> int:
         """ Convert the id of a single vector into a bitmap representation. """

clifford/test/test_bit_helpers.py

@@ -0,0 +1,16 @@
+""" Tests of clifford._bit_helpers """
+from clifford._bit_helpers import count_set_bits, set_bit_indices
+
+
+def test_count_bits():
+    assert count_set_bits(0b0) == 0
+    assert count_set_bits(0b1) == 1
+    assert count_set_bits(0b1001) == 2
+    assert count_set_bits(0b1111) == 4
+
+
+def test_bit_indices():
+    assert list(set_bit_indices(0b0)) == []
+    assert list(set_bit_indices(0b1)) == [0]
+    assert list(set_bit_indices(0b101)) == [0, 2]
+    assert list(set_bit_indices(0b101010)) == [1, 3, 5]