helpers: Infer bitmap size when possible for bitops

The Linux API for for_each_set_bit() involves a size, but when dealing
with complete array types (a relatively common case), we can infer the
size. Add this ability to the bitops helpers.

Signed-off-by: Stephen Brennan <[email protected]>

drgn/helpers/linux/bitops.py

@@ -9,9 +9,9 @@
 operations in the Linux kernel.
 """
 
-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,15 +20,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: pointer to, or array of, unsigned integer type,
         EG: ``unsigned long *``
-    :param size: Size of *bitmap* in bits.
+    :param size: Size of *bitmap* in bits. When *bitmap* is a sized array type
+        (EG: ``unsigned long[2]``), this can be omitted.
     """
-    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_()
@@ -37,15 +45,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: pointer to, or array of, unsigned integer type,
         EG: ``unsigned long *``
-    :param size: Size of *bitmap* in bits.
+    :param size: Size of *bitmap* in bits. When *bitmap* is a sized array type
+        (EG: ``unsigned long[2]``), this can be omitted.
     """
-    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_()

tests/linux_kernel/helpers/test_bitops.py

@@ -34,6 +34,7 @@ def test_for_each_set_bit(self):
         for type_ in self.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],
@@ -43,6 +44,7 @@ def test_for_each_clear_bit(self):
         for type_ in self.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],