boreholeio 0.1.0

use super::{
    utils::{convert, convert_vec},
    Error, TypedArray,
};
use ndarray::{ArrayViewD, IxDyn, ShapeBuilder, SliceInfoElem};
use ndarray_stats::QuantileExt;
use num::cast::AsPrimitive;
use std::{cmp::PartialOrd, default::Default, fmt::Debug};

/// `ndarray`'s view wrapper supporting a limited set of element types.
#[repr(u8)]
#[derive(Clone, Debug)]
pub enum ArrayProxy<'a> {
    I8(ArrayViewD<'a, i8>) = 0,
    U8(ArrayViewD<'a, u8>),
    I16(ArrayViewD<'a, i16>),
    U16(ArrayViewD<'a, u16>),
    I32(ArrayViewD<'a, i32>),
    U32(ArrayViewD<'a, u32>),
    I64(ArrayViewD<'a, i64>),
    U64(ArrayViewD<'a, u64>),
    F32(ArrayViewD<'a, f32>),
    F64(ArrayViewD<'a, f64>),
}

/// A wrapper for any generic function that needs to be applied to the
/// ArrayProxy without manually matching the exact value of the generic type to
/// the enum's variant.
pub trait ProxyFunctor<'a, DstT> {
    fn apply<T>(&self, view: &ArrayViewD<'a, T>) -> DstT
    where
        T: 'static + Debug + Copy + Default + PartialOrd + AsPrimitive<f64>;
}

/// A wrapper that automatically deducts the proper enum variant and creates
/// the `ArrayProxy` from the given view.
pub trait ProxyCreator<'a> {
    fn create<T: 'static>(&self, element_type: u8) -> Result<ArrayViewD<'a, T>, String>;
}

/// Enumeration of different mathematical statistics currently implemented for
/// `ArrayProxy`. They are calculated on the flattened data, that's why if one
/// needs to find, say, a coordinate-wise minimum of `N`-dimensional data, they
/// need to `slice` the data `N` times, call the `calc_stats` method on every
/// slice and manually concatenate the results into an `N`-dimensional point.
pub enum Statistics {
    /// Minimal value of the array.
    Min,
    /// Maximal value of the array.
    Max,
}

impl<'a> ArrayProxy<'a> {
    pub fn from_view<T: 'static>(view: ArrayViewD<'a, T>) -> Result<Self, Error> {
        struct Creator<'a, T> {
            view: ArrayViewD<'a, T>,
        }
        impl<'a, T: 'static> ProxyCreator<'a> for Creator<'a, T> {
            fn create<DstT: 'static>(
                &self,
                _element_type: u8,
            ) -> Result<ArrayViewD<'a, DstT>, Error> {
                reinterpret_cast(&self.view)
            }
        }
        Self::create(Creator { view })
    }

    pub fn from_mmap(
        mmap: &memmap2::Mmap,
        shape: Vec<u64>,
        strides_in_bytes: Vec<i64>,
        element_type: u8,
        data_start: u64,
    ) -> Result<Self, Error> {
        // Get the data bytes.
        let (span_min, span_max) =
            calculate_data_span(&shape, &strides_in_bytes, convert(data_start));
        let bytes = mmap[span_min..span_max].as_ptr();

        // Introduce a wrapper which converts dynamic element type to static.
        struct Creator {
            element_type: u8,
            bytes: *const u8,
            shape: Vec<u64>,
            strides_in_bytes: Vec<i64>,
        }
        impl<'a> ProxyCreator<'a> for Creator {
            fn create<DataType: 'static>(
                &self,
                element_type: u8,
            ) -> Result<ArrayViewD<'a, DataType>, Error> {
                // This method is run for all the supported types.
                if self.element_type != element_type {
                    return Err("Invalid element_type".to_string());
                }
                // Convert strides from bytes to array elements in accordance
                // with `ndarray`'s interfaces.
                let element_size = std::mem::size_of::<DataType>() as i64;
                let strides = bytes_to_elements(&self.strides_in_bytes, element_size);
                let strided_shape = IxDyn(&convert_vec(&self.shape)[..]).strides(IxDyn(&strides));
                // Reinterpret the bytes as array elements.
                unsafe {
                    Ok(ArrayViewD::from_shape_ptr(
                        strided_shape.clone(),
                        self.bytes as *const DataType,
                    ))
                }
            }
        }

        Self::create(Creator {
            element_type,
            bytes,
            shape,
            strides_in_bytes,
        })
    }

    pub fn shape(&self) -> Vec<usize> {
        struct ShapeGetter;
        impl<'a> ProxyFunctor<'a, Vec<usize>> for ShapeGetter {
            fn apply<T>(&self, view: &ArrayViewD<'a, T>) -> Vec<usize> {
                view.shape().to_vec()
            }
        }
        self.apply(ShapeGetter {})
    }

    pub fn strides(&self, in_bytes: bool) -> Vec<isize> {
        struct StridesGetter;
        impl<'a> ProxyFunctor<'a, Vec<isize>> for StridesGetter {
            fn apply<T>(&self, view: &ArrayViewD<'a, T>) -> Vec<isize> {
                view.strides().to_vec()
            }
        }
        let strides_in_elems = self.apply(StridesGetter {});
        let elem_size = if in_bytes {
            convert::<usize, isize>(self.element_size())
        } else {
            1
        };
        strides_in_elems
            .iter()
            .map(|s| s * elem_size)
            .collect::<Vec<isize>>()
    }

    pub fn dimensions(&self) -> usize {
        self.shape().len()
    }

    pub fn data(&self) -> *const u8 {
        struct BytesGetter;
        impl<'a> ProxyFunctor<'a, *const u8> for BytesGetter {
            fn apply<T>(&self, view: &ArrayViewD<'a, T>) -> *const u8 {
                view.as_ptr() as *const u8
            }
        }
        self.apply(BytesGetter {})
    }

    pub fn data_size(&self) -> usize {
        std::iter::zip(self.shape(), self.strides(true))
            .map(|(dim_shape, dim_stride_in_bytes)| {
                dim_shape * convert::<isize, usize>(dim_stride_in_bytes.abs())
            })
            .max()
            .unwrap()
    }

    pub fn element_type(&self) -> u8 {
        self.discriminant()
    }

    pub fn element_size(&self) -> usize {
        struct ElemSizeGetter;
        impl<'a> ProxyFunctor<'a, usize> for ElemSizeGetter {
            fn apply<T>(&self, _view: &ArrayViewD<'a, T>) -> usize {
                core::mem::size_of::<T>()
            }
        }
        self.apply(ElemSizeGetter {})
    }

    pub fn try_as<DstT: 'static>(&self) -> Result<ArrayViewD<'a, DstT>, Error> {
        type ResT<'a, T> = Result<ArrayViewD<'a, T>, Error>;
        struct Converter;
        impl<'a, DstT: 'static> ProxyFunctor<'a, ResT<'a, DstT>> for Converter {
            fn apply<T: 'static>(&self, view: &ArrayViewD<'a, T>) -> ResT<'a, DstT> {
                reinterpret_cast(view)
            }
        }
        self.apply(Converter {})
    }

    pub fn slice(&self, s: Vec<SliceInfoElem>) -> Result<Self, Error> {
        struct Getter {
            s: Vec<SliceInfoElem>,
        }
        type DstT<'b> = Result<ArrayProxy<'b>, Error>;
        impl<'b> ProxyFunctor<'b, DstT<'b>> for Getter {
            fn apply<T: 'static>(&self, view: &ArrayViewD<'b, T>) -> DstT<'b> {
                ArrayProxy::<'b>::from_view(view.clone().slice_move(self.s.as_slice()))
            }
        }
        self.apply(Getter { s })
    }

    /// Calculates the queried statistics.
    pub fn calc_stats(&self, s: Statistics) -> Result<f64, Error> {
        struct Calculator {
            stats: Statistics,
        }
        impl<'a> ProxyFunctor<'a, Result<f64, Error>> for Calculator {
            fn apply<T: Debug + PartialOrd + AsPrimitive<f64>>(
                &self,
                view: &ArrayViewD<'a, T>,
            ) -> Result<f64, Error> {
                let res = match self.stats {
                    Statistics::Min => view.min(),
                    Statistics::Max => view.max(),
                };
                let value = res.map_err(|e| e.to_string())?;
                Ok(value.as_())
            }
        }
        self.apply(Calculator { stats: s })
    }

    /// Returns a uniquely owned copy of the array.
    pub fn to_owned(&self) -> Result<TypedArray, Error> {
        struct Creator {}
        impl<'a> ProxyFunctor<'a, Result<TypedArray, Error>> for Creator {
            fn apply<T: 'static + Clone>(
                &self,
                view: &ArrayViewD<'a, T>,
            ) -> Result<TypedArray, Error> {
                TypedArray::from_array(view.to_owned())
            }
        }
        self.apply(Creator {})
    }

    pub fn apply<DstT, FunctorType: ProxyFunctor<'a, DstT>>(&self, functor: FunctorType) -> DstT {
        // Once a variant is added to / removed from the ArrayProxy enum, this
        // operator will generate a compilation error.
        match self {
            ArrayProxy::I8(view) => functor.apply(view),
            ArrayProxy::U8(view) => functor.apply(view),
            ArrayProxy::I16(view) => functor.apply(view),
            ArrayProxy::U16(view) => functor.apply(view),
            ArrayProxy::I32(view) => functor.apply(view),
            ArrayProxy::U32(view) => functor.apply(view),
            ArrayProxy::I64(view) => functor.apply(view),
            ArrayProxy::U64(view) => functor.apply(view),
            ArrayProxy::F32(view) => functor.apply(view),
            ArrayProxy::F64(view) => functor.apply(view),
        }
    }

    // This is one and the only place which needs to be manually synchronized
    // with the TypedArray enum variants.
    pub fn create<CreatorType: ProxyCreator<'a>>(creator: CreatorType) -> Result<Self, Error> {
        if let Ok(view) = creator.create(0) {
            return Ok(ArrayProxy::I8(view));
        }
        if let Ok(view) = creator.create(1) {
            return Ok(ArrayProxy::U8(view));
        }
        if let Ok(view) = creator.create(2) {
            return Ok(ArrayProxy::I16(view));
        }
        if let Ok(view) = creator.create(3) {
            return Ok(ArrayProxy::U16(view));
        }
        if let Ok(view) = creator.create(4) {
            return Ok(ArrayProxy::I32(view));
        }
        if let Ok(view) = creator.create(5) {
            return Ok(ArrayProxy::U32(view));
        }
        if let Ok(view) = creator.create(6) {
            return Ok(ArrayProxy::I64(view));
        }
        if let Ok(view) = creator.create(7) {
            return Ok(ArrayProxy::U64(view));
        }
        if let Ok(view) = creator.create(8) {
            return Ok(ArrayProxy::F32(view));
        }
        if let Ok(view) = creator.create(9) {
            return Ok(ArrayProxy::F64(view));
        }
        Err("Failed to create ArrayProxy from any of the supported types".to_string())
    }

    /// Returns `true` if `child`'s span is contained within the `self`'s one.
    pub(crate) fn contains(&self, child: ArrayProxy<'_>) -> Result<bool, Error> {
        struct Comparator<'c> {
            child: ArrayProxy<'c>,
        }
        impl<'f> ProxyFunctor<'f, Result<bool, Error>> for Comparator<'_> {
            fn apply<T: 'static>(&self, view: &ArrayViewD<'f, T>) -> Result<bool, Error> {
                let child = if let Ok(child) = self.child.try_as::<T>() {
                    child
                } else {
                    return Ok(false);
                };

                let (leftmost, rightmost) = view.get_span()?;
                let (child_leftmost, child_rightmost) = child.get_span()?;

                let mut is_inside = true;
                for child in [child_leftmost as usize, child_rightmost as usize] {
                    // `offset_from()` cannot be used here because it requires
                    // both pointers to originate from the same allocation. For
                    // more details please see:
                    // https://doc.rust-lang.org/std/primitive.pointer.html#safety-5
                    // https://stackoverflow.com/questions/68400492/why-must-pointers-used-by-offset-from-be-derived-from-a-pointer-to-the-same-ob
                    is_inside &= child >= leftmost as usize && child <= rightmost as usize;
                }

                Ok(is_inside)
            }
        }
        self.apply(Comparator { child })
    }

    // Each enum instance has a discriminant: an integer logically associated to
    // it that is used to determine which variant it holds.
    fn discriminant(&self) -> u8 {
        unsafe { *(self as *const Self as *const u8) }
    }
}

/// Given an array defined by its `shape`, `strides`, and starting offset
/// `data_start`, this method calculates the span of the array as `leftmost`
/// and `rightmost` offsets. Measurement units for the offsets are determined
/// by `data_start` and `strides` and can be either array elements or bytes.
/// Negative values for `strides` and `data_start` are supported.
fn calculate_data_span(shape: &[u64], strides: &[i64], data_start: i64) -> (usize, usize) {
    let mut stride_abs_min = i64::MAX;
    let mut leftmost = data_start;
    let mut rightmost = data_start;
    std::iter::zip(shape, strides).for_each(|(length, stride)| {
        let offset = convert::<u64, i64>(length - 1u64) * stride;
        if offset > 0 {
            rightmost += offset;
        } else {
            leftmost += offset;
        }
        stride_abs_min = stride_abs_min.min(stride.abs());
    });
    (convert(leftmost), convert(rightmost + stride_abs_min))
}

trait GetArrayViewSpan<T> {
    fn get_span(&self) -> Result<(*const T, *const T), Error>;
}

impl<T> GetArrayViewSpan<T> for ArrayViewD<'_, T> {
    /// Returns pointers to leftmost and rightmost array elements. Errors out
    /// if the input array is empty.
    ///
    /// This method works with a (third-party) `ArrayViewD`, whereas
    /// `calculate_data_span()` above runs earlier - before `ArrayViewD` is
    /// even created (moreover, its job is to help construct `ArrayViewD` from
    /// raw bytes).
    fn get_span(&self) -> Result<(*const T, *const T), Error> {
        if self.is_empty() {
            return Err("Cannot get span of an empty array view".into());
        }

        let first = self.first().unwrap() as *const T;
        let last = self.last().unwrap() as *const T;

        // Here it is safe to use the `offset_from` because both `first` and
        // `last` pointers are derived from the same allocation.
        if unsafe { last.offset_from(first) } > 0 {
            Ok((first, last))
        } else {
            Ok((last, first))
        }
    }
}

/// Given an array element size in bytes, this method converts the array's
/// strides from bytes to array elements.
///
/// The conversion math is a bit tricky because `ndarray` supports negative
/// strides in a funky way, for example:
///
/// https://github.com/rust-ndarray/ndarray/blob/master/tests/array.rs#L2039-L2061
///
/// let s = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
/// let mut answer = Array::from(s.to_vec())
///     .into_shape_with_order((2, 3, 2))
///     .unwrap();
/// let a = ArrayView::from_shape((2, 3, 2).strides((6, (-2isize) as usize, 1)), &s).unwrap();
/// answer.invert_axis(Axis(1));
/// assert_eq!(a, answer);
fn bytes_to_elements(strides_in_bytes: &[i64], element_size: i64) -> Vec<usize> {
    strides_in_bytes
        .iter()
        .map(|stride_in_bytes| {
            let stride_in_elements = *stride_in_bytes / element_size;
            if stride_in_elements > 0 {
                // Simple type cast.
                convert(stride_in_elements)
            } else {
                // Cast to `isize` first and to `usize` afterwards due to a
                // bizarre way of ndarray handling the negative strides:
                // https://github.com/rust-ndarray/ndarray/pull/901
                convert::<i64, isize>(stride_in_elements) as usize
            }
        })
        .collect::<Vec<usize>>()
}

fn reinterpret_cast<'a, SrcT: 'static, DstT: 'static>(
    view: &ArrayViewD<'a, SrcT>,
) -> Result<ArrayViewD<'a, DstT>, Error> {
    if std::any::TypeId::of::<SrcT>() == std::any::TypeId::of::<DstT>() {
        // As Rust does not support generics specialization, we make sure that
        // `SrcT == DstT` and perform the pointer's reinterpret cast
        let r_view = view as *const ArrayViewD<'a, SrcT> as *const ArrayViewD<'a, DstT>;
        return Ok(unsafe { (*r_view).clone() });
    }
    Err(format!(
        "Actual type {} differs from the queried one {}",
        std::any::type_name::<SrcT>(),
        std::any::type_name::<DstT>(),
    ))
}