use crate::reed_solomon::engine::{
CANTOR_BASIS, GF_BITS, GF_MODULUS, GF_ORDER, GF_POLYNOMIAL, GfElement, fwht, utils,
};
#[cfg(not(feature = "std"))]
use alloc::boxed::Box;
#[cfg(not(feature = "std"))]
use alloc::vec;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
#[cfg(not(feature = "std"))]
use once_cell::race::OnceBox;
#[cfg(feature = "std")]
use std::sync::{LazyLock, OnceLock};
pub type Exp = [GfElement; GF_ORDER];
pub type Log = [GfElement; GF_ORDER];
pub type Mul128 = [Multiply128lutT; GF_ORDER];
#[cfg(any(test, target_arch = "x86", target_arch = "x86_64"))]
pub(crate) type MulGfni = [MultiplyGfni; GF_ORDER];
#[cfg(any(test, target_arch = "x86", target_arch = "x86_64"))]
#[derive(Clone, Debug)]
pub(crate) struct MultiplyGfni {
pub(crate) low_from_low: u64,
pub(crate) low_from_high: u64,
pub(crate) high_from_low: u64,
pub(crate) high_from_high: u64,
}
#[derive(Clone, Debug)]
pub struct Multiply128lutT {
pub lo: [u128; 4],
pub hi: [u128; 4],
}
pub type LogWalsh = [GfElement; GF_ORDER];
pub type Mul16 = [[[GfElement; 16]; 4]; GF_ORDER];
pub type Skew = [GfElement; GF_MODULUS as usize];
pub struct ExpLog {
pub exp: Box<Exp>,
pub log: Box<Log>,
}
pub fn get_exp_log() -> &'static ExpLog {
#[cfg(feature = "std")]
{
static EXP_LOG: LazyLock<ExpLog> = LazyLock::new(initialize_exp_log);
&EXP_LOG
}
#[cfg(not(feature = "std"))]
{
static EXP_LOG: OnceBox<ExpLog> = OnceBox::new();
EXP_LOG.get_or_init(|| Box::new(initialize_exp_log()))
}
}
pub fn get_log_walsh() -> &'static LogWalsh {
#[cfg(feature = "std")]
{
static LOG_WALSH: LazyLock<Box<LogWalsh>> = LazyLock::new(initialize_log_walsh);
&LOG_WALSH
}
#[cfg(not(feature = "std"))]
{
static LOG_WALSH: OnceBox<LogWalsh> = OnceBox::new();
LOG_WALSH.get_or_init(initialize_log_walsh)
}
}
pub(crate) fn get_short_log_walsh(n: usize) -> &'static [GfElement] {
assert!(n.is_power_of_two() && n < GF_ORDER);
let level = n.trailing_zeros() as usize;
#[cfg(feature = "std")]
{
static KERNELS: [OnceLock<Vec<GfElement>>; GF_BITS] = [const { OnceLock::new() }; GF_BITS];
KERNELS[level].get_or_init(|| initialize_short_log_walsh(n))
}
#[cfg(not(feature = "std"))]
{
static KERNELS: [OnceBox<Vec<GfElement>>; GF_BITS] = [const { OnceBox::new() }; GF_BITS];
KERNELS[level].get_or_init(|| Box::new(initialize_short_log_walsh(n)))
}
}
pub fn get_mul16() -> &'static Mul16 {
#[cfg(feature = "std")]
{
static MUL16: LazyLock<Box<Mul16>> = LazyLock::new(initialize_mul16);
&MUL16
}
#[cfg(not(feature = "std"))]
{
static MUL16: OnceBox<Mul16> = OnceBox::new();
MUL16.get_or_init(initialize_mul16)
}
}
pub fn get_mul128() -> &'static Mul128 {
#[cfg(feature = "std")]
{
static MUL128: LazyLock<Box<Mul128>> = LazyLock::new(initialize_mul128);
&MUL128
}
#[cfg(not(feature = "std"))]
{
static MUL128: OnceBox<Mul128> = OnceBox::new();
MUL128.get_or_init(initialize_mul128)
}
}
#[cfg(any(test, target_arch = "x86", target_arch = "x86_64"))]
pub(crate) fn get_mul_gfni() -> &'static MulGfni {
#[cfg(feature = "std")]
{
static MUL_GFNI: LazyLock<Box<MulGfni>> = LazyLock::new(initialize_mul_gfni);
&MUL_GFNI
}
#[cfg(not(feature = "std"))]
{
static MUL_GFNI: OnceBox<MulGfni> = OnceBox::new();
MUL_GFNI.get_or_init(initialize_mul_gfni)
}
}
pub fn get_skew() -> &'static Skew {
#[cfg(feature = "std")]
{
static SKEW: LazyLock<Box<Skew>> = LazyLock::new(initialize_skew);
&SKEW
}
#[cfg(not(feature = "std"))]
{
static SKEW: OnceBox<Skew> = OnceBox::new();
SKEW.get_or_init(initialize_skew)
}
}
#[inline(always)]
pub fn mul(x: GfElement, log_m: GfElement, exp: &Exp, log: &Log) -> GfElement {
if x == 0 {
0
} else {
exp[utils::add_mod(log[x as usize], log_m) as usize]
}
}
fn initialize_exp_log() -> ExpLog {
let mut exp = Box::new([0; GF_ORDER]);
let mut log = Box::new([0; GF_ORDER]);
let mut state = 1;
for i in 0..GF_MODULUS {
exp[state] = i;
state <<= 1;
if state >= GF_ORDER {
state ^= GF_POLYNOMIAL;
}
}
exp[0] = GF_MODULUS;
log[0] = 0;
for (i, basis) in CANTOR_BASIS.iter().copied().enumerate().take(GF_BITS) {
let width = 1usize << i;
for j in 0..width {
log[j + width] = log[j] ^ basis;
}
}
for value in log.iter_mut() {
*value = exp[*value as usize];
}
for (i, value) in log.iter().copied().enumerate() {
exp[value as usize] = i as GfElement;
}
exp[GF_MODULUS as usize] = exp[0];
ExpLog { exp, log }
}
fn initialize_log_walsh() -> Box<LogWalsh> {
let log = get_exp_log().log.as_slice();
let mut log_walsh: Box<LogWalsh> = Box::new([0; GF_ORDER]);
log_walsh.copy_from_slice(log);
log_walsh[0] = 0;
fwht::fwht(log_walsh.as_mut(), GF_ORDER);
log_walsh
}
fn initialize_short_log_walsh(n: usize) -> Vec<GfElement> {
let log = &get_exp_log().log;
let mut kernel = log[..n].to_vec();
kernel[0] = 0;
fwht::fwht(&mut kernel, n);
for factor in &mut kernel {
*factor = factor.rotate_right(n.trailing_zeros());
}
kernel
}
fn initialize_mul16() -> Box<Mul16> {
let exp = &get_exp_log().exp;
let log = &get_exp_log().log;
let mut mul16 = vec![[[0; 16]; 4]; GF_ORDER];
for log_m in 0..=GF_MODULUS {
let lut = &mut mul16[log_m as usize];
let [row0, row1, row2, row3] = lut;
for (i, (((x0, x1), x2), x3)) in row0
.iter_mut()
.zip(row1.iter_mut())
.zip(row2.iter_mut())
.zip(row3.iter_mut())
.enumerate()
{
*x0 = mul(i as GfElement, log_m, exp, log);
*x1 = mul((i << 4) as GfElement, log_m, exp, log);
*x2 = mul((i << 8) as GfElement, log_m, exp, log);
*x3 = mul((i << 12) as GfElement, log_m, exp, log);
}
}
mul16.into_boxed_slice().try_into().unwrap()
}
fn initialize_mul128() -> Box<Mul128> {
let exp = &get_exp_log().exp;
let log = &get_exp_log().log;
let mut mul128 = vec![
Multiply128lutT {
lo: [0; 4],
hi: [0; 4],
};
GF_ORDER
];
for log_m in 0..=GF_MODULUS {
for i in 0..=3 {
let mut prod_lo = [0u8; 16];
let mut prod_hi = [0u8; 16];
for x in 0..16 {
let prod = mul((x << (i * 4)) as GfElement, log_m, exp, log);
prod_lo[x] = prod as u8;
prod_hi[x] = (prod >> 8) as u8;
}
mul128[log_m as usize].lo[i] = u128::from_ne_bytes(prod_lo);
mul128[log_m as usize].hi[i] = u128::from_ne_bytes(prod_hi);
}
}
mul128.into_boxed_slice().try_into().unwrap()
}
#[cfg(any(test, target_arch = "x86", target_arch = "x86_64"))]
fn initialize_mul_gfni() -> Box<MulGfni> {
let exp = &get_exp_log().exp;
let log = &get_exp_log().log;
let mut table = vec![
MultiplyGfni {
low_from_low: 0,
low_from_high: 0,
high_from_low: 0,
high_from_high: 0,
};
GF_ORDER
];
let mut basis_matrices = [[0u64; 4]; GF_BITS];
for (coefficient_bit, matrix) in basis_matrices.iter_mut().enumerate() {
let log_m = log[1 << coefficient_bit];
let mut rows = [[0u8; 8]; 4];
for input_bit in 0..16 {
let product = mul(1u16 << input_bit, log_m, exp, log);
for output_bit in 0..16 {
let block = (output_bit / 8) * 2 + input_bit / 8;
rows[block][output_bit % 8] |=
(((product >> output_bit) & 1) as u8) << (input_bit % 8);
}
}
*matrix = rows.map(u64::from_be_bytes);
}
let mut current = [0u64; 4];
for step in 1..GF_ORDER {
let toggled_bit = step.trailing_zeros() as usize;
for (block, basis) in current.iter_mut().zip(basis_matrices[toggled_bit]) {
*block ^= basis;
}
let coefficient = step ^ (step >> 1);
let [low_from_low, low_from_high, high_from_low, high_from_high] = current;
table[log[coefficient] as usize] = MultiplyGfni {
low_from_low,
low_from_high,
high_from_low,
high_from_high,
};
}
let identity = table[0].clone();
table[GF_MODULUS as usize] = identity;
table.into_boxed_slice().try_into().unwrap()
}
fn initialize_skew() -> Box<Skew> {
let exp = &get_exp_log().exp;
let log = &get_exp_log().log;
let mut skew = Box::new([0; GF_MODULUS as usize]);
let mut temp = [0; GF_BITS - 1];
for (i, value) in temp.iter_mut().enumerate() {
*value = 1 << (i + 1);
}
for m in 0..GF_BITS - 1 {
let step: usize = 1 << (m + 1);
skew[(1 << m) - 1] = 0;
for (i, temp_i) in temp.iter().copied().enumerate().skip(m) {
let s: usize = 1 << (i + 1);
let mut j = (1 << m) - 1;
while j < s {
skew[j + s] = skew[j] ^ temp_i;
j += step;
}
}
temp[m] = GF_MODULUS - log[mul(temp[m], log[(temp[m] ^ 1) as usize], exp, log) as usize];
for i in m + 1..GF_BITS - 1 {
let sum = utils::add_mod(log[(temp[i] ^ 1) as usize], temp[m]);
temp[i] = mul(temp[i], sum, exp, log);
}
}
for value in skew.iter_mut() {
*value = log[*value as usize];
}
skew
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mul128_byte_layout() {
let scalar = get_mul16();
for (vector, scalar) in get_mul128().iter().zip(scalar.iter()) {
for ((lo, hi), products) in vector.lo.iter().zip(&vector.hi).zip(scalar) {
let lo = lo.to_ne_bytes();
let hi = hi.to_ne_bytes();
for (index, product) in products.iter().enumerate() {
assert_eq!(lo[index], *product as u8);
assert_eq!(hi[index], (product >> 8) as u8);
}
}
}
}
fn affine(matrix: u64, input: u8) -> u8 {
let matrix = matrix.to_le_bytes();
let mut output = 0;
for output_bit in 0..8 {
output |= ((matrix[7 - output_bit] & input).count_ones() as u8 & 1) << output_bit;
}
output
}
#[test]
fn mul_gfni_matrix_semantics() {
assert_eq!(
core::mem::size_of::<MultiplyGfni>(),
4 * core::mem::size_of::<u64>()
);
assert_eq!(core::mem::size_of::<MulGfni>(), 2 * 1024 * 1024);
let exp_log = get_exp_log();
for (log_m, matrices) in get_mul_gfni().iter().enumerate() {
for input_bit in 0..16 {
let input = 1u16 << input_bit;
let input_lo = input as u8;
let input_hi = (input >> 8) as u8;
let output_lo = affine(matrices.low_from_low, input_lo)
^ affine(matrices.low_from_high, input_hi);
let output_hi = affine(matrices.high_from_low, input_lo)
^ affine(matrices.high_from_high, input_hi);
let actual = output_lo as u16 | (output_hi as u16) << 8;
let expected = mul(input, log_m as GfElement, &exp_log.exp, &exp_log.log);
assert_eq!(actual, expected, "log_m={log_m} input_bit={input_bit}");
}
}
for log_m in [0, GF_MODULUS as usize] {
let identity = &get_mul_gfni()[log_m];
assert_eq!(identity.low_from_low, 0x0102_0408_1020_4080);
assert_eq!(identity.low_from_high, 0);
assert_eq!(identity.high_from_low, 0);
assert_eq!(identity.high_from_high, 0x0102_0408_1020_4080);
}
}
}