executorch 0.9.0

use std::marker::PhantomData;
use std::pin::Pin;

use executorch_sys as sys;

use crate::memory::{Storable, Storage};
use crate::tensor::{
    DimOrderType, Scalar, ScalarType, SizesType, StridesType, TensorAccessor, TensorAccessorInner,
    TensorAccessorMut,
};
use crate::util::{Destroy, IntoCpp, IntoRust, NonTriviallyMovable, __ArrayRefImpl, c_new};
use crate::{Error, Result};

/// A raw tensor that does not own the underlying data.
///
/// This struct is a low-level match to the C++ `Tensor` class. The tensor does not own its data,
/// but rather point to a [`RawTensorImpl`]. A [`Tensor`](crate::tensor::Tensor) or any of its
/// variants (`TensorAny`, `TensorMut`, etc) is preferred for most use cases, but this struct
/// is exposed for low level users who need to avoid code size overhead (avoiding the regular
/// tensor generics).
///
/// The struct does not enforce any mutability rules, and the caller must ensure that the tensor
/// is used correctly according to its mutability.
pub struct RawTensor<'a>(
    NonTriviallyMovable<'a, sys::TensorStorage>,
    // phantom for the lifetime of the TensorImpl we depends on
    PhantomData<&'a ()>,
);
impl<'a> RawTensor<'a> {
    pub(crate) fn new_impl(tensor: NonTriviallyMovable<'a, sys::TensorStorage>) -> Self {
        Self(tensor, PhantomData)
    }

    /// Create a new tensor in a boxed heap memory.
    ///
    /// # Safety
    ///
    /// The returned tensor will allow to mutate the underlying data (which is owned by the
    /// `RawTensorImpl`), so the caller must ensure that the tensor is used correctly according to
    /// its mutability.
    #[cfg(feature = "alloc")]
    pub unsafe fn new(tensor_impl: &'a RawTensorImpl) -> Self {
        let impl_ = &tensor_impl.0 as *const sys::TensorImpl;
        let impl_ = impl_.cast_mut();
        // Safety: the closure init the pointer
        let tensor = unsafe {
            NonTriviallyMovable::new_boxed(|p: *mut sys::TensorStorage| {
                let p = sys::TensorRefMut { ptr: p as *mut _ };
                sys::executorch_Tensor_new(p, impl_)
            })
        };
        Self::new_impl(tensor)
    }

    /// Create a new tensor in the given storage.
    ///
    /// # Safety
    ///
    /// The returned tensor will allow to mutate the underlying data (which is owned by the
    /// `RawTensorImpl`), so the caller must ensure that the tensor is used correctly according to
    /// its mutability.
    pub unsafe fn new_in_storage(
        tensor_impl: &'a RawTensorImpl,
        storage: Pin<&'a mut Storage<RawTensor<'_>>>,
    ) -> Self {
        let impl_ = &tensor_impl.0 as *const sys::TensorImpl;
        let impl_ = impl_.cast_mut();
        // Safety: the closure init the pointer
        let tensor = unsafe {
            NonTriviallyMovable::new_in_storage(
                |p: *mut sys::TensorStorage| {
                    let p = sys::TensorRefMut { ptr: p as *mut _ };
                    sys::executorch_Tensor_new(p, impl_)
                },
                storage,
            )
        };
        Self::new_impl(tensor)
    }

    /// Create a new tensor from an immutable Cpp reference.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the given tensor is valid for the lifetime of the new tensor,
    /// and that the tensor is compatible with the data generic.
    /// The created tensor should not be modified as we take an immutable reference to the given
    /// Cpp tensor reference.
    pub(crate) unsafe fn from_inner_ref(tensor: sys::TensorRef) -> Self {
        debug_assert!(!tensor.ptr.is_null());
        let tensor = unsafe { &*(tensor.ptr as *const sys::TensorStorage) };
        Self::new_impl(NonTriviallyMovable::from_ref(tensor))
    }

    /// Create a new mutable tensor from a mutable Cpp reference.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the given tensor is valid for the lifetime of the new tensor.
    #[allow(unused)]
    pub(crate) unsafe fn from_inner_ref_mut(tensor: sys::TensorRefMut) -> Self {
        debug_assert!(!tensor.ptr.is_null());
        let tensor = unsafe { &mut *(tensor.ptr as *mut sys::TensorStorage) };
        Self::new_impl(NonTriviallyMovable::from_mut_ref(tensor))
    }

    /// Get the underlying Cpp tensor.
    pub(crate) fn as_cpp(&self) -> sys::TensorRef {
        sys::TensorRef {
            ptr: self.0.as_ref() as *const sys::TensorStorage as *const _,
        }
    }

    /// Get a mutable reference to the underlying Cpp tensor.
    ///
    /// # Safety
    ///
    /// The caller can not move out of the returned mut reference, and should use this function only
    /// if the tensor was created with a mutable tensor impl.
    pub(crate) unsafe fn as_cpp_mut(&mut self) -> Option<sys::TensorRefMut> {
        // Safety: the caller does not move out of the returned mut reference.
        Some(sys::TensorRefMut {
            ptr: unsafe { self.0.as_mut()? } as *mut sys::TensorStorage as *mut _,
        })
    }

    /// Returns the size of the tensor in bytes.
    ///
    /// NOTE: Only the alive space is returned not the total capacity of the
    /// underlying data blob.
    pub fn nbytes(&self) -> usize {
        unsafe { sys::executorch_Tensor_nbytes(self.as_cpp()) }
    }

    /// Returns the size of the tensor at the given dimension.
    ///
    /// NOTE: that size() intentionally does not return SizeType even though it
    /// returns an element of an array of SizeType. This is to help make calls of
    /// this method more compatible with at::Tensor, and more consistent with the
    /// rest of the methods on this class and in ETensor.
    pub fn size(&self, dim: usize) -> usize {
        unsafe { sys::executorch_Tensor_size(self.as_cpp(), dim) }
    }

    /// Returns the tensor's number of dimensions.
    pub fn dim(&self) -> usize {
        unsafe { sys::executorch_Tensor_dim(self.as_cpp()) }
    }

    /// Returns the number of elements in the tensor.
    pub fn numel(&self) -> usize {
        unsafe { sys::executorch_Tensor_numel(self.as_cpp()) }
    }

    /// Returns the type of the elements in the tensor (int32, float, bool, etc).
    pub fn scalar_type(&self) -> ScalarType {
        unsafe { sys::executorch_Tensor_scalar_type(self.as_cpp()) }.rs()
    }

    /// Returns the size in bytes of one element of the tensor.
    pub fn element_size(&self) -> usize {
        unsafe { sys::executorch_Tensor_element_size(self.as_cpp()) }
    }

    /// Returns the sizes of the tensor at each dimension.
    pub fn sizes(&self) -> &[SizesType] {
        unsafe {
            let arr = sys::executorch_Tensor_sizes(self.as_cpp());
            debug_assert!(!arr.data.is_null());
            std::slice::from_raw_parts(arr.data, arr.len)
        }
    }

    /// Returns the order the dimensions are laid out in memory.
    pub fn dim_order(&self) -> &[DimOrderType] {
        unsafe {
            let arr = sys::executorch_Tensor_dim_order(self.as_cpp());
            debug_assert!(!arr.data.is_null());
            std::slice::from_raw_parts(arr.data, arr.len)
        }
    }

    /// Returns the strides of the tensor at each dimension.
    ///
    /// Strides are in units of the elements size, not in bytes.
    pub fn strides(&self) -> &[StridesType] {
        unsafe {
            let arr = sys::executorch_Tensor_strides(self.as_cpp());
            debug_assert!(!arr.data.is_null());
            std::slice::from_raw_parts(arr.data, arr.len)
        }
    }

    /// Returns a pointer to the constant underlying data blob.
    ///
    /// # Safety
    ///
    /// The caller must access the values in the returned pointer according to the type, sizes, dim order and strides
    /// of the tensor.
    pub fn as_data_ptr(&self) -> *const () {
        let ptr = unsafe { sys::executorch_Tensor_const_data_ptr(self.as_cpp()) };
        debug_assert!(!ptr.is_null());
        ptr as *const ()
    }

    /// Returns a mutable pointer to the underlying data blob.
    ///
    /// # Returns
    ///
    /// Returns a mutable pointer to the underlying data blob, and sometimes `None` if the tensor is not mutable.
    /// `Some` may be returned in cases the tensor itself is mutable (owned or created by a mutable reference),
    /// but the underlying data is not mutable (e.g. a tensor created from a immutable tensor impl).
    /// In such cases, the caller should not call this function.
    /// This is why this function is marked as `unsafe`.
    ///
    /// # Safety
    ///
    /// The caller must access the values in the returned pointer according to the type, sizes, dim order and strides
    /// of the tensor.
    /// The caller should call this function only if the tensor was created with a mutable tensor impl.
    pub fn as_data_mut_ptr(&mut self) -> Option<*mut ()> {
        let tensor = unsafe { self.as_cpp_mut()? };
        let tensor = sys::TensorRef { ptr: tensor.ptr };
        let ptr = unsafe { sys::executorch_Tensor_mutable_data_ptr(tensor) };
        debug_assert!(!ptr.is_null());
        Some(ptr as *mut ())
    }

    fn coordinate_to_index(&self, coordinate: &[usize]) -> Option<usize> {
        let index = unsafe {
            sys::executorch_Tensor_coordinate_to_index(
                self.as_cpp(),
                sys::ArrayRefUsizeType::from_slice(coordinate),
            )
        };
        if index < 0 {
            None
        } else {
            Some(index as usize)
        }
    }
    unsafe fn coordinate_to_index_unchecked(&self, coordinate: &[usize]) -> usize {
        cfg_if::cfg_if! { if #[cfg(debug_assertions)] {
            let index = self.coordinate_to_index(coordinate);
            unsafe { index.unwrap_unchecked() }
        } else {
            let index = unsafe {
                sys::executorch_Tensor_coordinate_to_index_unchecked(
                    self.as_cpp(),
                    sys::ArrayRefUsizeType::from_slice(coordinate),
                )
            };
            index as usize
        } }
    }

    /// Get a reference to the element at `index`, without bounds checking.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the index is within bounds.
    pub unsafe fn get_unchecked<S: Scalar>(&self, index: &[usize]) -> &S {
        let index = unsafe { self.coordinate_to_index_unchecked(index) };
        let base_ptr = self.as_data_ptr() as *const S;
        debug_assert!(!base_ptr.is_null());
        unsafe { &*base_ptr.add(index) }
    }

    /// Get a mutable reference to the element at `index`, without bounds checking.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the index is within bounds.
    /// The caller must ensure that the tensor was created with a mutable tensor impl.
    pub unsafe fn get_unchecked_mut<S: Scalar>(&mut self, index: &[usize]) -> &mut S {
        let index = unsafe { self.coordinate_to_index_unchecked(index) };
        let base_ptr = unsafe { self.as_data_mut_ptr().unwrap_unchecked() } as *mut S;
        debug_assert!(!base_ptr.is_null());
        unsafe { &mut *base_ptr.add(index) }
    }

    /// Safety: the caller must ensure that type `S` is the correct scalar type of the tensor.
    pub(super) unsafe fn get_without_type_check<S: Scalar>(&self, index: &[usize]) -> Option<&S> {
        let index = self.coordinate_to_index(index)?;
        let base_ptr = self.as_data_ptr() as *const S;
        debug_assert!(!base_ptr.is_null());
        Some(unsafe { &*base_ptr.add(index) })
    }

    /// Safety: the caller must ensure that type `S` is the correct scalar type of the tensor.
    ///
    /// # Returns
    ///
    /// Returns a mutable reference to the element at `index`, or `None` if the index is out of bounds.
    /// The caller must ensure that the tensor was created with a mutable tensor impl.
    /// `Some` may be returned in cases the tensor itself is mutable (owned or created by a mutable reference),
    /// but the underlying data is not mutable (e.g. a tensor created from a immutable tensor impl).
    /// In such cases, the caller should not call this function.
    /// This is why this function is marked as `unsafe`.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the type `S` is the correct scalar type of the tensor,
    /// and that the tensor was created with a mutable tensor impl.
    pub(super) unsafe fn get_without_type_check_mut<S: Scalar>(
        &mut self,
        index: &[usize],
    ) -> Option<&mut S> {
        let index = self.coordinate_to_index(index)?;
        let base_ptr = self.as_data_mut_ptr()? as *mut S;
        debug_assert!(!base_ptr.is_null());
        Some(unsafe { &mut *base_ptr.add(index) })
    }

    /// Get a reference to the element at `index`, or `None` if the scalar type of the tensor does not
    /// match `S` or the index is out of bounds.
    pub fn get_as_typed<S: Scalar>(&self, index: &[usize]) -> Option<&S> {
        if self.scalar_type() == S::TYPE {
            // Safety: the scalar type is checked
            unsafe { self.get_without_type_check(index) }
        } else {
            None
        }
    }

    /// Get a mutable reference to the element at `index`, or `None` if the scalar type of the tensor does not
    /// match `S` or the index is out of bounds or the tensor is not mutable.
    ///
    /// `Some` may be returned in cases the tensor itself is mutable (owned or created by a mutable reference),
    /// but the underlying data is not mutable (e.g. a tensor created from a immutable tensor impl).
    ///
    /// # Safety
    ///
    /// The caller must ensure the tensor was created with a mutable tensor impl,
    pub unsafe fn get_as_typed_mut<S: Scalar>(&mut self, index: &[usize]) -> Option<&mut S> {
        if self.scalar_type() == S::TYPE {
            // Safety: the scalar type is checked
            unsafe { self.get_without_type_check_mut(index) }
        } else {
            None
        }
    }

    fn accessor_inner<S: Scalar, const N: usize>(&self) -> Option<TensorAccessorInner<'_, S, N>> {
        if self.scalar_type() != S::TYPE || self.dim() != N {
            return None;
        }
        if !self.dim_order().iter().map(|d| *d as usize).eq(0..N) {
            panic!("Non-default dim order is not supported for TensorAccessorInner");
        }
        let data = self.as_data_ptr() as *const S;
        let accessor = unsafe { TensorAccessorInner::new(data, self.sizes(), self.strides()) };
        Some(accessor)
    }

    /// Get an immutable accessor for the tensor.
    ///
    /// An accessor is a utility struct, templated over the type of the tensor elements and the number
    /// of dimensions, which make it very efficient to access tensor elements by index.
    /// See the [`TensorAccessor`] for more details.
    ///
    /// # Returns
    ///
    /// Returns an accessor if the scalar type of the tensor matches `S` and the number of dimensions
    /// matches `N`, otherwise returns `None`.
    pub fn accessor<S: Scalar, const N: usize>(&self) -> Option<TensorAccessor<'_, S, N>> {
        Some(TensorAccessor(self.accessor_inner()?))
    }

    /// Get a mutable accessor for the tensor.
    ///
    /// An accessor is a utility struct, templated over the type of the tensor elements and the number
    /// of dimensions, which make it very efficient to access tensor elements by index.
    /// See the [`TensorAccessorMut`] for more details.
    ///
    /// # Returns
    ///
    /// Returns an accessor if the scalar type of the tensor matches `S` and the number of dimensions
    /// matches `N`, otherwise returns `None`.
    ///
    /// # Safety
    ///
    /// The caller must ensure that the tensor was created with a mutable tensor impl.
    pub unsafe fn accessor_mut<S: Scalar, const N: usize>(
        &mut self,
    ) -> Option<TensorAccessorMut<'_, S, N>> {
        Some(TensorAccessorMut(self.accessor_inner()?))
    }
}
impl Destroy for sys::TensorStorage {
    unsafe fn destroy(&mut self) {
        unsafe {
            sys::executorch_Tensor_destructor(sys::TensorRefMut {
                ptr: self as *mut Self as *mut _,
            })
        }
    }
}
impl Storable for RawTensor<'_> {
    type __Storage = sys::TensorStorage;
}

/// A raw tensor implementation.
///
/// This struct is a low-level match to the C++ `TensorImpl` class.
/// A [`TensorImpl`](crate::tensor::TensorImpl) or any of its
/// variants (`TensorImplAny`, `TensorImplMut`, etc) is preferred for most uses, but this struct
/// is exposed for low level users who need to avoid code size overhead (avoiding the regular
/// tensor generics).
/// The struct does not enforce any mutability rules, and the caller must ensure that the tensor
/// is used correctly according to its mutability.
pub struct RawTensorImpl<'a>(sys::TensorImpl, PhantomData<&'a ()>);
impl<'a> RawTensorImpl<'a> {
    /// Create a new TensorImpl from a pointer to the data.
    ///
    /// # Arguments
    ///
    /// * `sizes` - The shape of the tensor.
    /// * `data` - A pointer to the data buffer.
    /// * `dim_order` - The order of the dimensions of the tensor, must have the same length as `sizes`.
    /// * `strides` - The strides of the tensor, in units of elements (not bytes), must have the same length as `sizes`.
    ///
    /// # Errors
    ///
    /// Returns an error if dim order is invalid, or if it doesn't match the strides, or if the strides are not dense,
    /// i.e. if the strides are not the standard layout strides of some permutation of the sizes.
    ///
    /// # Panics
    ///
    /// If the sizes, dim_order or strides slices are of different lengths.
    ///
    /// # Safety
    ///
    /// The caller must ensure elements in the data can be safely accessed according to the scalar type, sizes,
    /// dim order and strides of the tensor.
    /// The caller must ensure that the data is valid for the lifetime of the TensorImpl.
    pub unsafe fn from_ptr<S: Scalar>(
        sizes: &'a [SizesType],
        data: *mut S,
        dim_order: &'a [DimOrderType],
        strides: &'a [StridesType],
    ) -> Result<Self> {
        Self::from_ptr_impl(sizes, data, None, dim_order, strides, false)
    }

    pub(crate) unsafe fn from_ptr_impl<S: Scalar>(
        sizes: &'a [SizesType],
        data: *mut S,
        data_len: Option<usize>,
        dim_order: &'a [DimOrderType],
        strides: &'a [StridesType],
        mutable: bool,
    ) -> Result<Self> {
        Self::check_shape_strides(sizes, data_len, dim_order, strides, mutable)?;

        let dim = sizes.len();
        debug_assert_eq!(dim, dim_order.len());
        debug_assert_eq!(dim, strides.len());

        let sizes = sizes.as_ptr();
        let dim_order = dim_order.as_ptr();
        let strides = strides.as_ptr();
        debug_assert!(!data.is_null());

        let valid_strides = unsafe {
            sys::executorch_is_valid_dim_order_and_strides(dim, sizes, dim_order, strides)
        };
        if !valid_strides {
            crate::log::error!("Invalid strides");
            return Err(Error::InvalidArgument);
        }

        // Safety: sys::executorch_TensorImpl_new writes to the pointer.
        let impl_ = unsafe {
            c_new(|this| {
                sys::executorch_TensorImpl_new(
                    this,
                    S::TYPE.cpp(),
                    dim,
                    sizes as *mut SizesType,
                    data as *mut _,
                    dim_order as *mut DimOrderType,
                    strides as *mut StridesType,
                    sys::TensorShapeDynamism::TensorShapeDynamism_STATIC,
                )
            })
        };
        Ok(Self(impl_, PhantomData))
    }

    fn check_shape_strides(
        sizes: &[SizesType],
        data_len: Option<usize>,
        dim_order: &[DimOrderType],
        strides: &[StridesType],
        mutable: bool,
    ) -> Result<()> {
        enum TensorError {
            ShapeStridesDimOrderLenNotEq,
            Overflow,
            OutOfBounds,
            MultipleMutReferences,
            TooManyDimensions,
        }

        fn check_shape_strides_impl(
            sizes: &[SizesType],
            data_len: Option<usize>,
            dim_order: &[DimOrderType],
            strides: &[StridesType],
            mutable: bool,
        ) -> Result<(), TensorError> {
            // This function is based on code from the `ndarray` crate.

            if sizes.len() != strides.len() || sizes.len() != dim_order.len() {
                return Err(TensorError::ShapeStridesDimOrderLenNotEq);
            }

            // The product of non-zero axis lengths must not exceed `isize::MAX`.
            let size_nonzero = sizes
                .iter()
                .filter(|&&d| d != 0)
                .try_fold(1usize, |acc, &d| acc.checked_mul(d as usize))
                .ok_or(TensorError::Overflow)?;
            if size_nonzero > isize::MAX as usize {
                return Err(TensorError::Overflow);
            }

            // The absolute difference between least and greatest address accessible by moving along all axes.
            let max_offset: usize = sizes
                .iter()
                .zip(strides.iter())
                .try_fold(0usize, |acc, (&d, &s)| {
                    // Calculate maximum possible absolute movement along this axis.
                    let off = (d as usize)
                        .saturating_sub(1)
                        .checked_mul((s as isize).unsigned_abs())?;
                    acc.checked_add(off)
                })
                .ok_or(TensorError::Overflow)?;
            if max_offset > isize::MAX as usize {
                return Err(TensorError::Overflow);
            }

            // If the array will be empty (any axes are zero-length), the difference
            // between the least address and greatest address accessible by moving
            // along all axes must be ≤ `data.len()`. (It's fine in this case to move
            // one byte past the end of the slice since the pointers will be offset but
            // never dereferenced.)
            //
            // If the array will not be empty, the difference between the least address
            // and greatest address accessible by moving along all axes must be <
            // `data.len()`. This and #3 ensure that all dereferenceable pointers point
            // to elements within the slice.
            let is_empty = sizes.contains(&0);
            if let Some(data_len) = data_len {
                if is_empty && max_offset > data_len {
                    return Err(TensorError::OutOfBounds);
                }
                if !is_empty && max_offset >= data_len {
                    return Err(TensorError::OutOfBounds);
                }
            }

            // The strides must not allow any element to be referenced by two different indices.
            if !is_empty && mutable {
                // The axis ordering corresponding to the fastest variation (in ascending order).
                // Assumes that no stride value appears twice.
                const MAX_DIMS: usize = 16;
                if sizes.len() > MAX_DIMS {
                    // This is just a sanity check, the C++ code also limits the number of dimensions.
                    return Err(TensorError::TooManyDimensions);
                }
                let mut order_buf = [(0usize, 0isize); MAX_DIMS];
                let order_buf = &mut order_buf[..sizes.len()];
                for (i, stride) in strides.iter().enumerate() {
                    order_buf[i] = (i, *stride as isize);
                }
                order_buf.sort_by_key(|&(_, stride)| stride.abs());
                let order = order_buf.iter().take(sizes.len()).map(|&(i, _)| i);

                // There is overlap if, when iterating through the dimensions in order of
                // increasing stride, the current stride is less than or equal to the maximum
                // possible offset along the preceding axes. (Axes of length ≤1 are ignored.)
                let mut sum_prev_offsets = 0;
                for index in order {
                    match sizes[index] {
                        0 => unreachable!(), // !is_empty
                        1 => {}
                        d => {
                            let s = strides[index].unsigned_abs() as usize;
                            if s <= sum_prev_offsets {
                                return Err(TensorError::MultipleMutReferences);
                            }
                            sum_prev_offsets += (d as usize - 1) * s;
                        }
                    }
                }
            }

            Ok(())
        }

        if let Err(err) = check_shape_strides_impl(sizes, data_len, dim_order, strides, mutable) {
            let err_msg = match err {
                TensorError::ShapeStridesDimOrderLenNotEq => {
                                "Sizes, strides, and dim_order must have the same length"
                            },
                TensorError::Overflow => "shape product or strides max (abs) offset overflowed",
                TensorError::OutOfBounds => "shape and strides lead to out-of-bounds accesses",
                TensorError::MultipleMutReferences => {
                                "shape and strides allow to reference the same element with two different indices, invalid for mutable arrays"
                            }
                TensorError::TooManyDimensions => "too many dimensions, maximum supported is 16",
            };
            crate::log::error!("{err_msg}");
            return Err(Error::InvalidArgument);
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::tensor::{RawTensor, RawTensorImpl};
    use crate::tests::{check_send, check_sync};

    #[test]
    fn raw_tensor_send() {
        check_send::<RawTensor<'_>>();
    }
    #[test]
    fn raw_tensor_sync() {
        check_sync::<RawTensor<'_>>();
    }

    #[test]
    fn raw_tensor_impl_send() {
        check_send::<RawTensorImpl<'_>>();
    }
    #[test]
    fn raw_tensor_impl_sync() {
        check_sync::<RawTensorImpl<'_>>();
    }
}