Bitops improvements

drgn/helpers/linux/bitops.py

@@ -6,12 +6,17 @@
 --------------
 
 The ``drgn.helpers.linux.bitops`` module provides helpers for common bit
-operations in the Linux kernel.
+operations in the Linux kernel. These helpers operate on arrays of, or pointers
+to, ``unsigned long``, just as the kernel's helpers do. For little-endian
+architectures, they can operate equivalently if the underlying type is any other
+unsigned integer type. But for big-endian architectures, this does not hold
+true. For broadest compatibility, it is best to ensure the underlying type is
+``unsigned long``.
 """
 
-from typing import Iterator
+from typing import Iterator, Optional
 
-from drgn import IntegerLike, Object, sizeof
+from drgn import IntegerLike, Object, TypeKind, sizeof
 
 __all__ = (
     "for_each_clear_bit",
@@ -20,14 +25,23 @@
 )
 
 
-def for_each_set_bit(bitmap: Object, size: IntegerLike) -> Iterator[int]:
+def for_each_set_bit(
+    bitmap: Object, size: Optional[IntegerLike] = None
+) -> Iterator[int]:
     """
     Iterate over all set (one) bits in a bitmap.
 
-    :param bitmap: ``unsigned long *``
-    :param size: Size of *bitmap* in bits.
+    :param bitmap: pointer to, or array of, ``unsigned long``
+    :param size: Size of *bitmap* in bits. When *bitmap* is a sized array type
+        (EG: ``unsigned long[2]``), this value will default to the size of the
+        array in bits.
     """
-    size = int(size)
+    if size is not None:
+        size = int(size)
+    elif bitmap.type_.kind == TypeKind.ARRAY and bitmap.type_.length is not None:
+        size = 8 * sizeof(bitmap)
+    else:
+        raise ValueError("bitmap is not a complete array type, and size is not given")
     word_bits = 8 * sizeof(bitmap.type_.type)
     for i in range((size + word_bits - 1) // word_bits):
         word = bitmap[i].value_()
@@ -36,14 +50,23 @@ def for_each_set_bit(bitmap: Object, size: IntegerLike) -> Iterator[int]:
                 yield (word_bits * i) + j
 
 
-def for_each_clear_bit(bitmap: Object, size: IntegerLike) -> Iterator[int]:
+def for_each_clear_bit(
+    bitmap: Object, size: Optional[IntegerLike] = None
+) -> Iterator[int]:
     """
     Iterate over all clear (zero) bits in a bitmap.
 
-    :param bitmap: ``unsigned long *``
-    :param size: Size of *bitmap* in bits.
+    :param bitmap: pointer to, or array of, ``unsigned long``
+    :param size: Size of *bitmap* in bits. When *bitmap* is a sized array type
+        (EG: ``unsigned long[2]``), this value will default to the size of the
+        array in bits.
     """
-    size = int(size)
+    if size is not None:
+        size = int(size)
+    elif bitmap.type_.kind == TypeKind.ARRAY and bitmap.type_.length is not None:
+        size = 8 * sizeof(bitmap)
+    else:
+        raise ValueError("bitmap is not a complete array type, and size is not given")
     word_bits = 8 * sizeof(bitmap.type_.type)
     for i in range((size + word_bits - 1) // word_bits):
         word = bitmap[i].value_()
@@ -57,7 +80,7 @@ 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 *``
+    :param bitmap: pointer to, or array of, ``unsigned long``
     """
     nr = int(nr)
     word_bits = 8 * sizeof(bitmap.type_.type)

tests/linux_kernel/helpers/test_bitops.py

@@ -1,13 +1,13 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # SPDX-License-Identifier: LGPL-2.1-or-later
 
-from drgn import Object
+from drgn import Object, PlatformFlags
 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]
+    BITMAP = bytes.fromhex("80A6486698BC51B363BC805A61B6DD80")
     # fmt: off
     SET_BITS = [
         7, 9, 10, 13, 15, 19, 22, 25, 26, 29, 30, 35, 36, 39, 42, 43, 44, 45,
@@ -22,27 +22,44 @@ class TestBitOps(MockProgramTestCase):
         98, 99, 100, 103, 104, 107, 110, 113, 117, 120, 121, 122, 123, 124,
         125, 126,
     ]
+    TYPES = [
+        "unsigned long [2]",
+        "unsigned int [4]",
+        "unsigned short [8]",
+        "unsigned char [16]",
+    ]
     # fmt: on
 
+    def valid_integer_types(self):
+        if self.prog.platform.flags & PlatformFlags.IS_LITTLE_ENDIAN:
+            return self.TYPES
+        else:
+            return self.TYPES[:1]
+
     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],
-        )
+        for type_ in self.valid_integer_types():
+            bitmap = Object.from_bytes_(self.prog, type_, self.BITMAP)
+            self.assertEqual(list(for_each_set_bit(bitmap, 128)), self.SET_BITS)
+            self.assertEqual(list(for_each_set_bit(bitmap)), 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],
-        )
+        for type_ in self.valid_integer_types():
+            bitmap = Object.from_bytes_(self.prog, type_, self.BITMAP)
+            self.assertEqual(list(for_each_clear_bit(bitmap, 128)), self.CLEAR_BITS)
+            self.assertEqual(list(for_each_clear_bit(bitmap)), 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))
+        for type_ in self.valid_integer_types():
+            bitmap = Object.from_bytes_(self.prog, type_, 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))