use self::robj::{AsTypedSlice, Robj};
use super::*;
use crate::throw_r_error;
use extendr_ffi::{
Rf_GetArrayDimnames, Rf_GetColNames, Rf_GetRowNames, Rf_dimgets, Rf_dimnamesgets, Rf_namesgets,
TYPEOF,
};
use std::ops::{Index, IndexMut};
#[derive(Debug, PartialEq)]
pub struct RArray<T, const NDIM: usize> {
robj: Robj,
_data: std::marker::PhantomData<T>,
}
impl<T, const NDIM: usize> RArray<T, NDIM> {
pub fn get_dimnames(&self) -> List {
unsafe { List::try_from(Robj::from_sexp(Rf_GetArrayDimnames(self.get()))).unwrap() }
}
pub fn get_dim(&self) -> Vec<usize> {
self.robj
.get_attrib(wrapper::symbol::dim_symbol())
.unwrap()
.as_integer_vector()
.map(|vec| vec.into_iter().map(|x| x as usize).collect())
.unwrap()
}
pub fn set_names(&mut self, names: Strings) {
let _ = unsafe { Rf_namesgets(self.get_mut(), names.get()) };
}
pub fn set_dimnames(&mut self, dimnames: List) {
let _ = unsafe { Rf_dimnamesgets(self.get_mut(), dimnames.get()) };
}
pub fn set_dim(&mut self, dim: Robj) {
let _ = unsafe { Rf_dimgets(self.get_mut(), dim.get()) };
}
}
pub type RColumn<T> = RArray<T, 1>;
pub type RMatrix<T> = RArray<T, 2>;
pub type RMatrix3D<T> = RArray<T, 3>;
pub type RMatrix4D<T> = RArray<T, 4>;
pub type RMatrix5D<T> = RArray<T, 5>;
impl<T, const NDIM: usize> RArray<T, NDIM>
where
T: ToVectorValue,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new_array(dim: [usize; NDIM]) -> Self {
let sexptype = T::sexptype();
let len = dim.iter().product();
let mut robj = Robj::alloc_vector(sexptype, len);
robj.set_attrib(wrapper::symbol::dim_symbol(), dim).unwrap();
RArray::from_parts(robj)
}
}
impl<T> RMatrix<T>
where
T: ToVectorValue,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new(nrow: usize, ncol: usize) -> Self {
let sexptype = T::sexptype();
let matrix = Robj::alloc_matrix(sexptype, nrow as _, ncol as _);
RArray::from_parts(matrix)
}
}
impl<T> RMatrix<T>
where
T: ToVectorValue + CanBeNA,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new_with_na(nrow: usize, ncol: usize) -> Self {
let mut matrix = Self::new(nrow, ncol);
if nrow != 0 || ncol != 0 {
matrix
.as_typed_slice_mut()
.unwrap()
.iter_mut()
.for_each(|x| {
*x = T::na();
});
}
matrix
}
}
impl<T> RMatrix<T> {
pub fn get_colnames(&self) -> Option<Strings> {
unsafe {
let maybe_colnames = Rf_GetColNames(Rf_GetArrayDimnames(self.get()));
match TYPEOF(maybe_colnames) {
SEXPTYPE::NILSXP => None,
SEXPTYPE::STRSXP => {
let colnames = Robj::from_sexp(maybe_colnames);
Strings::try_from(colnames).ok()
}
_ => unreachable!(
"This should not have occurred. Please report an error at https://github.com/extendr/extendr/issues"
),
}
}
}
pub fn get_rownames(&self) -> Option<Strings> {
unsafe {
let maybe_rownames = Rf_GetRowNames(Rf_GetArrayDimnames(self.get()));
match TYPEOF(maybe_rownames) {
SEXPTYPE::NILSXP => None,
SEXPTYPE::STRSXP => {
let rownames = Robj::from_sexp(maybe_rownames);
Strings::try_from(rownames).ok()
}
_ => unreachable!(
"This should not have occurred. Please report an error at https://github.com/extendr/extendr/issues"
),
}
}
}
}
const BASE: usize = 0;
trait Offset<D> {
fn offset(&self, idx: D) -> usize;
}
impl<T, const NDIM: usize> Offset<[usize; NDIM]> for RArray<T, NDIM> {
fn offset(&self, index: [usize; NDIM]) -> usize {
let dims = self.get_dim();
if index.len() != dims.len() {
throw_r_error("array index: dimension mismatch");
}
index
.iter()
.zip(dims.iter())
.rev()
.enumerate()
.fold(0, |acc, (i, (&idx, &dim))| {
if idx - BASE >= dim {
let msg = format!("array index: dimension {} overflow (0-based dimension)", i);
throw_r_error(msg);
}
acc * dim + (idx - BASE)
})
}
}
impl<T, const NDIM: usize> RArray<T, NDIM>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn from_parts(robj: Robj) -> Self {
Self {
robj,
_data: std::marker::PhantomData,
}
}
pub fn data(&self) -> &[T] {
self.as_typed_slice().unwrap()
}
pub fn data_mut(&mut self) -> &mut [T] {
self.as_typed_slice_mut().unwrap()
}
pub fn ndim(&self) -> usize {
NDIM
}
pub fn dim(&self) -> Vec<usize> {
self.get_dim()
}
}
impl<T> RColumn<T>
where
T: ToVectorValue,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new_column<F: FnMut(usize) -> T>(nrows: usize, f: F) -> Self {
let mut robj = (0..nrows).map(f).collect_robj();
let dim = [nrows];
robj.set_attrib(wrapper::symbol::dim_symbol(), dim).unwrap();
RArray::from_parts(robj)
}
pub fn nrows(&self) -> usize {
self.get_dim()[0]
}
}
impl<T> RMatrix<T>
where
T: ToVectorValue,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new_matrix<F: Clone + FnMut(usize, usize) -> T>(
nrows: usize,
ncols: usize,
f: F,
) -> Self {
let mut robj = (0..ncols)
.flat_map(|c| {
let mut g = f.clone();
(0..nrows).map(move |r| g(r, c))
})
.collect_robj();
let dim = [nrows, ncols];
robj.set_attrib(wrapper::symbol::dim_symbol(), dim).unwrap();
RArray::from_parts(robj)
}
pub fn nrows(&self) -> usize {
self.get_dim()[0]
}
pub fn ncols(&self) -> usize {
self.get_dim()[1]
}
}
impl<T> RMatrix3D<T>
where
T: ToVectorValue,
Robj: for<'a> AsTypedSlice<'a, T>,
{
pub fn new_matrix3d<F: Clone + FnMut(usize, usize, usize) -> T>(
nrows: usize,
ncols: usize,
nmatrix: usize,
f: F,
) -> Self {
let mut robj = (0..nmatrix)
.flat_map(|m| {
let h = f.clone();
(0..ncols).flat_map(move |c| {
let mut g = h.clone();
(0..nrows).map(move |r| g(r, c, m))
})
})
.collect_robj();
let dim = [nrows, ncols, nmatrix];
robj.set_attrib(wrapper::symbol::dim_symbol(), dim).unwrap();
RArray::from_parts(robj)
}
pub fn nrows(&self) -> usize {
self.get_dim()[0]
}
pub fn ncols(&self) -> usize {
self.get_dim()[1]
}
pub fn nsub(&self) -> usize {
self.get_dim()[2]
}
}
impl<T> TryFrom<&Robj> for RColumn<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if let Some(_slice) = robj.as_typed_slice() {
Ok(RArray::from_parts(robj.clone()))
} else {
Err(Error::ExpectedVector(robj.clone()))
}
}
}
impl<T> TryFrom<&Robj> for RMatrix<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if !robj.is_matrix() {
Err(Error::ExpectedMatrix(robj.clone()))
} else if let Some(_slice) = robj.as_typed_slice() {
if let Some(dim) = robj.dim() {
let ndim = dim.len();
if ndim != 2 {
Err(Error::ExpectedMatrix(robj.clone()))
} else {
Ok(RArray::from_parts(robj.clone()))
}
} else {
Err(Error::ExpectedMatrix(robj.clone()))
}
} else {
Err(Error::TypeMismatch(robj.clone()))
}
}
}
impl<T> TryFrom<&Robj> for RMatrix3D<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if let Some(_slice) = robj.as_typed_slice() {
if let Some(dim) = robj.dim() {
if dim.len() != 3 {
Err(Error::ExpectedMatrix3D(robj.clone()))
} else {
Ok(RArray::from_parts(robj.clone()))
}
} else {
Err(Error::ExpectedMatrix3D(robj.clone()))
}
} else {
Err(Error::TypeMismatch(robj.clone()))
}
}
}
impl<T> TryFrom<&Robj> for RMatrix4D<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if let Some(_slice) = robj.as_typed_slice() {
if let Some(dim) = robj.dim() {
if dim.len() != 4 {
Err(Error::ExpectedMatrix4D(robj.clone()))
} else {
Ok(RArray::from_parts(robj.clone()))
}
} else {
Err(Error::ExpectedMatrix4D(robj.clone()))
}
} else {
Err(Error::TypeMismatch(robj.clone()))
}
}
}
impl<T> TryFrom<&Robj> for RMatrix5D<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if let Some(_slice) = robj.as_typed_slice() {
if let Some(dim) = robj.dim() {
if dim.len() != 5 {
Err(Error::ExpectedMatrix5D(robj.clone()))
} else {
Ok(RArray::from_parts(robj.clone()))
}
} else {
Err(Error::ExpectedMatrix5D(robj.clone()))
}
} else {
Err(Error::TypeMismatch(robj.clone()))
}
}
}
macro_rules! impl_try_from_robj_ref {
($($type : tt)*) => {
$(
impl<T> TryFrom<Robj> for $type<T>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: Robj) -> Result<Self> {
<$type<T>>::try_from(&robj)
}
}
impl<T> TryFrom<&Robj> for Option<$type<T>>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: &Robj) -> Result<Self> {
if robj.is_null() || robj.is_na() {
Ok(None)
} else {
Ok(Some(<$type<T>>::try_from(robj)?))
}
}
}
impl<T> TryFrom<Robj> for Option<$type<T>>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Error = Error;
fn try_from(robj: Robj) -> Result<Self> {
<Option::<$type<T>>>::try_from(&robj)
}
}
)*
}
}
impl_try_from_robj_ref!(
RMatrix
RColumn
RMatrix3D
RMatrix4D
RMatrix5D
);
impl<T, const DIM: usize> From<RArray<T, DIM>> for Robj {
fn from(array: RArray<T, DIM>) -> Self {
array.robj
}
}
pub trait MatrixConversions: GetSexp {
fn as_column<E>(&self) -> Option<RColumn<E>>
where
Robj: for<'a> AsTypedSlice<'a, E>,
{
<RColumn<E>>::try_from(self.as_robj()).ok()
}
fn as_matrix<E>(&self) -> Option<RMatrix<E>>
where
Robj: for<'a> AsTypedSlice<'a, E>,
{
<RMatrix<E>>::try_from(self.as_robj()).ok()
}
fn as_matrix3d<E>(&self) -> Option<RMatrix3D<E>>
where
Robj: for<'a> AsTypedSlice<'a, E>,
{
<RMatrix3D<E>>::try_from(self.as_robj()).ok()
}
}
impl MatrixConversions for Robj {}
impl<T, const NDIM: usize> Index<[usize; NDIM]> for RArray<T, NDIM>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
type Output = T;
fn index(&self, index: [usize; NDIM]) -> &Self::Output {
unsafe {
self.data()
.as_ptr()
.add(self.offset(index))
.as_ref()
.unwrap()
}
}
}
impl<T, const NDIM: usize> IndexMut<[usize; NDIM]> for RArray<T, NDIM>
where
Robj: for<'a> AsTypedSlice<'a, T>,
{
fn index_mut(&mut self, index: [usize; NDIM]) -> &mut Self::Output {
unsafe {
self.data_mut()
.as_mut_ptr()
.add(self.offset(index))
.as_mut()
.unwrap()
}
}
}
impl<T, const NDIM: usize> Deref for RArray<T, NDIM> {
type Target = Robj;
fn deref(&self) -> &Self::Target {
&self.robj
}
}
impl<T, const NDIM: usize> DerefMut for RArray<T, NDIM> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.robj
}
}
impl<T, const NDIM: usize> From<Option<RArray<T, NDIM>>> for Robj {
fn from(value: Option<RArray<T, NDIM>>) -> Self {
match value {
None => nil_value(),
Some(value) => value.into(),
}
}
}
#[cfg(test)]
mod tests {
use std::error::Error;
use super::*;
use crate as extendr_api;
use extendr_engine::with_r;
use extendr_ffi::Rf_PrintValue;
use prelude::{Rcplx, Rfloat, Rint};
#[test]
fn test_empty_matrix_new() -> std::result::Result<(), Box<dyn Error>> {
with_r(|| {
let m: RMatrix<Rbool> = RMatrix::new(5, 2);
unsafe { Rf_PrintValue(m.get()) };
let m: RMatrix<Rint> = RMatrix::new(5, 2);
unsafe { Rf_PrintValue(m.get()) };
let m: RMatrix<Rfloat> = RMatrix::new(5, 2);
unsafe { Rf_PrintValue(m.get()) };
let m: RMatrix<Rcplx> = RMatrix::new(5, 2);
unsafe { Rf_PrintValue(m.get()) };
rprintln!();
let m: RMatrix<Rbool> = RMatrix::new_with_na(10, 2);
assert_eq!(R!("matrix(NA, 10, 2)").unwrap(), m.into_robj());
let m: RMatrix<Rint> = RMatrix::new_with_na(10, 2);
assert_eq!(R!("matrix(NA_integer_, 10, 2)").unwrap(), m.into_robj());
let m: RMatrix<Rfloat> = RMatrix::new_with_na(10, 2);
assert_eq!(R!("matrix(NA_real_, 10, 2)").unwrap(), m.into_robj());
let m: RMatrix<Rcplx> = RMatrix::new_with_na(10, 2);
assert_eq!(R!("matrix(NA_complex_, 10, 2)").unwrap(), m.into_robj());
Ok(())
})
}
#[test]
fn matrix_ops() {
test! {
let vector = RColumn::new_column(3, |r| [1., 2., 3.][r]);
let robj = r!(vector);
assert_eq!(robj.is_vector(), true);
assert_eq!(robj.nrows(), 3);
let vector2 : RColumn<f64> = robj.as_column().ok_or("expected array")?;
assert_eq!(vector2.data().len(), 3);
assert_eq!(vector2.nrows(), 3);
let matrix = RMatrix::new_matrix(3, 2, |r, c| [
[1., 2., 3.],
[4., 5., 6.]][c][r]);
let robj = r!(matrix);
assert_eq!(robj.is_matrix(), true);
assert_eq!(robj.nrows(), 3);
assert_eq!(robj.ncols(), 2);
let matrix2 : RMatrix<f64> = robj.as_matrix().ok_or("expected matrix")?;
assert_eq!(matrix2.data().len(), 6);
assert_eq!(matrix2.nrows(), 3);
assert_eq!(matrix2.ncols(), 2);
let array = RMatrix3D::new_matrix3d(2, 2, 2, |r, c, m| [
[[1., 2.], [3., 4.]],
[[5., 6.], [7., 8.]]][m][c][r]);
let robj = r!(array);
assert_eq!(robj.is_array(), true);
assert_eq!(robj.nrows(), 2);
assert_eq!(robj.ncols(), 2);
let array2 : RMatrix3D<f64> = robj.as_matrix3d().ok_or("expected matrix3d")?;
assert_eq!(array2.data().len(), 8);
assert_eq!(array2.nrows(), 2);
assert_eq!(array2.ncols(), 2);
assert_eq!(array2.nsub(), 2);
}
}
#[test]
fn test_from_vec_doubles_to_matrix() {
test! {
let res: Vec<Doubles> = vec![
vec![17.0, 23.0, 4.0, 10.0, 11.0].try_into().unwrap(),
vec![24.0, 5.0, 6.0, 12.0, 18.0].try_into().unwrap(),
vec![1.0, 7.0, 13.0, 19.0, 25.0].try_into().unwrap(),
vec![8.0, 14.0, 20.0, 21.0, 2.0].try_into().unwrap(),
vec![15.0, 16.0, 22.0, 3.0, 9.0].try_into().unwrap(),
];
let (n_x, n_y) = (5, 5);
let _matrix = RMatrix::new_matrix(n_x, n_y, |r, c| res[c][r]);
}
}
#[test]
fn test_rmatrix_with_rstr() {
test! {
let string_data = ["a", "b", "c", "d", "e", "f"];
let matrix: RMatrix<Rstr> = RMatrix::new_matrix(2, 3, |r, c| {
let idx = c * 2 + r;
Rstr::from(string_data[idx])
});
assert_eq!(matrix.nrows(), 2);
assert_eq!(matrix.ncols(), 3);
let robj: Robj = matrix.into();
assert_eq!(robj.is_matrix(), true);
assert_eq!(robj.rtype(), Rtype::Strings);
}
}
}