#[allow(unused_imports)]
use crate::prelude::*;
use core::ops::{Add, Mul, Sub};
use num_traits::{One, Zero};
pub fn simd_poly_add<T>(a_coeffs: &[T], b_coeffs: &[T]) -> Vec<T>
where
T: Clone + Add<Output = T> + Zero,
{
let max_len = a_coeffs.len().max(b_coeffs.len());
let mut result = Vec::with_capacity(max_len);
for i in 0..max_len {
let a_val = a_coeffs.get(i).cloned().unwrap_or_else(T::zero);
let b_val = b_coeffs.get(i).cloned().unwrap_or_else(T::zero);
result.push(a_val + b_val);
}
while result.len() > 1 && result.last().is_some_and(|c| c.is_zero()) {
result.pop();
}
result
}
pub fn simd_poly_mul<T>(a_coeffs: &[T], b_coeffs: &[T]) -> Vec<T>
where
T: Clone + Add<Output = T> + Mul<Output = T> + Zero,
{
if a_coeffs.is_empty() || b_coeffs.is_empty() {
return vec![T::zero()];
}
let result_len = a_coeffs.len() + b_coeffs.len() - 1;
let mut result = vec![T::zero(); result_len];
for (i, a_coeff) in a_coeffs.iter().enumerate() {
for (j, b_coeff) in b_coeffs.iter().enumerate() {
result[i + j] = result[i + j].clone() + a_coeff.clone() * b_coeff.clone();
}
}
result
}
pub fn simd_poly_eval<T>(coeffs: &[T], x: &T) -> T
where
T: Clone + Add<Output = T> + Mul<Output = T> + Zero,
{
if coeffs.is_empty() {
return T::zero();
}
let mut result = coeffs
.last()
.expect("collection should not be empty")
.clone();
for coeff in coeffs.iter().rev().skip(1) {
result = coeff.clone() + result * x.clone();
}
result
}
pub fn simd_poly_eval_multi<T>(coeffs: &[T], points: &[T]) -> Vec<T>
where
T: Clone + Add<Output = T> + Mul<Output = T> + Zero + Send + Sync,
{
points.iter().map(|x| simd_poly_eval(coeffs, x)).collect()
}
pub fn simd_poly_derivative<T>(coeffs: &[T]) -> Vec<T>
where
T: Clone + Mul<Output = T> + From<usize> + Zero,
{
if coeffs.len() <= 1 {
return vec![T::zero()];
}
coeffs
.iter()
.enumerate()
.skip(1)
.map(|(i, c)| c.clone() * T::from(i))
.collect()
}
pub fn simd_poly_scale<T>(coeffs: &[T], scalar: &T) -> Vec<T>
where
T: Clone + Mul<Output = T>,
{
coeffs.iter().map(|c| c.clone() * scalar.clone()).collect()
}
pub fn simd_poly_compose<T>(p_coeffs: &[T], q_coeffs: &[T]) -> Vec<T>
where
T: Clone + Add<Output = T> + Mul<Output = T> + Zero,
{
if p_coeffs.is_empty() {
return vec![T::zero()];
}
let mut result = vec![p_coeffs[0].clone()];
let mut q_power = q_coeffs.to_vec();
for p_coeff in p_coeffs.iter().skip(1) {
let term = simd_poly_scale(&q_power, p_coeff);
result = simd_poly_add(&result, &term);
q_power = simd_poly_mul(&q_power, q_coeffs);
}
result
}
pub struct PolyEvaluator<T> {
coeffs: Vec<T>,
#[allow(dead_code)]
precomputed_powers: Vec<Vec<T>>,
}
impl<T> PolyEvaluator<T>
where
T: Clone + Add<Output = T> + Mul<Output = T> + Zero + One,
{
pub fn new(coeffs: Vec<T>) -> Self {
Self {
coeffs,
precomputed_powers: Vec::new(),
}
}
pub fn eval(&self, x: &T) -> T {
simd_poly_eval(&self.coeffs, x)
}
pub fn eval_batch(&self, points: &[T]) -> Vec<T>
where
T: Send + Sync,
{
simd_poly_eval_multi(&self.coeffs, points)
}
}
pub fn poly_add_coeffs(a: &[i64], b: &[i64]) -> Vec<i64> {
let len = a.len().max(b.len());
let mut result = Vec::with_capacity(len);
for i in 0..len {
let av = *a.get(i).unwrap_or(&0);
let bv = *b.get(i).unwrap_or(&0);
result.push(av + bv);
}
result
}
pub fn poly_sub_coeffs(a: &[i64], b: &[i64]) -> Vec<i64> {
let len = a.len().max(b.len());
let mut result = Vec::with_capacity(len);
for i in 0..len {
let av = *a.get(i).unwrap_or(&0);
let bv = *b.get(i).unwrap_or(&0);
result.push(av - bv);
}
result
}
pub fn poly_mul_scalar(coeffs: &[i64], scalar: i64) -> Vec<i64> {
coeffs.iter().map(|&c| c * scalar).collect()
}
pub fn poly_dot_product(a: &[i64], b: &[i64]) -> Result<i64, i128> {
assert_eq!(
a.len(),
b.len(),
"poly_dot_product: slices must have equal length"
);
let sum: i128 = a
.iter()
.zip(b.iter())
.map(|(&x, &y)| (x as i128) * (y as i128))
.sum();
i64::try_from(sum).map_err(|_| sum)
}
pub fn poly_eval_horner_i64(coeffs: &[i64], x: i64) -> i64 {
if coeffs.is_empty() {
return 0;
}
let mut result = *coeffs.last().expect("non-empty slice has a last element");
for &c in coeffs.iter().rev().skip(1) {
result = c + result * x;
}
result
}
pub fn simd_poly_gcd<T>(mut a: Vec<T>, mut b: Vec<T>) -> Vec<T>
where
T: Clone + Add<Output = T> + Sub<Output = T> + Mul<Output = T> + Zero + PartialEq,
{
while !is_zero_poly(&b) {
let (_, remainder) = poly_divide(&a, &b);
a = b;
b = remainder;
}
a
}
fn is_zero_poly<T>(coeffs: &[T]) -> bool
where
T: Zero + PartialEq,
{
coeffs.iter().all(|c| c.is_zero())
}
fn poly_divide<T>(dividend: &[T], _divisor: &[T]) -> (Vec<T>, Vec<T>)
where
T: Clone + Add<Output = T> + Sub<Output = T> + Mul<Output = T> + Zero,
{
(vec![T::zero()], dividend.to_vec())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_simd_poly_add() {
let a = vec![1.0, 2.0, 3.0]; let b = vec![4.0, 5.0]; let result = simd_poly_add(&a, &b);
assert_eq!(result, vec![5.0, 7.0, 3.0]);
}
#[test]
fn test_simd_poly_mul() {
let a = vec![1.0, 2.0]; let b = vec![3.0, 4.0]; let result = simd_poly_mul(&a, &b);
assert_eq!(result, vec![3.0, 10.0, 8.0]);
}
#[test]
fn test_simd_poly_eval() {
let coeffs = vec![1.0, 2.0, 3.0]; let result = simd_poly_eval(&coeffs, &2.0);
assert_eq!(result, 17.0);
}
#[test]
fn test_simd_poly_derivative() {
use num_bigint::BigInt;
let coeffs: Vec<BigInt> = vec![
BigInt::from(1),
BigInt::from(2),
BigInt::from(3),
BigInt::from(4),
]; let deriv = simd_poly_derivative(&coeffs);
assert_eq!(
deriv,
vec![BigInt::from(2), BigInt::from(6), BigInt::from(12)]
);
}
#[test]
fn test_simd_poly_scale() {
let coeffs = vec![1.0, 2.0, 3.0];
let result = simd_poly_scale(&coeffs, &2.0);
assert_eq!(result, vec![2.0, 4.0, 6.0]);
}
#[test]
fn test_poly_evaluator() {
let coeffs = vec![1.0, 2.0, 3.0];
let evaluator = PolyEvaluator::new(coeffs);
let result = evaluator.eval(&2.0);
assert_eq!(result, 17.0);
let batch = evaluator.eval_batch(&[1.0, 2.0, 3.0]);
assert_eq!(batch, vec![6.0, 17.0, 34.0]);
}
#[test]
fn test_simd_poly_compose() {
let p = vec![1.0, 0.0, 1.0]; let q = vec![0.0, 2.0]; let result = simd_poly_compose(&p, &q);
assert_eq!(result, vec![1.0, 0.0, 4.0]);
}
}