use std::env;
use std::fs::File;
use std::io::Write;
use std::path::Path;
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni",
feature = "simd-vsx",
))]
extern crate cc;
const FIELD_SIZE: usize = 256;
const GENERATING_POLYNOMIAL: usize = 29;
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni",
feature = "simd-vsx",
))]
#[derive(Copy, Clone)]
enum SimdCBuildTarget {
Baseline,
ExplicitArch,
}
fn gen_log_table(polynomial: usize) -> [u8; FIELD_SIZE] {
let mut result: [u8; FIELD_SIZE] = [0; FIELD_SIZE];
let mut b: usize = 1;
for log in 0..FIELD_SIZE - 1 {
result[b] = log as u8;
b <<= 1;
if FIELD_SIZE <= b {
b = (b - FIELD_SIZE) ^ polynomial;
}
}
result
}
const EXP_TABLE_SIZE: usize = FIELD_SIZE * 2 - 2;
fn gen_exp_table(log_table: &[u8; FIELD_SIZE]) -> [u8; EXP_TABLE_SIZE] {
let mut result: [u8; EXP_TABLE_SIZE] = [0; EXP_TABLE_SIZE];
for (i, &log_entry) in log_table.iter().enumerate().skip(1) {
let log = log_entry as usize;
result[log] = i as u8;
result[log + FIELD_SIZE - 1] = i as u8;
}
result
}
fn multiply(log_table: &[u8; FIELD_SIZE], exp_table: &[u8; EXP_TABLE_SIZE], a: u8, b: u8) -> u8 {
if a == 0 || b == 0 {
0
} else {
let log_a = log_table[a as usize];
let log_b = log_table[b as usize];
let log_result = log_a as usize + log_b as usize;
exp_table[log_result]
}
}
fn gen_mul_table(
log_table: &[u8; FIELD_SIZE],
exp_table: &[u8; EXP_TABLE_SIZE],
) -> [[u8; FIELD_SIZE]; FIELD_SIZE] {
let mut result: [[u8; FIELD_SIZE]; FIELD_SIZE] = [[0; 256]; 256];
for (a, row) in result.iter_mut().enumerate() {
for (b, cell) in row.iter_mut().enumerate() {
*cell = multiply(log_table, exp_table, a as u8, b as u8);
}
}
result
}
fn gen_mul_table_half(
log_table: &[u8; FIELD_SIZE],
exp_table: &[u8; EXP_TABLE_SIZE],
) -> ([[u8; 16]; FIELD_SIZE], [[u8; 16]; FIELD_SIZE]) {
let mut low: [[u8; 16]; FIELD_SIZE] = [[0; 16]; FIELD_SIZE];
let mut high: [[u8; 16]; FIELD_SIZE] = [[0; 16]; FIELD_SIZE];
for a in 0..low.len() {
for b in 0..low.len() {
let mut result = 0;
if !(a == 0 || b == 0) {
let log_a = log_table[a];
let log_b = log_table[b];
result = exp_table[log_a as usize + log_b as usize];
}
if (b & 0x0F) == b {
low[a][b] = result;
}
if (b & 0xF0) == b {
high[a][b >> 4] = result;
}
}
}
(low, high)
}
macro_rules! write_table {
(1D => $file:ident, $table:ident, $name:expr, $type:expr) => {{
let len = $table.len();
let mut table_str = String::from(format!("pub static {}: [{}; {}] = [", $name, $type, len));
for v in $table.iter() {
let str = format!("{}, ", v);
table_str.push_str(&str);
}
table_str.push_str("];\n");
$file.write_all(table_str.as_bytes()).unwrap();
}};
(2D => $file:ident, $table:ident, $name:expr, $type:expr) => {{
let rows = $table.len();
let cols = $table[0].len();
let mut table_str = String::from(format!(
"pub static {}: [[{}; {}]; {}] = [",
$name, $type, cols, rows
));
for a in $table.iter() {
table_str.push_str("[");
for b in a.iter() {
let str = format!("{}, ", b);
table_str.push_str(&str);
}
table_str.push_str("],\n");
}
table_str.push_str("];\n");
$file.write_all(table_str.as_bytes()).unwrap();
}};
}
fn write_tables() {
let log_table = gen_log_table(GENERATING_POLYNOMIAL);
let exp_table = gen_exp_table(&log_table);
let mul_table = gen_mul_table(&log_table, &exp_table);
let out_dir = env::var("OUT_DIR").unwrap();
let dest_path = Path::new(&out_dir).join("table.rs");
let mut f = File::create(&dest_path).unwrap();
write_table!(1D => f, log_table, "LOG_TABLE", "u8");
write_table!(1D => f, exp_table, "EXP_TABLE", "u8");
write_table!(2D => f, mul_table, "MUL_TABLE", "u8");
if cfg!(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni",
feature = "simd-vsx"
)) {
let (mul_table_low, mul_table_high) = gen_mul_table_half(&log_table, &exp_table);
write_table!(2D => f, mul_table_low, "MUL_TABLE_LOW", "u8");
write_table!(2D => f, mul_table_high, "MUL_TABLE_HIGH", "u8");
}
}
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni"
))]
fn target_cfg(name: &str) -> String {
env::var(name).unwrap_or_default()
}
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni"
))]
fn should_compile_simd_c_for_target() -> bool {
let target_arch = target_cfg("CARGO_CFG_TARGET_ARCH");
let target_env = target_cfg("CARGO_CFG_TARGET_ENV");
let target_os = target_cfg("CARGO_CFG_TARGET_OS");
let arch_supported = matches!(target_arch.as_str(), "x86_64" | "aarch64" | "powerpc64");
let env_supported = target_env != "msvc";
let os_supported = !matches!(target_os.as_str(), "android" | "ios");
arch_supported && env_supported && os_supported
}
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni"
))]
fn is_valid_march_value(arch: &str) -> bool {
!arch.is_empty()
&& arch
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b'+'))
}
#[cfg(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni"
))]
fn compile_simd_c() {
if !should_compile_simd_c_for_target() {
if let Ok(arch) = env::var("RUST_REED_SOLOMON_ERASURE_ARCH") {
println!(
"cargo:warning=ignoring RUST_REED_SOLOMON_ERASURE_ARCH={arch} because simd-c is disabled for this target"
);
}
return;
}
let mut build = cc::Build::new();
build.opt_level(3);
let mut build_target = SimdCBuildTarget::Baseline;
match env::var("RUST_REED_SOLOMON_ERASURE_ARCH") {
Ok(arch) => {
if is_valid_march_value(&arch) {
build.flag(format!("-march={arch}"));
println!("cargo:rustc-env=RSE_SIMD_C_ARCH={arch}");
println!("cargo:rustc-cfg=rse_simd_c_build_unknown");
build_target = SimdCBuildTarget::ExplicitArch;
} else {
println!(
"cargo:warning=invalid RUST_REED_SOLOMON_ERASURE_ARCH value '{arch}', expected [A-Za-z0-9_.+-]+; falling back to baseline simd-c build"
);
}
}
Err(_error) => {}
}
match build_target {
SimdCBuildTarget::Baseline => {
println!("cargo:rustc-cfg=rse_simd_c_build_baseline");
println!("cargo:rustc-env=RSE_SIMD_C_ARCH=baseline");
}
SimdCBuildTarget::ExplicitArch => {}
}
build
.flag("-std=c11")
.file("simd_c/reedsolomon.c")
.compile("reedsolomon");
}
#[cfg(not(any(
feature = "simd-neon",
feature = "simd-ssse3",
feature = "simd-avx2",
feature = "simd-avx512",
feature = "simd-gfni",
feature = "simd-vsx"
)))]
fn compile_simd_c() {}
fn generate_encode_codegen() {
let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default();
if target_arch != "x86_64" && target_arch != "aarch64" {
return;
}
let out_dir = env::var("OUT_DIR").unwrap();
let dest_path = Path::new(&out_dir).join("codegen_encode.rs");
let mut f = File::create(&dest_path).unwrap();
let configs: &[(usize, usize)] = &[(10, 4), (12, 4), (8, 3), (8, 4), (6, 3), (4, 2)];
if target_arch == "x86_64" {
generate_encode_codegen_avx2(&mut f, configs);
} else if target_arch == "aarch64" {
generate_encode_codegen_neon(&mut f, configs);
}
}
fn generate_encode_codegen_avx2(f: &mut File, configs: &[(usize, usize)]) {
use std::io::Write;
writeln!(
f,
"// Auto-generated AVX2 encode functions for common configurations."
)
.unwrap();
writeln!(f, "// Generated by build.rs — do not edit manually.").unwrap();
writeln!(f).unwrap();
writeln!(f, "#[cfg(all(").unwrap();
writeln!(f, " feature = \"simd-avx2\",").unwrap();
writeln!(f, " target_arch = \"x86_64\",").unwrap();
writeln!(f, " not(target_env = \"msvc\"),").unwrap();
writeln!(
f,
" not(any(target_os = \"android\", target_os = \"ios\"))"
)
.unwrap();
writeln!(f, "))]").unwrap();
writeln!(f, "mod avx2_impl {{").unwrap();
writeln!(f, " use core::arch::x86_64::*;").unwrap();
writeln!(f).unwrap();
for &(d, p) in configs {
generate_encode_fn_avx2(f, d, p);
}
writeln!(f, "}}").unwrap();
writeln!(f).unwrap();
writeln!(f, "#[cfg(all(").unwrap();
writeln!(f, " feature = \"simd-avx2\",").unwrap();
writeln!(f, " target_arch = \"x86_64\",").unwrap();
writeln!(f, " not(target_env = \"msvc\"),").unwrap();
writeln!(
f,
" not(any(target_os = \"android\", target_os = \"ios\"))"
)
.unwrap();
writeln!(f, "))]").unwrap();
writeln!(f, "pub(crate) fn try_encode_codegen_avx2(").unwrap();
writeln!(f, " data_shard_count: usize,").unwrap();
writeln!(f, " parity_shard_count: usize,").unwrap();
writeln!(f, " parity_rows: &[&[u8]],").unwrap();
writeln!(f, " data: &[&[u8]],").unwrap();
writeln!(f, " parity: &mut [&mut [u8]],").unwrap();
writeln!(f, " shard_len: usize,").unwrap();
writeln!(f, ") -> bool {{").unwrap();
writeln!(f, " match (data_shard_count, parity_shard_count) {{").unwrap();
for &(d, p) in configs {
writeln!(f, " ({d}, {p}) => unsafe {{").unwrap();
writeln!(
f,
" avx2_impl::encode_{d}x{p}_avx2(parity_rows, data, parity, shard_len);"
)
.unwrap();
writeln!(f, " true").unwrap();
writeln!(f, " }},").unwrap();
}
writeln!(f, " _ => false,").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f, "}}").unwrap();
}
fn generate_encode_fn_avx2(f: &mut File, d: usize, p: usize) {
use std::io::Write;
writeln!(
f,
" /// Specialized encode for {d} data + {p} parity shards using AVX2."
)
.unwrap();
writeln!(f, " ///").unwrap();
writeln!(f, " /// # Safety").unwrap();
writeln!(f, " ///").unwrap();
writeln!(
f,
" /// Caller must ensure AVX2 is available and all slices have length >= shard_len."
)
.unwrap();
writeln!(f, " #[target_feature(enable = \"avx2\")]").unwrap();
writeln!(f, " pub(super) unsafe fn encode_{d}x{p}_avx2(").unwrap();
writeln!(f, " parity_rows: &[&[u8]],").unwrap();
writeln!(f, " data: &[&[u8]],").unwrap();
writeln!(f, " parity: &mut [&mut [u8]],").unwrap();
writeln!(f, " shard_len: usize,").unwrap();
writeln!(f, " ) {{").unwrap();
writeln!(f, " unsafe {{").unwrap();
writeln!(
f,
" let nibble_mask: __m256i = _mm256_set1_epi8(0x0f);"
)
.unwrap();
writeln!(
f,
" // Load GF multiplication table halves for all coefficients."
)
.unwrap();
writeln!(
f,
" // Layout: coef_tables[p_idx][d_idx] = (low_tbl, high_tbl)"
)
.unwrap();
for pi in 0..p {
for di in 0..d {
writeln!(
f,
" let (coef_low_{pi}_{di}, coef_high_{pi}_{di}): (__m256i, __m256i) = {{"
)
.unwrap();
writeln!(f, " let c = parity_rows[{pi}][{di}];").unwrap();
writeln!(
f,
" let (lh, hh) = super::super::load_table_halves(c);"
)
.unwrap();
writeln!(
f,
" let low128: __m128i = _mm_loadu_si128(lh.as_ptr().cast());"
)
.unwrap();
writeln!(
f,
" let high128: __m128i = _mm_loadu_si128(hh.as_ptr().cast());"
)
.unwrap();
writeln!(f, " (_mm256_broadcastsi128_si256(low128), _mm256_broadcastsi128_si256(high128))").unwrap();
writeln!(f, " }};").unwrap();
}
}
writeln!(f).unwrap();
writeln!(f, " let bytes_done = shard_len & !31usize;").unwrap();
writeln!(f).unwrap();
writeln!(
f,
" // Main SIMD loop: process 32 bytes per iteration."
)
.unwrap();
writeln!(f, " let mut offset = 0usize;").unwrap();
writeln!(f, " while offset < bytes_done {{").unwrap();
writeln!(f, " // Load all {d} data shards.").unwrap();
for di in 0..d {
writeln!(f, " let d{di}: __m256i = _mm256_loadu_si256(data[{di}][offset..].as_ptr().cast());").unwrap();
}
writeln!(f).unwrap();
for pi in 0..p {
writeln!(f, " // Compute parity shard {pi}.").unwrap();
writeln!(
f,
" let low = _mm256_and_si256(d0, nibble_mask);"
)
.unwrap();
writeln!(
f,
" let high = _mm256_and_si256(_mm256_srli_epi64::<4>(d0), nibble_mask);"
)
.unwrap();
writeln!(
f,
" let mut acc_{pi}: __m256i = _mm256_xor_si256("
)
.unwrap();
writeln!(
f,
" _mm256_shuffle_epi8(coef_low_{pi}_0, low),"
)
.unwrap();
writeln!(
f,
" _mm256_shuffle_epi8(coef_high_{pi}_0, high),"
)
.unwrap();
writeln!(f, " );").unwrap();
for di in 1..d {
writeln!(
f,
" let low = _mm256_and_si256(d{di}, nibble_mask);"
)
.unwrap();
writeln!(f, " let high = _mm256_and_si256(_mm256_srli_epi64::<4>(d{di}), nibble_mask);").unwrap();
writeln!(
f,
" acc_{pi} = _mm256_xor_si256(acc_{pi}, _mm256_xor_si256("
)
.unwrap();
writeln!(
f,
" _mm256_shuffle_epi8(coef_low_{pi}_{di}, low),"
)
.unwrap();
writeln!(
f,
" _mm256_shuffle_epi8(coef_high_{pi}_{di}, high),"
)
.unwrap();
writeln!(f, " ));").unwrap();
}
writeln!(
f,
" _mm256_storeu_si256(parity[{pi}][offset..].as_mut_ptr().cast(), acc_{pi});"
)
.unwrap();
}
writeln!(f, " offset += 32;").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f).unwrap();
writeln!(f, " // Scalar tail for remaining bytes.").unwrap();
writeln!(f, " for i in bytes_done..shard_len {{").unwrap();
for pi in 0..p {
writeln!(f, " let mut acc: u8 = 0;").unwrap();
for di in 0..d {
writeln!(f, " acc ^= super::super::super::MUL_TABLE[parity_rows[{pi}][{di}] as usize][data[{di}][i] as usize];").unwrap();
}
writeln!(f, " parity[{pi}][i] = acc;").unwrap();
}
writeln!(f, " }}").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f).unwrap();
}
fn generate_encode_codegen_neon(f: &mut File, configs: &[(usize, usize)]) {
use std::io::Write;
writeln!(
f,
"// Auto-generated NEON encode functions for common configurations."
)
.unwrap();
writeln!(f, "// Generated by build.rs — do not edit manually.").unwrap();
writeln!(f).unwrap();
writeln!(f, "#[cfg(all(").unwrap();
writeln!(f, " feature = \"simd-neon\",").unwrap();
writeln!(f, " target_arch = \"aarch64\",").unwrap();
writeln!(f, " not(target_env = \"msvc\"),").unwrap();
writeln!(
f,
" not(any(target_os = \"android\", target_os = \"ios\"))"
)
.unwrap();
writeln!(f, "))]").unwrap();
writeln!(f, "mod neon_impl {{").unwrap();
writeln!(f, " use core::arch::aarch64::*;").unwrap();
writeln!(f).unwrap();
for &(d, p) in configs {
generate_encode_fn_neon(f, d, p);
}
writeln!(f, "}}").unwrap();
writeln!(f).unwrap();
writeln!(f, "#[cfg(all(").unwrap();
writeln!(f, " feature = \"simd-neon\",").unwrap();
writeln!(f, " target_arch = \"aarch64\",").unwrap();
writeln!(f, " not(target_env = \"msvc\"),").unwrap();
writeln!(
f,
" not(any(target_os = \"android\", target_os = \"ios\"))"
)
.unwrap();
writeln!(f, "))]").unwrap();
writeln!(f, "pub(crate) fn try_encode_codegen_neon(").unwrap();
writeln!(f, " data_shard_count: usize,").unwrap();
writeln!(f, " parity_shard_count: usize,").unwrap();
writeln!(f, " parity_rows: &[&[u8]],").unwrap();
writeln!(f, " data: &[&[u8]],").unwrap();
writeln!(f, " parity: &mut [&mut [u8]],").unwrap();
writeln!(f, " shard_len: usize,").unwrap();
writeln!(f, ") -> bool {{").unwrap();
writeln!(f, " match (data_shard_count, parity_shard_count) {{").unwrap();
for &(d, p) in configs {
writeln!(f, " ({d}, {p}) => unsafe {{").unwrap();
writeln!(
f,
" neon_impl::encode_{d}x{p}_neon(parity_rows, data, parity, shard_len);"
)
.unwrap();
writeln!(f, " true").unwrap();
writeln!(f, " }},").unwrap();
}
writeln!(f, " _ => false,").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f, "}}").unwrap();
}
fn generate_encode_fn_neon(f: &mut File, d: usize, p: usize) {
use std::io::Write;
writeln!(
f,
" /// Specialized encode for {d} data + {p} parity shards using NEON."
)
.unwrap();
writeln!(f, " ///").unwrap();
writeln!(f, " /// # Safety").unwrap();
writeln!(f, " ///").unwrap();
writeln!(
f,
" /// Caller must ensure NEON is available and all slices have length >= shard_len."
)
.unwrap();
writeln!(f, " #[target_feature(enable = \"neon\")]").unwrap();
writeln!(f, " pub(super) unsafe fn encode_{d}x{p}_neon(").unwrap();
writeln!(f, " parity_rows: &[&[u8]],").unwrap();
writeln!(f, " data: &[&[u8]],").unwrap();
writeln!(f, " parity: &mut [&mut [u8]],").unwrap();
writeln!(f, " shard_len: usize,").unwrap();
writeln!(f, " ) {{").unwrap();
writeln!(f, " let nibble_mask: uint8x16_t = vdupq_n_u8(0x0f);").unwrap();
writeln!(
f,
" // Load GF multiplication table halves for all coefficients."
)
.unwrap();
for pi in 0..p {
for di in 0..d {
writeln!(f, " let (coef_low_{pi}_{di}, coef_high_{pi}_{di}): (uint8x16_t, uint8x16_t) = unsafe {{").unwrap();
writeln!(f, " let c = parity_rows[{pi}][{di}];").unwrap();
writeln!(
f,
" (vld1q_u8(super::super::super::MUL_TABLE_LOW[c as usize].as_ptr()),"
)
.unwrap();
writeln!(
f,
" vld1q_u8(super::super::super::MUL_TABLE_HIGH[c as usize].as_ptr()))"
)
.unwrap();
writeln!(f, " }};").unwrap();
}
}
writeln!(f).unwrap();
writeln!(f, " let bytes_done = shard_len & !15usize;").unwrap();
writeln!(f).unwrap();
writeln!(
f,
" // Main SIMD loop: process 16 bytes per iteration."
)
.unwrap();
writeln!(f, " let mut offset = 0usize;").unwrap();
writeln!(f, " while offset < bytes_done {{").unwrap();
writeln!(f, " unsafe {{").unwrap();
for di in 0..d {
writeln!(
f,
" let d{di}: uint8x16_t = vld1q_u8(data[{di}][offset..].as_ptr());"
)
.unwrap();
}
writeln!(f).unwrap();
for pi in 0..p {
writeln!(f, " // Compute parity shard {pi}.").unwrap();
writeln!(f, " let low = vandq_u8(d0, nibble_mask);").unwrap();
writeln!(f, " let high = vshrq_n_u8::<4>(d0);").unwrap();
writeln!(f, " let mut acc_{pi}: uint8x16_t = veorq_u8(").unwrap();
writeln!(f, " vqtbl1q_u8(coef_low_{pi}_0, low),").unwrap();
writeln!(f, " vqtbl1q_u8(coef_high_{pi}_0, high),").unwrap();
writeln!(f, " );").unwrap();
for di in 1..d {
writeln!(f, " let low = vandq_u8(d{di}, nibble_mask);").unwrap();
writeln!(f, " let high = vshrq_n_u8::<4>(d{di});").unwrap();
writeln!(f, " acc_{pi} = veorq_u8(acc_{pi}, veorq_u8(").unwrap();
writeln!(f, " vqtbl1q_u8(coef_low_{pi}_{di}, low),").unwrap();
writeln!(f, " vqtbl1q_u8(coef_high_{pi}_{di}, high),").unwrap();
writeln!(f, " ));").unwrap();
}
writeln!(
f,
" vst1q_u8(parity[{pi}][offset..].as_mut_ptr(), acc_{pi});"
)
.unwrap();
}
writeln!(f, " }}").unwrap(); writeln!(f, " offset += 16;").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f).unwrap();
writeln!(f, " // Scalar tail for remaining bytes.").unwrap();
writeln!(f, " for i in bytes_done..shard_len {{").unwrap();
for pi in 0..p {
writeln!(f, " let mut acc: u8 = 0;").unwrap();
for di in 0..d {
writeln!(f, " acc ^= super::super::super::MUL_TABLE[parity_rows[{pi}][{di}] as usize][data[{di}][i] as usize];").unwrap();
}
writeln!(f, " parity[{pi}][i] = acc;").unwrap();
}
writeln!(f, " }}").unwrap();
writeln!(f, " }}").unwrap();
writeln!(f).unwrap();
}
fn main() {
println!("cargo:rerun-if-env-changed=RUST_REED_SOLOMON_ERASURE_ARCH");
println!("cargo:rerun-if-changed=simd_c/reedsolomon.c");
println!("cargo:rerun-if-changed=simd_c/reedsolomon.h");
println!("cargo:rustc-check-cfg=cfg(rse_simd_c_build_baseline)");
println!("cargo:rustc-check-cfg=cfg(rse_simd_c_build_unknown)");
compile_simd_c();
write_tables();
generate_encode_codegen();
}