diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index 920ecba6d60c..e1bf14bf2c34 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -168,6 +168,10 @@ pub(crate) fn read_up_to_byte_from_offset( /// * `len_in_bits` - Number of bits to process /// * `op` - Binary operation to apply (e.g., `|a, b| a & b`). Applied a word at a time /// +/// Only the bits in `left_offset_in_bits..left_offset_in_bits + len_in_bits` are +/// modified. Bits of `left` outside that range are left unchanged, including the +/// bits sharing a byte with either end of the range. +/// /// # Example: Modify entire buffer /// ``` /// # use arrow_buffer::MutableBuffer; @@ -247,6 +251,7 @@ pub fn apply_bitwise_binary_op( // Hope it gets inlined &mut |left| op(left, right_first_byte as u64), left_offset_in_bits, + bits_to_next_byte, ); } @@ -280,6 +285,10 @@ pub fn apply_bitwise_binary_op( /// * `len_in_bits` - Number of bits to process /// * `op` - Unary operation to apply (e.g., `|a| !a`). Applied a word at a time /// +/// Only the bits in `offset_in_bits..offset_in_bits + len_in_bits` are modified. +/// Bits outside that range are left unchanged, including the bits sharing a byte +/// with either end of the range. +/// /// # Example: Modify entire buffer /// ``` /// # use arrow_buffer::MutableBuffer; @@ -325,7 +334,7 @@ pub fn apply_bitwise_unary_op( if is_mutable_buffer_byte_aligned { byte_aligned_bitwise_unary_op_helper(buffer, offset_in_bits, len_in_bits, op); } else { - align_to_byte(buffer, &mut op, offset_in_bits); + align_to_byte(buffer, &mut op, offset_in_bits, len_in_bits); // If we are not byte aligned we will read the first few bits let bits_to_next_byte = 8 - left_bit_offset; @@ -449,13 +458,23 @@ fn byte_aligned_bitwise_unary_op_helper( /// * `op` - Unary operation to apply /// * `buffer` - The mutable buffer to modify /// * `offset_in_bits` - Starting bit offset (not byte-aligned) -fn align_to_byte(buffer: &mut [u8], op: &mut F, offset_in_bits: usize) -where +/// * `remaining_len_in_bits` - Number of bits still to process starting at `offset_in_bits`. +/// When this is smaller than the number of bits left in the byte, the trailing bits of +/// the byte are left untouched. +fn align_to_byte( + buffer: &mut [u8], + op: &mut F, + offset_in_bits: usize, + remaining_len_in_bits: usize, +) where F: FnMut(u64) -> u64, { let byte_offset = offset_in_bits / 8; let bit_offset = offset_in_bits % 8; + // Byte aligned offsets must take the byte aligned path instead + debug_assert_ne!(bit_offset, 0, "offset_in_bits must not be byte aligned"); + // 1. read the first byte from the buffer let first_byte: u8 = buffer[byte_offset]; @@ -468,12 +487,16 @@ where // 4. Shift back the result to the original position let result_first_byte = result_first_byte << bit_offset; - // 5. Mask the bits that are outside the relevant bits in the byte - // so the bits until bit_offset are 1 and the rest are 0 - let mask_for_first_bit_offset = (1 << bit_offset) - 1; + // 5. Mask in only the bits the caller asked to process, i.e. the bits in + // `bit_offset..bit_offset + bits_in_this_byte`. The request may end before the + // byte boundary, in which case the trailing bits must be preserved as well. + // + // `bits_in_this_byte + bit_offset <= 8`, so the mask always fits in a `u8`. The + // shift is done in `u16` only so the expression stays correct for a full byte. + let bits_in_this_byte = (8 - bit_offset).min(remaining_len_in_bits); + let write_mask = (((1u16 << bits_in_this_byte) - 1) << bit_offset) as u8; - let result_first_byte = - (first_byte & mask_for_first_bit_offset) | (result_first_byte & !mask_for_first_bit_offset); + let result_first_byte = (first_byte & !write_mask) | (result_first_byte & write_mask); // 6. write back the result to the buffer buffer[byte_offset] = result_first_byte; @@ -719,7 +742,9 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ // Unwrap as we already validated the slice is not empty .unwrap(); - let current = *current as u64; + // Shift the boundary byte to the position it occupies within `rem`, otherwise + // its bits would be compared against the wrong end of the mask below + let current = (*current as u64) << ((start_remainder_mut_slice.len() - 1) * 8); // Mask where the bits that are inside the remainder are 1 // and the bits outside the remainder are 0 @@ -740,7 +765,7 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ // Write back the result to the mutable slice { - let remainder_bytes = self::ceil(remainder_len, 8); + let remainder_bytes = start_remainder_mut_slice.len(); // we are counting starting from the least significant bit, so to_le_bytes should be correct let rem = &rem.to_le_bytes()[0..remainder_bytes]; @@ -1105,6 +1130,8 @@ mod tests { .map(|(l, r)| expected_op(*l, *r)) .collect(); + let before = left_buffer.as_slice().to_vec(); + apply_bitwise_binary_op( left_buffer.as_slice_mut(), left_offset_in_bits, @@ -1122,6 +1149,41 @@ mod tests { "Failed with left_offset={}, right_offset={}, len={}", left_offset_in_bits, right_offset_in_bits, len_in_bits ); + + assert_bits_outside_range_preserved( + &before, + left_buffer.as_slice(), + left_offset_in_bits, + len_in_bits, + &format!( + "left_offset={}, right_offset={}, len={}", + left_offset_in_bits, right_offset_in_bits, len_in_bits + ), + ); + } + + /// Asserts that every bit outside `offset_in_bits..offset_in_bits + len_in_bits` + /// is identical in `before` and `after`. + fn assert_bits_outside_range_preserved( + before: &[u8], + after: &[u8], + offset_in_bits: usize, + len_in_bits: usize, + context: &str, + ) { + assert_eq!(before.len(), after.len()); + for i in 0..before.len() * 8 { + if i >= offset_in_bits && i < offset_in_bits + len_in_bits { + continue; + } + assert_eq!( + get_bit(before, i), + get_bit(after, i), + "bit {} outside the requested range was modified ({})", + i, + context + ); + } } /// Verifies that a unary operation applied to a buffer using u64 chunks @@ -1146,6 +1208,8 @@ mod tests { .map(|b| expected_op(*b)) .collect(); + let before = buffer.as_slice().to_vec(); + apply_bitwise_unary_op(buffer.as_slice_mut(), offset_in_bits, len_in_bits, op); let result: Vec = @@ -1156,6 +1220,14 @@ mod tests { "Failed with offset={}, len={}", offset_in_bits, len_in_bits ); + + assert_bits_outside_range_preserved( + &before, + buffer.as_slice(), + offset_in_bits, + len_in_bits, + &format!("offset={}, len={}", offset_in_bits, len_in_bits), + ); } // Helper to create test data of specific length @@ -1424,6 +1496,66 @@ mod tests { ); } + /// Ranges that start and end inside the same non-byte-aligned byte must not + /// touch the trailing bits of that byte. + #[test] + fn test_ops_ending_inside_the_first_partial_byte() { + let (left, right) = create_test_data(32); + for offset in 1..8 { + // Inclusive so the range ending exactly on the byte boundary is covered too + for len in 1..=(8 - offset) { + test_all_binary_ops(&left, &right, offset, offset, len); + test_all_binary_ops(&left, &right, offset, (offset + 3) % 8, len); + test_mutable_buffer_unary_op_helper(&left, offset, len, |a| !a, |a| !a); + } + } + } + + #[test] + fn test_and_within_first_partial_byte_preserves_trailing_bits() { + let mut left = vec![0b11111111u8, 0b11111111u8]; + let right = vec![0b00000000u8, 0b00000000u8]; + // AND a single bit at bit offset 1: only bit 1 may be cleared + apply_bitwise_binary_op(&mut left, 1, &right, 0, 1, |a, b| a & b); + assert_eq!(left, vec![0b11111101u8, 0b11111111u8]); + } + + #[test] + fn test_not_within_first_partial_byte_preserves_trailing_bits() { + let mut buffer = vec![0b00000000u8]; + // NOT two bits at bit offset 3: only bits 3 and 4 may be flipped + apply_bitwise_unary_op(&mut buffer, 3, 2, |a| !a); + assert_eq!(buffer, vec![0b00011000u8]); + } + + /// When the remainder spans more than one byte, the byte holding the end of the + /// range is the *last* byte of the remainder, not the first. Its bits above the + /// remainder must survive. + #[test] + fn test_or_with_multi_byte_remainder_preserves_boundary_bits() { + let mut left = vec![0b00000000u8, 0b00000000u8, 0b11110000u8]; + let right = vec![0b11111111u8, 0b11111111u8, 0b11111111u8]; + // OR over 20 bits: bits 20..24 of `left` are outside the range and must stay set + apply_bitwise_binary_op(&mut left, 0, &right, 0, 20, |a, b| a | b); + assert_eq!( + left, + vec![0b11111111u8, 0b11111111u8, 0b11111111u8], + "the boundary byte lost its out-of-range bits" + ); + } + + #[test] + fn test_not_with_multi_byte_remainder_preserves_boundary_bits() { + let mut buffer = vec![0b00000000u8, 0b00000000u8, 0b11111111u8]; + // NOT over 20 bits: only bits 16..20 of the last byte may be flipped + apply_bitwise_unary_op(&mut buffer, 0, 20, |a| !a); + assert_eq!( + buffer, + vec![0b11111111u8, 0b11111111u8, 0b11110000u8], + "the boundary byte lost its out-of-range bits" + ); + } + #[test] fn test_bitwise_binary_op_offset_out_of_bounds() { let input = vec![0b10101010u8, 0b01010101u8];