Skip to main content

mzdata_bindata/
traits.rs

1use std::marker::PhantomData;
2use std::slice;
3use std::mem;
4use std::borrow::Cow;
5
6use bytemuck::Pod;
7use num_traits::{AsPrimitive, Num};
8use mzdata_param::Unit;
9
10use super::encodings::{ArrayRetrievalError, BinaryDataArrayType, Bytes};
11use super::ArrayType;
12
13/// A byte array that may be encoded and compressed as [`BinaryCompressionType`], or unpacked in
14/// memory and re-interpreted as native types appropriate to [`BinaryDataArrayType`].
15///
16/// The raw bytes must be made available through the [`ByteArrayView::view`] method,
17/// but all other behavior is added on top of that
18pub trait ByteArrayView<'transient, 'lifespan: 'transient> {
19    /// The core method that exposes the decoded byte array that the type *must* provide.
20    fn view(&'lifespan self) -> Result<Cow<'lifespan, [u8]>, ArrayRetrievalError>;
21
22    /// This is a helper method for the various `to_<type>` methods like [`Self::to_f64`]. It should not
23    /// need to be called directly. This method relies on [`bytemuck::try_cast_slice`] to do the heavy lifting.
24    fn coerce_from<T: Pod>(
25        buffer: Cow<'transient, [u8]>,
26    ) -> Result<Cow<'transient, [T]>, ArrayRetrievalError> {
27        let n = buffer.len();
28        if n == 0 {
29            return Ok(Cow::Owned(Vec::new()))
30        }
31        let z = mem::size_of::<T>();
32        if n % z != 0 {
33            return Err(ArrayRetrievalError::DataTypeSizeMismatch);
34        }
35        match buffer {
36            Cow::Borrowed(c) => {
37                Ok(Cow::Borrowed(bytemuck::try_cast_slice(c)?))
38            },
39            Cow::Owned(v) => {
40                let size_type = n / z;
41                let mut buf = Vec::with_capacity(size_type);
42                v.chunks_exact(z).try_for_each(|c| {
43                    buf.extend(bytemuck::try_cast_slice(c)?);
44                    Ok::<(), bytemuck::PodCastError>(())
45                })?;
46                Ok(Cow::Owned(buf))
47            },
48        }
49    }
50
51    /// This is a helper method for the various `to_<type>` methods like [`Self::to_f64`]. It should not
52    /// need to be called directly. This calls into [`Self::coerce_from`] internally after decoding.
53    fn coerce<T: Pod>(
54        &'lifespan self,
55    ) -> Result<Cow<'transient, [T]>, ArrayRetrievalError> {
56        match self.view() {
57            Ok(data) => {
58                #[cfg(target_endian = "big")]
59                {
60                    let mut data = Cow::Owned(data.to_vec());
61                    self.dtype().swap_bytes(&mut data)?;
62                    Self::coerce_from(data)
63                }
64                #[cfg(not(target_endian = "big"))]
65                {
66                    Self::coerce_from(data)
67                }
68            },
69            Err(err) => Err(err),
70        }
71    }
72
73    /// Decode the array, then copy it to a new array, converting each element from type `D` to to type `S`
74    ///
75    /// This makes a copy of the data, forcibly coercing the values using [`AsPrimitive`]
76    fn convert<S: Num + Clone + AsPrimitive<D> + Pod, D: Num + Clone + Copy + 'static>(
77        &'lifespan self,
78    ) -> Result<Cow<'transient, [D]>, ArrayRetrievalError> {
79        match self.coerce::<S>() {
80            Ok(view) => {
81                match view {
82                    Cow::Borrowed(view) => {
83                        Ok(Cow::Owned(view.iter().map(|a| a.as_()).collect()))
84                    }
85                    Cow::Owned(owned) => {
86                        let res = owned.iter().map(|a| a.as_()).collect();
87                        Ok(Cow::Owned(res))
88                    }
89                }
90            }
91            Err(err) => Err(err),
92        }
93    }
94
95    /// The kind of array this is
96    fn name(&self) -> &ArrayType;
97
98    /// The real data type encoded in bytes
99    fn dtype(&self) -> BinaryDataArrayType;
100
101    /// The unit of measurement each data point is in
102    fn unit(&self) -> Unit;
103
104    /// Get the identifier referencing a [`DataProcessing`](crate::meta::DataProcessing)
105    fn data_processing_reference(&self) -> Option<&str> {
106        None
107    }
108
109    /// Get a view of the data as [`f32`]. If the data are stored that way already, no
110    /// copy is required.
111    fn to_f32(&'lifespan self) -> Result<Cow<'transient, [f32]>, ArrayRetrievalError> {
112        type D = f32;
113        match self.dtype() {
114            BinaryDataArrayType::Float32 | BinaryDataArrayType::ASCII => self.coerce::<D>(),
115            BinaryDataArrayType::Float64 => {
116                type S = f64;
117                self.convert::<S, D>()
118            }
119            BinaryDataArrayType::Int32 => {
120                type S = i32;
121                self.convert::<S, D>()
122            }
123            BinaryDataArrayType::Int64 => {
124                type S = i64;
125                self.convert::<S, D>()
126            }
127            _ => Err(ArrayRetrievalError::DataTypeSizeMismatch),
128        }
129    }
130
131    /// Get a view of the data as [`f64`]. If the data are stored that way already, no
132    /// copy is required.
133    fn to_f64(&'lifespan self) -> Result<Cow<'transient, [f64]>, ArrayRetrievalError> {
134        type D = f64;
135        match self.dtype() {
136            BinaryDataArrayType::Float32 => {
137                type S = f32;
138                self.convert::<S, D>()
139            }
140            BinaryDataArrayType::Float64 | BinaryDataArrayType::ASCII => self.coerce(),
141            BinaryDataArrayType::Int32 => {
142                type S = i32;
143                self.convert::<S, D>()
144            }
145            BinaryDataArrayType::Int64 => {
146                type S = i64;
147                self.convert::<S, D>()
148            }
149            _ => Err(ArrayRetrievalError::DataTypeSizeMismatch),
150        }
151    }
152
153    /// Get a view of the data as [`i32`]. If the data are stored that way already, no
154    /// copy is required.
155    fn to_i32(&'lifespan self) -> Result<Cow<'transient, [i32]>, ArrayRetrievalError> {
156        type D = i32;
157        match self.dtype() {
158            BinaryDataArrayType::Float32 => {
159                type S = f32;
160                self.convert::<S, D>()
161            }
162            BinaryDataArrayType::Float64 => {
163                type S = f64;
164                self.convert::<S, D>()
165            }
166            BinaryDataArrayType::Int32 | BinaryDataArrayType::ASCII => self.coerce::<D>(),
167            BinaryDataArrayType::Int64 => {
168                type S = i64;
169                self.convert::<S, D>()
170            }
171            _ => Err(ArrayRetrievalError::DataTypeSizeMismatch),
172        }
173    }
174
175    /// Get a view of the data as [`i64`]. If the data are stored that way already, no
176    /// copy is required.
177    fn to_i64(&'lifespan self) -> Result<Cow<'transient, [i64]>, ArrayRetrievalError> {
178        type D = i64;
179        match self.dtype() {
180            BinaryDataArrayType::Float32 => {
181                type S = f32;
182                self.convert::<S, D>()
183            }
184            BinaryDataArrayType::Float64 => {
185                type S = f64;
186                self.convert::<S, D>()
187            }
188            BinaryDataArrayType::Int64 | BinaryDataArrayType::ASCII => self.coerce::<D>(),
189            BinaryDataArrayType::Int32 => {
190                type S = i32;
191                self.convert::<S, D>()
192            }
193            _ => Err(ArrayRetrievalError::DataTypeSizeMismatch),
194        }
195    }
196
197    /// The size of encoded array in terms of # of elements of the [`BinaryDataArrayType`] given by [`ByteArrayView::dtype`]
198    fn data_len(&'lifespan self) -> Result<usize, ArrayRetrievalError> {
199        let view = self.view()?;
200        let n = view.len();
201        Ok(n / self.dtype().size_of())
202    }
203
204    fn iter_type<T: Pod>(&'lifespan self) -> Result<DataSliceIter<'lifespan, T>, ArrayRetrievalError> {
205        Ok(DataSliceIter::new(self.view()?))
206    }
207
208    fn iter_u8(&'lifespan self) -> Result<DataSliceIter<'lifespan, u8>, ArrayRetrievalError> {
209        Ok(DataSliceIter::new(self.view()?))
210    }
211
212    fn iter_f32(&'lifespan self) -> Result<DataSliceIter<'lifespan, f32>, ArrayRetrievalError> {
213        Ok(DataSliceIter::new(self.view()?))
214    }
215
216    fn iter_f64(&'lifespan self) -> Result<DataSliceIter<'lifespan, f64>, ArrayRetrievalError> {
217        Ok(DataSliceIter::new(self.view()?))
218    }
219
220    fn iter_i32(&'lifespan self) -> Result<DataSliceIter<'lifespan, i32>, ArrayRetrievalError> {
221        Ok(DataSliceIter::new(self.view()?))
222    }
223
224    fn iter_i64(&'lifespan self) -> Result<DataSliceIter<'lifespan, i64>, ArrayRetrievalError> {
225        Ok(DataSliceIter::new(self.view()?))
226    }
227}
228
229/// A mutable byte array
230pub trait ByteArrayViewMut<'transient, 'lifespan: 'transient>:
231    ByteArrayView<'transient, 'lifespan>
232{
233
234    /// Specify the unit of the data array
235    fn unit_mut(&mut self) -> &mut Unit;
236
237    /// Get a mutable view of the bytes backing this data array.
238    ///
239    /// This is in turn used by [`ByteArrayViewMut::coerce_mut`] to produce a typed array
240    fn view_mut(&'transient mut self) -> Result<&'transient mut Bytes, ArrayRetrievalError>;
241
242    fn coerce_from_mut<T: Clone + Sized>(
243        buffer: &mut [u8],
244    ) -> Result<&'transient mut [T], ArrayRetrievalError> {
245        let n = buffer.len();
246        if n == 0 {
247            return Ok(&mut [])
248        }
249        let z = mem::size_of::<T>();
250        if n % z != 0 {
251            return Err(ArrayRetrievalError::DataTypeSizeMismatch);
252        }
253        let m = n / z;
254        unsafe { Ok(slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut T, m)) }
255    }
256
257    fn coerce_mut<T: Clone + Sized>(
258        &'lifespan mut self,
259    ) -> Result<&'transient mut [T], ArrayRetrievalError> {
260        let view = self.view_mut()?;
261        #[cfg(target_endian = "big")]
262        {
263            log::error!("Mutable view of raw bytes on big endian system is only partially supported.")
264        }
265        Self::coerce_from_mut(view)
266    }
267
268    #[allow(unused)]
269    /// Set the identifier referencing a [`DataProcessing`](crate::meta::DataProcessing)
270    fn set_data_processing_reference(&mut self, data_processing_reference: Option<Box<str>>) {}
271}
272
273#[derive(Debug)]
274pub struct DataSliceIter<'a, T: Pod> {
275    buffer: Cow<'a, [u8]>,
276    i: usize,
277    _t: PhantomData<T>
278}
279
280impl<T: Pod> ExactSizeIterator for DataSliceIter<'_, T> {
281    fn len(&self) -> usize {
282        let z = mem::size_of::<T>();
283        self.buffer.len() / z
284    }
285}
286
287impl<'a, T: Pod> DataSliceIter<'a, T> {
288    pub fn new(buffer: Cow<'a, [u8]>) -> Self {
289        Self { buffer, i: 0, _t: PhantomData }
290    }
291
292    pub fn next_value(&mut self) -> Option<T> {
293        let z = mem::size_of::<T>();
294        let offset = z * self.i;
295        if (offset + z) > self.buffer.len() {
296            None
297        } else {
298            let data = &self.buffer[offset..offset + z];
299            #[cfg(target_endian = "big")]
300            {
301                let mut data = data.to_vec();
302                data.reverse();
303                let val = bytemuck::from_bytes(data);
304                self.i += 1;
305                Some(*val)
306            }
307            #[cfg(not(target_endian = "big"))]
308            {
309                let val = bytemuck::from_bytes(data);
310                self.i += 1;
311                Some(*val)
312            }
313        }
314    }
315}
316
317impl<T: Pod> Iterator for DataSliceIter<'_, T> {
318    type Item = T;
319
320    fn next(&mut self) -> Option<Self::Item> {
321        self.next_value()
322    }
323}