kahlo 0.0.3

use std::arch::x86_64::{
    __m128i, __m256, __m256i, _mm256_add_epi16, _mm256_add_ps, _mm256_blend_epi16, _mm256_blend_ps,
    _mm256_castsi128_si256, _mm256_castsi256_si128, _mm256_cvtepi32_ps, _mm256_cvtepu16_epi32,
    _mm256_cvtepu8_epi16, _mm256_cvtepu8_epi32, _mm256_cvtps_epi32, _mm256_div_ps,
    _mm256_extracti128_si256, _mm256_fmadd_ps, _mm256_inserti128_si256, _mm256_loadu_epi8,
    _mm256_loadu_si256, _mm256_mul_epu32, _mm256_mul_ps, _mm256_mulhi_epu16, _mm256_mullo_epi16,
    _mm256_mullo_epi32, _mm256_or_si256, _mm256_set1_epi16, _mm256_set1_epi32, _mm256_set1_ps,
    _mm256_shuffle_epi32, _mm256_shuffle_epi8, _mm256_shuffle_ps, _mm256_shufflehi_epi16,
    _mm256_shufflelo_epi16, _mm256_srli_epi16, _mm256_srli_epi64, _mm256_storeu_si256,
    _mm256_sub_epi16, _mm256_sub_ps, _mm_cvtepi8_epi32, _mm_loadu_si128, _mm_loadu_si32,
    _mm_loadu_si64, _mm_or_si128, _mm_set1_epi32, _mm_slli_epi32, _mm_storeu_si128,
    _mm_storeu_si32, _mm_storeu_si64,
};

use crate::colour::Colour;

// multiply by 32897 (unsigned) and shift by 23 to divide by 255
const INV_255: u16 = 32897;

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

impl Ri8x16 {
    #[target_feature(enable = "avx2")]
    pub fn load4(data: &[u8; 4]) -> Self {
        Self(unsafe { _mm_loadu_si32(data.as_ptr()) })
    }

    #[target_feature(enable = "avx2")]
    pub fn load4_spaced_4_offset_3(data: &[u8; 4]) -> Self {
        let initial = unsafe { _mm_loadu_si32(data.as_ptr()) };
        // target is:   [0,0,0,a,0,0,0,b,0,0,0,c,0,0,0,d]
        // we now have: [a,b,c,d,0,0,0,0,0,0,0,0,0,0,0,0]
        let spread = _mm_cvtepi8_epi32(initial);
        // we now have: [a,0,0,0,b,0,0,0,c,0,0,0,d,0,0,0]
        // so shifting left by 24 bits gets the desired result
        Self(_mm_slli_epi32::<24>(spread))
    }

    #[target_feature(enable = "avx2")]
    pub fn widen_to_i16(self) -> Ri16x16 {
        Ri16x16(_mm256_cvtepu8_epi16(self.0))
    }
}

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

impl Ri8x32 {
    #[target_feature(enable = "avx2")]
    pub fn load32(data: &[u8; 32]) -> Self {
        Self(unsafe { _mm256_loadu_si256(data.as_ptr() as *const __m256i) })
    }

    #[target_feature(enable = "avx2")]
    pub fn store32(self, data: &mut [u8; 32]) {
        unsafe { _mm256_storeu_si256(data.as_ptr() as *mut __m256i, self.0) }
    }

    const ENDIAN_SHUFFLE: [i8; 32] = [
        // low lane, first pixel
        3, 2, 1, 0, // low lane, second pixel
        7, 6, 5, 4, // low lane, third pixel
        11, 10, 9, 8, // low lane, fourth pixel
        15, 14, 13, 12, // high lane, first pixel
        3, 2, 1, 0, // high lane, second pixel
        7, 6, 5, 4, // high lane, third pixel
        11, 10, 9, 8, // high lane, fourth pixel
        15, 14, 13, 12,
    ];

    #[target_feature(enable = "avx2")]
    pub fn swap_endianness(self) -> Self {
        Self(_mm256_shuffle_epi8(self.0, unsafe {
            _mm256_loadu_epi8(Self::ENDIAN_SHUFFLE.as_ptr())
        }))
    }
}

/// standard format for pixel manipulation: each pixel is represented by 64 bits, r/g/b/a with a 16-bit value each
#[derive(Clone, Copy)]
pub struct Ri16x16(pub __m256i);

#[allow(unused)]
impl Ri16x16 {
    pub fn unwrap(self) -> __m256i {
        self.0
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn load8_and_expand(val: &[u8; 16]) -> Self {
        Self(unsafe { _mm256_cvtepu8_epi16(_mm_loadu_si128(val.as_ptr() as *const __m128i)) })
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn load_from(val: &[u8; 32]) -> Self {
        Self(unsafe { _mm256_loadu_epi8(val.as_ptr() as *const i8) })
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_constant(val: u16) -> Self {
        Self(_mm256_set1_epi16(val as i16))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn shrink_and_store8(self, into: &mut [u8; 16]) {
        // since we don't have AVX512, there's no convenient _mm256_cvtepi16_epi8
        // available. hence we byte-shuffle the LSBs to be contiguous with zeroes
        // in the upper/lower half of the 128-bit lanes, then OR the two together.

        // shuffle vector
        const SHUFFLE: [i8; 32] = [
            0, 2, 4, 6, 8, 10, 12, 14, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, 0, 2, 4, 6, 8, 10, 12, 14,
        ];

        let shuffled = _mm256_shuffle_epi8(self.0, unsafe { _mm256_loadu_epi8(SHUFFLE.as_ptr()) });

        let recombined = _mm_or_si128(
            _mm256_castsi256_si128(shuffled),
            _mm256_extracti128_si256::<1>(shuffled),
        );
        unsafe { _mm_storeu_si128(into.as_mut_ptr() as *mut __m128i, recombined) };
    }

    /// extracts alpha of each pixel and broadcasts its value to the r/g/b channels
    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn extract_alpha_rgba(self) -> Self {
        Self(_mm256_shufflelo_epi16::<0b11_11_11_11>(
            _mm256_shufflehi_epi16::<0b11_11_11_11>(self.0),
        ))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn div_255(self) -> Self {
        Ri16x16(_mm256_srli_epi16::<7>(_mm256_mulhi_epu16(
            _mm256_add_epi16(self.0, _mm256_set1_epi16(128)),
            _mm256_set1_epi16(INV_255 as i16),
        )))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn to_i32x8(self) -> (Ri32x8, Ri32x8) {
        (
            Ri32x8(_mm256_cvtepu16_epi32(_mm256_castsi256_si128(self.0))),
            Ri32x8(_mm256_cvtepu16_epi32(_mm256_extracti128_si256::<1>(self.0))),
        )
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_i32x8(lo: Ri32x8, hi: Ri32x8) -> Ri16x16 {
        // we don't have _mm256_cvtepi32_epi16, as nice as it would be, since that's AVX-512.
        // so instead, do a byte-level shuffle to shove the two least-significant bytes to one
        // half of the 128-bit lane
        const SHUFFLE: [i8; 32] = [
            // two least-significant bytes for each of the first four i32s
            0, 1, 4, 5, 8, 9, 12, 13, // high 64 bits of the first 128-bit lane
            -1, -1, -1, -1, -1, -1, -1, -1, // low 64 bits of the second 128-bit lane
            -1, -1, -1, -1, -1, -1, -1, -1,
            // two-least significant bytes for each of the second four i32s
            0, 1, 4, 5, 8, 9, 12, 13,
        ];
        let shuffle = unsafe { _mm256_loadu_epi8(SHUFFLE.as_ptr()) };
        let lo_shuf = _mm256_shuffle_epi8(lo.0, shuffle);
        let hi_shuf = _mm256_shuffle_epi8(hi.0, shuffle);

        // now combine the low and high halves of each shuffled __mm256i
        let lo16 = _mm_or_si128(
            _mm256_castsi256_si128(lo_shuf),
            _mm256_extracti128_si256::<1>(lo_shuf),
        );

        let hi16 = _mm_or_si128(
            _mm256_castsi256_si128(hi_shuf),
            _mm256_extracti128_si256::<1>(hi_shuf),
        );

        Ri16x16(_mm256_inserti128_si256::<1>(
            _mm256_castsi128_si256(lo16),
            hi16,
        ))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn mix(&self, other: Ri16x16, amount: Ri16x16) -> Ri16x16 {
        ((*self * amount) + (other * (Ri16x16::from_constant(255) - amount))).div_255()
    }

    /// Combines the alpha values from another Ri16x16 and the rgb values from this Ri16x16 and
    /// returns the result.
    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn copy_alpha_rgba(&self, other: Ri16x16) -> Ri16x16 {
        Self(_mm256_blend_epi16::<0b_1000_1000>(self.0, other.0))
    }
}

impl std::fmt::Debug for Ri16x16 {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let a = unsafe { std::mem::transmute::<_, [u16; 16]>(self.0) };

        write!(
            fmt,
            "Ri16x16 {{ {} | {} | {} | {} }}",
            format!("{:04x} {:04x} {:04x} {:04x}", a[0], a[1], a[2], a[3]),
            format!("{:04x} {:04x} {:04x} {:04x}", a[4], a[5], a[6], a[7]),
            format!("{:04x} {:04x} {:04x} {:04x}", a[8], a[9], a[10], a[11]),
            format!("{:04x} {:04x} {:04x} {:04x}", a[12], a[13], a[14], a[15]),
        )
    }
}

impl std::ops::Add for Ri16x16 {
    type Output = Ri16x16;
    fn add(self, r: Self) -> Self {
        Self(unsafe { _mm256_add_epi16(self.0, r.0) })
    }
}

impl std::ops::Sub for Ri16x16 {
    type Output = Ri16x16;
    fn sub(self, r: Self) -> Self {
        Self(unsafe { _mm256_sub_epi16(self.0, r.0) })
    }
}

impl std::ops::Mul for Ri16x16 {
    type Output = Ri16x16;
    fn mul(self, r: Self) -> Self {
        Self(unsafe { _mm256_mullo_epi16(self.0, r.0) })
    }
}

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

#[allow(unused)]
impl Ri32x4 {
    pub fn unwrap(self) -> __m128i {
        self.0
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_constant(val: u32) -> Self {
        Self(_mm_set1_epi32(val as i32))
    }

    pub fn cast_to_i8x16(self) -> Ri8x16 {
        Ri8x16(self.0)
    }
}

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

#[allow(unused)]
impl Ri32x8 {
    pub fn unwrap(self) -> __m256i {
        self.0
    }

    pub fn cast_to_i16x16(self) -> Ri16x16 {
        Ri16x16(self.0)
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_constant(val: u32) -> Self {
        Self(_mm256_set1_epi32(val as i32))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn div_255(&self) -> Self {
        // first add 128 for rounding
        let withround = _mm256_add_epi16(self.0, _mm256_set1_epi32(128));

        // the multiplication by INV_255 can and will result in a value larger than 32 bits, so we
        // need to use the 64-bit multiplication result instead.
        let loprod = _mm256_mul_epu32(self.0, _mm256_set1_epi32(INV_255 as i32));
        // need to swap the high/low 64-bit groups in each lane and then multiply
        let hiprod = _mm256_mul_epu32(
            _mm256_shuffle_epi32::<0b_10_11_00_01>(self.0),
            _mm256_set1_epi32(INV_255 as i32),
        );

        // at this point, loprod has (v*INV_255).aceg and hiprod has (v*INV_255).bdfh as u64s

        // do the post-multiplication reduction...
        let loshift = _mm256_srli_epi64::<23>(loprod);
        let hishift = _mm256_srli_epi64::<23>(hiprod);

        // at this point, loprod has (v*INV_255 >> 23).aceg and hiprod has (v*INV_255 >> 23).bdfh as u64s,
        // so we want to convert down to u32s. we want:
        // - loprod to a_c_e_g_
        // - hiprod to _b_d_f_h
        //
        // but loprod is already in this form, so we just have to shuffle hiprod around a bit
        let hiswizz = _mm256_shuffle_epi32::<0b_10_11_00_01>(hishift);

        Ri32x8(_mm256_or_si256(loshift, hiswizz))
    }
}

impl std::fmt::Debug for Ri32x8 {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let a = unsafe { std::mem::transmute::<_, [u32; 8]>(self.0) };

        write!(
            fmt,
            "Ri32x8 {{ {} | {} }}",
            format!("{:08x} {:08x} {:08x} {:08x}", a[0], a[1], a[2], a[3]),
            format!("{:08x} {:08x} {:08x} {:08x}", a[4], a[5], a[6], a[7]),
        )
    }
}

impl std::ops::Add for Ri32x8 {
    type Output = Self;
    fn add(self, r: Self) -> Self {
        Self(unsafe { _mm256_add_epi16(self.0, r.0) })
    }
}

impl std::ops::Sub for Ri32x8 {
    type Output = Self;
    fn sub(self, r: Self) -> Self {
        Self(unsafe { _mm256_sub_epi16(self.0, r.0) })
    }
}

impl std::ops::Mul for Ri32x8 {
    type Output = Self;
    fn mul(self, r: Self) -> Self {
        Self(unsafe { _mm256_mullo_epi32(self.0, r.0) })
    }
}

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

impl std::fmt::Debug for Ri64x4 {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let a = unsafe { std::mem::transmute::<_, [u64; 4]>(self.0) };

        write!(
            fmt,
            "Ri32x8 {{ {} }}",
            format!("{:016x} {:016x} | {:016x} {:016x}", a[0], a[1], a[2], a[3]),
        )
    }
}

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

impl Rf32x8 {
    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn load8_and_expand(val: &[u8; 8]) -> Self {
        // _mm_loadu_si128(val.as_ptr() as *const __m128i)
        Self(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(unsafe {
            _mm_loadu_si64(val.as_ptr())
        })))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn load4_and_expand(val: &[u8; 4]) -> Self {
        Self(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(unsafe {
            _mm_loadu_si32(val.as_ptr())
        })))
    }

    const SHRINK_SHUFFLE: [i8; 32] = [
        // least-significant byte for the first four i32s
        0, 4, 8, 12, // high 96 bits of the low lane
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        // skip the low 32 bits of the high lane
        -1, -1, -1, -1, // least-significant byte for the next four i32s
        0, 4, 8, 12, // high 64 bits of the high lane
        -1, -1, -1, -1, -1, -1, -1, -1,
    ];

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn shrink_and_store8(self, val: &mut [u8; 8]) {
        // really would love to use _mm256_cvtepi32_epi8, but that's a avx512f mnemonic, so...
        // shuffling it is.

        let shuffle = unsafe { _mm256_loadu_epi8(Self::SHRINK_SHUFFLE.as_ptr()) };
        let shuffled = _mm256_shuffle_epi8(_mm256_cvtps_epi32(self.0), shuffle);
        let combined = _mm_or_si128(
            _mm256_castsi256_si128(shuffled),
            _mm256_extracti128_si256::<1>(shuffled),
        );

        unsafe { _mm_storeu_si64(val.as_mut_ptr(), combined) }
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn shrink_and_store4(self, val: &mut [u8; 4]) {
        // really would love to use _mm256_cvtepi32_epi8, but that's a avx512f mnemonic, so...
        // shuffling it is.

        let shuffle = unsafe { _mm256_loadu_epi8(Self::SHRINK_SHUFFLE.as_ptr()) };
        let shuffled = _mm256_shuffle_epi8(_mm256_cvtps_epi32(self.0), shuffle);
        // unlike the store8 case, here we have no upper lane to extract bits from, so we can skip
        // straight to the store.
        unsafe { _mm_storeu_si32(val.as_mut_ptr(), _mm256_castsi256_si128(shuffled)) }
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_constant(val: f32) -> Self {
        Self(_mm256_set1_ps(val))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn from_colours_rgba(lo: Colour, hi: Colour) -> Self {
        let data = [lo.r, lo.g, lo.b, lo.a, hi.r, hi.g, hi.b, hi.a];
        Self(_mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(unsafe {
            _mm_loadu_si64(data.as_ptr())
        })))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn extract_alpha_rgba(self) -> Self {
        Self(_mm256_shuffle_ps::<0b_11_11_11_11>(self.0, self.0))
    }

    #[target_feature(enable = "avx2")]
    #[inline]
    pub fn copy_alpha_from(self, alpha: Self) -> Self {
        Self(_mm256_blend_ps::<0b_1000_1000>(self.0, alpha.0))
    }

    #[target_feature(enable = "fma")]
    #[inline]
    pub fn fma(self, m: Self, a: Self) -> Self {
        Self(_mm256_fmadd_ps(self.0, m.0, a.0))
    }
}

impl std::fmt::Debug for Rf32x8 {
    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
        let a = unsafe { std::mem::transmute::<_, [f32; 8]>(self.0) };

        write!(
            fmt,
            "Rf32x8 {{ {} | {} }}",
            format_args!("{:.4} {:.4} {:.4} {:.4}", a[0], a[1], a[2], a[3]),
            format_args!("{:.4} {:.4} {:.4} {:.4}", a[4], a[5], a[6], a[7]),
        )
    }
}

impl std::ops::Add for Rf32x8 {
    type Output = Self;
    fn add(self, r: Self) -> Self {
        Self(unsafe { _mm256_add_ps(self.0, r.0) })
    }
}

impl std::ops::Sub for Rf32x8 {
    type Output = Self;
    fn sub(self, r: Self) -> Self {
        Self(unsafe { _mm256_sub_ps(self.0, r.0) })
    }
}

impl std::ops::Mul for Rf32x8 {
    type Output = Self;
    fn mul(self, r: Self) -> Self {
        Self(unsafe { _mm256_mul_ps(self.0, r.0) })
    }
}

impl std::ops::Div for Rf32x8 {
    type Output = Self;
    fn div(self, r: Self) -> Self {
        Self(unsafe { _mm256_div_ps(self.0, r.0) })
    }
}