diff --git a/vortex-buffer/src/bit/ops.rs b/vortex-buffer/src/bit/ops.rs index be1401d912a..5006fbcd4cc 100644 --- a/vortex-buffer/src/bit/ops.rs +++ b/vortex-buffer/src/bit/ops.rs @@ -5,9 +5,8 @@ use std::mem::MaybeUninit; use crate::BitBuffer; use crate::BitBufferMut; -use crate::Buffer; +use crate::BufferMut; use crate::ByteBufferMut; -use crate::trusted_len::TrustedLenExt; /// Read up to 8 bytes as a little-endian `u64`, zero-padding the high bytes when fewer than 8 are /// supplied. Using [`u64::from_le_bytes`] keeps the bit-numbering identical on little- and @@ -188,36 +187,52 @@ pub(super) fn bitwise_binary_op u64>( mut op: F, ) -> BitBuffer { assert_eq!(left.len(), right.len()); - - // If the buffers are aligned, we can use the fast path. - if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) { - let left_chunks = left.unaligned_chunks(); - let right_chunks = right.unaligned_chunks(); - if left_chunks.lead_padding() == 0 - && left_chunks.trailing_padding() == 0 - && right_chunks.lead_padding() == 0 - && right_chunks.trailing_padding() == 0 - { - let iter = left_chunks - .iter() - .zip(right_chunks.iter()) - .map(|(l, r)| op(l, r)); - let iter = unsafe { iter.trusted_len() }; - let result = Buffer::::from_trusted_len_iter(iter).into_byte_buffer(); - return BitBuffer::new(result, left.len()); - } + let len = left.len(); + if len == 0 { + return BitBuffer::empty(); } - let iter = left - .chunks() - .iter_padded() - .zip(right.chunks().iter_padded()) - .map(|(l, r)| op(l, r)); - let iter = unsafe { iter.trusted_len() }; + let n_bytes = len.div_ceil(8); + let out = if left.offset().is_multiple_of(8) && right.offset().is_multiple_of(8) { + // Byte-aligned operands: logical bits map onto physical `u64` words, so read the backing + // bytes straight as words and build the result from a `TrustedLen` iterator. + 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 result = Buffer::::from_trusted_len_iter(iter).into_byte_buffer(); + let (lhs_words, lhs_tail) = lhs.as_chunks::<8>(); + let (rhs_words, rhs_tail) = rhs.as_chunks::<8>(); + + let mut out = BufferMut::::from_trusted_len_iter( + lhs_words + .iter() + .zip(rhs_words) + .map(|(l, r)| op(u64::from_le_bytes(*l), u64::from_le_bytes(*r))), + ); + if !lhs_tail.is_empty() { + out.push(op(read_u64_le(lhs_tail), read_u64_le(rhs_tail))); + } + out + } else { + // Sub-byte offset: `iter_padded` realigns the bits and appends one pad word, so take + // exactly `ceil(len / 64)` words. + let n_words = len.div_ceil(64); + let mut out = BufferMut::::with_capacity(n_words); + for (l, r) in left + .chunks() + .iter_padded() + .zip(right.chunks().iter_padded()) + .take(n_words) + { + out.push(op(l, r)); + } + out + }; - BitBuffer::new(result, left.len()) + let mut bytes = out.into_byte_buffer(); + bytes.truncate(n_bytes); + BitBuffer::new(bytes.freeze(), len) } #[cfg(test)] @@ -317,6 +332,41 @@ mod tests { assert_eq!(result, bitbuffer![false, true, true, false]); } + /// `bitwise_binary_op` must match a naive per-bit reference for every op, offset and length, + /// independent of the chunked kernels. + #[rstest] + #[case::aligned(0, 0)] + #[case::byte_aligned(8, 16)] + #[case::byte_aligned_mismatch(16, 0)] + #[case::sub_byte(3, 3)] + #[case::sub_byte_mismatch(0, 5)] + fn binary_op_matches_naive(#[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, 200, 256] { + let left = make(left_offset, len, 0xC3); + let right = make(right_offset, len, 0x5A); + for op in ops { + let got = bitwise_binary_op(&left, &right, op); + let expected: BitBuffer = (0..len) + .map(|i| op(u64::from(left.value(i)), u64::from(right.value(i))) & 1 == 1) + .collect(); + assert_eq!( + got, expected, + "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