add into_mut methods on mask and buffer

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-buffer/src/bit/buf.rs

@@ -14,6 +14,9 @@ use crate::{Alignment, BitBufferMut, Buffer, BufferMut, ByteBuffer, buffer};
 #[derive(Debug, Clone, Eq)]
 pub struct BitBuffer {
     buffer: ByteBuffer,
+    /// Represents the offset of the bit buffer into the first byte.
+    ///
+    /// This is always less than 8 (for when the bit buffer is not aligned to a byte).
     offset: usize,
     len: usize,
 }

vortex-buffer/src/bit/buf_mut.rs

@@ -66,6 +66,15 @@ impl BitBufferMut {
         }
     }
 
+    /// Creates a `BitBufferMut` from a [`BitBuffer`] by copying all of the data over.
+    pub fn copy_from(bit_buffer: &BitBuffer) -> Self {
+        Self {
+            buffer: ByteBufferMut::copy_from(bit_buffer.inner()),
+            offset: bit_buffer.offset(),
+            len: bit_buffer.len(),
+        }
+    }
+
     /// Create a new empty mutable bit buffer with requested capacity (in bits).
     pub fn with_capacity(capacity: usize) -> Self {
         Self {

vortex-buffer/src/buffer.rs

@@ -9,7 +9,7 @@ use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
 use std::ops::{Deref, RangeBounds};
 
-use bytes::{Buf, Bytes};
+use bytes::{Buf, Bytes, BytesMut};
 use vortex_error::{VortexExpect, vortex_panic};
 
 use crate::debug::TruncatedDebug;
@@ -399,12 +399,6 @@ impl<T> Buffer<T> {
         }
     }
 
-    /// Convert self into `BufferMut<T>`, copying if there are multiple strong references.
-    pub fn into_mut(self) -> BufferMut<T> {
-        self.try_into_mut()
-            .unwrap_or_else(|buffer| BufferMut::<T>::copy_from(&buffer))
-    }
-
     /// Try to convert self into `BufferMut<T>` if there is only a single strong reference.
     pub fn try_into_mut(self) -> Result<BufferMut<T>, Self> {
         self.bytes
@@ -423,6 +417,18 @@ impl<T> Buffer<T> {
             })
     }
 
+    /// Convert self into `BufferMut<T>`, cloning the data if there are multiple strong references.
+    pub fn into_mut(self) -> BufferMut<T> {
+        let new_bytes = self.bytes.try_into_mut().unwrap_or_else(BytesMut::from);
+
+        BufferMut {
+            bytes: new_bytes,
+            length: self.length,
+            alignment: self.alignment,
+            _marker: Default::default(),
+        }
+    }
+
     /// Returns whether a `Buffer<T>` is aligned to the given alignment.
     pub fn is_aligned(&self, alignment: Alignment) -> bool {
         self.bytes.as_ptr().align_offset(*alignment) == 0

vortex-mask/src/mask_mut.rs

@@ -300,7 +300,8 @@ impl MaskMut {
 }
 
 impl Mask {
-    /// Attempts to convert an immutable mask into a mutable one.
+    /// Attempts to convert an immutable mask into a mutable one, returning an error of `Self` if
+    /// the underlying [`BitBuffer`] data if there are any other references.
     pub fn try_into_mut(self) -> Result<MaskMut, Self> {
         match self {
             Mask::AllTrue(len) => Ok(MaskMut::new_true(len)),
@@ -316,6 +317,29 @@ impl Mask {
             }
         }
     }
+
+    /// Convert an immutable mask into a mutable one, cloning the underlying [`BitBuffer`] data if
+    /// there are any other references.
+    pub fn into_mut(self) -> MaskMut {
+        match self {
+            Mask::AllTrue(len) => MaskMut::new_true(len),
+            Mask::AllFalse(len) => MaskMut::new_false(len),
+            Mask::Values(values) => {
+                let bit_buffer_mut = match Arc::try_unwrap(values) {
+                    Ok(mask_values) => {
+                        let bit_buffer = mask_values.into_buffer();
+                        bit_buffer.into_mut()
+                    }
+                    Err(arc_mask_values) => {
+                        let bit_buffer = arc_mask_values.bit_buffer();
+                        BitBufferMut::copy_from(bit_buffer)
+                    }
+                };
+
+                MaskMut(Inner::Builder(bit_buffer_mut))
+            }
+        }
+    }
 }
 
 #[cfg(test)]