Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Improved API of getting mutable from Buffer #1399

Merged
merged 2 commits into from
Feb 13, 2023
Merged

Improved API of getting mutable from Buffer #1399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e95700c
Removed potential source of UB
jorgecarleitao Feb 12, 2023
6e69ce6
Improved API
jorgecarleitao Feb 12, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 21 additions & 29 deletions src/array/binary/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,8 @@ impl<O: Offset> BinaryArray<O> {
}

/// Try to convert this `BinaryArray` to a `MutableBinaryArray`
pub fn into_mut(mut self) -> Either<Self, MutableBinaryArray<O>> {
#[must_use]
pub fn into_mut(self) -> Either<Self, MutableBinaryArray<O>> {
use Either::*;
if let Some(bitmap) = self.validity {
match bitmap.into_mut() {
Expand All @@ -251,29 +252,26 @@ impl<O: Offset> BinaryArray<O> {
self.values,
Some(bitmap),
)),
Right(mutable_bitmap) => match (
self.values.get_mut().map(std::mem::take),
self.offsets.get_mut(),
) {
(None, None) => Left(BinaryArray::new(
Right(mutable_bitmap) => match (self.values.into_mut(), self.offsets.into_mut()) {
(Left(values), Left(offsets)) => Left(BinaryArray::new(
self.data_type,
self.offsets,
self.values,
offsets,
values,
Some(mutable_bitmap.into()),
)),
(None, Some(offsets)) => Left(BinaryArray::new(
(Left(values), Right(offsets)) => Left(BinaryArray::new(
self.data_type,
offsets.into(),
self.values,
values,
Some(mutable_bitmap.into()),
)),
(Some(mutable_values), None) => Left(BinaryArray::new(
(Right(values), Left(offsets)) => Left(BinaryArray::new(
self.data_type,
self.offsets,
mutable_values.into(),
offsets,
values.into(),
Some(mutable_bitmap.into()),
)),
(Some(values), Some(offsets)) => Right(
(Right(values), Right(offsets)) => Right(
MutableBinaryArray::try_new(
self.data_type,
offsets,
Expand All @@ -285,29 +283,23 @@ impl<O: Offset> BinaryArray<O> {
},
}
} else {
match (
self.values.get_mut().map(std::mem::take),
self.offsets.get_mut(),
) {
(None, None) => Left(BinaryArray::new(
self.data_type,
self.offsets,
self.values,
None,
)),
(None, Some(offsets)) => Left(BinaryArray::new(
match (self.values.into_mut(), self.offsets.into_mut()) {
(Left(values), Left(offsets)) => {
Left(BinaryArray::new(self.data_type, offsets, values, None))
}
(Left(values), Right(offsets)) => Left(BinaryArray::new(
self.data_type,
offsets.into(),
self.values,
values,
None,
)),
(Some(values), None) => Left(BinaryArray::new(
(Right(values), Left(offsets)) => Left(BinaryArray::new(
self.data_type,
self.offsets,
offsets,
values.into(),
None,
)),
(Some(values), Some(offsets)) => Right(
(Right(values), Right(offsets)) => Right(
MutableBinaryArray::try_new(self.data_type, offsets, values, None).unwrap(),
),
}
Expand Down
18 changes: 9 additions & 9 deletions src/array/primitive/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ impl<T: NativeType> PrimitiveArray<T> {

/// Returns an option of a mutable reference to the values of this [`PrimitiveArray`].
pub fn get_mut_values(&mut self) -> Option<&mut [T]> {
self.values.get_mut().map(|x| x.as_mut())
self.values.get_mut_slice()
}

/// Returns its internal representation
Expand All @@ -316,7 +316,7 @@ impl<T: NativeType> PrimitiveArray<T> {
///
/// This function is primarily used to re-use memory regions.
#[must_use]
pub fn into_mut(mut self) -> Either<Self, MutablePrimitiveArray<T>> {
pub fn into_mut(self) -> Either<Self, MutablePrimitiveArray<T>> {
use Either::*;

if let Some(bitmap) = self.validity {
Expand All @@ -326,28 +326,28 @@ impl<T: NativeType> PrimitiveArray<T> {
self.values,
Some(bitmap),
)),
Right(mutable_bitmap) => match self.values.get_mut().map(std::mem::take) {
Some(values) => Right(
Right(mutable_bitmap) => match self.values.into_mut() {
Right(values) => Right(
MutablePrimitiveArray::try_new(
self.data_type,
values,
Some(mutable_bitmap),
)
.unwrap(),
),
None => Left(PrimitiveArray::new(
Left(values) => Left(PrimitiveArray::new(
self.data_type,
self.values,
values,
Some(mutable_bitmap.into()),
)),
},
}
} else {
match self.values.get_mut().map(std::mem::take) {
Some(values) => {
match self.values.into_mut() {
Right(values) => {
Right(MutablePrimitiveArray::try_new(self.data_type, values, None).unwrap())
}
None => Left(PrimitiveArray::new(self.data_type, self.values, None)),
Left(values) => Left(PrimitiveArray::new(self.data_type, values, None)),
}
}
}
Expand Down
47 changes: 21 additions & 26 deletions src/array/utf8/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,8 @@ impl<O: Offset> Utf8Array<O> {
}

/// Try to convert this `Utf8Array` to a `MutableUtf8Array`
pub fn into_mut(mut self) -> Either<Self, MutableUtf8Array<O>> {
#[must_use]
pub fn into_mut(self) -> Either<Self, MutableUtf8Array<O>> {
use Either::*;
if let Some(bitmap) = self.validity {
match bitmap.into_mut() {
Expand All @@ -274,44 +275,41 @@ impl<O: Offset> Utf8Array<O> {
Some(bitmap),
)
}),
Right(mutable_bitmap) => match (
self.values.get_mut().map(std::mem::take),
self.offsets.get_mut(),
) {
(None, None) => {
Right(mutable_bitmap) => match (self.values.into_mut(), self.offsets.into_mut()) {
(Left(values), Left(offsets)) => {
// Safety: invariants are preserved
Left(unsafe {
Utf8Array::new_unchecked(
self.data_type,
self.offsets,
self.values,
offsets,
values,
Some(mutable_bitmap.into()),
)
})
}
(None, Some(offsets)) => {
(Left(values), Right(offsets)) => {
// Safety: invariants are preserved
Left(unsafe {
Utf8Array::new_unchecked(
self.data_type,
offsets.into(),
self.values,
values,
Some(mutable_bitmap.into()),
)
})
}
(Some(mutable_values), None) => {
(Right(values), Left(offsets)) => {
// Safety: invariants are preserved
Left(unsafe {
Utf8Array::new_unchecked(
self.data_type,
self.offsets,
mutable_values.into(),
offsets,
values.into(),
Some(mutable_bitmap.into()),
)
})
}
(Some(values), Some(offsets)) => Right(unsafe {
(Right(values), Right(offsets)) => Right(unsafe {
MutableUtf8Array::new_unchecked(
self.data_type,
offsets,
Expand All @@ -322,20 +320,17 @@ impl<O: Offset> Utf8Array<O> {
},
}
} else {
match (
self.values.get_mut().map(std::mem::take),
self.offsets.get_mut(),
) {
(None, None) => Left(unsafe {
Utf8Array::new_unchecked(self.data_type, self.offsets, self.values, None)
}),
(None, Some(offsets)) => Left(unsafe {
Utf8Array::new_unchecked(self.data_type, offsets.into(), self.values, None)
match (self.values.into_mut(), self.offsets.into_mut()) {
(Left(values), Left(offsets)) => {
Left(unsafe { Utf8Array::new_unchecked(self.data_type, offsets, values, None) })
}
(Left(values), Right(offsets)) => Left(unsafe {
Utf8Array::new_unchecked(self.data_type, offsets.into(), values, None)
}),
(Some(values), None) => Left(unsafe {
Utf8Array::new_unchecked(self.data_type, self.offsets, values.into(), None)
(Right(values), Left(offsets)) => Left(unsafe {
Utf8Array::new_unchecked(self.data_type, offsets, values.into(), None)
}),
(Some(values), Some(offsets)) => Right(unsafe {
(Right(values), Right(offsets)) => Right(unsafe {
MutableUtf8Array::new_unchecked(self.data_type, offsets, values, None)
}),
}
Expand Down
51 changes: 41 additions & 10 deletions src/buffer/immutable.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
use std::{iter::FromIterator, ops::Deref, sync::Arc, usize};

use either::Either;

use super::Bytes;
use super::IntoIter;

Expand All @@ -21,15 +23,15 @@ use super::IntoIter;
/// assert_eq!(buffer.as_ref(), [1, 2, 3].as_ref());
///
/// // it supports copy-on-write semantics (i.e. back to a `Vec`)
/// let vec: &mut [u32] = buffer.get_mut().unwrap();
/// assert_eq!(vec, &mut [1, 2, 3]);
/// let vec: Vec<u32> = buffer.into_mut().right().unwrap();
/// assert_eq!(vec, vec![1, 2, 3]);
///
/// // cloning and slicing is `O(1)` (data is shared)
/// let mut buffer: Buffer<u32> = vec![1, 2, 3].into();
/// let slice = buffer.clone().slice(1, 1);
/// assert_eq!(slice.as_ref(), [2].as_ref());
/// // but cloning forbids getting mut since `slice` and `buffer` now share data
/// assert_eq!(buffer.get_mut(), None);
/// assert_eq!(buffer.get_mut_slice(), None);
/// ```
#[derive(Clone)]
pub struct Buffer<T> {
Expand Down Expand Up @@ -155,18 +157,47 @@ impl<T> Buffer<T> {

/// Returns a mutable reference to its underlying [`Vec`], if possible.
///
/// This operation returns [`Some`] iff this [`Buffer`]:
/// * has not been sliced with an offset
/// This operation returns [`Either::Right`] iff this [`Buffer`]:
/// * has not been cloned (i.e. [`Arc`]`::get_mut` yields [`Some`])
/// * has not been imported from the c data interface (FFI)
pub fn get_mut(&mut self) -> Option<&mut Vec<T>> {
if self.offset != 0 {
None
} else {
Arc::get_mut(&mut self.data).and_then(|b| b.get_vec())
#[inline]
pub fn into_mut(mut self) -> Either<Self, Vec<T>> {
match Arc::get_mut(&mut self.data)
.and_then(|b| b.get_vec())
.map(std::mem::take)
{
Some(inner) => Either::Right(inner),
None => Either::Left(self),
}
}

/// Returns a mutable reference to its underlying `Vec`, if possible.
/// Note that only `[self.offset(), self.offset() + self.len()[` in this vector is visible
/// by this buffer.
///
/// This operation returns [`Some`] iff this [`Buffer`]:
/// * has not been cloned (i.e. [`Arc`]`::get_mut` yields [`Some`])
/// * has not been imported from the c data interface (FFI)
/// # Safety
/// The caller must ensure that the vector in the mutable reference keeps a length of at least `self.offset() + self.len() - 1`.
#[inline]
pub unsafe fn get_mut(&mut self) -> Option<&mut Vec<T>> {
Arc::get_mut(&mut self.data).and_then(|b| b.get_vec())
}

/// Returns a mutable reference to its slice, if possible.
///
/// This operation returns [`Some`] iff this [`Buffer`]:
/// * has not been cloned (i.e. [`Arc`]`::get_mut` yields [`Some`])
/// * has not been imported from the c data interface (FFI)
#[inline]
pub fn get_mut_slice(&mut self) -> Option<&mut [T]> {
Arc::get_mut(&mut self.data)
.and_then(|b| b.get_vec())
// Safety: the invariant of this struct
.map(|x| unsafe { x.get_unchecked_mut(self.offset..self.offset + self.length) })
}

/// Get the strong count of underlying `Arc` data buffer.
pub fn shared_count_strong(&self) -> usize {
Arc::strong_count(&self.data)
Expand Down
12 changes: 4 additions & 8 deletions src/offset.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,16 +363,12 @@ impl<O: Offset> OffsetsBuffer<O> {

/// Copy-on-write API to convert [`OffsetsBuffer`] into [`Offsets`].
#[inline]
pub fn get_mut(&mut self) -> Option<Offsets<O>> {
pub fn into_mut(self) -> either::Either<Self, Offsets<O>> {
self.0
.get_mut()
.map(|x| {
let mut new = vec![O::zero()];
std::mem::swap(x, &mut new);
new
})
.into_mut()
// Safety: Offsets and OffsetsBuffer share invariants
.map(|offsets| unsafe { Offsets::new_unchecked(offsets) })
.map_right(|offsets| unsafe { Offsets::new_unchecked(offsets) })
.map_left(Self)
}

/// Returns a reference to its internal [`Buffer`].
Expand Down