Compiler Explorer

Source code

#![feature(stdsimd)]

extern crate core;

pub fn haraka256_5(dst: &mut [u8; 32], src: &[u8; 32]) {
    haraka256::<5>(dst, src)
}

#[macro_export]
macro_rules! array_ref {
    ($arr:expr, $offset:expr, $len:expr) => {{
        {
            #[inline]
            unsafe fn as_array<T>(slice: &[T]) -> &[T; $len] {
                &*(slice.as_ptr() as *const [_; $len])
            }
            let offset = $offset;
            let slice = & $arr[offset..offset + $len];
            #[allow(unused_unsafe)]
            unsafe {
                as_array(slice)
            }
        }
    }}
}

#[macro_export]
macro_rules! array_mut_ref {
    ($arr:expr, $offset:expr, $len:expr) => {{
        {
            #[inline]
            unsafe fn as_array<T>(slice: &mut [T]) -> &mut [T; $len] {
                &mut *(slice.as_mut_ptr() as *mut [_; $len])
            }
            let offset = $offset;
            let slice = &mut $arr[offset..offset + $len];
            #[allow(unused_unsafe)]
            unsafe {
                as_array(slice)
            }
        }
    }}
}

#[inline(always)]
fn aes2(s0: &mut Simd128, s1: &mut Simd128, rci: usize) {
    Simd128::aesenc(s0, &HARAKA_CONSTANTS[rci]);
    Simd128::aesenc(s1, &HARAKA_CONSTANTS[rci + 1]);
    Simd128::aesenc(s0, &HARAKA_CONSTANTS[rci + 2]);
    Simd128::aesenc(s1, &HARAKA_CONSTANTS[rci + 3]);
}

#[inline(always)]
fn mix2(s0: &mut Simd128, s1: &mut Simd128) {
    let mut tmp = *s0;
    Simd128::unpackhi_epi32(&mut tmp, s1);
    Simd128::unpacklo_epi32(s0, s1);
    *s1 = tmp;
}

#[inline(always)]
fn aes_mix2(s0: &mut Simd128, s1: &mut Simd128, rci: usize) {
    aes2(s0, s1, rci);
    mix2(s0, s1);
}

pub fn haraka256<const N_ROUNDS: usize>(dst: &mut [u8; 32], src: &[u8; 32]) {
    let mut s0 = Simd128::read(array_ref![src, 0, 16]);
    let mut s1 = Simd128::read(array_ref![src, 16, 16]);

for i in 0..N_ROUNDS {
        aes_mix2(&mut s0, &mut s1, 4 * i);
    }

let t0 = Simd128::read(array_ref![src, 0, 16]);
    let t1 = Simd128::read(array_ref![src, 16, 16]);
    Simd128::pxor(&mut s0, &t0);
    Simd128::pxor(&mut s1, &t1);

s0.write(array_mut_ref![dst, 0, 16]);
    s1.write(array_mut_ref![dst, 16, 16]);
}

#[cfg(target_arch = "aarch64")]
use core::arch::aarch64::{
    uint8x16_t, vaeseq_u8, vaesmcq_u8, vdupq_n_u8, veorq_u8, vld1q_u8, vreinterpretq_u32_u8,
    vreinterpretq_u8_u32, vst1q_u8, vzip1q_u32, vzip2q_u32,
};
#[cfg(target_arch = "arm")]
use core::arch::arm::{
    uint8x16_t, vaeseq_u8, vaesmcq_u8, vdupq_n_u8, veorq_u8, vld1q_u8, vreinterpretq_u32_u8,
    vreinterpretq_u8_u32, vst1q_u8, vzipq_u32,
};
#[cfg(target_arch = "x86")]
use std::arch::x86::{
    __m128i, _mm_aesenc_si128, _mm_loadu_si128, _mm_storeu_si128, _mm_unpackhi_epi32,
    _mm_unpacklo_epi32, _mm_xor_si128,
};
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::{
    __m128i, _mm_aesenc_si128, _mm_loadu_si128, _mm_storeu_si128, _mm_unpackhi_epi32,
    _mm_unpacklo_epi32, _mm_xor_si128,
};
use std::mem::transmute;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[derive(Clone, Copy)]
pub(crate) struct Simd128(__m128i);

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
impl Simd128 {
    pub const fn from(x: u128) -> Self {
        Self(unsafe { transmute(x) })
    }

/// Read from array pointer (potentially unaligned)
    #[inline(always)]
    pub fn read(src: &[u8; 16]) -> Self {
        let x = unsafe { _mm_loadu_si128(src.as_ptr() as *const _ as *const __m128i) };
        Self(x)
    }

/// Write into array pointer (potentially unaligned)
    #[inline(always)]
    pub fn write(self, dst: &mut [u8; 16]) {
        unsafe {
            _mm_storeu_si128(dst.as_mut_ptr() as *mut _ as *mut __m128i, self.0);
        }
    }

#[inline(always)]
    pub(crate) fn aesenc(block: &mut Self, key: &Self) {
        unsafe {
            block.0 = _mm_aesenc_si128(block.0, key.0);
        }
    }

#[inline(always)]
    pub(crate) fn pxor(dst: &mut Self, src: &Self) {
        unsafe {
            dst.0 = _mm_xor_si128(dst.0, src.0);
        }
    }

#[inline(always)]
    pub(crate) fn unpacklo_epi32(dst: &mut Self, src: &Self) {
        unsafe {
            dst.0 = _mm_unpacklo_epi32(dst.0, src.0);
        }
    }

#[inline(always)]
    pub(crate) fn unpackhi_epi32(dst: &mut Self, src: &Self) {
        unsafe {
            dst.0 = _mm_unpackhi_epi32(dst.0, src.0);
        }
    }
}

#[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
#[derive(Clone, Copy)]
pub(crate) struct Simd128(uint8x16_t);

#[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
impl Simd128 {
    pub const fn from(x: u128) -> Self {
        Self(unsafe { transmute(x) })
    }

/// Read from array pointer (potentially unaligned)
    #[inline(always)]
    pub fn read(src: &[u8; 16]) -> Self {
        let x = unsafe { vld1q_u8(src.as_ptr() as *const _ as *const u8) };
        Self(x)
    }

/// Write into array pointer (potentially unaligned)
    #[inline(always)]
    pub fn write(self, dst: &mut [u8; 16]) {
        unsafe {
            vst1q_u8(dst.as_mut_ptr() as *mut _ as *mut u8, self.0);
        }
    }

#[inline(always)]
    pub(crate) fn aesenc(block: &mut Self, key: &Self) {
        unsafe {
            let zero = vdupq_n_u8(0);
            let x = vaeseq_u8(block.0, zero);
            let y = vaesmcq_u8(x);
            block.0 = veorq_u8(y, key.0);
        }
    }

#[inline(always)]
    pub(crate) fn pxor(dst: &mut Self, src: &Self) {
        unsafe {
            dst.0 = veorq_u8(dst.0, src.0);
        }
    }

#[inline(always)]
    pub(crate) fn unpacklo_epi32(dst: &mut Self, src: &Self) {
        unsafe {
            let a = vreinterpretq_u32_u8(dst.0);
            let b = vreinterpretq_u32_u8(src.0);
            // TODO: vzip1q_u32 is missing from core::arch::arm.
            #[cfg(target_arch = "arm")]
            let x = vzipq_u32(a, b).0;
            #[cfg(target_arch = "aarch64")]
            let x = vzip1q_u32(a, b);
            dst.0 = vreinterpretq_u8_u32(x);
        }
    }

#[inline(always)]
    pub(crate) fn unpackhi_epi32(dst: &mut Self, src: &Self) {
        unsafe {
            let a = vreinterpretq_u32_u8(dst.0);
            let b = vreinterpretq_u32_u8(src.0);
            // TODO: vzip2q_u32 is missing from core::arch::arm.
            #[cfg(target_arch = "arm")]
            let x = vzipq_u32(a, b).1;
            #[cfg(target_arch = "aarch64")]
            let x = vzip2q_u32(a, b);
            dst.0 = vreinterpretq_u8_u32(x);
        }
    }
}

static HARAKA_CONSTANTS: &[Simd128; 48] = &[
    Simd128::from(0x0684704ce620c00ab2c5fef075817b9d),
    Simd128::from(0x8b66b4e188f3a06b640f6ba42f08f717),
    Simd128::from(0x3402de2d53f28498cf029d609f029114),
    Simd128::from(0x0ed6eae62e7b4f08bbf3bcaffd5b4f79),
    Simd128::from(0xcbcfb0cb4872448b79eecd1cbe397044),
    Simd128::from(0x7eeacdee6e9032b78d5335ed2b8a057b),
    Simd128::from(0x67c28f435e2e7cd0e2412761da4fef1b),
    Simd128::from(0x2924d9b0afcacc07675ffde21fc70b3b),
    Simd128::from(0xab4d63f1e6867fe9ecdb8fcab9d465ee),
    Simd128::from(0x1c30bf84d4b7cd645b2a404fad037e33),
    Simd128::from(0xb2cc0bb9941723bf69028b2e8df69800),
    Simd128::from(0xfa0478a6de6f55724aaa9ec85c9d2d8a),
    Simd128::from(0xdfb49f2b6b772a120efa4f2e29129fd4),
    Simd128::from(0x1ea10344f449a23632d611aebb6a12ee),
    Simd128::from(0xaf0449884b0500845f9600c99ca8eca6),
    Simd128::from(0x21025ed89d199c4f78a2c7e327e593ec),
    Simd128::from(0xbf3aaaf8a759c9b7b9282ecd82d40173),
    Simd128::from(0x6260700d6186b01737f2efd910307d6b),
    Simd128::from(0x5aca45c22130044381c29153f6fc9ac6),
    Simd128::from(0x9223973c226b68bb2caf92e836d1943a),
    Simd128::from(0xd3bf9238225886eb6cbab958e51071b4),
    Simd128::from(0xdb863ce5aef0c677933dfddd24e1128d),
    Simd128::from(0xbb606268ffeba09c83e48de3cb2212b1),
    Simd128::from(0x734bd3dce2e4d19c2db91a4ec72bf77d),
    Simd128::from(0x43bb47c361301b434b1415c42cb3924e),
    Simd128::from(0xdba775a8e707eff603b231dd16eb6899),
    Simd128::from(0x6df3614b3c7559778e5e23027eca472c),
    Simd128::from(0xcda75a17d6de7d776d1be5b9b88617f9),
    Simd128::from(0xec6b43f06ba8e9aa9d6c069da946ee5d),
    Simd128::from(0xcb1e6950f957332ba25311593bf327c1),
    Simd128::from(0x2cee0c7500da619ce4ed0353600ed0d9),
    Simd128::from(0xf0b1a5a196e90cab80bbbabc63a4a350),
    Simd128::from(0xae3db1025e962988ab0dde30938dca39),
    Simd128::from(0x17bb8f38d554a40b8814f3a82e75b442),
    Simd128::from(0x34bb8a5b5f427fd7aeb6b779360a16f6),
    Simd128::from(0x26f65241cbe5543843ce5918ffbaafde),
    Simd128::from(0x4ce99a54b9f3026aa2ca9cf7839ec978),
    Simd128::from(0xae51a51a1bdff7be40c06e2822901235),
    Simd128::from(0xa0c1613cba7ed22bc173bc0f48a659cf),
    Simd128::from(0x756acc03022882884ad6bdfde9c59da1),
    Simd128::from(0x2ff372380de7d31e367e4778848f2ad2),
    Simd128::from(0x08d95c6acf74be8bee36b135b73bd58f),
    Simd128::from(0x5880f434c9d6ee9866ae1838a3743e4a),
    Simd128::from(0x593023f0aefabd99d0fdf4c79a9369bd),
    Simd128::from(0x329ae3d1eb606e6fa5cc637b6f1ecb2a),
    Simd128::from(0xe00207eb49e01594a4dc93d6cb7594ab),
    Simd128::from(0x1caa0c4ff751c880942366a665208ef8),
    Simd128::from(0x02f7f57fdb2dc1ddbd03239fe3e67e4a),
];