helpers: add bit operation helpers

Extract for_each_set_bit() that was added internally for the cpumask and
nodemask helpers, and add for_each_clear_bit() and test_bit() to go with
it.

Signed-off-by: Omar Sandoval <osandov@osandov.com>

drgn/helpers/linux/bitops.py

@@ -0,0 +1,64 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+"""
+Bit Operations
+--------------
+
+The ``drgn.helpers.linux.bitops`` module provides helpers for common bit
+operations in the Linux kernel.
+"""
+
+from typing import Iterator
+
+from drgn import IntegerLike, Object, sizeof
+
+__all__ = (
+    "for_each_clear_bit",
+    "for_each_set_bit",
+    "test_bit",
+)
+
+
+def for_each_set_bit(bitmap: Object, size: IntegerLike) -> Iterator[int]:
+    """
+    Iterate over all set (one) bits in a bitmap.
+
+    :param bitmap: ``unsigned long *``
+    :param size: Size of *bitmap* in bits.
+    """
+    size = int(size)
+    word_bits = 8 * sizeof(bitmap.type_.type)
+    for i in range((size + word_bits - 1) // word_bits):
+        word = bitmap[i].value_()
+        for j in range(min(word_bits, size - word_bits * i)):
+            if word & (1 << j):
+                yield (word_bits * i) + j
+
+
+def for_each_clear_bit(bitmap: Object, size: IntegerLike) -> Iterator[int]:
+    """
+    Iterate over all clear (zero) bits in a bitmap.
+
+    :param bitmap: ``unsigned long *``
+    :param size: Size of *bitmap* in bits.
+    """
+    size = int(size)
+    word_bits = 8 * sizeof(bitmap.type_.type)
+    for i in range((size + word_bits - 1) // word_bits):
+        word = bitmap[i].value_()
+        for j in range(min(word_bits, size - word_bits * i)):
+            if not (word & (1 << j)):
+                yield (word_bits * i) + j
+
+
+def test_bit(nr: IntegerLike, bitmap: Object) -> bool:
+    """
+    Return whether a bit in a bitmap is set.
+
+    :param nr: Bit number.
+    :param bitmap: ``unsigned long *``
+    """
+    nr = int(nr)
+    word_bits = 8 * sizeof(bitmap.type_.type)
+    return ((bitmap[nr // word_bits].value_() >> (nr & (word_bits - 1))) & 1) != 0

drgn/helpers/linux/cpumask.py

@@ -11,7 +11,8 @@ masks from :linux:`include/linux/cpumask.h`.
 
 from typing import Iterator
 
-from drgn import Object, Program, sizeof
+from drgn import Object, Program
+from drgn.helpers.linux.bitops import for_each_set_bit
 
 __all__ = (
     "for_each_cpu",
@@ -21,15 +22,6 @@ __all__ = (
 )
 
 
-def _for_each_set_bit(bitmap: Object, size: int) -> Iterator[int]:
-    word_bits = 8 * sizeof(bitmap.type_.type)
-    for i in range((size + word_bits - 1) // word_bits):
-        word = bitmap[i].value_()
-        for j in range(min(word_bits, size - word_bits * i)):
-            if word & (1 << j):
-                yield (word_bits * i) + j
-
-
 def for_each_cpu(mask: Object) -> Iterator[int]:
     """
     Iterate over all of the CPUs in the given mask.
@@ -40,7 +32,7 @@ def for_each_cpu(mask: Object) -> Iterator[int]:
         nr_cpu_ids = mask.prog_["nr_cpu_ids"].value_()
     except KeyError:
         nr_cpu_ids = 1
-    return _for_each_set_bit(mask.bits, nr_cpu_ids)
+    return for_each_set_bit(mask.bits, nr_cpu_ids)
 
 
 def _for_each_cpu_mask(prog: Program, name: str) -> Iterator[int]:

drgn/helpers/linux/nodemask.py

@@ -12,7 +12,7 @@ NUMA node masks from :linux:`include/linux/nodemask.h`.
 from typing import Iterator
 
 from drgn import IntegerLike, Object, Program
-from drgn.helpers.linux.cpumask import _for_each_set_bit
+from drgn.helpers.linux.bitops import for_each_set_bit
 
 __all__ = (
     "for_each_node",
@@ -32,7 +32,7 @@ def for_each_node_mask(mask: Object) -> Iterator[int]:
         nr_node_ids = mask.prog_["nr_node_ids"].value_()
     except KeyError:
         nr_node_ids = 1
-    return _for_each_set_bit(mask.bits, nr_node_ids)
+    return for_each_set_bit(mask.bits, nr_node_ids)
 
 
 def for_each_node_state(prog: Program, state: IntegerLike) -> Iterator[int]:

tests/helpers/linux/test_bitops.py

@@ -0,0 +1,48 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from drgn import Object
+from drgn.helpers.linux.bitops import for_each_clear_bit, for_each_set_bit, test_bit
+from tests import MockProgramTestCase
+
+
+class TestBitOps(MockProgramTestCase):
+    BITMAP = [0xB351BC986648A680, 0x80DDB6615A80BC63]
+    # fmt: off
+    SET_BITS = [
+        7, 9, 10, 13, 15, 19, 22, 25, 26, 29, 30, 35, 36, 39, 42, 43, 44, 45,
+        47, 48, 52, 54, 56, 57, 60, 61, 63, 64, 65, 69, 70, 74, 75, 76, 77, 79,
+        87, 89, 91, 92, 94, 96, 101, 102, 105, 106, 108, 109, 111, 112, 114,
+        115, 116, 118, 119, 127,
+    ]
+    CLEAR_BITS = [
+        0, 1, 2, 3, 4, 5, 6, 8, 11, 12, 14, 16, 17, 18, 20, 21, 23, 24, 27, 28,
+        31, 32, 33, 34, 37, 38, 40, 41, 46, 49, 50, 51, 53, 55, 58, 59, 62, 66,
+        67, 68, 71, 72, 73, 78, 80, 81, 82, 83, 84, 85, 86, 88, 90, 93, 95, 97,
+        98, 99, 100, 103, 104, 107, 110, 113, 117, 120, 121, 122, 123, 124,
+        125, 126,
+    ]
+    # fmt: on
+
+    def test_for_each_set_bit(self):
+        bitmap = Object(self.prog, "unsigned long [2]", self.BITMAP)
+        self.assertEqual(list(for_each_set_bit(bitmap, 128)), self.SET_BITS)
+        self.assertEqual(
+            list(for_each_set_bit(bitmap, 101)),
+            [bit for bit in self.SET_BITS if bit < 101],
+        )
+
+    def test_for_each_clear_bit(self):
+        bitmap = Object(self.prog, "unsigned long [2]", self.BITMAP)
+        self.assertEqual(list(for_each_clear_bit(bitmap, 128)), self.CLEAR_BITS)
+        self.assertEqual(
+            list(for_each_clear_bit(bitmap, 100)),
+            [bit for bit in self.CLEAR_BITS if bit < 100],
+        )
+
+    def test_test_bit(self):
+        bitmap = Object(self.prog, "unsigned long [2]", self.BITMAP)
+        for bit in self.SET_BITS:
+            self.assertTrue(test_bit(bit, bitmap))
+        for bit in self.CLEAR_BITS:
+            self.assertFalse(test_bit(bit, bitmap))