rlst 0.10.1

//! Basic traits for elements of vector spaces.

use crate::{
    Array, MutableArrayImpl, RawAccessMut, base_types::MemoryLayout, dense::array::StaticArray,
};

use super::UnsafeRandom1DAccessByValue;

/// Define a basic item type associated with an object.
pub trait BaseItem {
    /// Item type.
    type Item: Copy + Default;
}

/// Allow shared access guarded by runtime borrow or lock checking.
///
/// Implementations choose the guard type: for example, shared arrays use either a `RefCell`
/// borrow guard or an `RwLock` result.
pub trait BorrowAtRuntime {
    /// The result of the borrow
    type Target<'a>
    where
        Self: 'a;

    /// Return a read-only borrow.
    fn borrow(&self) -> Self::Target<'_>;
}

/// Allow exclusive access guarded by runtime borrow or lock checking.
pub trait BorrowAtRuntimeMut {
    /// The mutable guard or lock-result type.
    type Target<'a>
    where
        Self: 'a;

    /// Return a read-write borrow.
    fn borrow_mut(&self) -> Self::Target<'_>;
}

/// Associate a shape with a given object.
pub trait Shape<const NDIM: usize> {
    /// Return the shape of the object.
    fn shape(&self) -> [usize; NDIM];

    /// Return the length, which for n dimension is the product of the dimensions.
    #[inline(always)]
    fn len(&self) -> usize {
        self.shape().iter().product()
    }

    /// Return true if the array is empty.
    ///
    /// An array is empty if at least one dimension is zero.
    #[inline(always)]
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Associate a stride with a given object.
///
/// A stride describes the memory layout of an array.
/// If the stride is e.g. `[1, 5]`, this means
/// moving from one to the next element in the first
/// dimension jumps one memory entry and in the second
/// dimension jumps 5 memory entries.
pub trait Stride<const NDIM: usize> {
    /// Return the stride of the object.
    fn stride(&self) -> [usize; NDIM];

    /// Return the memory layout.
    ///
    /// A column-major layout has entries continuous in memory
    /// starting with the first outer dimension. A row-major layout
    /// has entries continuous in memory starting with the last outer dimension.
    /// 1d contiguous layouts are always returned as column major.
    /// Possible return values are [MemoryLayout::ColumnMajor], [MemoryLayout::RowMajor],
    /// and [MemoryLayout::Unknown].
    fn memory_layout(&self) -> MemoryLayout
    where
        Self: Shape<NDIM>,
    {
        let stride = self.stride();
        let shape = self.shape();
        let mut factor = 1;
        if stride[0] == 1 {
            // Contiguous 1d arrays are always column-major.
            if NDIM == 1 {
                return MemoryLayout::ColumnMajor;
            }
            // Check if strides are column-major
            for i in 1..NDIM {
                factor *= shape[i - 1];
                if stride[i] != factor {
                    return MemoryLayout::Unknown;
                }
            }
            MemoryLayout::ColumnMajor
        } else if stride[NDIM - 1] == 1 {
            // Check if strides are row-major
            for i in 2..NDIM + 1 {
                factor *= shape[NDIM - i + 1];
                if stride[NDIM - i] != factor {
                    return MemoryLayout::Unknown;
                }
            }
            MemoryLayout::RowMajor
        } else {
            MemoryLayout::Unknown
        }
    }

    /// Return true if the array is contiguous in memory.
    ///
    /// This is the case if the data is either row-major or column-major.
    /// Other gapless permutations are reported as non-contiguous.
    fn is_contiguous(&self) -> bool
    where
        Self: Shape<NDIM>,
    {
        self.memory_layout() != MemoryLayout::Unknown
    }
}

/// Number of elements.
pub trait NumberOfElements {
    /// Return the number of elements.
    fn number_of_elements(&self) -> usize;
}

/// Resize in place.
pub trait ResizeInPlace<const NDIM: usize> {
    /// Resize an object in place.
    fn resize_in_place(&mut self, shape: [usize; NDIM]);
}

/// Fill an object with values from another object.
///
/// Individual implementations define how incompatible shapes are handled. RLST dense arrays
/// require equal shapes and panic otherwise.
pub trait FillFrom<Other> {
    /// Fill `self` with values from `other`.
    fn fill_from(&mut self, other: &Other);
}

/// Fill an array with values from an iterator.
///
/// For RLST dense arrays, values are consumed in logical column-major order until either side is
/// exhausted; a short iterator leaves trailing elements unchanged.
pub trait FillFromIter<Iter: Iterator> {
    /// Fill `self` with values from `iter`.
    fn fill_from_iter(&mut self, iter: Iter);
}

/// Convert a scalar into a one-element array.
pub trait ScalarAsArray: Sized + Copy + Default {
    /// Convert scalar value into a 1-element array of given dimension `NDIM`.
    fn as_array<const NDIM: usize>(&self) -> StaticArray<Self, NDIM, 1>;
}

/// Fill from another object and resize if necessary.
pub trait FillFromResize<Other> {
    /// Fill `self` with values from `other` and resize if necessary.
    fn fill_from_resize(&mut self, other: &Other);
}

/// Sum into current object from another object.
pub trait SumFrom<Other> {
    /// Sum values from `other` into `self`.
    fn sum_from(&mut self, other: &Other);
}

/// Componentwise Multiply other array into this array.
pub trait CmpMulFrom<Other> {
    /// Componentwise multiply the elements of `self` with `other` and store in `self`.
    fn cmp_mul_from(&mut self, other: &Other);
}

/// Multiply with a scalar.
pub trait ScalarMul<Scalar> {
    /// Output of multiplication with a scalar.
    type Output;
    /// Multiply `self` with `scalar`.
    fn scalar_mul(self, scalar: Scalar) -> Self::Output;
}

/// Componentwise form `Self = Self * Other1 + Other2`.
pub trait CmpMulAddFrom<Other1, Other2> {
    /// Componentwise form `Self = Self * Other1 + Other2`.
    fn cmp_mul_add_from(&mut self, other1: &Other1, other2: &Other2);
}

/// Fill an object with a specific value.
pub trait FillWithValue: BaseItem {
    /// Fill `self` with `value`.
    fn fill_with_value(&mut self, value: Self::Item);
}

/// Set elements current object to zero.
pub trait SetZero {
    /// Set all elements of `self` to `zero`.
    fn set_zero(&mut self);
}

/// Set elements of current object to one.
pub trait SetOne {
    /// Set all elements of `self` to one.
    fn set_one(&mut self);
}

/// Set the current object to the identity element.
pub trait SetIdentity {
    /// Set `self` to be the identity.
    fn set_identity(&mut self);
}

/// Scale in place.
pub trait ScaleInPlace: BaseItem {
    /// Scale `self` by `alpha`.
    fn scale_inplace(&mut self, alpha: Self::Item);
}

/// Compute the trace of an object.
pub trait Trace: BaseItem {
    /// Return the trace of `self`.
    fn trace(&self) -> Self::Item;
}

/// Sum all elements.
pub trait Sum: BaseItem {
    /// Compute the sum of all elements of `self`.
    fn sum(&self) -> Self::Item;
}

/// Compute the length of an object.
///
/// For multi-dimensional array the length is the product of the dimensions.
pub trait Len {
    /// Return the length of `self`.
    fn len(&self) -> usize;

    /// Return true if `self` is empty.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Return the conjugate object.
///
/// For an array this is the array of conjugate entries.
pub trait ConjObject {
    /// The output type.
    type Output;

    /// Return the conjugation of `self`.
    fn conj(self) -> Self::Output;
}

/// Evaluate an array or lazy expression into a newly created column-major array.
///
/// Evaluation never aliases the source. The source's container type controls whether the output
/// uses dynamic storage or a statically sized container.
pub trait EvaluateArray<Item, const NDIM: usize>
where
    Item: Copy + Default,
{
    /// The output implementation type of the evaluated object.
    type OutputImpl: MutableArrayImpl<Item, NDIM> + RawAccessMut<Item = Item> + Stride<NDIM>;

    /// Evaluate the object.
    fn eval(&self) -> Array<Self::OutputImpl, NDIM>;
}

/// Evaluate an array or lazy expression into a newly created row-major array.
///
/// Logical values are preserved while physical storage order changes.
pub trait EvaluateRowMajorArray<Item, const NDIM: usize>
where
    Item: Copy + Default,
{
    /// The output implementation type of the evaluated array.
    type OutputImpl: MutableArrayImpl<Item, NDIM> + RawAccessMut<Item = Item> + Stride<NDIM>;

    /// Evaluate the array into a new row-major array.
    fn eval_row_major(&self) -> Array<Self::OutputImpl, NDIM>;
}

/// Dispatch the evaluation of an array to an actual implementation.
pub trait DispatchEval<const NDIM: usize> {
    /// The output type.
    type Output;

    /// The implementation type.
    type ArrayImpl: Shape<NDIM> + UnsafeRandom1DAccessByValue;

    /// Dispatch the evaluation of the array to an actual implementation.
    fn dispatch(&self, arr: &Array<Self::ArrayImpl, NDIM>) -> Self::Output;
}

/// Dispatch the row-major evaluation of an array to an actual implementation.
pub trait DispatchEvalRowMajor<const NDIM: usize> {
    /// The output type of the evaluated array.
    type Output;

    /// The implementation type of the array.
    type ArrayImpl: Shape<NDIM> + UnsafeRandom1DAccessByValue;

    /// Dispatch the evaluation of the array to an actual implementation.
    fn dispatch(&self, arr: &Array<Self::ArrayImpl, NDIM>) -> Self::Output;
}

/// Convert to a new type.
///
/// This trait is used to convert an object of one type into an object of another type.
pub trait ToType<T> {
    /// The element type of the array.
    type Item;
    /// The output type of the array.
    type Output;

    /// Convert the array into a new array.
    fn into_type(self) -> Self::Output
    where
        Self::Item: Into<T>;
}

/// Fill an array with equidistant values.
pub trait Linspace<Item> {
    /// Fill an array with equally spaced points, including first value `a` and last value `b`.
    ///
    /// # Panics
    ///
    /// Panics if `b <= a` or the array has fewer than two elements.
    fn linspace(&mut self, a: Item, b: Item);
}