// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{
    arch::x86_64::*,
    ops::{Add, AddAssign, BitAnd, BitOr, BitOrAssign, BitXor, Div, Mul, MulAssign, Neg, Not, Sub},
    ptr,
};

#[derive(Clone, Copy, Debug)]
pub struct m8x16(__m128i);

impl m8x16 {
    pub fn all(self) -> bool {
        unsafe { _mm_movemask_epi8(_mm_cmpeq_epi8(self.0, _mm_setzero_si128())) == 0 }
    }
}

impl Default for m8x16 {
    fn default() -> Self {
        Self(unsafe { _mm_setzero_si128() })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct m32x4(__m128i);

#[derive(Clone, Copy, Debug)]
pub struct m32x8(__m256i);

impl m32x8 {
    pub fn all(self) -> bool {
        unsafe { _mm256_movemask_epi8(_mm256_cmpeq_epi32(self.0, _mm256_setzero_si256())) == 0 }
    }

    pub fn any(self) -> bool {
        unsafe { _mm256_movemask_epi8(_mm256_cmpeq_epi32(self.0, _mm256_setzero_si256())) != -1 }
    }
}

impl Not for m32x8 {
    type Output = Self;

    fn not(self) -> Self::Output {
        Self(unsafe { _mm256_xor_si256(self.0, _mm256_cmpeq_epi32(self.0, self.0)) })
    }
}

impl BitOr for m32x8 {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_or_si256(self.0, rhs.0) })
    }
}

impl BitOrAssign for m32x8 {
    fn bitor_assign(&mut self, rhs: Self) {
        *self = *self | rhs;
    }
}

impl BitXor for m32x8 {
    type Output = Self;

    fn bitxor(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_xor_si256(self.0, rhs.0) })
    }
}

impl Default for m32x8 {
    fn default() -> Self {
        Self(unsafe { _mm256_setzero_si256() })
    }
}

#[allow(dead_code)] // TODO(https://fxbug.dev/318827209)
#[derive(Clone, Copy, Debug)]
pub struct u8x32(__m256i);

impl u8x32 {
    pub fn splat(val: u8) -> Self {
        Self(unsafe { _mm256_set1_epi8(i8::from_ne_bytes(val.to_ne_bytes())) })
    }

    pub fn from_u32_interleaved(vals: [u32x8; 4]) -> Self {
        unsafe {
            let mask = _mm256_set1_epi32(0xFF);

            let bytes = _mm256_packus_epi16(
                _mm256_packus_epi32(
                    _mm256_and_si256(vals[0].0, mask),
                    _mm256_and_si256(vals[1].0, mask),
                ),
                _mm256_packus_epi32(
                    _mm256_and_si256(vals[2].0, mask),
                    _mm256_and_si256(vals[3].0, mask),
                ),
            );

            let shuffle = _mm256_set_epi8(
                15, 11, 7, 3, 14, 10, 6, 2, 13, 9, 5, 1, 12, 8, 4, 0, 15, 11, 7, 3, 14, 10, 6, 2,
                13, 9, 5, 1, 12, 8, 4, 0,
            );

            Self(_mm256_shuffle_epi8(bytes, shuffle))
        }
    }
}

impl Default for u8x32 {
    fn default() -> Self {
        Self(unsafe { _mm256_setzero_si256() })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct i8x16(__m128i);

impl i8x16 {
    #[cfg(test)]
    pub fn as_mut_array(&mut self) -> &mut [i8; 16] {
        unsafe { std::mem::transmute(&mut self.0) }
    }

    pub fn splat(val: i8) -> Self {
        Self(unsafe { _mm_set1_epi8(val) })
    }

    pub fn eq(self, other: Self) -> m8x16 {
        m8x16(unsafe { _mm_cmpeq_epi8(self.0, other.0) })
    }

    pub fn abs(self) -> Self {
        Self(unsafe { _mm_abs_epi8(self.0) })
    }
}

impl Default for i8x16 {
    fn default() -> Self {
        Self(unsafe { _mm_setzero_si128() })
    }
}

impl Add for i8x16 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm_add_epi8(self.0, rhs.0) })
    }
}

impl AddAssign for i8x16 {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl BitAnd for i8x16 {
    type Output = Self;

    fn bitand(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm_and_si128(self.0, rhs.0) })
    }
}

impl From<i8x16> for [i32x8; 2] {
    fn from(val: i8x16) -> Self {
        unsafe {
            let _i8x16_lo = _mm_unpacklo_epi64(val.0, _mm_setzero_si128());
            let _i8x16_hi = _mm_unpackhi_epi64(val.0, _mm_setzero_si128());

            [i32x8(_mm256_cvtepi8_epi32(_i8x16_lo)), i32x8(_mm256_cvtepi8_epi32(_i8x16_hi))]
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct i16x16(__m256i);

impl i16x16 {
    pub fn splat(val: i16) -> Self {
        Self(unsafe { _mm256_set1_epi16(val) })
    }
}

impl Default for i16x16 {
    fn default() -> Self {
        Self(unsafe { _mm256_setzero_si256() })
    }
}

impl From<i16x16> for [i32x8; 2] {
    fn from(val: i16x16) -> Self {
        unsafe {
            let mut _i16x8_lo = _mm_undefined_si128();
            let mut _i16x8_hi = _mm_undefined_si128();

            _mm256_storeu2_m128i(ptr::addr_of_mut!(_i16x8_hi), ptr::addr_of_mut!(_i16x8_lo), val.0);

            [i32x8(_mm256_cvtepi16_epi32(_i16x8_lo)), i32x8(_mm256_cvtepi16_epi32(_i16x8_hi))]
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct i32x8(__m256i);

impl i32x8 {
    pub fn splat(val: i32) -> Self {
        Self(unsafe { _mm256_set1_epi32(val) })
    }

    pub fn eq(self, other: Self) -> m32x8 {
        m32x8(unsafe { _mm256_cmpeq_epi32(self.0, other.0) })
    }

    pub fn shr<const N: i32>(self) -> Self {
        Self(unsafe { _mm256_srav_epi32(self.0, _mm256_set1_epi32(N)) })
    }

    pub fn abs(self) -> Self {
        Self(unsafe { _mm256_abs_epi32(self.0) })
    }
}

impl Default for i32x8 {
    fn default() -> Self {
        Self(unsafe { _mm256_setzero_si256() })
    }
}

impl Add for i32x8 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_add_epi32(self.0, rhs.0) })
    }
}

impl Sub for i32x8 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_sub_epi32(self.0, rhs.0) })
    }
}

impl Mul for i32x8 {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_mullo_epi32(self.0, rhs.0) })
    }
}

impl BitAnd for i32x8 {
    type Output = Self;

    fn bitand(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_and_si256(self.0, rhs.0) })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct u32x4(__m128i);

impl u32x4 {
    pub fn splat(val: u32) -> Self {
        Self(unsafe { _mm_set1_epi32(val as i32) })
    }
}

impl From<u32x4> for [u8; 4] {
    fn from(val: u32x4) -> Self {
        unsafe {
            let mask = _mm_set1_epi32(0xFF);
            let val = _mm_and_si128(val.0, mask);

            let val = _mm_packus_epi32(val, val);
            let val = _mm_packus_epi16(val, val);

            _mm_cvtsi128_si32(val).to_ne_bytes()
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct u32x8(__m256i);

impl u32x8 {
    pub fn splat(val: u32) -> Self {
        Self(unsafe { _mm256_set1_epi32(val as i32) })
    }

    pub fn mul_add(self, a: Self, b: Self) -> Self {
        Self(unsafe { _mm256_add_epi32(_mm256_mullo_epi32(self.0, a.0), b.0) })
    }

    pub fn to_array(self) -> [u32; 8] {
        zerocopy::transmute!(self.0)
    }
}

impl From<f32x8> for u32x8 {
    fn from(val: f32x8) -> Self {
        // Sets negative value to 0 to prevent _mm256_cvttps_epi32 from
        // returning negative values.
        Self(unsafe { _mm256_cvttps_epi32(val.max(f32x8::splat(0.0)).0) })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct f32x4(__m128);

impl f32x4 {
    pub fn new(vals: [f32; 4]) -> Self {
        Self(unsafe { _mm_set_ps(vals[3], vals[2], vals[1], vals[0]) })
    }

    pub fn splat(val: f32) -> Self {
        Self(unsafe { _mm_set1_ps(val) })
    }

    pub fn from_bits(val: u32x4) -> Self {
        Self(unsafe { _mm_castsi128_ps(val.0) })
    }

    pub fn to_bits(self) -> u32x4 {
        u32x4(unsafe { _mm_castps_si128(self.0) })
    }

    pub fn set<const INDEX: i32>(self, val: f32) -> Self {
        Self(unsafe {
            _mm_castsi128_ps(_mm_insert_epi32::<INDEX>(
                _mm_castps_si128(self.0),
                val.to_bits() as i32,
            ))
        })
    }

    pub fn le(self, other: Self) -> m32x4 {
        m32x4(unsafe { _mm_castps_si128(_mm_cmp_ps(self.0, other.0, _CMP_LE_OQ)) })
    }

    pub fn select(self, other: Self, mask: m32x4) -> Self {
        Self(unsafe { _mm_blendv_ps(other.0, self.0, _mm_castsi128_ps(mask.0)) })
    }

    pub fn clamp(self, min: Self, max: Self) -> Self {
        Self(unsafe { _mm_min_ps(_mm_max_ps(self.0, min.0), max.0) })
    }

    pub fn sqrt(self) -> Self {
        Self(unsafe { _mm_sqrt_ps(self.0) })
    }

    pub fn mul_add(self, a: Self, b: Self) -> Self {
        Self(unsafe { _mm_fmadd_ps(self.0, a.0, b.0) })
    }
}

impl Add for f32x4 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm_add_ps(self.0, rhs.0) })
    }
}

impl Mul for f32x4 {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm_mul_ps(self.0, rhs.0) })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct f32x8(__m256);

impl f32x8 {
    pub fn splat(val: f32) -> Self {
        Self(unsafe { _mm256_set1_ps(val) })
    }

    pub fn indexed() -> Self {
        Self(unsafe { _mm256_set_ps(7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0, 0.0) })
    }

    pub fn from_bits(val: u32x8) -> Self {
        Self(unsafe { _mm256_castsi256_ps(val.0) })
    }

    pub fn to_bits(self) -> u32x8 {
        u32x8(unsafe { _mm256_castps_si256(self.0) })
    }

    pub fn from_array(val: [f32; 8]) -> Self {
        Self(unsafe { _mm256_loadu_ps(val.as_ptr()) })
    }

    #[cfg(test)]
    pub fn to_array(self) -> [f32; 8] {
        zerocopy::transmute!(self.0)
    }

    pub fn eq(self, other: Self) -> m32x8 {
        m32x8(unsafe { _mm256_castps_si256(_mm256_cmp_ps(self.0, other.0, _CMP_EQ_OQ)) })
    }

    pub fn lt(self, other: Self) -> m32x8 {
        m32x8(unsafe { _mm256_castps_si256(_mm256_cmp_ps(self.0, other.0, _CMP_LT_OQ)) })
    }

    pub fn le(self, other: Self) -> m32x8 {
        m32x8(unsafe { _mm256_castps_si256(_mm256_cmp_ps(self.0, other.0, _CMP_LE_OQ)) })
    }

    pub fn select(self, other: Self, mask: m32x8) -> Self {
        Self(unsafe { _mm256_blendv_ps(other.0, self.0, _mm256_castsi256_ps(mask.0)) })
    }

    pub fn abs(self) -> Self {
        Self(unsafe { _mm256_andnot_ps(_mm256_set1_ps(-0.0), self.0) })
    }

    pub fn min(self, other: Self) -> Self {
        Self(unsafe { _mm256_min_ps(self.0, other.0) })
    }

    pub fn max(self, other: Self) -> Self {
        Self(unsafe { _mm256_max_ps(self.0, other.0) })
    }

    pub fn clamp(self, min: Self, max: Self) -> Self {
        self.min(max).max(min)
    }

    pub fn sqrt(self) -> Self {
        Self(unsafe { _mm256_sqrt_ps(self.0) })
    }

    pub fn recip(self) -> Self {
        Self(unsafe { _mm256_rcp_ps(self.0) })
    }

    pub fn mul_add(self, a: Self, b: Self) -> Self {
        Self(unsafe { _mm256_fmadd_ps(self.0, a.0, b.0) })
    }
}

impl Default for f32x8 {
    fn default() -> Self {
        Self(unsafe { _mm256_setzero_ps() })
    }
}

impl Add for f32x8 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_add_ps(self.0, rhs.0) })
    }
}

impl AddAssign for f32x8 {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl Sub for f32x8 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_sub_ps(self.0, rhs.0) })
    }
}

impl Mul for f32x8 {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_mul_ps(self.0, rhs.0) })
    }
}

impl MulAssign for f32x8 {
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl Div for f32x8 {
    type Output = Self;

    fn div(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_div_ps(self.0, rhs.0) })
    }
}

impl Neg for f32x8 {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(unsafe { _mm256_xor_ps(self.0, _mm256_set1_ps(-0.0)) })
    }
}

impl BitOr for f32x8 {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        Self(unsafe { _mm256_or_ps(self.0, rhs.0) })
    }
}

impl BitOrAssign for f32x8 {
    fn bitor_assign(&mut self, rhs: Self) {
        *self = *self | rhs;
    }
}

impl From<i32x8> for f32x8 {
    fn from(val: i32x8) -> Self {
        Self(unsafe { _mm256_cvtepi32_ps(val.0) })
    }
}