diff --git a/arrow-buffer/src/buffer/mutable_ops.rs b/arrow-buffer/src/buffer/mutable_ops.rs index 0784672bedaf..e3e567b35496 100644 --- a/arrow-buffer/src/buffer/mutable_ops.rs +++ b/arrow-buffer/src/buffer/mutable_ops.rs @@ -52,41 +52,6 @@ impl BufferSupportedRhs for BooleanBufferBuilder { } } -/// Trait that will be operated on as the left-hand side (LHS) buffer in mutable operations. -/// -/// This consumer of the trait must satisfies the following guarantees: -/// 1. It will not change the length of the buffer. -/// -/// # Implementation notes -/// -/// ## Why is this trait `pub(crate)`? -/// Because we don't wanna expose the inner mutable buffer to the public. -/// as this is the choice of the implementor of the trait and sometimes it is not desirable -/// (e.g. `BooleanBufferBuilder`). -/// -/// ## Why this trait is needed, can't we just use `MutableBuffer` directly? -/// Sometimes we don't want to expose the inner `MutableBuffer` -/// so it can't be misused. -/// -/// For example, [`BooleanBufferBuilder`] does not expose the inner `MutableBuffer` -/// as exposing it will allow the user to change the length of the buffer that will make the -/// `BooleanBufferBuilder` invalid. -/// -pub(crate) trait MutableOpsBufferSupportedLhs { - /// Get a mutable reference to the inner `MutableBuffer`. - /// - /// This is used to perform in-place operations on the buffer. - /// - /// the caller must ensure that the length of the buffer is not changed. - fn inner_mutable_buffer(&mut self) -> &mut MutableBuffer; -} - -impl MutableOpsBufferSupportedLhs for MutableBuffer { - fn inner_mutable_buffer(&mut self) -> &mut MutableBuffer { - self - } -} - /// Apply a binary bitwise operation to two bit-packed buffers. /// /// This is the main entry point for binary operations. It handles both byte-aligned @@ -106,7 +71,7 @@ impl MutableOpsBufferSupportedLhs for MutableBuffer { reason = "MutableOpsBufferSupportedLhs and BufferSupportedRhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_bitwise_bin_op_helper( - left: &mut impl MutableOpsBufferSupportedLhs, + left: &mut MutableBuffer, left_offset_in_bits: usize, right: &impl BufferSupportedRhs, right_offset_in_bits: usize, @@ -119,7 +84,7 @@ pub fn mutable_bitwise_bin_op_helper( return; } - let mutable_buffer = left.inner_mutable_buffer(); + let mutable_buffer = left; let mutable_buffer_len = mutable_buffer.len(); let mutable_buffer_cap = mutable_buffer.capacity(); @@ -734,7 +699,7 @@ fn handle_mutable_buffer_remainder_unary( reason = "MutableOpsBufferSupportedLhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_bitwise_unary_op_helper( - buffer: &mut impl MutableOpsBufferSupportedLhs, + buffer: &mut MutableBuffer, offset_in_bits: usize, len_in_bits: usize, mut op: F, @@ -745,7 +710,7 @@ pub fn mutable_bitwise_unary_op_helper( return; } - let mutable_buffer = buffer.inner_mutable_buffer(); + let mutable_buffer = buffer; let mutable_buffer_len = mutable_buffer.len(); let mutable_buffer_cap = mutable_buffer.capacity(); @@ -829,7 +794,7 @@ pub fn mutable_bitwise_unary_op_helper( reason = "MutableOpsBufferSupportedLhs and BufferSupportedRhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_buffer_bin_and( - left: &mut impl MutableOpsBufferSupportedLhs, + left: &mut MutableBuffer, left_offset_in_bits: usize, right: &impl BufferSupportedRhs, right_offset_in_bits: usize, @@ -863,7 +828,7 @@ pub fn mutable_buffer_bin_and( reason = "MutableOpsBufferSupportedLhs and BufferSupportedRhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_buffer_bin_or( - left: &mut impl MutableOpsBufferSupportedLhs, + left: &mut MutableBuffer, left_offset_in_bits: usize, right: &impl BufferSupportedRhs, right_offset_in_bits: usize, @@ -897,7 +862,7 @@ pub fn mutable_buffer_bin_or( reason = "MutableOpsBufferSupportedLhs and BufferSupportedRhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_buffer_bin_xor( - left: &mut impl MutableOpsBufferSupportedLhs, + left: &mut MutableBuffer, left_offset_in_bits: usize, right: &impl BufferSupportedRhs, right_offset_in_bits: usize, @@ -929,7 +894,7 @@ pub fn mutable_buffer_bin_xor( reason = "MutableOpsBufferSupportedLhs exposes the inner internals which is the implementor choice and we dont want to leak internals" )] pub fn mutable_buffer_unary_not( - buffer: &mut impl MutableOpsBufferSupportedLhs, + buffer: &mut MutableBuffer, offset_in_bits: usize, len_in_bits: usize, ) { @@ -964,14 +929,16 @@ mod tests { .map(|(l, r)| expected_op(*l, *r)) .collect(); - super::mutable_bitwise_bin_op_helper( - &mut left_buffer, - left_offset_in_bits, - right_buffer.inner(), - right_offset_in_bits, - len_in_bits, - op, - ); + unsafe { + super::mutable_bitwise_bin_op_helper( + left_buffer.mutable_buffer(), + left_offset_in_bits, + right_buffer.inner(), + right_offset_in_bits, + len_in_bits, + op, + ); + } let result: Vec = BitIterator::new(left_buffer.as_slice(), left_offset_in_bits, len_in_bits).collect(); @@ -1002,7 +969,14 @@ mod tests { .map(|b| expected_op(*b)) .collect(); - super::mutable_bitwise_unary_op_helper(&mut buffer, offset_in_bits, len_in_bits, op); + unsafe { + super::mutable_bitwise_unary_op_helper( + buffer.mutable_buffer(), + offset_in_bits, + len_in_bits, + op, + ); + } let result: Vec = BitIterator::new(buffer.as_slice(), offset_in_bits, len_in_bits).collect(); diff --git a/arrow-buffer/src/builder/boolean.rs b/arrow-buffer/src/builder/boolean.rs index ff8a1f66a303..77dc17b96073 100644 --- a/arrow-buffer/src/builder/boolean.rs +++ b/arrow-buffer/src/builder/boolean.rs @@ -16,9 +16,8 @@ // under the License. use crate::{ - BooleanBuffer, Buffer, MutableBuffer, MutableOpsBufferSupportedLhs, bit_mask, bit_util, - mutable_buffer_bin_and, mutable_buffer_bin_or, mutable_buffer_bin_xor, - mutable_buffer_unary_not, + BooleanBuffer, Buffer, MutableBuffer, bit_mask, bit_util, mutable_buffer_bin_and, + mutable_buffer_bin_or, mutable_buffer_bin_xor, mutable_buffer_unary_not, }; use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Range}; @@ -260,11 +259,14 @@ impl BooleanBufferBuilder { pub fn finish_cloned(&self) -> BooleanBuffer { BooleanBuffer::new(Buffer::from_slice_ref(self.as_slice()), 0, self.len) } -} -/// This trait is not public API so it does not leak the inner mutable buffer -impl MutableOpsBufferSupportedLhs for BooleanBufferBuilder { - fn inner_mutable_buffer(&mut self) -> &mut MutableBuffer { + /// Return a mutable reference to the internal buffer + /// + /// # Safety + /// The caller must ensure that any modifications to the buffer maintain the invariant + /// `self.len < buffer.len() / 8` (that is that the buffer has enough capacity to hold `self.len` bits). + #[inline] + pub unsafe fn mutable_buffer(&mut self) -> &mut MutableBuffer { &mut self.buffer } }