perf(mask): fuse true_count into boolean operators

vortex-buffer/Cargo.toml

@@ -48,3 +48,7 @@ harness = false
 [[bench]]
 name = "vortex_bitbuffer"
 harness = false
+
+[[bench]]
+name = "bitbuffer_count"
+harness = false

vortex-buffer/benches/bitbuffer_count.rs

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Benchmarks comparing strategies for computing a bitwise-op result *and* its `true_count`:
+//!
+//! * `*_two_pass`: current behaviour, materialise the result buffer then run the (vectorised)
+//!   `count_ones` kernel as a separate pass (`&a & &b` followed by `.true_count()`).
+//! * `*_fused`: single pass that combines the words and accumulates `count_ones` together
+//!   (`BitBuffer::bitand_with_true_count`).
+
+#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
+
+use divan::Bencher;
+use divan::black_box;
+use vortex_buffer::BitBuffer;
+
+fn main() {
+    divan::main();
+}
+
+// Lengths in bits, from L1-resident up to L2/L3.
+const SIZES: &[usize] = &[
+    1_024,   // 128 B
+    8_192,   // 1 KB
+    65_536,  // 8 KB
+    524_288, // 64 KB
+];
+
+fn inputs(len: usize) -> (BitBuffer, BitBuffer) {
+    let lhs =
+        BitBuffer::from_iter((0..len).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) & 1 == 0));
+    let rhs = BitBuffer::from_iter((0..len).map(|i| (i.wrapping_mul(40_503) >> 7) & 1 == 0));
+    (lhs, rhs)
+}
+
+#[divan::bench(args = SIZES)]
+fn and_two_pass(bencher: Bencher, len: usize) {
+    let (lhs, rhs) = inputs(len);
+    bencher.bench_local(|| {
+        let result = black_box(&lhs) & black_box(&rhs);
+        let count = result.true_count();
+        (result, count)
+    });
+}
+
+#[divan::bench(args = SIZES)]
+fn and_fused(bencher: Bencher, len: usize) {
+    let (lhs, rhs) = inputs(len);
+    bencher.bench_local(|| black_box(&lhs).bitand_with_true_count(black_box(&rhs)));
+}
+
+// Owned-LHS scenario: the current code can combine in-place (no allocation) then count
+// separately; the fused path always allocates but does a single pass. A fresh owned, uniquely
+// owned LHS is cloned per iteration so the in-place fast path is actually taken.
+#[divan::bench(args = SIZES)]
+fn and_owned_two_pass(bencher: Bencher, len: usize) {
+    let (lhs, rhs) = inputs(len);
+    bencher
+        .with_inputs(|| lhs.clone())
+        .bench_local_values(|lhs| {
+            let result = lhs & black_box(&rhs);
+            let count = result.true_count();
+            (result, count)
+        });
+}
+
+#[divan::bench(args = SIZES)]
+fn and_owned_fused(bencher: Bencher, len: usize) {
+    let (lhs, rhs) = inputs(len);
+    bencher
+        .with_inputs(|| lhs.clone())
+        .bench_local_values(|lhs| lhs.bitand_with_true_count(black_box(&rhs)));
+}

vortex-buffer/src/bit/buf.rs

@@ -25,6 +25,7 @@ use crate::bit::collect_bool_word;
 use crate::bit::count_ones::count_ones;
 use crate::bit::get_bit_unchecked;
 use crate::bit::ops::bitwise_binary_op;
+use crate::bit::ops::bitwise_binary_op_counted;
 use crate::bit::ops::bitwise_binary_op_lhs_owned;
 use crate::bit::ops::bitwise_unary_op;
 use crate::bit::ops::bitwise_unary_op_copy;
@@ -541,6 +542,21 @@ impl BitBuffer {
         bitwise_binary_op_lhs_owned(self, rhs, |a, b| a & !b)
     }
 
+    /// Bitwise AND of two buffers, returning the result and its true count from a single pass.
+    pub fn bitand_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) {
+        bitwise_binary_op_counted(self, rhs, |a, b| a & b)
+    }
+
+    /// Bitwise OR of two buffers, returning the result and its true count from a single pass.
+    pub fn bitor_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) {
+        bitwise_binary_op_counted(self, rhs, |a, b| a | b)
+    }
+
+    /// Bitwise AND NOT (`self & !rhs`), returning the result and its true count from a single pass.
+    pub fn bitand_not_with_true_count(&self, rhs: &BitBuffer) -> (BitBuffer, usize) {
+        bitwise_binary_op_counted(self, rhs, |a, b| a & !b)
+    }
+
     /// Iterate through bits in a buffer.
     ///
     /// # Arguments

vortex-buffer/src/bit/ops.rs

@@ -235,6 +235,113 @@ pub(super) fn bitwise_binary_op<F: FnMut(u64, u64) -> u64>(
     BitBuffer::new(bytes.freeze(), len)
 }
 
+/// Combine `left` and `right` word-by-word with `op`, calling `comb` on each result word so a
+/// caller can fold a reduction (e.g. `count_ones`) into the same pass. The word given to `comb`
+/// has bits past `len` cleared; the buffer keeps `op`'s raw output in those positions.
+pub(super) fn bitwise_binary_op_with<F, C>(
+    left: &BitBuffer,
+    right: &BitBuffer,
+    op: F,
+    comb: C,
+) -> BitBuffer
+where
+    F: FnMut(u64, u64) -> u64,
+    C: FnMut(u64),
+{
+    assert_eq!(left.len(), right.len());
+    let len = left.len();
+    if len == 0 {
+        return BitBuffer::empty();
+    }
+
+    // Byte-aligned operands: logical bits map onto physical `u64` words, so we can read the backing
+    // bytes straight as words instead of shifting through `iter_padded`.
+    if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) {
+        let n_bytes = len.div_ceil(8);
+        let l_start = left.offset() / 8;
+        let r_start = right.offset() / 8;
+        let lhs = &left.inner().as_slice()[l_start..l_start + n_bytes];
+        let rhs = &right.inner().as_slice()[r_start..r_start + n_bytes];
+
+        let (lhs_words, lhs_tail) = lhs.as_chunks::<8>();
+        let (rhs_words, rhs_tail) = rhs.as_chunks::<8>();
+
+        let words = lhs_words
+            .iter()
+            .zip(rhs_words)
+            .map(|(l, r)| (u64::from_le_bytes(*l), u64::from_le_bytes(*r)))
+            .chain((!lhs_tail.is_empty()).then(|| (read_u64_le(lhs_tail), read_u64_le(rhs_tail))));
+        return combine_words(words, len, op, comb);
+    }
+
+    // Sub-byte offset: `iter_padded` realigns the bits and zero-pads the final word.
+    let words = left
+        .chunks()
+        .iter_padded()
+        .zip(right.chunks().iter_padded());
+    combine_words(words, len, op, comb)
+}
+
+/// Combine an iterator of `(left, right)` words via `op` into a `len`-bit [`BitBuffer`], folding
+/// `comb` over each in-range result word. Consumes exactly `ceil(len / 64)` words (any extra the
+/// producer yields, e.g. `iter_padded`'s trailing pad word, is dropped). Only the final word can
+/// hold bits past `len`, so it alone is masked for `comb`; the buffer keeps `op`'s raw output.
+#[inline]
+fn combine_words<I, F, C>(words: I, len: usize, mut op: F, mut comb: C) -> BitBuffer
+where
+    I: Iterator<Item = (u64, u64)>,
+    F: FnMut(u64, u64) -> u64,
+    C: FnMut(u64),
+{
+    let n_words = len.div_ceil(64);
+    let mut out: BufferMut<u64> = BufferMut::with_capacity(n_words);
+    let mut words = words.take(n_words);
+
+    // Words fully within `len` need no masking; handling the final word separately keeps this loop
+    // uniform, with no per-word branch.
+    for (l, r) in words.by_ref().take(n_words - 1) {
+        let word = op(l, r);
+        comb(word);
+        out.push(word);
+    }
+    // Final word: clear bits past `len` for `comb` only.
+    if let Some((l, r)) = words.next() {
+        let word = op(l, r);
+        comb(word & last_word_mask(len));
+        out.push(word);
+    }
+
+    let mut bytes = out.into_byte_buffer();
+    bytes.truncate(len.div_ceil(8));
+    BitBuffer::new(bytes.freeze(), len)
+}
+
+/// Mask of the in-range bits in the final word: the low `len % 64` bits, or all 64 when `len` is
+/// a multiple of 64.
+#[inline]
+fn last_word_mask(len: usize) -> u64 {
+    let rem = len % 64;
+    if rem == 0 {
+        u64::MAX
+    } else {
+        (1u64 << rem) - 1
+    }
+}
+
+/// Like [`bitwise_binary_op`], but also returns the result's set-bit count, folded into the same
+/// pass instead of re-scanning the buffer.
+pub(super) fn bitwise_binary_op_counted<F: FnMut(u64, u64) -> u64>(
+    left: &BitBuffer,
+    right: &BitBuffer,
+    op: F,
+) -> (BitBuffer, usize) {
+    let mut count = 0usize;
+    let buffer = bitwise_binary_op_with(left, right, op, |word| {
+        count += word.count_ones() as usize;
+    });
+    (buffer, count)
+}
+
 #[cfg(test)]
 mod tests {
     use std::ops::Not;
@@ -367,6 +474,44 @@ mod tests {
         }
     }
 
+    /// The fused path must yield the same buffer and true count as the reference
+    /// (`bitwise_binary_op` + a separate `true_count`), across all ops, offsets, and lengths,
+    /// including lengths like 127 whose final word holds out-of-range bits.
+    #[rstest]
+    #[case::aligned(0, 0)]
+    #[case::equal_nonzero(3, 3)]
+    #[case::byte_aligned(8, 16)]
+    #[case::mismatch(0, 5)]
+    fn counted_matches_reference(#[case] left_offset: usize, #[case] right_offset: usize) {
+        #[allow(clippy::cast_possible_truncation)]
+        let make = |offset: usize, len: usize, salt: u8| -> BitBuffer {
+            let bytes: ByteBufferMut = (0..(offset + len).div_ceil(8).max(1))
+                .map(|i| (i as u8).wrapping_mul(31).wrapping_add(salt))
+                .collect();
+            BitBufferMut::from_buffer(bytes, offset, len).freeze()
+        };
+        let ops: [fn(u64, u64) -> u64; 4] =
+            [|a, b| a & b, |a, b| a | b, |a, b| a ^ b, |a, b| a & !b];
+
+        for len in [1usize, 5, 8, 63, 64, 65, 127, 128, 129, 200, 256] {
+            let left = make(left_offset, len, 0xC3);
+            let right = make(right_offset, len, 0x5A);
+            for op in ops {
+                let (got_buf, got_count) = bitwise_binary_op_counted(&left, &right, op);
+                let expected_buf = bitwise_binary_op(&left, &right, op);
+                assert_eq!(
+                    got_buf, expected_buf,
+                    "buffer mismatch loff={left_offset} roff={right_offset} len={len}"
+                );
+                assert_eq!(
+                    got_count,
+                    expected_buf.true_count(),
+                    "count mismatch loff={left_offset} roff={right_offset} len={len}"
+                );
+            }
+        }
+    }
+
     /// Regression test for a bug where [`bitwise_unary_op`] produced corrupt results when
     /// the [`BitBuffer`]'s underlying byte pointer was not u64-aligned. Slicing a buffer by
     /// a non-multiple-of-8 number of bytes can cause this misalignment. The bug only