zarrs: bump unsafe_cell_slice to 0.2.0

LDeakin · Oct 3, 2024 · d267659 · d267659
1 parent 94bbe1e
commit d267659
Show file tree

Hide file tree

Showing 11 changed files with 81 additions and 66 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+- Bump `unsafe_cell_slice` to 0.2.0
+- **Breaking**: Change `output` parameter of `decode_into` codec trait methods to `&UnsafeCellSlice<u8>`
+
 ## [0.17.0] - 2024-10-02
 
 ### Highlights / Major Changes

diff --git a/zarrs/Cargo.toml b/zarrs/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "zarrs"
-version = "0.17.0"
+version = "0.17.1"
 authors = ["Lachlan Deakin <ljdgit@gmail.com>"]
 edition = "2021"
 rust-version = "1.76"
@@ -62,7 +62,7 @@ serde = { version = "1.0.184", features = ["derive"] }
 serde_json = { version = "1.0.71", features = ["float_roundtrip", "preserve_order"] }
 thiserror = "1.0.61"
 thread_local = "1.1.8"
-unsafe_cell_slice = "0.1.0"
+unsafe_cell_slice = "0.2.0"
 zarrs_filesystem = { workspace = true, optional = true }
 zarrs_metadata = { workspace = true }
 zarrs_storage = { workspace = true }

diff --git a/zarrs/src/array/array_async_readable.rs b/zarrs/src/array/array_async_readable.rs
@@ -348,7 +348,7 @@ impl<TStorage: ?Sized + AsyncReadableStorageTraits + 'static> Array<TStorage> {
     async unsafe fn async_retrieve_chunk_into(
         &self,
         chunk_indices: &[u64],
-        output: &mut [u8],
+        output: &UnsafeCellSlice<'_, u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,
@@ -662,13 +662,12 @@ impl<TStorage: ?Sized + AsyncReadableStorageTraits + 'static> Array<TStorage> {
                                     let chunk_subset = self.chunk_subset(&chunk_indices)?;
                                     let chunk_subset_overlap =
                                         chunk_subset.overlap(array_subset)?;
-                                    let output = unsafe { output.as_mut_slice() };
                                     unsafe {
                                         self.async_retrieve_chunk_subset_into(
                                             &chunk_indices,
                                             &chunk_subset_overlap
                                                 .relative_to(chunk_subset.start())?,
-                                            output,
+                                            &output,
                                             array_subset.shape(),
                                             &chunk_subset_overlap
                                                 .relative_to(array_subset.start())?,
@@ -790,7 +789,7 @@ impl<TStorage: ?Sized + AsyncReadableStorageTraits + 'static> Array<TStorage> {
         &self,
         chunk_indices: &[u64],
         chunk_subset: &ArraySubset,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<'_, u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,

diff --git a/zarrs/src/array/array_bytes.rs b/zarrs/src/array/array_bytes.rs
@@ -2,6 +2,7 @@ use std::borrow::Cow;
 
 use itertools::Itertools;
 use thiserror::Error;
+use unsafe_cell_slice::UnsafeCellSlice;
 
 use crate::{
     array_subset::{ArraySubset, IncompatibleArraySubsetAndShapeError},
@@ -261,7 +262,7 @@ fn validate_bytes(
 /// This function is used internally by various array/codec methods to write the bytes of a chunk subset into an output with an associated array subset.
 /// This approach only works for fixed length data types.
 pub fn update_bytes_flen(
-    output_bytes: &mut [u8],
+    output_bytes: &UnsafeCellSlice<u8>,
     output_shape: &[u64],
     subset_bytes: &RawBytes,
     subset: &ArraySubset,
@@ -285,8 +286,11 @@ pub fn update_bytes_flen(
         let output_offset = usize::try_from(array_subset_element_index).unwrap() * data_type_size;
         debug_assert!((output_offset + length) <= output_bytes.len());
         debug_assert!((decoded_offset + length) <= subset_bytes.len());
-        output_bytes[output_offset..output_offset + length]
-            .copy_from_slice(&subset_bytes[decoded_offset..decoded_offset + length]);
+        unsafe {
+            output_bytes
+                .index_mut(output_offset..output_offset + length)
+                .copy_from_slice(&subset_bytes[decoded_offset..decoded_offset + length]);
+        }
         decoded_offset += length;
     }
 }
@@ -382,13 +386,16 @@ pub fn update_array_bytes<'a>(
             DataTypeSize::Fixed(data_type_size),
         ) => {
             let mut chunk_bytes = chunk_bytes.into_owned();
-            update_bytes_flen(
-                &mut chunk_bytes,
-                &output_shape,
-                &chunk_subset_bytes,
-                subset,
-                data_type_size,
-            );
+            {
+                let chunk_bytes = UnsafeCellSlice::new(&mut chunk_bytes);
+                update_bytes_flen(
+                    &chunk_bytes,
+                    &output_shape,
+                    &chunk_subset_bytes,
+                    subset,
+                    data_type_size,
+                );
+            }
             ArrayBytes::new_flen(chunk_bytes)
         }
         (_, _, _) => {
@@ -583,21 +590,24 @@ mod tests {
     #[test]
     fn test_flen_update_subset() {
         let mut bytes_array = vec![0u8; 4 * 4];
-        update_bytes_flen(
-            &mut bytes_array,
-            &vec![4, 4],
-            &vec![1u8, 2].into(),
-            &ArraySubset::new_with_ranges(&[1..2, 1..3]),
-            1,
-        );
+        {
+            let bytes_array = UnsafeCellSlice::new(&mut bytes_array);
+            update_bytes_flen(
+                &bytes_array,
+                &vec![4, 4],
+                &vec![1u8, 2].into(),
+                &ArraySubset::new_with_ranges(&[1..2, 1..3]),
+                1,
+            );
 
-        update_bytes_flen(
-            &mut bytes_array,
-            &vec![4, 4],
-            &vec![3u8, 4].into(),
-            &ArraySubset::new_with_ranges(&[3..4, 0..2]),
-            1,
-        );
+            update_bytes_flen(
+                &bytes_array,
+                &vec![4, 4],
+                &vec![3u8, 4].into(),
+                &ArraySubset::new_with_ranges(&[3..4, 0..2]),
+                1,
+            );
+        }
 
         debug_assert_eq!(
             bytes_array,

diff --git a/zarrs/src/array/array_sync_readable.rs b/zarrs/src/array/array_sync_readable.rs
@@ -471,7 +471,7 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> Array<TStorage> {
     unsafe fn retrieve_chunk_into(
         &self,
         chunk_indices: &[u64],
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,
@@ -733,12 +733,11 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> Array<TStorage> {
                             let retrieve_chunk = |chunk_indices: Vec<u64>| {
                                 let chunk_subset = self.chunk_subset(&chunk_indices)?;
                                 let chunk_subset_overlap = chunk_subset.overlap(array_subset)?;
-                                let output = unsafe { output.as_mut_slice() };
                                 unsafe {
                                     self.retrieve_chunk_subset_into(
                                         &chunk_indices,
                                         &chunk_subset_overlap.relative_to(chunk_subset.start())?,
-                                        output,
+                                        &output,
                                         array_subset.shape(),
                                         &chunk_subset_overlap.relative_to(array_subset.start())?,
                                         &options,
@@ -750,7 +749,7 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> Array<TStorage> {
                                 //     &options,
                                 // )?;
                                 // update_bytes_flen(
-                                //     unsafe { output.as_mut_slice() },
+                                //     &output,
                                 //     array_subset.shape(),
                                 //     &chunk_subset_bytes.into_fixed()?,
                                 //     &chunk_subset_overlap.relative_to(array_subset.start())?,
@@ -845,7 +844,7 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> Array<TStorage> {
         &self,
         chunk_indices: &[u64],
         chunk_subset: &ArraySubset,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,

diff --git a/zarrs/src/array/array_sync_sharded_readable_ext.rs b/zarrs/src/array/array_sync_sharded_readable_ext.rs
@@ -418,7 +418,7 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> ArrayShardedReadableExt
                                     .remove(0)
                                     .into_owned();
                                 update_bytes_flen(
-                                    unsafe { output.as_mut_slice() },
+                                    &output,
                                     array_subset.shape(),
                                     &bytes.into_fixed()?,
                                     &shard_subset_overlap.relative_to(array_subset.start())?,

diff --git a/zarrs/src/array/chunk_cache/array_chunk_cache_ext_sync.rs b/zarrs/src/array/chunk_cache/array_chunk_cache_ext_sync.rs
@@ -404,7 +404,7 @@ impl<TStorage: ?Sized + ReadableStorageTraits + 'static> ArrayChunkCacheExt<TSto
                                 };
 
                                 update_bytes_flen(
-                                    unsafe { output.as_mut_slice() },
+                                    &output,
                                     array_subset.shape(),
                                     fixed,
                                     &chunk_subset_overlap.relative_to(array_subset.start())?,

diff --git a/zarrs/src/array/codec.rs b/zarrs/src/array/codec.rs
@@ -67,6 +67,7 @@ pub use byte_interval_partial_decoder::ByteIntervalPartialDecoder;
 
 #[cfg(feature = "async")]
 pub use byte_interval_partial_decoder::AsyncByteIntervalPartialDecoder;
+use unsafe_cell_slice::UnsafeCellSlice;
 
 use crate::{
     array_subset::{ArraySubset, IncompatibleArraySubsetAndShapeError},
@@ -364,7 +365,7 @@ pub trait ArrayPartialDecoderTraits: Send + Sync {
     unsafe fn partial_decode_into(
         &self,
         array_subset: &ArraySubset,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,
@@ -420,7 +421,7 @@ pub trait AsyncArrayPartialDecoderTraits: Send + Sync {
     async unsafe fn partial_decode_into(
         &self,
         array_subset: &ArraySubset,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,
@@ -633,7 +634,7 @@ pub trait ArrayToBytesCodecTraits: ArrayCodecTraits + core::fmt::Debug {
         &self,
         bytes: RawBytes<'_>,
         decoded_representation: &ChunkRepresentation,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,

diff --git a/zarrs/src/array/codec/array_to_bytes/codec_chain.rs b/zarrs/src/array/codec/array_to_bytes/codec_chain.rs
@@ -2,6 +2,8 @@
 
 use std::sync::Arc;
 
+use unsafe_cell_slice::UnsafeCellSlice;
+
 use crate::{
     array::{
         array_bytes::update_bytes_flen,
@@ -307,7 +309,7 @@ impl ArrayToBytesCodecTraits for CodecChain {
         &self,
         mut bytes: RawBytes<'_>,
         decoded_representation: &ChunkRepresentation,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,

diff --git a/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs b/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs
@@ -281,7 +281,6 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
                     let decode_chunk = |chunk_index: usize| {
                         let chunk_subset = self
                             .chunk_index_to_subset(chunk_index as u64, chunks_per_shard.as_slice());
-                        let output = unsafe { output.as_mut_slice() };
 
                         // Read the offset/size
                         let offset = shard_index[chunk_index * 2];
@@ -300,8 +299,11 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
                                     );
                                     let shard_offset =
                                         usize::try_from(index * data_type_size as u64).unwrap();
-                                    output[shard_offset..shard_offset + fv.len()]
-                                        .copy_from_slice(fv);
+                                    unsafe {
+                                        output
+                                            .index_mut(shard_offset..shard_offset + fv.len())
+                                            .copy_from_slice(fv);
+                                    }
                                 }
                             } else {
                                 unreachable!();
@@ -321,7 +323,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
                                 &options,
                             )?;
                             update_bytes_flen(
-                                output,
+                                &output,
                                 &shard_representation.shape_u64(),
                                 &decoded_chunk.into_fixed()?,
                                 &chunk_subset,
@@ -350,7 +352,7 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
         &self,
         encoded_shard: RawBytes<'_>,
         shard_representation: &ChunkRepresentation,
-        output: &mut [u8],
+        output: &UnsafeCellSlice<u8>,
         output_shape: &[u64],
         output_subset: &ArraySubset,
         options: &CodecOptions,
@@ -408,11 +410,9 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
                 };
 
                 {
-                    let output = UnsafeCellSlice::new(output);
                     let decode_chunk = |chunk_index: usize| {
                         let chunk_subset = self
                             .chunk_index_to_subset(chunk_index as u64, chunks_per_shard.as_slice());
-                        let output = unsafe { output.as_mut_slice() };
 
                         let output_subset_chunk = ArraySubset::new_with_start_shape(
                             std::iter::zip(output_subset.start(), chunk_subset.start())
@@ -439,7 +439,8 @@ impl ArrayToBytesCodecTraits for ShardingCodec {
                                     );
                                     let shard_offset =
                                         usize::try_from(index * data_type_size as u64).unwrap();
-                                    output[shard_offset..shard_offset + fv.len()]
+                                    output
+                                        .index_mut(shard_offset..shard_offset + fv.len())
                                         .copy_from_slice(fv);
                                 }
                             } else {
@@ -680,14 +681,13 @@ impl ShardingCodec {
                         }
 
                         unsafe {
-                            let shard_index_unsafe = shard_index_slice.as_mut_slice();
-                            shard_index_unsafe[chunk_index * 2] =
-                                u64::try_from(chunk_offset).unwrap();
-                            shard_index_unsafe[chunk_index * 2 + 1] =
-                                u64::try_from(chunk_encoded.len()).unwrap();
-
-                            let shard_unsafe = shard_slice.as_mut_slice();
-                            shard_unsafe[chunk_offset..chunk_offset + chunk_encoded.len()]
+                            let shard_index_unsafe =
+                                shard_index_slice.index_mut(chunk_index * 2..chunk_index * 2 + 2);
+                            shard_index_unsafe[0] = u64::try_from(chunk_offset).unwrap();
+                            shard_index_unsafe[1] = u64::try_from(chunk_encoded.len()).unwrap();
+
+                            shard_slice
+                                .index_mut(chunk_offset..chunk_offset + chunk_encoded.len())
                                 .copy_from_slice(&chunk_encoded);
                         }
                     }
@@ -837,13 +837,13 @@ impl ShardingCodec {
                     let chunk_offset = encoded_shard_offset
                         .fetch_add(chunk_encoded.len(), std::sync::atomic::Ordering::Relaxed);
                     unsafe {
-                        let shard_index_unsafe = shard_index_slice.as_mut_slice();
-                        shard_index_unsafe[chunk_index * 2] = u64::try_from(chunk_offset).unwrap();
-                        shard_index_unsafe[chunk_index * 2 + 1] =
-                            u64::try_from(chunk_encoded.len()).unwrap();
+                        let shard_index_unsafe =
+                            shard_index_slice.index_mut(chunk_index * 2..chunk_index * 2 + 2);
+                        shard_index_unsafe[0] = u64::try_from(chunk_offset).unwrap();
+                        shard_index_unsafe[1] = u64::try_from(chunk_encoded.len()).unwrap();
 
-                        let shard_unsafe = shard_slice.as_mut_slice();
-                        shard_unsafe[chunk_offset..chunk_offset + chunk_encoded.len()]
+                        shard_slice
+                            .index_mut(chunk_offset..chunk_offset + chunk_encoded.len())
                             .copy_from_slice(&chunk_encoded);
                     }
                 }

diff --git a/zarrs/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs b/zarrs/src/array/codec/array_to_bytes/sharding/sharding_partial_decoder.rs
@@ -265,7 +265,7 @@ impl ArrayPartialDecoderTraits for ShardingPartialDecoder {
                         };
                         let decoded_bytes = decoded_bytes.into_fixed()?;
                         update_bytes_flen(
-                            unsafe { out_array_subset_slice.as_mut_slice() },
+                            &out_array_subset_slice,
                             array_subset.shape(),
                             &decoded_bytes,
                             &chunk_subset_overlap
@@ -557,7 +557,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncShardingPartialDecoder {
                                     ArraySubset,
                                 ) = subset_and_decoded_chunk?;
                                 update_bytes_flen(
-                                    unsafe { shard_slice.as_mut_slice() },
+                                    &shard_slice,
                                     array_subset.shape(),
                                     &chunk_subset_bytes.into(),
                                     &chunk_subset_overlap
@@ -596,7 +596,7 @@ impl AsyncArrayPartialDecoderTraits for AsyncShardingPartialDecoder {
                                     .as_ne_bytes()
                                     .repeat(chunk_subset_overlap.num_elements_usize());
                                 update_bytes_flen(
-                                    unsafe { shard_slice.as_mut_slice() },
+                                    &shard_slice,
                                     array_subset.shape(),
                                     &filled_chunk.into(),
                                     &chunk_subset_overlap