move fn to BitBuffer

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

Cargo.lock

@@ -42,6 +42,7 @@ workspace = true
 [dev-dependencies]
 arrow-buffer = { workspace = true }
 divan = { workspace = true }
+rstest = { workspace = true }
 
 [[bench]]
 name = "vortex_buffer"

vortex-buffer/Cargo.toml

@@ -406,6 +406,108 @@ impl BitBuffer {
     pub fn bitand_not(&self, rhs: &BitBuffer) -> BitBuffer {
         bitwise_binary_op(self, rhs, |a, b| a & !b)
     }
+
+    /// Iterate through bits in a buffer.
+    ///
+    /// # Arguments
+    ///
+    /// * `range` - Bit range to iterate through
+    /// * `f` - Callback function taking (bit_index, is_set)
+    ///
+    /// # Panics
+    ///
+    /// Panics if the range is outside valid bounds of the buffer.
+    #[inline]
+    pub fn iter_bits<F>(&self, range: std::ops::Range<usize>, mut f: F)
+    where
+        F: FnMut(usize, bool),
+    {
+        let start = range.start;
+        let end = range.end;
+
+        assert!(start <= end);
+        assert!(end <= self.len);
+
+        let buffer_ptr = self.buffer.as_ptr();
+        let offset = self.offset;
+
+        let full_bytes = (end - start) / 8;
+        let remaining_bits = (end - start) % 8;
+
+        for byte_idx in 0..full_bytes {
+            let bit_offset = offset + start + byte_idx * 8;
+            let byte_offset = bit_offset / 8;
+            let byte = unsafe { *buffer_ptr.add(byte_offset) };
+
+            for bit_idx in 0..8 {
+                let is_set = (byte & (1 << bit_idx)) != 0;
+                f(start + byte_idx * 8 + bit_idx, is_set);
+            }
+        }
+
+        if remaining_bits > 0 {
+            let bit_idx_start = start + full_bytes * 8;
+            let bit_offset = offset + bit_idx_start;
+            let byte_offset = bit_offset / 8;
+            let byte = unsafe { *buffer_ptr.add(byte_offset) };
+
+            for i in 0..remaining_bits {
+                let is_set = (byte & (1 << i)) != 0;
+                f(bit_idx_start + i, is_set);
+            }
+        }
+    }
+
+    /// Iterate through bits in a buffer in reverse.
+    ///
+    /// # Arguments
+    ///
+    /// * `range` - Bit range to iterate through in reverse (start inclusive, end exclusive)
+    /// * `f` - Callback function taking (bit_index, is_set)
+    ///
+    /// # Panics
+    ///
+    /// Panics if the range is outside valid bounds of the buffer.
+    #[inline]
+    pub fn iter_bits_reverse<F>(&self, range: std::ops::Range<usize>, mut f: F)
+    where
+        F: FnMut(usize, bool),
+    {
+        let start = range.start;
+        let end = range.end;
+
+        assert!(start <= end);
+        assert!(end <= self.len);
+
+        let buffer_ptr = self.buffer.as_ptr();
+        let offset = self.offset;
+
+        let full_bytes = (end - start) / 8;
+        let remaining_bits = (end - start) % 8;
+
+        if remaining_bits > 0 {
+            let bit_idx_start = start + full_bytes * 8;
+            let bit_offset = offset + bit_idx_start;
+            let byte_offset = bit_offset / 8;
+            let byte = unsafe { *buffer_ptr.add(byte_offset) };
+
+            for bit_idx in (0..remaining_bits).rev() {
+                let is_set = (byte & (1 << bit_idx)) != 0;
+                f(bit_idx_start + bit_idx, is_set);
+            }
+        }
+
+        for byte_idx in (0..full_bytes).rev() {
+            let bit_offset = offset + start + byte_idx * 8;
+            let byte_offset = bit_offset / 8;
+            let byte = unsafe { *buffer_ptr.add(byte_offset) };
+
+            for bit_idx in (0..8).rev() {
+                let is_set = (byte & (1 << bit_idx)) != 0;
+                f(start + byte_idx * 8 + bit_idx, is_set);
+            }
+        }
+    }
 }
 
 impl<'a> IntoIterator for &'a BitBuffer {
@@ -419,6 +521,8 @@ impl<'a> IntoIterator for &'a BitBuffer {
 
 #[cfg(test)]
 mod tests {
+    use rstest::rstest;
+
     use crate::bit::BitBuffer;
     use crate::{ByteBuffer, buffer};
 
@@ -488,4 +592,52 @@ mod tests {
             "Buffer slices with different bits should not be equal (`PartialEq` needs `iter_padded()`)"
         );
     }
+
+    #[rstest]
+    #[case(5)]
+    #[case(8)]
+    #[case(10)]
+    #[case(13)]
+    #[case(16)]
+    #[case(23)]
+    #[case(100)]
+    fn test_iter_bits(#[case] len: usize) {
+        let buf = BitBuffer::collect_bool(len, |i| i % 2 == 0);
+
+        let mut collected = Vec::new();
+        buf.iter_bits(0..len, |idx, is_set| {
+            collected.push((idx, is_set));
+        });
+
+        assert_eq!(collected.len(), len);
+
+        for (idx, is_set) in collected {
+            assert_eq!(is_set, idx % 2 == 0);
+        }
+    }
+
+    #[rstest]
+    #[case(5)]
+    #[case(8)]
+    #[case(10)]
+    #[case(13)]
+    #[case(16)]
+    #[case(23)]
+    #[case(100)]
+    fn test_iter_bits_reverse(#[case] len: usize) {
+        let buf = BitBuffer::collect_bool(len, |i| i % 2 == 0);
+
+        let mut collected = Vec::new();
+        buf.iter_bits_reverse(0..len, |idx, is_set| {
+            collected.push((idx, is_set));
+        });
+
+        assert_eq!(collected.len(), len);
+
+        for (i, (idx, is_set)) in collected.iter().enumerate() {
+            let expected_idx = len - 1 - i;
+            assert_eq!(*idx, expected_idx);
+            assert_eq!(*is_set, expected_idx % 2 == 0);
+        }
+    }
 }

vortex-buffer/src/bit/buf.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::ops::Range;
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskValues};
 
@@ -100,7 +99,7 @@ fn expand_inplace<T: Copy>(buf_mut: &mut BufferMut<T>, mask_values: &MaskValues)
     let bit_buffer = mask_values.bit_buffer();
 
     // Iterate backwards through the mask to avoid overwriting unprocessed elements.
-    iter_bits_reverse(bit_buffer, 0..mask_len, |idx, is_valid| {
+    bit_buffer.iter_bits_reverse(0..mask_len, |idx, is_valid| {
         if is_valid {
             element_idx -= 1;
             unsafe { *buf_slice.get_unchecked_mut(idx) = buf_slice[element_idx] };
@@ -134,7 +133,7 @@ fn expand_copy<T: Copy>(src: &[T], mask_values: &MaskValues) -> Buffer<T> {
 
     let bit_buffer = mask_values.bit_buffer();
 
-    iter_bits(bit_buffer, 0..mask_len, |mask_idx, is_valid| {
+    bit_buffer.iter_bits(0..mask_len, |mask_idx, is_valid| {
         if is_valid {
             unsafe {
                 target_slice
@@ -157,110 +156,6 @@ fn expand_copy<T: Copy>(src: &[T], mask_values: &MaskValues) -> Buffer<T> {
     target_buf.freeze()
 }
 
-/// Iterate through bits in a buffer.
-///
-/// # Arguments
-///
-/// * `bit_buffer` - The bit buffer to iterate through
-/// * `range` - Bit range to iterate through
-/// * `f` - Callback function taking (bit_index, is_set)
-///
-/// # Safety
-///
-/// The caller must ensure that the range is within valid bounds of the bit buffer.
-#[inline]
-fn iter_bits<F>(bit_buffer: &vortex_buffer::BitBuffer, range: Range<usize>, mut f: F)
-where
-    F: FnMut(usize, bool),
-{
-    let start = range.start;
-    let end = range.end;
-
-    assert!(start <= end);
-    assert!(end <= bit_buffer.len());
-
-    let buffer_ptr = bit_buffer.inner().as_ptr();
-    let offset = bit_buffer.offset();
-
-    let full_bytes = (end - start) / 8;
-    let remaining_bits = (end - start) % 8;
-
-    for byte_idx in 0..full_bytes {
-        let bit_offset = offset + start + byte_idx * 8;
-        let byte_offset = bit_offset / 8;
-        let byte = unsafe { *buffer_ptr.add(byte_offset) };
-
-        for bit_idx in 0..8 {
-            let is_set = (byte & (1 << bit_idx)) != 0;
-            f(start + byte_idx * 8 + bit_idx, is_set);
-        }
-    }
-
-    if remaining_bits > 0 {
-        let bit_idx_start = start + full_bytes * 8;
-        let bit_offset = offset + bit_idx_start;
-        let byte_offset = bit_offset / 8;
-        let byte = unsafe { *buffer_ptr.add(byte_offset) };
-
-        for i in 0..remaining_bits {
-            let is_set = (byte & (1 << i)) != 0;
-            f(bit_idx_start + i, is_set);
-        }
-    }
-}
-
-/// Iterate through bits in a buffer in reverse.
-///
-/// # Arguments
-///
-/// * `bit_buffer` - The bit buffer to iterate through
-/// * `range` - Bit range to iterate through in reverse (start inclusive, end exclusive)
-/// * `f` - Callback function taking (bit_index, is_set)
-///
-/// # Safety
-///
-/// The caller must ensure that the range is within valid bounds of the bit buffer.
-#[inline]
-fn iter_bits_reverse<F>(bit_buffer: &vortex_buffer::BitBuffer, range: Range<usize>, mut f: F)
-where
-    F: FnMut(usize, bool),
-{
-    let start = range.start;
-    let end = range.end;
-
-    assert!(start <= end);
-    assert!(end <= bit_buffer.len());
-
-    let buffer_ptr = bit_buffer.inner().as_ptr();
-    let offset = bit_buffer.offset();
-
-    let full_bytes = (end - start) / 8;
-    let remaining_bits = (end - start) % 8;
-
-    if remaining_bits > 0 {
-        let bit_idx_start = start + full_bytes * 8;
-        let bit_offset = offset + bit_idx_start;
-        let byte_offset = bit_offset / 8;
-        let byte = unsafe { *buffer_ptr.add(byte_offset) };
-
-        for bit_idx in (0..remaining_bits).rev() {
-            let is_set = (byte & (1 << bit_idx)) != 0;
-            f(bit_idx_start + bit_idx, is_set);
-        }
-    }
-
-    for byte_idx in (0..full_bytes).rev() {
-        let bit_offset = offset + start + byte_idx * 8;
-        let byte_offset = bit_offset / 8;
-        let byte = unsafe { *buffer_ptr.add(byte_offset) };
-
-        for bit_idx in (0..8).rev() {
-            let is_set = (byte & (1 << bit_idx)) != 0;
-            f(start + byte_idx * 8 + bit_idx, is_set);
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_buffer::{buffer, buffer_mut};