rustmex_separated_complex/

lib.rs

#![allow(unused_unsafe)]
use std::slice::{from_raw_parts_mut, from_raw_parts};
use core::ops::DerefMut;
use rustmex_core::{
	mxArray,

	num_struct,
	data_or_dangling,
	create_uninit_numeric_array,
	Complexity,

	numeric::{NumericArray, ComplexNumericArray, MutNumericArray, MatlabNumber},
	pointers::{MatlabPtr, MutMatlabPtr, MxArray},
	mappable::{*},
	convert::{ToMatlabResult, DataShapeMismatch, FromMatlabError},
	shim::rustmex_create_uninit_numeric_array,
	raw::{
		mxClassID,
		mwSize,
	},

	MatlabClass,
	MutMatlabClass,
	OwnedMatlabClass,
	NewEmpty,
	TakeData,
};

use num_complex::Complex;

use ndarray::{
	Array,
	Dimension,
	ArrayD,
	ArrayViewD,
	ArrayViewMutD,
	Zip,
};

use core::ffi::c_void;

// If these symbols don't exist in the backend, this crate does not compile. This was the
// best way I could of to ensure the interleaved and separated implementations don't get
// mixed up.
extern "Rust" {
	fn rustmex_get_separated_complex(mx: *const mxArray) -> Complex<*mut c_void>;
	fn rustmex_set_separated_complex(mx: *mut mxArray, newdata: Complex<*mut c_void>);
}

num_struct!(SeparatedComplexArray, Complexity::Complex);

macro_rules! data_access {
	($s:ident, $t1:ty, $builder:ident) => {{
		let numel = $s.array.numel();
		unsafe { rustmex_get_separated_complex($s.array.deref()) }
			.map(|elem| data_or_dangling!(elem, $t1))
			.map(|elem| unsafe { $builder(elem, numel) } )
	}}
}

impl<'p, T, P> NumericArray<'p> for SeparatedComplexArray<T, P> where
	T: MatlabNumber + 'p,
	P: MatlabPtr + 'p,
{
	type Data = Complex<&'p [T]>;
	fn data(&self) -> Self::Data {
		data_access!(self, *const T, from_raw_parts)
	}
}

impl<'p, T, P> MutNumericArray<'p> for SeparatedComplexArray<T, P> where
	T: MatlabNumber + 'p,
	P: MutMatlabPtr + 'p,
{
	type MutData = Complex<&'p mut [T]>;
	fn mut_data(&mut self) -> Self::MutData {
		data_access!(self, *mut T, from_raw_parts_mut)
	}
}

impl<T, P> TakeData<P> for SeparatedComplexArray<T, P> where
	T: MatlabNumber,
	P: MutMatlabPtr,
{
	type OwnedData = Complex<Box<[T]>>;
	fn take_data(&mut self) -> Self::OwnedData {

		// use mut_data to get the data slice, since that method implements that
		// correctly
		let data = self.mut_data();

		// SAFETY: Omg this is such a mess
		// Technically, Vec wants to know the _original_ size of the allocation,
		// which is something we can't get out of matlab. So there are two
		// options: either do a complicated and expensive realloc, or (possibly)
		// lie to Vec and Box about the size of the allocation.
		//
		// For now, we're doing the latter. The Matlab allocator does not care
		// about the allocated size, after all, it only cares about the pointer
		// it gets (it probably tracks the size internally).
		let data = data.map(|part| {
			unsafe { Vec::from_raw_parts(part.as_mut_ptr(), part.len(), part.len()) }
				.into_boxed_slice()
		});

		// SAFETY: Set the data pointer of the array to NULL, so we don't have an
		// aliased pointer lying around.
		unsafe { rustmex_set_separated_complex(self.deref_mut(),
			Complex::init_value(core::ptr::null_mut())) };

		data
	}
}

impl<'p, T, P> ComplexNumericArray<'p> for SeparatedComplexArray<T, P> where
	T: MatlabNumber + 'p,
	P: MatlabPtr + 'p,
{
	type CDT = T;
}

macro_rules! copy_complex {
	($data:ident, $f:expr) => {{
		let mut v = Vec::with_capacity($data.len());
		v.extend($data.iter().map($f));
		v.into_boxed_slice()
	}}
}

impl<T: MatlabNumber + Copy> SeparatedComplexArray<T, MxArray> {
	pub fn new(data: Box<[Complex<T>]>, shape: &[usize]) -> ToMatlabResult<Self, Box<[Complex<T>]>> {
		let dimprod = shape.iter().product::<usize>();
		if dimprod != data.len() {
			return Err(DataShapeMismatch::because_numel_shape(data));
		}
		let split = Complex {
			re: copy_complex!(data, |x|x.re),
			im: copy_complex!(data, |x|x.im)
		};
		Ok(Self::new_separated_unchecked(split, shape))
	}
}

impl<T: MatlabNumber> SeparatedComplexArray<T, MxArray> {
	pub fn new_separated(data: Complex<Box<[T]>>, shape: &[usize]) -> ToMatlabResult<Self, Complex<Box<[T]>>> {
		if data.re.len() != data.im.len() {
			return Err(DataShapeMismatch::because_im_re(data));
		}

		if shape.iter().product::<usize>() != data.re.len() {
			return Err(DataShapeMismatch::because_numel_shape(data));
		}

		Ok(Self::new_separated_unchecked(data, shape))
	}

	fn new_separated_unchecked(data: Complex<Box<[T]>>, shape: &[usize]) -> Self {
		let mx = create_uninit_numeric_array!(shape, T, Complexity::Real);

		unsafe { rustmex_set_separated_complex(mx, data.map(|x| Box::into_raw(x) as *mut core::ffi::c_void)); }

		Self::construct(unsafe { MxArray::assume_responsibility_ptr(mx)})
	}
}

impl<'a, 'b, T, P> From<&'b SeparatedComplexArray<T, P>> for Complex<ArrayViewD<'a, T>> where
	T: MatlabNumber,
	P: MatlabPtr + 'a,
	'b: 'a,
{
	fn from(num: &SeparatedComplexArray<T, P>) -> Self {
		let data = num.data();
		let dimensions = num.dimensions();

		data.map(|part| rustmex_core::from_num_to_ndarray!(part, dimensions, ArrayViewD<T>))
	}
}

impl<'a, T> From<SeparatedComplexArray<T, &'a mxArray>> for Complex<ArrayViewD<'a, T>> where
	T: MatlabNumber,
{
	fn from(arr: SeparatedComplexArray<T, &'a mxArray>) -> Self {
		arr.into()
	}
}

impl<'a, 'b, T, P> From<&'b mut SeparatedComplexArray<T, P>> for Complex<ArrayViewMutD<'a, T>> where
	T: MatlabNumber,
	P: MutMatlabPtr + 'a,
	'b: 'a,
{
	fn from(num: &mut SeparatedComplexArray<T, P>) -> Self {
		let data: Complex<&mut[T]> = num.mut_data();
		let dimensions = num.dimensions();

		data.map(|part| rustmex_core::from_num_to_ndarray!(part, dimensions, ArrayViewMutD<T>))
	}
}

impl<'a, T> From<SeparatedComplexArray<T, &'a mut mxArray>> for Complex<ArrayViewMutD<'a, T>> where
	T: MatlabNumber,
{
	fn from(arr: SeparatedComplexArray<T, &'a mut mxArray>) -> Self {
		arr.into()
	}
}


impl<T, P> From<SeparatedComplexArray<T, P>> for ArrayD<Complex<T>> where
	T: MatlabNumber,
	P: MatlabPtr,
{
	fn from(arr: SeparatedComplexArray<T, P>) -> Self {
		// Technically one could construct the Array just from the two slices,
		// but that is massively combersome, so just create two views into the
		// real and imaginary part, and then zip them together. The cost of
		// creating two views will in general be overshadowed by the copy anyway,
		// but now we don't duplicate the array construction code.
		let arr: Complex<ArrayViewD<T>> = (&arr).into();
		Zip::from(&arr.re).and(&arr.im).map_collect(|&re, &im| Complex { re, im })
	}
}

impl<T> From<Complex<T>> for SeparatedComplexArray<T, MxArray> where T: MatlabNumber {
	fn from(c: Complex<T>) -> Self {
		let alloced = c.map(|val| Box::from([val]));
		Self::new_separated_unchecked(alloced, &[1,1])
	}
}

impl<T> From<[Complex<T>;0]> for SeparatedComplexArray<T, MxArray> where T: MatlabNumber {
	fn from(_empty: [Complex<T>;0]) -> Self {
		Self::new_empty()
	}
}

impl<'a, T, D> From<Array<Complex<T>, D>> for SeparatedComplexArray<T, MxArray> where
	T: MatlabNumber + Copy,
	D: Dimension,
	SeparatedComplexArray<T, MxArray>: TryFrom<Complex<Array<T, D>>>
{
	// This method is implemented separately from the Complex<Array> case. Even
	// though this method does split the Array<Complex> into a Complex<Array>, we
	// know that the sizes of the imaginary and real parts match, so we can skip the
	// extra check.
	fn from(arr: Array<Complex<T>, D>) -> Self {
		let arr = arr.view().split_complex();
		let arr = <Complex<_> as Mappable<_, _>>::map(arr, |component| {
			let mut component = component.into_owned();
			rustmex_core::from_ndarray_to_num!(@shuffle component)
		});
		let dim = arr.re.raw_dim();
		let dimview = dim.as_array_view();
		let shape = dimview.to_slice().unwrap();
		let arr = arr.map(|x|x.into_raw_vec().into_boxed_slice());
		Self::new_separated_unchecked(arr, &shape)
	}
}

impl<'a, T, D> TryFrom<Complex<Array<T, D>>> for SeparatedComplexArray<T, MxArray> where
	T: MatlabNumber,
	D: Dimension,
{
	type Error = DataShapeMismatch<Complex<Array<T,D>>>;
	fn try_from(arr: Complex<Array<T, D>>) -> Result<Self, Self::Error> {
		// Check whether the shapes of both arrays match; hereby also inferring
		// that the number of elements is correct (we're assuming that that is
		// handled correctly by NDarray).
		if arr.re.shape() != arr.im.shape() {
			return Err(DataShapeMismatch::because_im_re(arr));
		}

		let arr = arr.map(|mut arr| {
			rustmex_core::from_ndarray_to_num!(@shuffle arr)
		});

		// At this point we know that the shapes match, so it does not matter
		// which one we pick
		//
		// We do this dance with the raw_dim(), because NDarray knows more about
		// the layout of the dimension than we do. It is probably more efficient
		// than always copying the .shape() into a Vec.
		let dim = arr.re.raw_dim();
		let dimview = dim.as_array_view();
		let shape = dimview.to_slice().unwrap();

		Ok(Self::new_separated_unchecked(arr.map(|x| {
			x.into_raw_vec()
			 .into_boxed_slice()
		}), &shape))
	}
}