npsimd 0.3.0

use core::fmt;
use core::mem::size_of;

use super::*;

#[doc(inline)]
pub use crate::intel_vector as vector;

/// A SIMD vector.
#[repr(transparent)]
pub struct Vector<T: Array, G = EmptyGroup, FS = ()>
where FS: Features<G> {
    /// The underlying primitive vector type.
    pub(crate) primitive: T::Primitive,

    /// A set of features known to be available.
    pub(crate) features: FeatureSet<G, FS>,
}

/// Construct a compile-time vector constant.
///
/// This macro provides an array construction syntax for initializing arbitrary
/// vector types.  It can only be used at compile-time because constructing a
/// vector from arbitrary runtime values is not efficient.  The actual vector
/// type is left implicit; it must be indicated by the surrounding context.
///
/// # Examples
///
/// ```
/// use npsimd::intel::{Vector, vector};
///
/// // 'vector!' can be used in constants.
/// const ALL_ONES: Vector<[u8; 16]> = vector![0xFF; 16];
///
/// // The macro also works for regular `let` bindings.
/// let iota: Vector<[u32; 4]> = vector![0, 1, 2, 3];
///
/// // The values can be `const` expressions.
/// let complex: Vector<[u64; 2]> = vector![0xDEAD + 0xBEEF, 0xABCD * 2];
///
/// // The type can be inferred from the element expressions.
/// let iota16 = vector![0u16, 1, 2, 3, 4, 5, 6, 7];
/// ```
#[doc(hidden)]
#[macro_export]
macro_rules! intel_vector {
    [$($elem:expr),* $(,)?] => {
        const {
            const fn transmute<E: $crate::intel::Element<LEN>, const LEN: usize>
                    (array: [E; LEN]) -> $crate::intel::Vector<[E; LEN]> {
                // `transmute()` doesn't work because we can't show the compiler
                // that `V` and `V::Array` are always of the same size, even if
                // it's guaranteed by the safety requirements for `Vector`.
                let p = unsafe { $crate::__core::mem::transmute_copy(&array) };
                $crate::intel::Vector::from_primitive(p)
            }

            transmute([$(const { $elem }),*])
        }
    };

    [$elem:expr; $size:expr] => {
        const {
            const fn transmute<E, const LEN: usize>(elem: E)
                -> $crate::intel::Vector<[E; LEN]>
            where E: $crate::intel::Element<LEN> {
                let array = [elem; LEN];

                // `transmute()` doesn't work because we can't show the compiler
                // that `V` and `V::Array` are always of the same size, even if
                // it's guaranteed by the safety requirements for `Vector`.
                let p = unsafe { $crate::__core::mem::transmute_copy(&array) };
                $crate::intel::Vector::from_primitive(p)
            }

            transmute::<_, $size>($elem)
        }
    };
}

impl<E: Element<LEN>, const LEN: usize> Vector<[E; LEN]> {
    /// Construct a new [`Vector`] from the given primitive.
    pub const fn from_primitive(primitive: E::Primitive) -> Self {
        Self { primitive, features: FeatureSet::new(()) }
    }
}

impl<E, G, FS, const LEN: usize> Vector<[E; LEN], G, FS>
where E: Element<LEN>, FS: Features<G> {
    /// The feature set of this vector.
    pub const fn features(self) -> FeatureSet<G, FS> {
        self.features
    }

    /// Replace the feature set of this vector with a different one.
    pub const fn with_features<NG, NFS: Features<NG>>
            (self, features: FeatureSet<NG, NFS>)
            -> Vector<[E; LEN], NG, NFS> {
        Vector { primitive: self.primitive, features }
    }

    /// The primitive vector underlying this.
    pub const fn into_primitive(self) -> E::Primitive {
        self.primitive
    }

    /// Convert this into the underlying array.
    ///
    /// This function should only be used for debugging purposes: the cost of
    /// this conversion is unspecified and will likely be high.
    pub const fn as_array(&self) -> &[E; LEN] {
        unsafe { &*(&self.primitive as *const _ as *const [E; LEN]) }
    }

    /// Convert a slice of vectors into one of elements.
    pub const fn slice_as_array(this: &[Self]) -> &[E] {
        unsafe {
            core::slice::from_raw_parts(
                this.as_ptr() as *const E, this.len() * LEN)
        }
    }

    /// Truncate this vector to a shorter one.
    ///
    /// Truncation is a zero-cost operation.  The same vector register may be
    /// used for the value, with the truncated elements left undefined.
    #[inline]
    pub const fn truncate<const NLEN: usize>(self) -> Vector<[E; NLEN], G, FS>
    where E: Element<NLEN> {
        const_assert!(NLEN <= LEN);
        // TODO: Verify this does not cause a load/store.
        use core::mem::transmute_copy;
        Vector::from_primitive(unsafe {
            // SAFETY:
            // - We know that 'NLEN <= LEN'.
            // - Thus 'size_of::<[E; NLEN]>() <= size_of::<[E; LEN]()'.
            // - Since the element type is identical, transmutation cannot
            //   result in an invalid element.
            transmute_copy(&self.primitive)
        }).with_features(self.features)
    }

    /// Cast the type of this vector to another of the same size.
    ///
    /// Casting is a zero-cost operation.
    #[inline]
    pub const fn cast<NE, const NLEN: usize>(self) -> Vector<[NE; NLEN], G, FS>
    where NE: Element<NLEN> {
        const_assert!(size_of::<[E; LEN]>() == size_of::<[NE; NLEN]>());
        // TODO: Verify this does not cause a load/store.
        use core::mem::transmute_copy;
        Vector::from_primitive(unsafe {
            // SAFETY:
            // - We know 'size_of::<[E; LEN]>() == size_of::<[NE; NLEN]>()'.
            // - Since the element type is identical, transmutation cannot
            //   result in an invalid element.
            transmute_copy(&self.primitive)
        }).with_features(self.features)
    }
    
}

impl<E, G, FS, const LEN: usize> Clone for Vector<[E; LEN], G, FS>
where E: Element<LEN>, FS: Features<G> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<E, G, FS, const LEN: usize> Copy for Vector<[E; LEN], G, FS>
where E: Element<LEN>, FS: Features<G> {}

impl<E, G, FS, const LEN: usize> fmt::Debug for Vector<[E; LEN], G, FS>
where E: Element<LEN> + fmt::Debug, FS: Features<G> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.as_array().fmt(f)
    }
}

#[cfg(feature = "proptest")]
mod props {
    use super::*;

    use proptest::arbitrary::Arbitrary;
    use proptest::strategy::{*, Strategy};
    use proptest::test_runner::TestRunner;

    impl<E, const LEN: usize> Arbitrary for Vector<[E; LEN]>
    where E: Element<LEN> + Arbitrary {
        type Parameters = ();
        type Strategy = VectorStrategy<<[E; LEN] as Arbitrary>::Strategy>;

        fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
            VectorStrategy(<[E; LEN]>::arbitrary())
        }
    }

    #[derive(Clone, Debug)]
    pub struct VectorStrategy<S>(S);

    impl<S, E, const LEN: usize> Strategy for VectorStrategy<S>
    where S: Strategy<Value = [E; LEN]>, E: Element<LEN> + fmt::Debug {
        type Tree = VectorStrategy<S::Tree>;
        type Value = Vector<[E; LEN]>;

        fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
            self.0.new_tree(runner).map(VectorStrategy)
        }
    }

    impl<S, E, const LEN: usize> ValueTree for VectorStrategy<S>
    where S: ValueTree<Value = [E; LEN]>, E: Element<LEN> + fmt::Debug {
        type Value = Vector<[E; LEN]>;

        fn current(&self) -> Self::Value {
            let array = self.0.current();
            Vector::from_primitive(unsafe {
                core::mem::transmute_copy(&array)
            })
        }

        fn simplify(&mut self) -> bool {
            self.0.simplify()
        }

        fn complicate(&mut self) -> bool {
            self.0.complicate()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn with_features() {
        use feats::{Features, Feature, RuntimeSupport};

        #[derive(Copy, Clone)]
        struct Feat;
        struct FeatGroup;

        impl Features<FeatGroup> for Feat {
            fn query(_: &RuntimeSupport) -> Option<Self> {
                Some(Self)
            }
        }

        impl Feature<FeatGroup> for Feat {}

        let vector = vector![0u8; 16];
        let _: Vector<_, FeatGroup, (Feat, ())>
            = vector.with_features(FeatureSet::new((Feat, ())));
    }

    #[test]
    fn truncate() {
        let vector = vector![
            0u8, 7, 1, 15, 4, 2, 8, 11, 5, 0, 6, 14, 12, 13, 3, 9
        ];
        let shorter: Vector<[u8; 8]> = vector.truncate();
        assert_eq!(shorter.as_array(), &[0u8, 7, 1, 15, 4, 2, 8, 11]);
    }

    #[test]
    fn cast() {
        let vector = vector![
            0u8, 7, 1, 15, 4, 2, 8, 11, 5, 0, 6, 14, 12, 13, 3, 9
        ];
        let casted: Vector<[u16; 8]> = vector.cast();
        assert_eq!(casted.as_array(), &[
            1792u16, 3841, 516, 2824, 5, 3590, 3340, 2307
        ]);
    }
}