use std::sync::Arc;
use zarrs_storage::StorageError;
use zarrs_storage::byte_range::{ByteLength, ByteOffset, ByteRange, ByteRangeIterator};
use crate::{ArrayBytesRaw, CodecError, CodecOptions};
#[cfg(feature = "async")]
use crate::AsyncBytesPartialDecoderTraits;
use crate::BytesPartialDecoderTraits;
pub struct ByteIntervalPartialDecoder {
input_handle: Arc<dyn BytesPartialDecoderTraits>,
byte_offset: ByteOffset,
byte_length: ByteLength,
}
impl ByteIntervalPartialDecoder {
pub fn new(
input_handle: Arc<dyn BytesPartialDecoderTraits>,
byte_offset: ByteOffset,
byte_length: ByteLength,
) -> Self {
Self {
input_handle,
byte_offset,
byte_length,
}
}
}
impl BytesPartialDecoderTraits for ByteIntervalPartialDecoder {
fn exists(&self) -> Result<bool, StorageError> {
self.input_handle.exists()
}
fn size_held(&self) -> usize {
self.input_handle.size_held()
}
fn partial_decode_many(
&self,
byte_ranges: ByteRangeIterator,
options: &CodecOptions,
) -> Result<Option<Vec<ArrayBytesRaw<'_>>>, CodecError> {
let byte_ranges = byte_ranges.map(|byte_range| match byte_range {
ByteRange::FromStart(offset, None) => {
ByteRange::FromStart(self.byte_offset + offset, Some(self.byte_length))
}
ByteRange::FromStart(offset, Some(length)) => {
ByteRange::FromStart(self.byte_offset + offset, Some(length))
}
ByteRange::Suffix(length) => {
ByteRange::FromStart(self.byte_offset + self.byte_length - length, Some(length))
}
});
self.input_handle
.partial_decode_many(Box::new(byte_ranges), options)
}
fn supports_partial_decode(&self) -> bool {
self.input_handle.supports_partial_decode()
}
}
#[cfg(feature = "async")]
pub struct AsyncByteIntervalPartialDecoder {
input_handle: Arc<dyn AsyncBytesPartialDecoderTraits>,
byte_offset: ByteOffset,
byte_length: ByteLength,
}
#[cfg(feature = "async")]
impl AsyncByteIntervalPartialDecoder {
pub fn new(
input_handle: Arc<dyn AsyncBytesPartialDecoderTraits>,
byte_offset: ByteOffset,
byte_length: ByteLength,
) -> Self {
Self {
input_handle,
byte_offset,
byte_length,
}
}
}
#[cfg(feature = "async")]
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
impl AsyncBytesPartialDecoderTraits for AsyncByteIntervalPartialDecoder {
async fn exists(&self) -> Result<bool, StorageError> {
self.input_handle.exists().await
}
fn size_held(&self) -> usize {
self.input_handle.size_held()
}
async fn partial_decode_many<'a>(
&'a self,
byte_ranges: ByteRangeIterator<'a>,
options: &CodecOptions,
) -> Result<Option<Vec<ArrayBytesRaw<'a>>>, CodecError> {
let byte_ranges = byte_ranges.map(|byte_range| match byte_range {
ByteRange::FromStart(offset, None) => {
ByteRange::FromStart(self.byte_offset + offset, Some(self.byte_length))
}
ByteRange::FromStart(offset, Some(length)) => {
ByteRange::FromStart(self.byte_offset + offset, Some(length))
}
ByteRange::Suffix(length) => {
ByteRange::FromStart(self.byte_offset + self.byte_length - length, Some(length))
}
});
self.input_handle
.partial_decode_many(Box::new(byte_ranges), options)
.await
}
fn supports_partial_decode(&self) -> bool {
self.input_handle.supports_partial_decode()
}
}