macro_rules! impl_word {
($name:ident, $prim:ty, $bits:literal) => {
#[doc = concat!("A ", $bits, "-bit wrapping integer: the ring ℤ/2^", $bits, "ℤ, where every operation is total and wraps modulo 2^", $bits, ".")]
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
#[repr(transparent)]
pub struct $name(pub $prim);
impl $name {
pub const BITS: u32 = <$prim>::BITS;
pub const ZERO: Self = Self(0);
pub const ONE: Self = Self(1);
pub const MAX: Self = Self(<$prim>::MAX);
#[inline]
pub const fn get(self) -> $prim {
self.0
}
#[inline]
pub const fn add(self, o: Self) -> Self {
Self(self.0.wrapping_add(o.0))
}
#[inline]
pub const fn sub(self, o: Self) -> Self {
Self(self.0.wrapping_sub(o.0))
}
#[inline]
pub const fn mul(self, o: Self) -> Self {
Self(self.0.wrapping_mul(o.0))
}
#[inline]
pub const fn bitand(self, o: Self) -> Self {
Self(self.0 & o.0)
}
#[inline]
pub const fn bitor(self, o: Self) -> Self {
Self(self.0 | o.0)
}
#[inline]
pub const fn bitxor(self, o: Self) -> Self {
Self(self.0 ^ o.0)
}
#[inline]
pub const fn not(self) -> Self {
Self(!self.0)
}
#[inline]
pub const fn shl(self, n: u32) -> Self {
Self(self.0.wrapping_shl(n))
}
#[inline]
pub const fn shr(self, n: u32) -> Self {
Self(self.0.wrapping_shr(n))
}
#[inline]
pub const fn rotl(self, n: u32) -> Self {
Self(self.0.rotate_left(n))
}
#[inline]
pub const fn rotr(self, n: u32) -> Self {
Self(self.0.rotate_right(n))
}
}
impl ::core::ops::Add for $name {
type Output = Self;
#[inline]
fn add(self, o: Self) -> Self { Self(self.0.wrapping_add(o.0)) }
}
impl ::core::ops::Sub for $name {
type Output = Self;
#[inline]
fn sub(self, o: Self) -> Self { Self(self.0.wrapping_sub(o.0)) }
}
impl ::core::ops::Mul for $name {
type Output = Self;
#[inline]
fn mul(self, o: Self) -> Self { Self(self.0.wrapping_mul(o.0)) }
}
impl ::core::ops::BitAnd for $name {
type Output = Self;
#[inline]
fn bitand(self, o: Self) -> Self { Self(self.0 & o.0) }
}
impl ::core::ops::BitOr for $name {
type Output = Self;
#[inline]
fn bitor(self, o: Self) -> Self { Self(self.0 | o.0) }
}
impl ::core::ops::BitXor for $name {
type Output = Self;
#[inline]
fn bitxor(self, o: Self) -> Self { Self(self.0 ^ o.0) }
}
impl ::core::ops::Not for $name {
type Output = Self;
#[inline]
fn not(self) -> Self { Self(!self.0) }
}
impl ::core::ops::Div for $name {
type Output = Self;
#[inline]
fn div(self, o: Self) -> Self { Self(self.0 / o.0) }
}
impl ::core::ops::Rem for $name {
type Output = Self;
#[inline]
fn rem(self, o: Self) -> Self { Self(self.0 % o.0) }
}
impl ::core::ops::Shl<u32> for $name {
type Output = Self;
#[inline]
fn shl(self, n: u32) -> Self { Self(self.0.wrapping_shl(n)) }
}
impl ::core::ops::Shr<u32> for $name {
type Output = Self;
#[inline]
fn shr(self, n: u32) -> Self { Self(self.0.wrapping_shr(n)) }
}
impl ::core::fmt::Display for $name {
#[inline]
fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
::core::write!(f, "{}", self.0)
}
}
};
}
impl_word!(Word8, u8, "8");
impl_word!(Word16, u16, "16");
impl_word!(Word32, u32, "32");
impl_word!(Word64, u64, "64");
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum WordVal {
W32(Word32),
W64(Word64),
}
impl WordVal {
#[inline]
pub const fn width(self) -> u32 {
match self {
WordVal::W32(_) => 32,
WordVal::W64(_) => 64,
}
}
#[inline]
pub const fn to_u64(self) -> u64 {
match self {
WordVal::W32(w) => w.0 as u64,
WordVal::W64(w) => w.0,
}
}
#[inline]
pub const fn from_u64(width: u32, bits: u64) -> Option<Self> {
match width {
32 => Some(WordVal::W32(Word32(bits as u32))),
64 => Some(WordVal::W64(Word64(bits))),
_ => None,
}
}
#[inline]
fn zip(
self,
o: Self,
f32: impl FnOnce(Word32, Word32) -> Word32,
f64: impl FnOnce(Word64, Word64) -> Word64,
) -> Option<Self> {
match (self, o) {
(WordVal::W32(a), WordVal::W32(b)) => Some(WordVal::W32(f32(a, b))),
(WordVal::W64(a), WordVal::W64(b)) => Some(WordVal::W64(f64(a, b))),
_ => None,
}
}
#[inline]
pub fn add(self, o: Self) -> Option<Self> {
self.zip(o, Word32::add, Word64::add)
}
#[inline]
pub fn sub(self, o: Self) -> Option<Self> {
self.zip(o, Word32::sub, Word64::sub)
}
#[inline]
pub fn mul(self, o: Self) -> Option<Self> {
self.zip(o, Word32::mul, Word64::mul)
}
#[inline]
pub fn bitand(self, o: Self) -> Option<Self> {
self.zip(o, Word32::bitand, Word64::bitand)
}
#[inline]
pub fn bitor(self, o: Self) -> Option<Self> {
self.zip(o, Word32::bitor, Word64::bitor)
}
#[inline]
pub fn bitxor(self, o: Self) -> Option<Self> {
self.zip(o, Word32::bitxor, Word64::bitxor)
}
#[inline]
pub const fn not(self) -> Self {
match self {
WordVal::W32(w) => WordVal::W32(w.not()),
WordVal::W64(w) => WordVal::W64(w.not()),
}
}
#[inline]
pub const fn shl(self, n: u32) -> Self {
match self {
WordVal::W32(w) => WordVal::W32(w.shl(n)),
WordVal::W64(w) => WordVal::W64(w.shl(n)),
}
}
#[inline]
pub const fn shr(self, n: u32) -> Self {
match self {
WordVal::W32(w) => WordVal::W32(w.shr(n)),
WordVal::W64(w) => WordVal::W64(w.shr(n)),
}
}
#[inline]
pub const fn rotl(self, n: u32) -> Self {
match self {
WordVal::W32(w) => WordVal::W32(w.rotl(n)),
WordVal::W64(w) => WordVal::W64(w.rotl(n)),
}
}
#[inline]
pub const fn rotr(self, n: u32) -> Self {
match self {
WordVal::W32(w) => WordVal::W32(w.rotr(n)),
WordVal::W64(w) => WordVal::W64(w.rotr(n)),
}
}
}
impl std::fmt::Display for WordVal {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_u64())
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[repr(C, align(32))]
pub struct Lanes8Word32(pub [u32; 8]);
impl Lanes8Word32 {
pub const LANES: usize = 8;
#[inline]
pub const fn splat(x: u32) -> Self {
Self([x; 8])
}
#[inline]
pub fn from_words(s: &[Word32]) -> Self {
let mut a = [0u32; 8];
for (i, w) in s.iter().take(8).enumerate() {
a[i] = w.0;
}
Self(a)
}
#[inline]
pub fn to_words(self) -> [Word32; 8] {
self.0.map(Word32)
}
#[inline]
pub fn lane(self, i: usize) -> Word32 {
Word32(self.0[i])
}
#[inline(always)]
pub fn bitxor(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_xor_si256(a, b));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i] ^ o.0[i];
}
Self(r)
}
}
#[cfg(test)]
#[inline]
fn bitxor_scalar(self, o: Self) -> Self {
Self(core::array::from_fn(|i| self.0[i] ^ o.0[i]))
}
#[cfg(test)]
#[inline]
fn add_scalar(self, o: Self) -> Self {
Self(core::array::from_fn(|i| self.0[i].wrapping_add(o.0[i])))
}
#[cfg(test)]
#[inline]
fn sub_scalar(self, o: Self) -> Self {
Self(core::array::from_fn(|i| self.0[i].wrapping_sub(o.0[i])))
}
#[cfg(test)]
#[inline]
fn rotl_scalar(self, n: u32) -> Self {
Self(core::array::from_fn(|i| self.0[i].rotate_left(n)))
}
#[cfg(test)]
#[inline]
fn montmul32_scalar(self, b: Self, q: Self, qinv: Self) -> Self {
Self(core::array::from_fn(|i| {
let p = (self.0[i] as i32 as i64) * (b.0[i] as i32 as i64);
let t = (p as i32).wrapping_mul(qinv.0[i] as i32) as i64;
(((p - t * (q.0[i] as i32 as i64)) >> 32) as i32) as u32
}))
}
#[inline]
pub fn bitand(self, o: Self) -> Self {
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i] & o.0[i];
}
Self(r)
}
#[inline]
pub fn bitor(self, o: Self) -> Self {
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i] | o.0[i];
}
Self(r)
}
#[inline]
pub fn not(self) -> Self {
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = !self.0[i];
}
Self(r)
}
#[inline(always)]
pub fn add(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_add_epi32(a, b));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i].wrapping_add(o.0[i]);
}
Self(r)
}
}
#[inline(always)]
pub fn rotl(self, n: u32) -> Self {
let n = n % 32;
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let x = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let l = _mm256_sll_epi32(x, _mm_cvtsi32_si128(n as i32));
let r_sh = _mm256_srl_epi32(x, _mm_cvtsi32_si128((32 - n) as i32));
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_or_si256(l, r_sh));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i].rotate_left(n);
}
Self(r)
}
}
#[inline(always)]
pub fn sub(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_sub_epi32(a, b));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u32; 8];
for i in 0..8 {
r[i] = self.0[i].wrapping_sub(o.0[i]);
}
Self(r)
}
}
#[inline(always)]
pub fn montmul32(self, b: Self, q: Self, qinv: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let bb = _mm256_loadu_si256(b.0.as_ptr() as *const __m256i);
let qv = _mm256_loadu_si256(q.0.as_ptr() as *const __m256i);
let qiv = _mm256_loadu_si256(qinv.0.as_ptr() as *const __m256i);
let pe = _mm256_mul_epi32(a, bb);
let po = _mm256_mul_epi32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(bb, 32));
let te = _mm256_mul_epi32(pe, qiv);
let re = _mm256_sub_epi64(pe, _mm256_mul_epi32(te, qv));
let to = _mm256_mul_epi32(po, qiv);
let ro = _mm256_sub_epi64(po, _mm256_mul_epi32(to, qv));
let res = _mm256_blend_epi32(_mm256_srli_epi64(re, 32), ro, 0xAA);
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, res);
return Self(r);
}
#[allow(unreachable_code)]
{
let qq = q.0[0] as i32 as i64;
let qi = qinv.0[0] as i32;
let mut r = [0u32; 8];
for i in 0..8 {
let p = (self.0[i] as i32 as i64) * (b.0[i] as i32 as i64);
let t = (p as i32).wrapping_mul(qi) as i64;
r[i] = (((p - t * qq) >> 32) as i32) as u32;
}
Self(r)
}
}
#[inline(always)]
pub fn ntt_bcast_lo(self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let v = match h {
4 => _mm256_permute2x128_si256::<0x00>(a, a),
2 => _mm256_shuffle_epi32::<0x44>(a),
1 => _mm256_shuffle_epi32::<0xA0>(a),
_ => unreachable!("i32 within-vector NTT stride is 4/2/1"),
};
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| self.0[(i / (2 * h)) * (2 * h) + (i % (2 * h)) % h]))
}
#[inline(always)]
pub fn ntt_bcast_hi(self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let v = match h {
4 => _mm256_permute2x128_si256::<0x11>(a, a),
2 => _mm256_shuffle_epi32::<0xEE>(a),
1 => _mm256_shuffle_epi32::<0xF5>(a),
_ => unreachable!("i32 within-vector NTT stride is 4/2/1"),
};
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| self.0[(i / (2 * h)) * (2 * h) + h + (i % (2 * h)) % h]))
}
#[inline(always)]
pub fn ntt_blend(self, o: Self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let v = match h {
4 => _mm256_permute2x128_si256::<0x30>(a, b),
2 => _mm256_blend_epi32::<0xCC>(a, b),
1 => _mm256_blend_epi32::<0xAA>(a, b),
_ => unreachable!("i32 within-vector NTT stride is 4/2/1"),
};
let mut r = [0u32; 8];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| if (i % (2 * h)) < h { self.0[i] } else { o.0[i] }))
}
}
impl ::core::ops::BitXor for Lanes8Word32 {
type Output = Self;
#[inline]
fn bitxor(self, o: Self) -> Self {
Lanes8Word32::bitxor(self, o)
}
}
impl ::core::ops::Add for Lanes8Word32 {
type Output = Self;
#[inline]
fn add(self, o: Self) -> Self {
Lanes8Word32::add(self, o)
}
}
impl ::core::ops::Sub for Lanes8Word32 {
type Output = Self;
#[inline]
fn sub(self, o: Self) -> Self {
Lanes8Word32::sub(self, o)
}
}
impl ::core::ops::BitAnd for Lanes8Word32 {
type Output = Self;
#[inline]
fn bitand(self, o: Self) -> Self {
Lanes8Word32::bitand(self, o)
}
}
impl ::core::ops::BitOr for Lanes8Word32 {
type Output = Self;
#[inline]
fn bitor(self, o: Self) -> Self {
Lanes8Word32::bitor(self, o)
}
}
impl ::core::ops::Not for Lanes8Word32 {
type Output = Self;
#[inline]
fn not(self) -> Self {
Lanes8Word32::not(self)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[repr(C, align(16))]
pub struct Lanes4Word32(pub [u32; 4]);
impl Lanes4Word32 {
pub const LANES: usize = 4;
#[inline]
pub const fn splat(x: u32) -> Self {
Self([x; 4])
}
#[inline]
pub fn from_words(s: &[Word32]) -> Self {
let mut a = [0u32; 4];
for (i, w) in s.iter().take(4).enumerate() {
a[i] = w.0;
}
Self(a)
}
#[inline]
pub fn to_words(self) -> [Word32; 4] {
self.0.map(Word32)
}
#[inline]
pub fn lane(self, i: usize) -> Word32 {
Word32(self.0[i])
}
#[inline(always)]
pub fn sha1rnds4(self, msg: Self, func: u32) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "sha"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(msg.0);
let r = match func & 3 {
0 => _mm_sha1rnds4_epu32(a, b, 0),
1 => _mm_sha1rnds4_epu32(a, b, 1),
2 => _mm_sha1rnds4_epu32(a, b, 2),
_ => _mm_sha1rnds4_epu32(a, b, 3),
};
Self(core::mem::transmute(r))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "sha")))]
{
#[cfg(target_arch = "x86_64")]
{
if shani_available() {
return unsafe { self.sha1rnds4_hw(msg, func) };
}
}
Self(crate::sha_ops::sha1rnds4(self.0, msg.0, func))
}
}
#[inline(always)]
pub fn sha1msg1(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "sha"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(o.0);
Self(core::mem::transmute(_mm_sha1msg1_epu32(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "sha")))]
{
#[cfg(target_arch = "x86_64")]
{
if shani_available() {
return unsafe { self.sha1msg1_hw(o) };
}
}
Self(crate::sha_ops::sha1msg1(self.0, o.0))
}
}
#[inline(always)]
pub fn sha1msg2(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "sha"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(o.0);
Self(core::mem::transmute(_mm_sha1msg2_epu32(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "sha")))]
{
#[cfg(target_arch = "x86_64")]
{
if shani_available() {
return unsafe { self.sha1msg2_hw(o) };
}
}
Self(crate::sha_ops::sha1msg2(self.0, o.0))
}
}
#[inline(always)]
pub fn sha1nexte(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "sha"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(o.0);
Self(core::mem::transmute(_mm_sha1nexte_epu32(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "sha")))]
{
#[cfg(target_arch = "x86_64")]
{
if shani_available() {
return unsafe { self.sha1nexte_hw(o) };
}
}
Self(crate::sha_ops::sha1nexte(self.0, o.0))
}
}
#[cfg(all(target_arch = "x86_64", not(target_feature = "sha")))]
#[target_feature(enable = "sha,sse2,sse4.1")]
unsafe fn sha1rnds4_hw(self, msg: Self, func: u32) -> Self {
use std::arch::x86_64::*;
let a = _mm_loadu_si128(self.0.as_ptr() as *const __m128i);
let b = _mm_loadu_si128(msg.0.as_ptr() as *const __m128i);
let r = match func & 3 {
0 => _mm_sha1rnds4_epu32(a, b, 0),
1 => _mm_sha1rnds4_epu32(a, b, 1),
2 => _mm_sha1rnds4_epu32(a, b, 2),
_ => _mm_sha1rnds4_epu32(a, b, 3),
};
let mut out = [0u32; 4];
_mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, r);
Self(out)
}
#[cfg(all(target_arch = "x86_64", not(target_feature = "sha")))]
#[target_feature(enable = "sha,sse2,ssse3,sse4.1")]
unsafe fn sha1msg1_hw(self, o: Self) -> Self {
use std::arch::x86_64::*;
let a = _mm_loadu_si128(self.0.as_ptr() as *const __m128i);
let b = _mm_loadu_si128(o.0.as_ptr() as *const __m128i);
let mut out = [0u32; 4];
_mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm_sha1msg1_epu32(a, b));
Self(out)
}
#[cfg(all(target_arch = "x86_64", not(target_feature = "sha")))]
#[target_feature(enable = "sha,sse2,ssse3,sse4.1")]
unsafe fn sha1msg2_hw(self, o: Self) -> Self {
use std::arch::x86_64::*;
let a = _mm_loadu_si128(self.0.as_ptr() as *const __m128i);
let b = _mm_loadu_si128(o.0.as_ptr() as *const __m128i);
let mut out = [0u32; 4];
_mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm_sha1msg2_epu32(a, b));
Self(out)
}
#[cfg(all(target_arch = "x86_64", not(target_feature = "sha")))]
#[target_feature(enable = "sha,sse2,sse4.1")]
unsafe fn sha1nexte_hw(self, o: Self) -> Self {
use std::arch::x86_64::*;
let a = _mm_loadu_si128(self.0.as_ptr() as *const __m128i);
let b = _mm_loadu_si128(o.0.as_ptr() as *const __m128i);
let mut out = [0u32; 4];
_mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm_sha1nexte_epu32(a, b));
Self(out)
}
#[inline]
pub fn add(self, o: Self) -> Self {
let mut out = [0u32; 4];
for i in 0..4 {
out[i] = self.0[i].wrapping_add(o.0[i]);
}
Self(out)
}
#[inline]
pub fn bitxor(self, o: Self) -> Self {
let mut out = [0u32; 4];
for i in 0..4 {
out[i] = self.0[i] ^ o.0[i];
}
Self(out)
}
}
impl ::core::ops::Add for Lanes4Word32 {
type Output = Self;
#[inline]
fn add(self, o: Self) -> Self {
Lanes4Word32::add(self, o)
}
}
impl ::core::ops::BitXor for Lanes4Word32 {
type Output = Self;
#[inline]
fn bitxor(self, o: Self) -> Self {
Lanes4Word32::bitxor(self, o)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[repr(C, align(16))]
pub struct Lanes16Word8(pub [u8; 16]);
impl Lanes16Word8 {
pub const LANES: usize = 16;
#[inline]
pub const fn splat(x: u8) -> Self {
Self([x; 16])
}
#[inline]
pub fn from_bytes(s: &[u8]) -> Self {
let mut a = [0u8; 16];
for (i, b) in s.iter().take(16).enumerate() {
a[i] = *b;
}
Self(a)
}
#[inline]
pub fn to_bytes(self) -> [u8; 16] {
self.0
}
#[inline]
pub fn lane(self, i: usize) -> u8 {
self.0[i]
}
#[inline]
pub fn shuffle(self, idx: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "ssse3"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(idx.0);
Self(core::mem::transmute(_mm_shuffle_epi8(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "ssse3")))]
{
let mut r = [0u8; 16];
for i in 0..16 {
let x = idx.0[i];
r[i] = if x & 0x80 != 0 { 0 } else { self.0[(x & 0x0f) as usize] };
}
Self(r)
}
}
#[inline]
pub fn bitand(self, o: Self) -> Self {
let mut r = [0u8; 16];
for i in 0..16 {
r[i] = self.0[i] & o.0[i];
}
Self(r)
}
#[inline]
pub fn shr_bytes(self, n: u32) -> Self {
let mut r = [0u8; 16];
for i in 0..16 {
r[i] = self.0[i] >> n;
}
Self(r)
}
#[inline]
pub fn interleave_lo(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "ssse3"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(o.0);
Self(core::mem::transmute(_mm_unpacklo_epi8(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "ssse3")))]
{
let mut r = [0u8; 16];
for i in 0..8 {
r[2 * i] = self.0[i];
r[2 * i + 1] = o.0[i];
}
Self(r)
}
}
#[inline]
pub fn interleave_hi(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "ssse3"))]
unsafe {
use core::arch::x86_64::*;
let a: __m128i = core::mem::transmute(self.0);
let b: __m128i = core::mem::transmute(o.0);
Self(core::mem::transmute(_mm_unpackhi_epi8(a, b)))
}
#[cfg(not(all(target_arch = "x86_64", target_feature = "ssse3")))]
{
let mut r = [0u8; 16];
for i in 0..8 {
r[2 * i] = self.0[8 + i];
r[2 * i + 1] = o.0[8 + i];
}
Self(r)
}
}
#[inline]
pub fn byte_add(self, o: Self) -> Self {
let mut r = [0u8; 16];
for i in 0..16 {
r[i] = self.0[i].wrapping_add(o.0[i]);
}
Self(r)
}
#[inline]
pub fn maddubs(self, o: Self) -> Self {
let mut r = [0u8; 16];
for i in 0..8 {
let a0 = self.0[2 * i] as i32;
let a1 = self.0[2 * i + 1] as i32;
let b0 = o.0[2 * i] as i8 as i32;
let b1 = o.0[2 * i + 1] as i8 as i32;
let s = (a0 * b0 + a1 * b1).clamp(-32768, 32767) as i16 as u16;
r[2 * i] = (s & 0xff) as u8;
r[2 * i + 1] = (s >> 8) as u8;
}
Self(r)
}
#[inline]
pub fn packus(self, o: Self) -> Self {
let mut r = [0u8; 16];
for i in 0..8 {
let a = (self.0[2 * i] as u16 | (self.0[2 * i + 1] as u16) << 8) as i16;
r[i] = a.clamp(0, 255) as u8;
let b = (o.0[2 * i] as u16 | (o.0[2 * i + 1] as u16) << 8) as i16;
r[8 + i] = b.clamp(0, 255) as u8;
}
Self(r)
}
}
impl ::core::ops::BitAnd for Lanes16Word8 {
type Output = Self;
#[inline]
fn bitand(self, o: Self) -> Self {
Lanes16Word8::bitand(self, o)
}
}
#[cfg(all(target_arch = "x86_64", not(target_feature = "sha")))]
#[inline]
fn shani_available() -> bool {
std::is_x86_feature_detected!("sha")
&& std::is_x86_feature_detected!("ssse3")
&& std::is_x86_feature_detected!("sse4.1")
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[repr(C, align(32))]
pub struct Lanes4Word64(pub [u64; 4]);
impl Lanes4Word64 {
pub const LANES: usize = 4;
#[inline]
pub fn from_words(s: &[Word64]) -> Self {
let mut a = [0u64; 4];
for (i, w) in s.iter().take(4).enumerate() {
a[i] = w.0;
}
Self(a)
}
#[inline]
pub fn to_words(self) -> [Word64; 4] {
self.0.map(Word64)
}
#[inline]
pub fn lane(self, i: usize) -> Word64 {
Word64(self.0[i])
}
#[inline]
pub fn hsum(self) -> u64 {
self.0[0]
.wrapping_add(self.0[1])
.wrapping_add(self.0[2])
.wrapping_add(self.0[3])
}
#[inline(always)]
pub fn add(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_add_epi64(a, b));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u64; 4];
for i in 0..4 {
r[i] = self.0[i].wrapping_add(o.0[i]);
}
Self(r)
}
}
#[cfg(test)]
#[inline]
fn add_scalar(self, o: Self) -> Self {
let mut r = [0u64; 4];
for i in 0..4 {
r[i] = self.0[i].wrapping_add(o.0[i]);
}
Self(r)
}
#[inline(always)]
pub fn mul_lo32_wide(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_mul_epu32(a, b));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u64; 4];
for i in 0..4 {
r[i] = (self.0[i] & 0xffff_ffff) * (o.0[i] & 0xffff_ffff);
}
Self(r)
}
}
#[cfg(test)]
#[inline]
fn mul_lo32_wide_scalar(self, o: Self) -> Self {
let mut r = [0u64; 4];
for i in 0..4 {
r[i] = (self.0[i] & 0xffff_ffff) * (o.0[i] & 0xffff_ffff);
}
Self(r)
}
#[inline(always)]
pub fn splat(x: u64) -> Self {
Self([x; 4])
}
#[inline(always)]
pub fn bitxor(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_xor_si256(a, b));
return Self(r);
}
#[allow(unreachable_code)]
Self([self.0[0] ^ o.0[0], self.0[1] ^ o.0[1], self.0[2] ^ o.0[2], self.0[3] ^ o.0[3]])
}
#[inline(always)]
pub fn and(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_and_si256(a, b));
return Self(r);
}
#[allow(unreachable_code)]
Self([self.0[0] & o.0[0], self.0[1] & o.0[1], self.0[2] & o.0[2], self.0[3] & o.0[3]])
}
#[inline(always)]
pub fn andnot(self, o: Self) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_andnot_si256(a, b));
return Self(r);
}
#[allow(unreachable_code)]
Self([!self.0[0] & o.0[0], !self.0[1] & o.0[1], !self.0[2] & o.0[2], !self.0[3] & o.0[3]])
}
#[inline(always)]
pub fn rotl(self, n: u32) -> Self {
let n = n % 64;
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
unsafe {
use std::arch::x86_64::*;
if n == 0 {
return self;
}
let x = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let l = _mm256_sll_epi64(x, _mm_cvtsi32_si128(n as i32));
let r_sh = _mm256_srl_epi64(x, _mm_cvtsi32_si128((64 - n) as i32));
let mut r = [0u64; 4];
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, _mm256_or_si256(l, r_sh));
return Self(r);
}
#[allow(unreachable_code)]
{
let mut r = [0u64; 4];
for i in 0..4 {
r[i] = self.0[i].rotate_left(n);
}
Self(r)
}
}
}
impl ::core::ops::Add for Lanes4Word64 {
type Output = Self;
#[inline]
fn add(self, o: Self) -> Self {
Lanes4Word64::add(self, o)
}
}
impl ::core::ops::BitXor for Lanes4Word64 {
type Output = Self;
#[inline]
fn bitxor(self, o: Self) -> Self {
Lanes4Word64::bitxor(self, o)
}
}
impl ::core::ops::BitAnd for Lanes4Word64 {
type Output = Self;
#[inline]
fn bitand(self, o: Self) -> Self {
Lanes4Word64::and(self, o)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
#[repr(C, align(32))]
pub struct Lanes16Word16(pub [u16; 16]);
#[cfg(target_arch = "x86_64")]
type Avx256 = std::arch::x86_64::__m256i;
#[cfg(not(target_arch = "x86_64"))]
type Avx256 = [u16; 16];
#[cfg(target_arch = "x86_64")]
#[inline]
fn simd_add(a: Avx256, b: Avx256) -> Avx256 {
unsafe { std::arch::x86_64::_mm256_add_epi16(a, b) }
}
#[cfg(target_arch = "x86_64")]
#[inline]
fn simd_sub(a: Avx256, b: Avx256) -> Avx256 {
unsafe { std::arch::x86_64::_mm256_sub_epi16(a, b) }
}
#[cfg(target_arch = "x86_64")]
#[inline]
fn simd_mullo(a: Avx256, b: Avx256) -> Avx256 {
unsafe { std::arch::x86_64::_mm256_mullo_epi16(a, b) }
}
#[cfg(target_arch = "x86_64")]
#[inline]
fn simd_mulhi(a: Avx256, b: Avx256) -> Avx256 {
unsafe { std::arch::x86_64::_mm256_mulhi_epi16(a, b) }
}
#[cfg(not(target_arch = "x86_64"))]
#[inline]
fn simd_add(_a: Avx256, _b: Avx256) -> Avx256 {
unreachable!("SIMD lane path is x86_64-only")
}
#[cfg(not(target_arch = "x86_64"))]
#[inline]
fn simd_sub(_a: Avx256, _b: Avx256) -> Avx256 {
unreachable!("SIMD lane path is x86_64-only")
}
#[cfg(not(target_arch = "x86_64"))]
#[inline]
fn simd_mullo(_a: Avx256, _b: Avx256) -> Avx256 {
unreachable!("SIMD lane path is x86_64-only")
}
#[cfg(not(target_arch = "x86_64"))]
#[inline]
fn simd_mulhi(_a: Avx256, _b: Avx256) -> Avx256 {
unreachable!("SIMD lane path is x86_64-only")
}
impl Lanes16Word16 {
pub const LANES: usize = 16;
#[inline]
pub const fn splat(x: u16) -> Self {
Self([x; 16])
}
#[inline]
pub fn from_words(s: &[Word16]) -> Self {
let mut a = [0u16; 16];
for (i, w) in s.iter().take(16).enumerate() {
a[i] = w.0;
}
Self(a)
}
#[inline]
pub fn to_words(self) -> [Word16; 16] {
self.0.map(Word16)
}
#[inline]
pub fn lane(self, i: usize) -> Word16 {
Word16(self.0[i])
}
#[inline(always)]
pub fn add(self, o: Self) -> Self {
self.binop(o, |a, b| a.wrapping_add(b), |a, b| simd_add(a, b))
}
#[inline(always)]
pub fn sub(self, o: Self) -> Self {
self.binop(o, |a, b| a.wrapping_sub(b), |a, b| simd_sub(a, b))
}
#[inline(always)]
pub fn mullo(self, o: Self) -> Self {
self.binop(o, |a, b| a.wrapping_mul(b), |a, b| simd_mullo(a, b))
}
#[inline(always)]
pub fn mulhi(self, o: Self) -> Self {
self.binop(
o,
|a, b| (((a as i16 as i32) * (b as i16 as i32)) >> 16) as u16,
|a, b| simd_mulhi(a, b),
)
}
#[inline(always)]
pub fn ntt_bcast_lo(self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
{
use std::arch::x86_64::*;
let mut r = [0u16; 16];
unsafe {
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let v = match h {
8 => _mm256_permute2x128_si256::<0x00>(a, a),
4 => _mm256_shuffle_epi32::<0x44>(a),
2 => _mm256_shuffle_epi32::<0xA0>(a),
_ => unreachable!("within-vector NTT stride is 8/4/2"),
};
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
}
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| self.0[(i / (2 * h)) * (2 * h) + (i % (2 * h)) % h]))
}
#[inline(always)]
pub fn ntt_bcast_hi(self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
{
use std::arch::x86_64::*;
let mut r = [0u16; 16];
unsafe {
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let v = match h {
8 => _mm256_permute2x128_si256::<0x11>(a, a),
4 => _mm256_shuffle_epi32::<0xEE>(a),
2 => _mm256_shuffle_epi32::<0xF5>(a),
_ => unreachable!("within-vector NTT stride is 8/4/2"),
};
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
}
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| self.0[(i / (2 * h)) * (2 * h) + h + (i % (2 * h)) % h]))
}
#[inline(always)]
pub fn ntt_blend(self, o: Self, h: usize) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
{
use std::arch::x86_64::*;
let mut r = [0u16; 16];
unsafe {
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
let v = match h {
8 => _mm256_permute2x128_si256::<0x30>(a, b),
4 => _mm256_blend_epi32::<0xCC>(a, b),
2 => _mm256_blend_epi32::<0xAA>(a, b),
_ => unreachable!("within-vector NTT stride is 8/4/2"),
};
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, v);
}
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| if (i % (2 * h)) < h { self.0[i] } else { o.0[i] }))
}
#[inline(always)]
fn binop(
self,
o: Self,
scalar: impl Fn(u16, u16) -> u16,
#[allow(unused_variables)] vector: impl Fn(Avx256, Avx256) -> Avx256,
) -> Self {
#[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
{
use std::arch::x86_64::*;
let mut r = [0u16; 16];
unsafe {
let a = _mm256_loadu_si256(self.0.as_ptr() as *const __m256i);
let b = _mm256_loadu_si256(o.0.as_ptr() as *const __m256i);
_mm256_storeu_si256(r.as_mut_ptr() as *mut __m256i, vector(a, b));
}
return Self(r);
}
#[allow(unreachable_code)]
Self(core::array::from_fn(|i| scalar(self.0[i], o.0[i])))
}
}
macro_rules! lanes16_op {
($trait:ident, $tm:ident, $m:ident) => {
impl ::core::ops::$trait for Lanes16Word16 {
type Output = Self;
#[inline]
fn $tm(self, o: Self) -> Self {
Lanes16Word16::$m(self, o)
}
}
};
}
lanes16_op!(Add, add, add);
lanes16_op!(Sub, sub, sub);
lanes16_op!(Mul, mul, mullo);
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum LanesVal {
L8W32(Lanes8Word32),
L4W64(Lanes4Word64),
L16W16(Lanes16Word16),
L4W32(Lanes4Word32),
L16W8(Lanes16Word8),
}
impl LanesVal {
#[inline]
pub const fn lanes(self) -> usize {
match self {
LanesVal::L8W32(_) => 8,
LanesVal::L4W64(_) => 4,
LanesVal::L16W16(_) => 16,
LanesVal::L4W32(_) => 4,
LanesVal::L16W8(_) => 16,
}
}
#[inline]
pub const fn elem_bits(self) -> u32 {
match self {
LanesVal::L8W32(_) => 32,
LanesVal::L4W64(_) => 64,
LanesVal::L16W16(_) => 16,
LanesVal::L4W32(_) => 32,
LanesVal::L16W8(_) => 8,
}
}
#[inline]
pub const fn type_name(self) -> &'static str {
match self {
LanesVal::L8W32(_) => "Lanes8Word32",
LanesVal::L4W64(_) => "Lanes4Word64",
LanesVal::L16W16(_) => "Lanes16Word16",
LanesVal::L4W32(_) => "Lanes4Word32",
LanesVal::L16W8(_) => "Lanes16Word8",
}
}
#[inline]
pub fn lane(self, i: usize) -> u64 {
match self {
LanesVal::L8W32(v) => v.lane(i).0 as u64,
LanesVal::L4W64(v) => v.lane(i).0,
LanesVal::L16W16(v) => v.lane(i).0 as u64,
LanesVal::L4W32(v) => v.lane(i).0 as u64,
LanesVal::L16W8(v) => v.lane(i) as u64,
}
}
#[inline]
pub fn shuffle(self, idx: Self) -> Option<Self> {
match (self, idx) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.shuffle(b))),
_ => None,
}
}
#[inline]
pub fn byte_and(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.bitand(b))),
_ => None,
}
}
#[inline]
pub fn shr_bytes(self, n: u32) -> Option<Self> {
match self {
LanesVal::L16W8(a) => Some(LanesVal::L16W8(a.shr_bytes(n))),
_ => None,
}
}
#[inline]
pub fn interleave_lo(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.interleave_lo(b))),
_ => None,
}
}
#[inline]
pub fn interleave_hi(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.interleave_hi(b))),
_ => None,
}
}
#[inline]
pub fn byte_add(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.byte_add(b))),
_ => None,
}
}
#[inline]
pub fn maddubs_bytes(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.maddubs(b))),
_ => None,
}
}
#[inline]
pub fn packus_bytes(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.packus(b))),
_ => None,
}
}
#[inline]
pub fn sha1rnds4(self, msg: Self, func: u32) -> Option<Self> {
match (self, msg) {
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.sha1rnds4(b, func))),
_ => None,
}
}
#[inline]
pub fn sha1msg1(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.sha1msg1(b))),
_ => None,
}
}
#[inline]
pub fn sha1msg2(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.sha1msg2(b))),
_ => None,
}
}
#[inline]
pub fn sha1nexte(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.sha1nexte(b))),
_ => None,
}
}
#[inline]
pub fn bitxor(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.bitxor(b))),
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.bitxor(b))),
_ => None,
}
}
#[inline]
pub fn bitand(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.bitand(b))),
(LanesVal::L16W8(a), LanesVal::L16W8(b)) => Some(LanesVal::L16W8(a.bitand(b))),
_ => None,
}
}
#[inline]
pub fn bitor(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.bitor(b))),
_ => None,
}
}
#[inline]
pub fn lane_not(self) -> Option<Self> {
match self {
LanesVal::L8W32(a) => Some(LanesVal::L8W32(a.not())),
_ => None,
}
}
#[inline]
pub fn add(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.add(b))),
(LanesVal::L4W64(a), LanesVal::L4W64(b)) => Some(LanesVal::L4W64(a.add(b))),
(LanesVal::L16W16(a), LanesVal::L16W16(b)) => Some(LanesVal::L16W16(a.add(b))),
(LanesVal::L4W32(a), LanesVal::L4W32(b)) => Some(LanesVal::L4W32(a.add(b))),
_ => None,
}
}
#[inline]
pub fn sub(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W16(a), LanesVal::L16W16(b)) => Some(LanesVal::L16W16(a.sub(b))),
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.sub(b))),
_ => None,
}
}
#[inline]
pub fn mullo(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W16(a), LanesVal::L16W16(b)) => Some(LanesVal::L16W16(a.mullo(b))),
_ => None,
}
}
#[inline]
pub fn mulhi(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L16W16(a), LanesVal::L16W16(b)) => Some(LanesVal::L16W16(a.mulhi(b))),
_ => None,
}
}
#[inline]
pub fn ntt_bcast_lo(self, h: usize) -> Option<Self> {
match self {
LanesVal::L16W16(a) => Some(LanesVal::L16W16(a.ntt_bcast_lo(h))),
LanesVal::L8W32(a) => Some(LanesVal::L8W32(a.ntt_bcast_lo(h))),
_ => None,
}
}
#[inline]
pub fn ntt_bcast_hi(self, h: usize) -> Option<Self> {
match self {
LanesVal::L16W16(a) => Some(LanesVal::L16W16(a.ntt_bcast_hi(h))),
LanesVal::L8W32(a) => Some(LanesVal::L8W32(a.ntt_bcast_hi(h))),
_ => None,
}
}
#[inline]
pub fn ntt_blend(self, o: Self, h: usize) -> Option<Self> {
match (self, o) {
(LanesVal::L16W16(a), LanesVal::L16W16(b)) => Some(LanesVal::L16W16(a.ntt_blend(b, h))),
(LanesVal::L8W32(a), LanesVal::L8W32(b)) => Some(LanesVal::L8W32(a.ntt_blend(b, h))),
_ => None,
}
}
#[inline]
pub fn rotl(self, n: u32) -> Option<Self> {
match self {
LanesVal::L8W32(a) => Some(LanesVal::L8W32(a.rotl(n))),
_ => None,
}
}
#[inline]
pub fn mul_lo32_wide(self, o: Self) -> Option<Self> {
match (self, o) {
(LanesVal::L4W64(a), LanesVal::L4W64(b)) => Some(LanesVal::L4W64(a.mul_lo32_wide(b))),
_ => None,
}
}
#[inline]
pub fn montmul32(self, b: Self, q: Self, qinv: Self) -> Option<Self> {
match (self, b, q, qinv) {
(LanesVal::L8W32(a), LanesVal::L8W32(b), LanesVal::L8W32(q), LanesVal::L8W32(qi)) => {
Some(LanesVal::L8W32(a.montmul32(b, q, qi)))
}
_ => None,
}
}
#[inline]
pub fn hsum(self) -> i64 {
match self {
LanesVal::L8W32(a) => (0..8).map(|i| a.lane(i).0 as i64).sum(),
LanesVal::L4W64(a) => a.hsum() as i64,
LanesVal::L16W16(a) => (0..16).map(|i| a.lane(i).0 as i64).sum(),
LanesVal::L4W32(a) => (0..4).map(|i| a.lane(i).0 as i64).sum(),
LanesVal::L16W8(a) => (0..16).map(|i| a.lane(i) as i64).sum(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn word32_add_wraps_not_promotes() {
assert_eq!(Word32::MAX.add(Word32::ONE), Word32::ZERO);
assert_eq!(Word32(0xFFFF_FFFF).add(Word32(2)), Word32(1));
}
#[test]
fn word64_add_wraps_not_promotes() {
assert_eq!(Word64::MAX.add(Word64::ONE), Word64::ZERO);
assert_eq!(Word64(u64::MAX).add(Word64(5)), Word64(4));
}
#[test]
fn div_rem_shl_shr_operators_are_unsigned_primitive() {
use core::ops::{Div, Rem, Shl, Shr};
assert_eq!(Word32(0x1234_5678).rem(Word32(65536)), Word32(0x5678), "% 2^16 masks the low bits");
assert_eq!(Word32(0x1234_5678).div(Word32(65536)), Word32(0x1234), "/ 2^16 shifts right 16");
assert_eq!(Word32(7).div(Word32(2)), Word32(3), "unsigned truncating division");
assert_eq!(Word32(7).rem(Word32(3)), Word32(1), "unsigned remainder");
assert_eq!(Word32(1).shl(31), Word32(0x8000_0000), "<< 31 sets the top bit");
assert_eq!(Word32(0x8000_0000).shr(31), Word32(1), ">> 31 is a logical (unsigned) shift");
assert_eq!(Word64(1).shl(63).shr(63), Word64(1), "Word64 logical shift round-trips the top bit");
assert_eq!(Word8(200).div(Word8(8)), Word8(25), "Word8 unsigned division (200 ÷ 8)");
}
#[test]
fn rotl_is_cyclic_and_inverse_of_rotr() {
let x = Word32(0x1234_5678);
assert_eq!(x.rotl(8), Word32(0x3456_7812), "rotl by 8 is the documented permutation");
assert_eq!(x.rotl(0), x, "rotate by 0 is identity");
assert_eq!(x.rotl(32), x, "a full rotation is identity");
for n in 0..Word32::BITS {
assert_eq!(x.rotl(n).rotr(n), x, "rotr undoes rotl by {n}");
}
}
#[test]
fn agrees_with_native_wrapping_ops_over_fuzz() {
let mut s = 0xDEAD_BEEF_CAFE_F00Du64;
let mut next = || {
s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407);
s
};
for _ in 0..5000 {
let a = next() as u32;
let b = next() as u32;
let n = (next() % 32) as u32;
assert_eq!(Word32(a).add(Word32(b)).get(), a.wrapping_add(b), "add");
assert_eq!(Word32(a).sub(Word32(b)).get(), a.wrapping_sub(b), "sub");
assert_eq!(Word32(a).mul(Word32(b)).get(), a.wrapping_mul(b), "mul");
assert_eq!(Word32(a).bitxor(Word32(b)).get(), a ^ b, "xor");
assert_eq!(Word32(a).bitand(Word32(b)).get(), a & b, "and");
assert_eq!(Word32(a).bitor(Word32(b)).get(), a | b, "or");
assert_eq!(Word32(a).not().get(), !a, "not");
assert_eq!(Word32(a).shl(n).get(), a.wrapping_shl(n), "shl");
assert_eq!(Word32(a).shr(n).get(), a.wrapping_shr(n), "shr");
assert_eq!(Word32(a).rotl(n).get(), a.rotate_left(n), "rotl");
assert_eq!(Word32(a).rotr(n).get(), a.rotate_right(n), "rotr");
}
}
#[test]
fn lanes8word32_avx2_xor_equals_scalar_lanes() {
let mut s = 0x0BAD_F00D_1234_5678u64;
let mut next = || {
s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
s
};
for _ in 0..2000 {
let a = Lanes8Word32(core::array::from_fn(|_| next() as u32));
let b = Lanes8Word32(core::array::from_fn(|_| next() as u32));
let n = (next() % 64) as u32; let xor_exp = Lanes8Word32(core::array::from_fn(|i| a.0[i] ^ b.0[i]));
assert_eq!(a.bitxor_scalar(b), xor_exp, "scalar lane xor is the spec");
assert_eq!(a.bitxor(b), xor_exp, "dispatched (AVX2) xor == spec");
let add_exp = Lanes8Word32(core::array::from_fn(|i| a.0[i].wrapping_add(b.0[i])));
assert_eq!(a.add_scalar(b), add_exp, "scalar lane add is the spec");
assert_eq!(a.add(b), add_exp, "dispatched (AVX2) add == spec");
let rot_exp = Lanes8Word32(core::array::from_fn(|i| a.0[i].rotate_left(n % 32)));
assert_eq!(a.rotl_scalar(n % 32), rot_exp, "scalar lane rotl is the spec");
assert_eq!(a.rotl(n), rot_exp, "dispatched (AVX2) rotl == spec (n mod 32)");
let sub_exp = Lanes8Word32(core::array::from_fn(|i| a.0[i].wrapping_sub(b.0[i])));
assert_eq!(a.sub_scalar(b), sub_exp, "scalar lane sub32 is the spec");
assert_eq!(a.sub(b), sub_exp, "dispatched (AVX2) sub32 == spec");
const Q: i32 = 8_380_417;
const QINV: i32 = 58_728_449;
let qv = Lanes8Word32::splat(Q as u32);
let qiv = Lanes8Word32::splat(QINV as u32);
let ar = Lanes8Word32(core::array::from_fn(|i| (a.0[i] as i32 % Q) as u32));
let br = Lanes8Word32(core::array::from_fn(|i| (b.0[i] as i32 % Q) as u32));
let mont_exp = Lanes8Word32(core::array::from_fn(|i| {
let p = (ar.0[i] as i32 as i64) * (br.0[i] as i32 as i64);
let t = (p as i32).wrapping_mul(QINV) as i64;
(((p - t * Q as i64) >> 32) as i32) as u32
}));
assert_eq!(ar.montmul32_scalar(br, qv, qiv), mont_exp, "scalar montmul32 is the spec");
assert_eq!(ar.montmul32(br, qv, qiv), mont_exp, "dispatched (AVX2) montmul32 == spec");
for &h in &[4usize, 2, 1] {
let bl = Lanes8Word32(core::array::from_fn(|i| a.0[(i / (2 * h)) * (2 * h) + (i % (2 * h)) % h]));
let bh = Lanes8Word32(core::array::from_fn(|i| a.0[(i / (2 * h)) * (2 * h) + h + (i % (2 * h)) % h]));
let bd = Lanes8Word32(core::array::from_fn(|i| if (i % (2 * h)) < h { a.0[i] } else { b.0[i] }));
assert_eq!(a.ntt_bcast_lo(h), bl, "i32 ntt_bcast_lo({h}) (AVX2) == scalar gather");
assert_eq!(a.ntt_bcast_hi(h), bh, "i32 ntt_bcast_hi({h}) (AVX2) == scalar gather");
assert_eq!(a.ntt_blend(b, h), bd, "i32 ntt_blend({h}) (AVX2) == scalar gather");
}
}
let v = Lanes8Word32::splat(0xABCD_1234);
assert_eq!(v.lane(3), Word32(0xABCD_1234));
assert_eq!(v.to_words()[7], Word32(0xABCD_1234));
let packed = Lanes8Word32::from_words(&[Word32(1), Word32(2), Word32(3)]);
assert_eq!(packed.0, [1, 2, 3, 0, 0, 0, 0, 0], "short slice zero-fills");
}
#[test]
fn lanes4word64_avx2_add_and_widemul_equal_scalar_lanes() {
let mut s = 0x1357_9bdf_2468_ace0u64;
let mut next = || {
s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
s
};
for _ in 0..2000 {
let a = Lanes4Word64(core::array::from_fn(|_| next()));
let b = Lanes4Word64(core::array::from_fn(|_| next()));
let add_exp = Lanes4Word64(core::array::from_fn(|i| a.0[i].wrapping_add(b.0[i])));
assert_eq!(a.add_scalar(b), add_exp, "scalar lane add64 is the spec");
assert_eq!(a.add(b), add_exp, "dispatched (AVX2) add64 == spec");
let mul_exp =
Lanes4Word64(core::array::from_fn(|i| (a.0[i] & 0xffff_ffff) * (b.0[i] & 0xffff_ffff)));
assert_eq!(a.mul_lo32_wide_scalar(b), mul_exp, "scalar widening mul is the spec");
assert_eq!(a.mul_lo32_wide(b), mul_exp, "dispatched (AVX2) vpmuludq == spec");
let hs = a.0[0].wrapping_add(a.0[1]).wrapping_add(a.0[2]).wrapping_add(a.0[3]);
assert_eq!(a.hsum(), hs, "horizontal sum");
}
let v = Lanes4Word64::from_words(&[Word64(5), Word64(9)]);
assert_eq!(v.0, [5, 9, 0, 0], "short slice zero-fills");
assert_eq!(v.lane(1), Word64(9));
assert_eq!(LanesVal::L4W64(v).hsum(), 14, "5 + 9 + 0 + 0");
assert_eq!(LanesVal::L4W64(v).type_name(), "Lanes4Word64");
}
#[test]
fn lanes16word16_avx2_ntt_ops_equal_scalar_lanes() {
let mut s = 0x2468_ace0_1357_9bdfu64;
let mut next = || {
s = s.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1);
(s >> 40) as u16
};
for _ in 0..3000 {
let a = Lanes16Word16(core::array::from_fn(|_| next()));
let b = Lanes16Word16(core::array::from_fn(|_| next()));
let add_exp = Lanes16Word16(core::array::from_fn(|i| a.0[i].wrapping_add(b.0[i])));
let sub_exp = Lanes16Word16(core::array::from_fn(|i| a.0[i].wrapping_sub(b.0[i])));
let lo_exp = Lanes16Word16(core::array::from_fn(|i| a.0[i].wrapping_mul(b.0[i])));
let hi_exp = Lanes16Word16(core::array::from_fn(|i| {
(((a.0[i] as i16 as i32) * (b.0[i] as i16 as i32)) >> 16) as u16
}));
assert_eq!(a.add(b), add_exp, "lane add16 == spec");
assert_eq!(a.sub(b), sub_exp, "lane sub16 == spec");
assert_eq!(a.mullo(b), lo_exp, "lane mullo16 == spec");
assert_eq!(a.mulhi(b), hi_exp, "lane (signed) mulhi16 == spec");
for &h in &[8usize, 4, 2] {
let bl = Lanes16Word16(core::array::from_fn(|i| a.0[(i / (2 * h)) * (2 * h) + (i % (2 * h)) % h]));
let bh = Lanes16Word16(core::array::from_fn(|i| a.0[(i / (2 * h)) * (2 * h) + h + (i % (2 * h)) % h]));
let bd = Lanes16Word16(core::array::from_fn(|i| if (i % (2 * h)) < h { a.0[i] } else { b.0[i] }));
assert_eq!(a.ntt_bcast_lo(h), bl, "ntt_bcast_lo({h}) (AVX2) == scalar gather");
assert_eq!(a.ntt_bcast_hi(h), bh, "ntt_bcast_hi({h}) (AVX2) == scalar gather");
assert_eq!(a.ntt_blend(b, h), bd, "ntt_blend({h}) (AVX2) == scalar gather");
}
}
assert_eq!(Lanes16Word16::splat(7).lane(9), Word16(7));
let p = Lanes16Word16::from_words(&[Word16(3), Word16(5)]);
assert_eq!(p.0[0..3], [3, 5, 0], "short slice zero-fills");
assert_eq!(LanesVal::L16W16(p).type_name(), "Lanes16Word16");
assert_eq!(LanesVal::L16W16(p).lanes(), 16);
}
#[test]
fn wordval_dispatches_by_width_and_rejects_mismatch() {
let a = WordVal::W32(Word32(0xFFFF_FFFF));
let b = WordVal::W32(Word32(1));
assert_eq!(a.add(b), Some(WordVal::W32(Word32(0))), "same-width add wraps");
assert_eq!(a.width(), 32);
assert_eq!(a.to_u64(), 0xFFFF_FFFF);
let c = WordVal::W64(Word64(1));
assert_eq!(a.add(c), None, "mixed-width add is a type error");
assert_eq!(a.bitxor(c), None, "mixed-width xor is a type error");
assert_eq!(WordVal::W64(Word64(0)).not(), WordVal::W64(Word64(u64::MAX)));
assert_eq!(WordVal::W32(Word32(0x1234_5678)).rotl(8), WordVal::W32(Word32(0x3456_7812)));
assert_eq!(WordVal::from_u64(32, 0x1_0000_0005), Some(WordVal::W32(Word32(5))), "32 truncates");
assert_eq!(WordVal::from_u64(64, 5), Some(WordVal::W64(Word64(5))));
assert_eq!(WordVal::from_u64(16, 5), None, "only 32/64 are valid widths");
assert_eq!(WordVal::W32(Word32(42)).to_string(), "42");
}
}