plat-core 0.1.1

//! Polynomial arithmetic in the ring Z_q[X]/(X^N + 1).
//!
//! All operations are modular arithmetic modulo a prime q,
//! in the negacyclic polynomial ring of dimension N.

use crate::modular::{mod_add, mod_mul, mod_neg, mod_sub};
use crate::params::Params;

/// A polynomial in Z_q[X]/(X^N + 1), stored as coefficient vector.
#[derive(Clone, Debug)]
pub struct Poly {
    /// Coefficients in [0, q). Length must equal params.n.
    pub coeffs: Vec<u64>,
}

impl Poly {
    /// Create the zero polynomial of dimension n.
    pub fn zero(n: usize) -> Self {
        Self {
            coeffs: vec![0u64; n],
        }
    }

    /// Create a polynomial from coefficients (clamped to n, zero-padded if short).
    pub fn from_coeffs(coeffs: &[u64], n: usize) -> Self {
        let mut c = vec![0u64; n];
        let len = coeffs.len().min(n);
        c[..len].copy_from_slice(&coeffs[..len]);
        Self { coeffs: c }
    }

    /// Create a constant polynomial (value in slot 0).
    pub fn constant(value: u64, n: usize) -> Self {
        let mut c = vec![0u64; n];
        c[0] = value;
        Self { coeffs: c }
    }

    /// Polynomial addition: (a + b) mod q, coefficient-wise.
    pub fn add(&self, other: &Poly, params: &Params) -> Poly {
        debug_assert_eq!(self.coeffs.len(), other.coeffs.len());
        let coeffs = self
            .coeffs
            .iter()
            .zip(&other.coeffs)
            .map(|(&a, &b)| mod_add(a, b, params.q))
            .collect();
        Poly { coeffs }
    }

    /// Polynomial subtraction: (a - b) mod q.
    pub fn sub(&self, other: &Poly, params: &Params) -> Poly {
        debug_assert_eq!(self.coeffs.len(), other.coeffs.len());
        let coeffs = self
            .coeffs
            .iter()
            .zip(&other.coeffs)
            .map(|(&a, &b)| mod_sub(a, b, params.q))
            .collect();
        Poly { coeffs }
    }

    /// Polynomial negation: (-a) mod q.
    pub fn neg(&self, params: &Params) -> Poly {
        let coeffs = self.coeffs.iter().map(|&a| mod_neg(a, params.q)).collect();
        Poly { coeffs }
    }

    /// Scalar multiplication: (a * scalar) mod q.
    pub fn scalar_mul(&self, scalar: u64, params: &Params) -> Poly {
        let coeffs = self
            .coeffs
            .iter()
            .map(|&a| mod_mul(a, scalar, params.q))
            .collect();
        Poly { coeffs }
    }

    /// Schoolbook polynomial multiplication in Z_q[X]/(X^N + 1).
    /// This is O(N^2) — used for correctness testing and small N.
    pub fn mul_schoolbook(&self, other: &Poly, params: &Params) -> Poly {
        let n = params.n;
        let q = params.q;
        let mut result = vec![0u64; n];

        for i in 0..n {
            if self.coeffs[i] == 0 {
                continue;
            }
            for j in 0..n {
                if other.coeffs[j] == 0 {
                    continue;
                }
                let prod = mod_mul(self.coeffs[i], other.coeffs[j], q);
                let idx = i + j;
                if idx < n {
                    result[idx] = mod_add(result[idx], prod, q);
                } else {
                    // X^N ≡ -X^0 in negacyclic ring
                    let idx = idx - n;
                    result[idx] = mod_sub(result[idx], prod, q);
                }
            }
        }

        Poly { coeffs: result }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn tiny() -> Params {
        Params::test_tiny()
    }

    #[test]
    fn test_add_sub() {
        let p = tiny();
        let a = Poly::from_coeffs(&[1, 2, 3], p.n);
        let b = Poly::from_coeffs(&[4, 5, 6], p.n);
        let sum = a.add(&b, &p);
        assert_eq!(sum.coeffs[0], 5);
        assert_eq!(sum.coeffs[1], 7);
        assert_eq!(sum.coeffs[2], 9);

        let diff = sum.sub(&b, &p);
        assert_eq!(diff.coeffs[0], 1);
        assert_eq!(diff.coeffs[1], 2);
        assert_eq!(diff.coeffs[2], 3);
    }

    #[test]
    fn test_negacyclic_mul() {
        let p = tiny();
        // (1 + X) * (1 + X) = 1 + 2X + X^2 in normal ring
        // but X^N wraps, so for small coefficients this is fine
        let a = Poly::from_coeffs(&[1, 1], p.n);
        let b = Poly::from_coeffs(&[1, 1], p.n);
        let c = a.mul_schoolbook(&b, &p);
        assert_eq!(c.coeffs[0], 1);
        assert_eq!(c.coeffs[1], 2);
        assert_eq!(c.coeffs[2], 1);
    }

    #[test]
    fn test_negacyclic_wraparound() {
        // In Z_q[X]/(X^N + 1), X^{N-1} * X = X^N = -1
        let p = tiny();
        let n = p.n;
        // a = X^{N-1}
        let mut a_coeffs = vec![0u64; n];
        a_coeffs[n - 1] = 1;
        let a = Poly { coeffs: a_coeffs };
        // b = X
        let b = Poly::from_coeffs(&[0, 1], n);
        let c = a.mul_schoolbook(&b, &p);
        // X^N = -1 mod (X^N + 1), so constant term should be q-1 (-1 mod q)
        assert_eq!(c.coeffs[0], p.q - 1);
        for i in 1..n {
            assert_eq!(c.coeffs[i], 0);
        }
    }

    #[test]
    fn test_scalar_mul() {
        let p = tiny();
        let a = Poly::from_coeffs(&[3, 5, 7], p.n);
        let b = a.scalar_mul(10, &p);
        assert_eq!(b.coeffs[0], 30);
        assert_eq!(b.coeffs[1], 50);
        assert_eq!(b.coeffs[2], 70);
    }
}