Skip to content

Commit

Permalink
fix: enforce RawBytesOffsets to be non-empty and add .last()
Browse files Browse the repository at this point in the history
  • Loading branch information
LDeakin committed Jan 22, 2025
1 parent e5df5d4 commit 6582e14
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 17 deletions.
24 changes: 10 additions & 14 deletions zarrs/src/array/array_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,18 +76,13 @@ impl<'a> ArrayBytes<'a> {
offsets: RawBytesOffsets<'a>,
) -> Result<Self, RawBytesOffsetsOutOfBoundsError> {
let bytes = bytes.into();
match offsets.last() {
Some(&last) => {
if last <= bytes.len() {
Ok(Self::Variable(bytes, offsets))
} else {
Err(RawBytesOffsetsOutOfBoundsError {
offset: last,
len: bytes.len(),
})
}
}
None => Err(RawBytesOffsetsOutOfBoundsError { offset: 0, len: 0 }),
if offsets.last() <= bytes.len() {
Ok(Self::Variable(bytes, offsets))
} else {
Err(RawBytesOffsetsOutOfBoundsError {
offset: offsets.last(),
len: bytes.len(),
})
}
}

Expand All @@ -100,7 +95,7 @@ impl<'a> ArrayBytes<'a> {
offsets: RawBytesOffsets<'a>,
) -> Self {
let bytes = bytes.into();
debug_assert!(offsets.last().is_some_and(|&last| last <= bytes.len()));
debug_assert!(offsets.last() <= bytes.len());
Self::Variable(bytes, offsets)
}

Expand Down Expand Up @@ -501,7 +496,7 @@ pub(crate) fn merge_chunks_vlen<'a>(

// Write bytes
// TODO: Go parallel
let mut bytes = vec![0; *offsets.last().unwrap()];
let mut bytes = vec![0; offsets.last()];
for (chunk_bytes, chunk_subset) in chunk_bytes_and_subsets {
let (chunk_bytes, chunk_offsets) = chunk_bytes.into_variable()?;
let indices = chunk_subset.linearised_indices(array_shape).unwrap();
Expand Down Expand Up @@ -664,6 +659,7 @@ mod tests {
#[test]
fn array_bytes_vlen() {
let data = [0u8, 1, 2, 3, 4];
assert!(ArrayBytes::new_vlen(&data, vec![0].try_into().unwrap()).is_ok());
assert!(ArrayBytes::new_vlen(&data, vec![0, 5].try_into().unwrap()).is_ok());
assert!(ArrayBytes::new_vlen(&data, vec![0, 5, 5].try_into().unwrap()).is_ok());
assert!(ArrayBytes::new_vlen(&data, vec![0, 5, 6].try_into().unwrap()).is_err());
Expand Down
28 changes: 25 additions & 3 deletions zarrs/src/array/array_bytes/raw_bytes_offsets.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,14 @@ impl Deref for RawBytesOffsets<'_> {
///
/// This error occurs when the offsets are not monotonically increasing.
#[derive(Debug, Error, Display)]
pub struct RawBytesOffsetsCreateError;
pub enum RawBytesOffsetsCreateError {
/// The length must be greater than zero.
#[display("length must be greater than zero")]
ZeroLength,
/// The offsets are not monotonically increasing.
#[display("offsets are not monotonically increasing")]
NotMonotonicallyIncreasing,
}

impl<'a> RawBytesOffsets<'a> {
/// Creates a new `RawBytesOffsets`.
Expand All @@ -30,10 +37,12 @@ impl<'a> RawBytesOffsets<'a> {
/// Returns an error if the offsets are not monotonically increasing.
pub fn new(offsets: impl Into<Cow<'a, [usize]>>) -> Result<Self, RawBytesOffsetsCreateError> {
let offsets = offsets.into();
if offsets.windows(2).all(|w| w[1] >= w[0]) {
if offsets.is_empty() {
Err(RawBytesOffsetsCreateError::ZeroLength)
} else if offsets.windows(2).all(|w| w[1] >= w[0]) {
Ok(Self(offsets))
} else {
Err(RawBytesOffsetsCreateError)
Err(RawBytesOffsetsCreateError::NotMonotonicallyIncreasing)
}
}

Expand All @@ -44,6 +53,7 @@ impl<'a> RawBytesOffsets<'a> {
#[must_use]
pub unsafe fn new_unchecked(offsets: impl Into<Cow<'a, [usize]>>) -> Self {
let offsets = offsets.into();
debug_assert!(!offsets.is_empty());
debug_assert!(offsets.windows(2).all(|w| w[1] >= w[0]));
Self(offsets)
}
Expand All @@ -53,6 +63,15 @@ impl<'a> RawBytesOffsets<'a> {
pub fn into_owned(self) -> RawBytesOffsets<'static> {
RawBytesOffsets(self.0.into_owned().into())
}

/// Returns the last offset.
#[must_use]
pub fn last(&self) -> usize {
unsafe {
// SAFETY: The offsets cannot be empty.
*self.0.last().unwrap_unchecked()
}
}
}

impl<'a> TryFrom<Cow<'a, [usize]>> for RawBytesOffsets<'a> {
Expand Down Expand Up @@ -95,6 +114,9 @@ mod tests {
fn raw_bytes_offsets() {
let offsets = RawBytesOffsets::new(vec![0, 1, 2, 3]).unwrap();
assert_eq!(&*offsets, &[0, 1, 2, 3]);
assert!(RawBytesOffsets::new(vec![]).is_err());
assert!(RawBytesOffsets::new(vec![0]).is_ok());
assert!(RawBytesOffsets::new(vec![10]).is_ok()); // nonsense, but not invalid
assert!(RawBytesOffsets::new(vec![0, 1, 1]).is_ok());
assert!(RawBytesOffsets::new(vec![0, 1, 0]).is_err());
assert!(RawBytesOffsets::try_from(vec![0, 1, 2]).is_ok());
Expand Down

0 comments on commit 6582e14

Please sign in to comment.