diff --git a/drgn/helpers/linux/bitops.py b/drgn/helpers/linux/bitops.py index 521a61417..fe3fd147c 100644 --- a/drgn/helpers/linux/bitops.py +++ b/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_() diff --git a/tests/linux_kernel/helpers/test_bitops.py b/tests/linux_kernel/helpers/test_bitops.py index 39dbbdc1d..29bb74fca 100644 --- a/tests/linux_kernel/helpers/test_bitops.py +++ b/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],