ndarray_npy/npy/elements/
complex.rs

1use super::check_for_extra_bytes;
2use crate::{ReadDataError, ReadableElement};
3use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
4use num_complex_0_4::Complex;
5use py_literal::Value as PyValue;
6use std::io;
7use std::mem;
8use std::slice;
9
10/// Casts the slice of complex numbers to a slice of the inner values.
11///
12/// The resulting slice alternates between the real and imaginary parts of
13/// consecutive complex elements.
14///
15/// # Panics
16///
17/// Panics if `T` is zero-sized.
18fn complex_slice_as_inner_slice_mut<T>(slice: &mut [Complex<T>]) -> &mut [T] {
19    assert!(mem::size_of::<T>() > 0);
20
21    // These assertions should always pass, since `Complex` is `repr(C)` and
22    // has only two fields, both of type `T`.
23    assert_eq!(
24        mem::size_of::<Complex<T>>(),
25        mem::size_of::<T>().checked_mul(2).unwrap()
26    );
27    assert_eq!(mem::align_of::<Complex<T>>(), mem::align_of::<T>());
28
29    // This should never panic, because we've checked that `T` is not
30    // zero-sized, and slices are guaranteed to contain no more than
31    // `isize::MAX` bytes. (See the docs of `std::slice::from_raw_parts`.)
32    let inner_len = slice.len().checked_mul(2).unwrap();
33
34    // This is sound because:
35    //
36    // - Since `slice` is an existing slice, its pointer is non-null, it's
37    //   properly aligned, its length is correct, the elements are initialized,
38    //   and the lifetime is correct.
39    //
40    // - We've checked above that `inner_len` will be the correct length and
41    //   the alignment will be the same.
42    //
43    // - The fields of `Complex` are public, so we aren't violating any
44    //   visibility constraints.
45    unsafe { slice::from_raw_parts_mut(slice.as_mut_ptr().cast(), inner_len) }
46}
47
48macro_rules! impl_readable_complex_multi_byte {
49    ($elem:ty, $little_desc:expr, $big_desc:expr, $zero:expr, $inner_read_into:ident) => {
50        impl ReadableElement for $elem {
51            fn read_to_end_exact_vec<R: io::Read>(
52                mut reader: R,
53                type_desc: &PyValue,
54                len: usize,
55            ) -> Result<Vec<Self>, ReadDataError> {
56                let mut out = vec![$zero; len];
57                let inner_slice = complex_slice_as_inner_slice_mut(&mut out);
58                match *type_desc {
59                    PyValue::String(ref s) if s == $little_desc => {
60                        reader.$inner_read_into::<LittleEndian>(inner_slice)?;
61                    }
62                    PyValue::String(ref s) if s == $big_desc => {
63                        reader.$inner_read_into::<BigEndian>(inner_slice)?;
64                    }
65                    ref other => {
66                        return Err(ReadDataError::WrongDescriptor(other.clone()));
67                    }
68                }
69                check_for_extra_bytes(&mut reader)?;
70                Ok(out)
71            }
72        }
73    };
74}
75
76macro_rules! impl_complex_multi_byte {
77    ($elem:ty, $little_desc:expr, $big_desc:expr, $zero:expr, $inner_read_into:ident) => {
78        impl_writable_element_always_valid_cast!($elem, $little_desc, $big_desc);
79        impl_readable_complex_multi_byte!($elem, $little_desc, $big_desc, $zero, $inner_read_into);
80        #[cfg(target_endian = "little")]
81        impl_view_and_view_mut_always_valid_cast_multi_byte!($elem, $little_desc, $big_desc);
82        #[cfg(target_endian = "big")]
83        impl_view_and_view_mut_always_valid_cast_multi_byte!($elem, $big_desc, $little_desc);
84    };
85}
86
87impl_complex_multi_byte!(
88    Complex<f32>,
89    "<c8",
90    ">c8",
91    Complex::new(0., 0.),
92    read_f32_into
93);
94impl_complex_multi_byte!(
95    Complex<f64>,
96    "<c16",
97    ">c16",
98    Complex::new(0., 0.),
99    read_f64_into
100);